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

import argparse
import json
import os
from pathlib import Path

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


def json_mat_to_numpy(json_mat: dict):
    numpy_mat = np.array(
        [
            [
                json_mat["e00"],
                json_mat["e01"],
                json_mat["e02"],
                json_mat["e03"],
            ],
            [
                json_mat["e10"],
                json_mat["e11"],
                json_mat["e12"],
                json_mat["e13"],
            ],
            [
                json_mat["e20"],
                json_mat["e21"],
                json_mat["e22"],
                json_mat["e23"],
            ],
            [
                json_mat["e30"],
                json_mat["e31"],
                json_mat["e32"],
                json_mat["e33"],
            ],
        ]
    )
    return numpy_mat


def get_params(camera_params_file: Path, is_interdigital: bool):

    camera_params_per_frame = {}
    camera_params_per_frame_np = {}
    with open(camera_params_file) as file:
        for idx, line in enumerate(file):
            if idx == 0 and is_interdigital:
                assert "Version" in line, "We should not skip this line!"
                continue
            camera_params = json.loads(line)
            if is_interdigital:
                frame = camera_params["frame"]
            else:
                frame = camera_params["frames"]
            camera_params_per_frame[frame] = camera_params
            projectionMatrix = camera_params["projectionMatrix"]
            worldToCameraMatrix = camera_params["worldToCameraMatrix"]
            camera_params_per_frame_np[frame] = {}
            if is_interdigital:
                camera_params_per_frame_np[frame]["projectionMatrix"] = (
                    json_mat_to_numpy(projectionMatrix)
                )
                camera_params_per_frame_np[frame]["worldToCameraMatrix"] = (
                    json_mat_to_numpy(worldToCameraMatrix)
                )
            else:
                camera_params_per_frame_np[frame]["projectionMatrix"] = (
                    json_mat_to_numpy(projectionMatrix).T
                )
                camera_params_per_frame_np[frame]["worldToCameraMatrix"] = (
                    json_mat_to_numpy(worldToCameraMatrix)
                ).T

    return camera_params_per_frame, camera_params_per_frame_np


# 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"
    camera_params_per_frame, camera_params_per_frame_np = get_params(
        Path(args.cam_params), is_interdigital
    )

    Z_far = camera_params_per_frame[99]["farClipPlane"]
    Z_near = camera_params_per_frame[99]["nearClipPlane"]

    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,
                )

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

                # for now only support constant intrinsics!
                assert (
                    camera_params_per_frame_np[ref_frame]["projectionMatrix"]
                    == camera_params_per_frame_np[pred_frame]["projectionMatrix"]
                ).all()
                # assert (camera_params_per_frame_np[ref_frame]['InvProjectionMatrix'] == camera_params_per_frame_np[pred_frame]['InvProjectionMatrix']).all()
                K = camera_params_per_frame_np[ref_frame]["projectionMatrix"]
                Kinv = np.linalg.inv(K)

                if is_interdigital:
                    a = -(2 * Z_far * Z_near) / (Z_far - Z_near)
                    c = -(Z_far + Z_near) / (Z_far - Z_near)
                    Zndc = 2 * Z_buf_engine - 1
                    Z = a / (Zndc - c)
                    Zndc = -Zndc
                else:
                    a = Z_near
                    c = 0
                    Zndc = Z_buf_engine
                    Z = a / (Zndc - c)

                points2D = np.mgrid[0:im_height, 0:im_width]
                points2DandZndc_h = np.concatenate(
                    (points2D, Zndc[None, :], np.ones_like(Zndc[None, :])), axis=0
                )
                # construction avoided loops, but x and y are in wrong order, need to swap them:
                points2DandZndc_h[[0, 1], :] = points2DandZndc_h[[1, 0], :]

                # adjust x and y to be in NDC coords -1 to 1
                points2DandZndc_h[0, :] = (
                    points2DandZndc_h[0, :] + 0.5
                ) * 2 / im_width - 1
                points2DandZndc_h[1, :] = (
                    points2DandZndc_h[1, :] + 0.5
                ) * 2 / im_height - 1
                # invert y-axis
                points2DandZndc_h[1, :] = -1 * points2DandZndc_h[1, :]

                points3D_h = np.einsum("ij,jkl->ikl", Kinv, points2DandZndc_h)
                points3D_h_n = points3D_h / points3D_h[3, :]  # normalized

                assert np.allclose(
                    points3D_h_n[2, ...].flatten(), Z.flatten()
                ), "Mismatch between depth calculated using Znear/Zfar and from matrix"

                p1 = camera_params_per_frame_np[pred_frame]["worldToCameraMatrix"]
                p1inv = np.linalg.inv(p1)
                p2 = camera_params_per_frame_np[ref_frame]["worldToCameraMatrix"]
                p1top2 = p2 @ p1inv
                points3D_h_n_other_frame = np.einsum(
                    "ij,jkl->ikl", p1top2, points3D_h_n
                )
                points2DandZndc_h_other_frame = np.einsum(
                    "ij,jkl->ikl", K, points3D_h_n_other_frame
                )
                points2DandZndc_h_other_frame_n = (
                    points2DandZndc_h_other_frame / points2DandZndc_h_other_frame[3]
                )

                gridY = (
                    points2DandZndc_h_other_frame_n[: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(
        "--cam_params",
        required=True,
        help="""Path of the camara parameters file""",
    )
    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)
