#!/usr/bin/env python
# coding: utf-8

import argparse
import os
from pathlib import Path

os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"
import cv2
import numpy as np
import yuvio


# grid sample for numpy, from
# https://stackoverflow.com/questions/61574143/sampling-numpy-array-with-float-indices-similar-to-pytorch-grid-sample


def grid_sample(tensor, grid):
    """Given an input and a flow-field grid, computes the output using input
    values and pixel locations from grid.

    Args:
        tensor: (N, C, H_in, W_in) tensor
        grid: (N, H_out, W_out, 2) tensor in the range of [-1, 1]

    Returns:
        (N, C, H_out, W_out) tensor

    See `torch.nn.functional.grid_sample`.
    """
    b, c, h, w = tensor.shape
    b_, h_out, w_out, w_ = grid.shape
    assert b == b_ and w_ == 2
    out = []
    for t, g in zip(tensor, grid):
        x_ = 0.5 * (w - 1) * (g[..., 0].reshape(-1) + 1)
        y_ = 0.5 * (h - 1) * (g[..., 1].reshape(-1) + 1)
        ix = np.floor(x_).astype(np.int32).clip(0, w - 2)
        iy = np.floor(y_).astype(np.int32).clip(0, h - 2)
        dx = x_ - ix
        dy = y_ - iy
        out.append(
            (1 - dx) * (1 - dy) * t[..., iy, ix]
            + dx * (1 - dy) * t[..., iy, ix + 1]
            + (1 - dx) * dy * t[..., iy + 1, ix]
            + dx * dy * t[..., iy + 1, ix + 1]
        )
    return np.concatenate(out, axis=0).reshape(b, c, h_out, w_out)


# DarkTree


def main(args: argparse.Namespace):

    example_exr_data = cv2.imread(
        args.first_exr_file,
        cv2.IMREAD_UNCHANGED,
    )

    frame_height, frame_width, _ = example_exr_data.shape

    reader = yuvio.get_reader(
        args.input_yuv,
        frame_width,
        frame_height,
        "yuv420p10le",
    )

    example_exr_file = Path(args.first_exr_file)
    exr_dir = example_exr_file.parent

    is_interdigital = args.type == "interdigital"

    with open(
        args.output_yuv,
        "wb",
    ) as add_ref_file:

        for frame_idx in range(600):
            print(f"Generating frame {frame_idx}")

            ref_frame = frame_idx - 1
            pred_frame = frame_idx

            yuv_frame = reader.read(pred_frame, 1)[0]
            pred_Y = yuv_frame.y
            pred_U = yuv_frame.u
            pred_V = yuv_frame.v

            if frame_idx > 0:

                yuv_frame = reader.read(ref_frame, 1)[0]
                ref_Y = yuv_frame.y
                ref_U = yuv_frame.u
                ref_V = yuv_frame.v

                luma_height, luma_width = ref_Y.shape
                chroma_height, chroma_width = ref_U.shape

                exr_file_this_frame = exr_dir / (
                    f"{example_exr_file.name.split('frame')[0]}frame{pred_frame}.exr"
                )

                exr_data = cv2.imread(
                    exr_file_this_frame.as_posix(),
                    cv2.IMREAD_UNCHANGED,
                )

                mv = exr_data[:, :, :2]

                Z_buf_engine = exr_data[:, :, 2]
                im_height, im_width = Z_buf_engine.shape
                assert (luma_height, luma_width) == (im_height, im_width)

                points2D = np.mgrid[0:im_height, 0:im_width]

                # construction avoided loops, but x and y are in wrong order, need to swap them:
                points2D[[0, 1], :] = points2D[[1, 0], :]

                mv = mv.transpose(2, 0, 1)
                mv[[0, 1], :] = mv[[1, 0], :]
                points2D = points2D + 0.5
                if is_interdigital:
                    mv[0, ...] = -1 * mv[0, ...]
                    points2D_other_frame = points2D + mv
                    points2D_other_frame[0, ...] = (
                        points2D_other_frame[0, ...] / im_width
                    )
                    points2D_other_frame[1, ...] = (
                        points2D_other_frame[1, ...] / im_height
                    )
                else:
                    mv[0, ...] = mv[0, ...] * im_width
                    mv[1, ...] = mv[1, ...] * im_height

                    points2D_other_frame = points2D - mv

                    points2D_other_frame[0, ...] = (
                        points2D_other_frame[0, ...] / im_width
                    )
                    points2D_other_frame[1, ...] = (
                        points2D_other_frame[1, ...] / im_height
                    )

                # from 0 ... 1 to -1 ... 1
                points2D_other_frame = points2D_other_frame * 2 - 1
                gridY = (
                    points2D_other_frame[:2, :]
                    .copy()
                    .reshape(2, im_height, im_width, 1)
                    .copy()
                )
                gridY = gridY.transpose(3, 1, 2, 0).copy()

                # clip stuff outside -1...1
                gridY = np.clip(gridY, -1, 1)
                # invert y-axis back
                # gridY[..., 1] = -1 * gridY[..., 1]
                gridUV = gridY[:, ::2, ::2, :]

                new_Y = grid_sample(
                    ref_Y.reshape(1, 1, luma_height, luma_width), gridY
                ).reshape(luma_height, luma_width)
                new_U = grid_sample(
                    ref_U.reshape(1, 1, chroma_height, chroma_width), gridUV
                ).reshape(chroma_height, chroma_width)
                new_V = grid_sample(
                    ref_V.reshape(1, 1, chroma_height, chroma_width), gridUV
                ).reshape(chroma_height, chroma_width)

            if frame_idx == 0:
                new_Y = np.ones_like(pred_Y) * 512
                new_U = np.ones_like(pred_U) * 512
                new_V = np.ones_like(pred_V) * 512
                new_frame = yuvio.frame((new_Y, new_U, new_V), "yuv420p10le")
            else:
                new_frame = yuvio.frame((new_Y, new_U, new_V), "yuv420p10le")
            yuvio.imwrite(add_ref_file, new_frame)


if __name__ == "__main__":

    parser = argparse.ArgumentParser()
    parser.add_argument("--input_yuv", required=True)
    parser.add_argument(
        "--first_exr_file",
        required=True,
        help="""Path on one of the exr files. Will be modified to take exr 
        of correct frame, i.e. replacing POC in filename given here""",
    )
    parser.add_argument(
        "--output_yuv",
        required=True,
        help="Path of generated yuv file.",
    )
    parser.add_argument(
        "--type",
        required=True,
        choices=["huawei", "interdigital"],
        help="which auxiliary data format to use.",
    )

    args = parser.parse_args()

    main(args)
