import math
import functools
import numpy as np
from skimage import io
from skimage import util
from skimage import color
from skimage import transform
import matplotlib.pyplot as plt


def normalize(v):
    return v / np.linalg.norm(v)


def world_to_camera_with_pose(view_pose):
    lookat_pose = view_pose[3:]
    camera_pose = view_pose[:3]
    up = np.array([0, 1, 0])
    R = np.diag(np.ones(4))
    R[2, :3] = normalize(lookat_pose - camera_pose)
    R[0, :3] = normalize(np.cross(R[2, :3], up))
    R[1, :3] = -normalize(np.cross(R[0, :3], R[2, :3]))
    T = np.diag(np.ones(4))
    T[:3, 3] = -camera_pose
    return R.dot(T)


def camera_intrinsic_transform(vfov=45, hfov=60, pixel_width=320, pixel_height=240):
    camera_intrinsics = np.zeros((4, 4))
    camera_intrinsics[2, 2] = 1
    camera_intrinsics[0, 0] = (pixel_width / 2.0) / math.tan(math.radians(hfov / 2.0))
    camera_intrinsics[0, 2] = pixel_width / 2.0
    camera_intrinsics[1, 1] = (pixel_height / 2.0) / math.tan(math.radians(vfov / 2.0))
    camera_intrinsics[1, 2] = pixel_height / 2.0
    camera_intrinsics[3, 3] = 1
    return camera_intrinsics


if __name__ == "__main__":
    H, W = 240, 320
    K = camera_intrinsic_transform(pixel_width=W, pixel_height=H)

    dst_vec = np.loadtxt("data/125.txt")  # [6]
    dst_Rt = world_to_camera_with_pose(dst_vec)
    dst_P = K @ dst_Rt
    dst_img = util.img_as_float(io.imread("data/125.jpg"))  # [H, W, 3]
    dst_dep = io.imread("data/125.png").astype(np.float32)  # [H, W]
    dst_dep = 1 / (dst_dep * 0.001 + 1e-8)  # depth -> disparity

    src_vec = np.loadtxt("data/0.txt")
    src_Rt = world_to_camera_with_pose(src_vec)
    src_P = K @ src_Rt
    src_img = util.img_as_float(io.imread("data/0.jpg"))  # [H, W, 3]
    src_dep = io.imread("data/0.png").astype(np.float32)  # [H, W]
    src_dep = 1 / (src_dep * 0.001 + 1e-8)  # depth -> disparity

    M = src_P @ np.linalg.pinv(dst_P)

    dst_xx, dst_yy = np.meshgrid(np.arange(W), np.arange(H))
    dst_xx, dst_yy = dst_xx.ravel(), dst_yy.ravel()  # [H * W], [H * W]
    dst_coords = np.stack(
        [dst_xx, dst_yy, np.ones_like(dst_xx), dst_dep[dst_yy, dst_xx]]
    )  # [4, W * H]
    src_coords = M @ dst_coords # [4, W * H]
    src_coords[:2] = src_coords[:2] / (src_coords[2] + 1e-8) # homogeneous normalize
    src_xx, src_yy = src_coords[:2]

    coords = np.stack([src_yy, src_xx], axis=0) # [2, H * W]
    coords = coords.reshape(2, H, W) # coords[:, out_r, out_c] = (inp_r, inp_c)
    out_img = np.stack([
        transform.warp(src_img[..., 0], coords), # red
        transform.warp(src_img[..., 1], coords), # green
        transform.warp(src_img[..., 2], coords), # blue
    ], axis=2)

    fig, ax = plt.subplots(2, 2)
    ax = ax.ravel()
    ax[0].set_title('dst')
    ax[1].set_title('src')
    ax[2].set_title('out')
    ax[3].set_title('diff')
    ax[0].imshow(dst_img)
    ax[1].imshow(src_img)
    ax[2].imshow(out_img)
    ax[3].imshow(np.max(dst_img - out_img, axis=2))
    plt.show()