import os
import torch
import json
import open3d
import toolz
import numpy as np
from PIL import Image
from tqdm import tqdm, trange
from scipy.ndimage import map_coordinates

from misc.post_proc import np_coor2xy, np_coorx2u, np_coory2v
from misc.panostretch import pano_connect_points
from eval_general import layout_2_depth
import glob
import tqdm
import kornia
import cv2

import math
import numpy.matlib as matlib
import typing


dst = torch.Tensor([[
    [-1, 1],
    [-1, -1],
    [1, -1],
    [1, 1], 
]]).float()

colors = [
    [255, 0, 0],    # blue - floor
    [0, 255, 0],    # green - ceil
    [0, 0, 255],    # red - proj floor
    [255, 0, 255],  # purple - proj ceil
    [255, 255, 0],  # cyan - proj
    [55, 255, 100],  # ??? - ceil svd
    [255, 25, 120],  # ??? - floor svd
    [55, 25, 120],  # ??? - joint svd
]

def floor_xyz(
    top_corners: torch.Tensor, # in [-1, 1]
    bottom_corners: torch.Tensor, # in [-1, 1]
    floor_z: float=-1.6,
):
    b, N, _ = top_corners.size()
    u = bottom_corners[:, :, 0] * np.pi
    v = bottom_corners[:, :, 1] * -0.5 * np.pi
    c = floor_z / torch.tan(v)
    x = c * torch.sin(u)
    y = -c * torch.cos(u)
    floor_xy = torch.stack([x, y], dim=-1)
    c = torch.linalg.norm(floor_xy, ord=2, dim=-1)
    v = top_corners[:, :, 1] * -0.5 * np.pi
    ceil_z = (c * torch.tan(v)).mean(dim=1, keepdim=True)
    ceil_z = ceil_z.unsqueeze(1).expand(b, 4, 1).contiguous()
    bottom_points = torch.cat([floor_xy, floor_z * torch.ones_like(c.unsqueeze(-1))], dim=-1)
    top_points = torch.cat([floor_xy, ceil_z], dim=-1)
    return top_points, bottom_points

def ceil_xyz(
    top_corners: torch.Tensor, # in [-1, 1]
    bottom_corners: torch.Tensor, # in [-1, 1]
    ceil_z: float=1.6,
):
    b, N, _ = top_corners.size()
    u = top_corners[:, :, 0] * np.pi
    v = top_corners[:, :, 1] * -0.5 * np.pi
    c = ceil_z / torch.tan(v)
    x = c * torch.sin(u)
    y = -c * torch.cos(u)
    ceil_xy = torch.stack([x, y], dim=-1)
    c = torch.linalg.norm(ceil_xy, ord=2, dim=-1)
    v = bottom_corners[:, :, 1] * -0.5 * np.pi
    floor_z = (c * torch.tan(v)).mean(dim=1, keepdim=True)
    fix_ceil = -ceil_z / floor_z
    floor_z = -1.6
    ceil_z_fix = ceil_z * fix_ceil
    ceil_z_fix = ceil_z_fix.unsqueeze(1).expand(b, 4, 1).contiguous()
    bottom_points = torch.cat([ceil_xy, floor_z * torch.ones_like(c.unsqueeze(-1))], dim=-1)
    top_points = torch.cat([ceil_xy, ceil_z_fix], dim=-1)
    return top_points, bottom_points

def get_scale_123(coords: torch.Tensor) -> torch.Tensor:
    return torch.pairwise_distance(coords[:, :2, :], coords[:, 1:3, :])

def get_scale_all(coords: torch.Tensor, eps: float=1e-12) -> torch.Tensor:
    a_x1 = torch.linalg.norm(coords[:, 0, :] - coords[:, 1, :], ord=2, dim=1)
    a_y1 = torch.linalg.norm(coords[:, 1, :] - coords[:, 2, :], ord=2, dim=1)
    a_x2 = torch.linalg.norm(coords[:, 2, :] - coords[:, 3, :], ord=2, dim=1)
    a_y2 = torch.linalg.norm(coords[:, 3, :] - coords[:, 0, :], ord=2, dim=1)
    a_x = 0.5 * (a_x1 + a_x2)
    a_y = 0.5 * (a_y1 + a_y2)
    # return torch.stack([a_x, a_y], dim=1)
    return torch.stack([a_y, a_x], dim=1)

def get_angle1(coords: torch.Tensor, epsilon: float=1e-12) -> torch.Tensor:
        normalized = torch.nn.functional.normalize(coords, p=2, dim=-1)
        angle = torch.atan2(normalized[:, 0, 1], normalized[:, 0, 0] + epsilon)
        return angle  / np.pi

def get_rot2d(angle: torch.Tensor) -> torch.Tensor:
    def get_mat(a: torch.Tensor) -> torch.Tensor:
        cos = torch.cos(a)
        sin = torch.sin(a)
        return torch.Tensor([
            [cos, -sin],
            [sin, cos],
        ]).to(a)
    return torch.stack([get_mat(a) for a in angle])

def homography_floor_aligned(
    top_corners: torch.Tensor, # in [-1, 1]
    bottom_corners: torch.Tensor, # in [-1, 1]
    floor_z: float=-1.6,
):
    b, N, _ = top_corners.size()
    u = bottom_corners[:, :, 0] * np.pi
    v = bottom_corners[:, :, 1] * (-0.5 * np.pi)
    c = floor_z / torch.tan(v)
    x = c * torch.sin(u)
    y = -c * torch.cos(u)
    floor_xy = torch.stack([x, y], dim=-1)
    # angle = get_angle1(floor_xy)
    # rot = get_rot2d(-1.0 * angle)
    # floor_xy = (rot @ floor_xy.transpose(-2, -1)).transpose(-2, -1)
    scale = get_scale_all(floor_xy)
    scale = scale / 2.0 # scale.flip(-1) / 2.0
    centroid = floor_xy.mean(dim=1)
    c = torch.linalg.norm(floor_xy, ord=2, dim=-1)
    v = top_corners[:, :, 1] * (-0.5 * np.pi)
    ceil_z = (c * torch.tan(v)).mean(dim=1, keepdim=True)
    ceil_z = ceil_z.unsqueeze(1).expand(b, 4, 1).contiguous()
    floor_xy = floor_xy - centroid.unsqueeze(1)
    inds = torch.sort(torch.atan2(floor_xy[..., 0], floor_xy[..., 1] + 1e-12))[1]
    axes = dst[:, inds.squeeze(), :]
    homography = kornia.get_perspective_transform(floor_xy, axes)
    # homogeneous = torch.cat([floor_xy + centroid.unsqueeze(1), torch.ones_like(floor_xy[..., -1:])], dim=2)
    homogeneous = torch.cat([floor_xy, torch.ones_like(floor_xy[..., -1:])], dim=2)
    xformed = (homography @ homogeneous.transpose(1, 2)).transpose(1, 2)
    xformed = xformed[:, :, :2] / xformed[:, :, 2].unsqueeze(-1)
    # scale = torch.pairwise_distance(floor_xy[:, :2, :], floor_xy[:, 1:3, :])

    # for i in range(b):
    #     trace = torch.trace(homography[i])
    #     print(f"i tr(h) = {trace}")
    #     print(f"{centroid[if expression:
    #         pass]}")
    #     if trace < 1.3:
    #         print("error")
    # centroid = centroid
    rect_floor_xy = xformed * scale.unsqueeze(1) + centroid.unsqueeze(1)
    # rect_floor_xy = xformed * scale.unsqueeze(1)
    bottom_points = torch.cat([rect_floor_xy, floor_z * torch.ones_like(c.unsqueeze(-1))], dim=-1)
    top_points = torch.cat([rect_floor_xy, ceil_z], dim=-1)
    return top_points, bottom_points

def homography_floor_based(
    top_corners: torch.Tensor, # in [-1, 1]
    bottom_corners: torch.Tensor, # in [-1, 1]
    floor_z: float=-1.6,
):
    b, N, _ = top_corners.size()
    u = bottom_corners[:, :, 0] * np.pi
    v = bottom_corners[:, :, 1] * -0.5 * np.pi
    c = floor_z / torch.tan(v)
    x = c * torch.sin(u)
    y = -c * torch.cos(u)
    floor_xy = torch.stack([x, y], dim=-1)
    c = torch.linalg.norm(floor_xy, ord=2, dim=-1)
    v = top_corners[:, :, 1] * -0.5 * np.pi
    ceil_z = (c * torch.tan(v)).mean(dim=1, keepdim=True)
    ceil_z = ceil_z.unsqueeze(1).expand(b, 4, 1).contiguous()    
    homography = kornia.get_perspective_transform(floor_xy, dst.expand(b, N, 2))
    homogeneous = torch.cat([floor_xy, torch.ones(b, N, 1)], dim=2)
    for i in range (b):
        h = homography[i]
        print(f"Homography: {h}, det = {torch.det(h)}, trace = {torch.trace(h)}")
    xformed = (homography @ homogeneous.transpose(1, 2)).transpose(1, 2)
    xformed = xformed[:, :, :2] / xformed[:, :, 2].unsqueeze(-1)
    # circ_shifted = torch.cat([floor_xy[:, 1:, :], floor_xy[:, :1, :]], dim=1)
    # scales = torch.pairwise_distance(floor_xy.transpose(-1, -2), circ_shifted.transpose(-1, -2))
    # scales = torch.pairwise_distance(floor_xy, circ_shifted)
    # x_scale = torch.mean(scales[:, ::2], dim=1)
    # y_scale = torch.mean(scales[:, 1::2], dim=1)
    # scale = torch.stack([x_scale, y_scale], dim=1)
    # scale = torch.pairwise_distance(floor_xy[:, :2, :], circ_shifted[:, :2, :])
    # scale = torch.pairwise_distance(floor_xy[:, :2, :], floor_xy[:, 1:3, :])
    # scale = get_scale_123(floor_xy)
    scale = get_scale_all(floor_xy)
    scale = scale / 2.0 # scale.flip(-1) / 2.0
    centroid = floor_xy.mean(dim=1)
    rect_floor_xy = xformed * scale.unsqueeze(1) + centroid.unsqueeze(1)
    bottom_points = torch.cat([rect_floor_xy, floor_z * torch.ones_like(c.unsqueeze(-1))], dim=-1)
    top_points = torch.cat([rect_floor_xy, ceil_z], dim=-1)
    return top_points, bottom_points

def svd(
        points1: torch.Tensor,
        points2: torch.Tensor
    ) -> torch.Tensor:
        '''
            Computes a similarity transform (sR, t) that takes
            a set of 3D points S1 (3 x N) closest to a set of 3D points S2,
            where R is an 3x3 rotation matrix, t 3x1 translation, s scale.
            i.e. solves the orthogonal Procrutes problem.
        '''
        b, _, c = points1.shape
        # 1. Remove mean.
        points1 = torch.transpose(points1, -2, -1)
        points2 = torch.transpose(points2, -2, -1)
        centroid1 = points1.mean(dim=-1, keepdims=True)
        centroid2 = points1.mean(dim=-1, keepdims=True)
        centered1 = points1 - centroid1
        centered2 = points2 - centroid2
        # 2. Compute variance of X1 used for scale.
        variance = torch.sum(centered1 ** 2, dim=[1, 2])
        # 3. The outer product of X1 and X2.    
        K = centered1 @ torch.transpose(centered2, -2, -1)
        # 4. Solution that Maximizes trace(R'K) is R=U*V', where U, V are singular vectors of K.
        U, s, V = torch.svd(K)
        # Construct Z that fixes the orientation of R to get det(R)=1.
        Z = torch.eye(c).to(U).unsqueeze(0).repeat(b, 1, 1)
        Z[:,-1, -1] *= torch.sign(torch.det(U @ torch.transpose(V, -2, -1)))
        # Construct R.
        rotation = V @ (Z @ torch.transpose(U, -2, -1))
        # 5. Recover scale.
        scale = torch.cat([torch.trace(x).unsqueeze(0) for x in (rotation @ K)]) / variance
        # 6. Recover translation.
        scale = scale.unsqueeze(-1).unsqueeze(-1)
        translation = centroid2 - (scale * (rotation @ centroid1))
        return rotation, translation, scale

def transform_points(
        points: torch.Tensor,
        rotation: torch.Tensor,
        translation: torch.Tensor,
        scale: torch.Tensor,
    ) -> torch.Tensor:
        xformed = scale * (rotation @ torch.transpose(points, -2, -1)) + translation
        return torch.transpose(xformed, -2, -1)

def homography_floor_svd(
        top_corners: torch.Tensor, # in [-1, 1]
        bottom_corners: torch.Tensor, # in [-1, 1]
        floor_z: float=-1.6,
    ):
        b, N, _ = top_corners.size()
        u = bottom_corners[:, :, 0] * np.pi
        v = bottom_corners[:, :, 1] * (-0.5 * np.pi)
        c = floor_z / torch.tan(v)
        x = c * torch.sin(u)
        y = -c * torch.cos(u)
        floor_xy = torch.stack([x, y], dim=-1)
        scale = get_scale_all(floor_xy)
        scale = scale / 2.0
        centroid = floor_xy.mean(dim=1)
        c = torch.linalg.norm(floor_xy, ord=2, dim=-1)
        v = top_corners[:, :, 1] * (-0.5 * np.pi)
        ceil_z = (c * torch.tan(v)).mean(dim=1, keepdim=True)
        ceil_z = ceil_z.unsqueeze(1).expand(b, 4, 1).contiguous()
        floor_xy = floor_xy - centroid.unsqueeze(1)
        inds = torch.sort(torch.atan2(floor_xy[..., 0], floor_xy[..., 1] + 1e-12))[1]
        # axes = self.cuboid_axes[:, inds.squeeze(), :]
        # homography = kornia.get_perspective_transform(floor_xy, axes)    
        homography = kornia.get_perspective_transform(floor_xy, dst.expand(b, N, 2))
        homogeneous = torch.cat([floor_xy, torch.ones_like(floor_xy[..., -1:])], dim=2)
        xformed = (homography @ homogeneous.transpose(1, 2)).transpose(1, 2)
        xformed = xformed[:, :, :2] / xformed[:, :, 2].unsqueeze(-1)
        rect_floor_xy = xformed * scale.unsqueeze(1) + centroid.unsqueeze(1)
        # rect_floor_xy = xformed + centroid.unsqueeze(1)
        original_xy = floor_xy + centroid.unsqueeze(1)
        R, t, s = svd(rect_floor_xy, original_xy[:, inds.squeeze(), :])
        rect_floor_xy = transform_points(rect_floor_xy, R, t, s)
        bottom_points = torch.cat([rect_floor_xy, floor_z * torch.ones_like(c.unsqueeze(-1))], dim=-1)
        top_points = torch.cat([rect_floor_xy, ceil_z], dim=-1)
        return top_points, bottom_points

def homography_ceil_svd(
        top_corners: torch.Tensor, # in [-1, 1]
        bottom_corners: torch.Tensor, # in [-1, 1]
        ceil_z: float=1.6,
    ):
        b, N, _ = top_corners.size()
        u_t = top_corners[:, :, 0] * np.pi
        v_t = top_corners[:, :, 1] * (-0.5 * np.pi)
        c = ceil_z / torch.tan(v_t)
        x = c * torch.sin(u_t)
        y = -c * torch.cos(u_t)
        ceil_xy = torch.stack([x, y], dim=-1)
        scale = get_scale_all(ceil_xy)        
        c = torch.linalg.norm(ceil_xy, ord=2, dim=-1)
        v_b = bottom_corners[:, :, 1] * (-0.5 * np.pi)
        floor_z = (c * torch.tan(v_b)).mean(dim=1, keepdim=True)
        fix_ceil = -ceil_z / floor_z
        floor_z = -ceil_z
        ceil_z_fix = ceil_z * fix_ceil
        ceil_z_fix = ceil_z_fix.unsqueeze(1).expand(b, 4, 1).contiguous()
        c = ceil_z_fix.squeeze(-1) / torch.tan(v_t)
        x = c * torch.sin(u_t)
        y = -c * torch.cos(u_t)
        ceil_xy = torch.stack([x, y], dim=-1)
        # scale = get_scale_all(ceil_xy) #NOTE: PREVIOUS SCALE (after floor rescaling)

        # scale = self._get_scale_123(ceil_xy)
        # scale = scale.flip(-1) / 2.0
        scale = scale / 2.0
        centroid = ceil_xy.mean(dim=1)
        ceil_xy = ceil_xy - centroid.unsqueeze(1)
        inds = torch.sort(torch.atan2(ceil_xy[..., 0], ceil_xy[..., 1] + 1e-12))[1]
        # axes = dst[:, inds.squeeze(), :]
        # homography = kornia.get_perspective_transform(ceil_xy, torch.sign(ceil_xy))        
        # homography = kornia.get_perspective_transform(ceil_xy, axes)
        homography = kornia.get_perspective_transform(ceil_xy, dst.expand(b, N, 2))
        homogeneous = torch.cat([ceil_xy, torch.ones_like(ceil_xy[..., -1:])], dim=2)
        # homogeneous = torch.cat([ceil_xy + centroid.unsqueeze(1), torch.ones_like(ceil_xy[..., -1:])], dim=2)
        xformed = (homography @ homogeneous.transpose(1, 2)).transpose(1, 2)
        xformed = xformed[:, :, :2] / xformed[:, :, 2].unsqueeze(-1)
        # scale = torch.pairwise_distance(ceil_xy[:, :2, :], ceil_xy[:, 1:3, :])
        rect_ceil_xy = xformed * scale.unsqueeze(1) + centroid.unsqueeze(1)
        
        original_xy = ceil_xy + centroid.unsqueeze(1)
        R, t, s = svd(rect_ceil_xy, original_xy[:, inds.squeeze(), :])
        rect_ceil_xy = transform_points(rect_ceil_xy, R, t, s)

        bottom_points = torch.cat([rect_ceil_xy, floor_z * torch.ones_like(c.unsqueeze(-1))], dim=-1)
        top_points = torch.cat([rect_ceil_xy, ceil_z_fix], dim=-1)
        return top_points, bottom_points

def homography_joint_svd(
        top_corners: torch.Tensor, # in [-1, 1]
        bottom_corners: torch.Tensor, # in [-1, 1]
        floor_z: float=-1.6,
        ceil_z: float=1.6,
    ):
        b, N, _ = top_corners.size()
        floor_u = bottom_corners[:, :, 0] * np.pi
        floor_v = bottom_corners[:, :, 1] * (-0.5 * np.pi)
        floor_c = floor_z / torch.tan(floor_v)
        floor_x = floor_c * torch.sin(floor_u)
        floor_y = -floor_c * torch.cos(floor_u)
        floor_xy = torch.stack([floor_x, floor_y], dim=-1)
        floor_scale = get_scale_all(floor_xy)
        floor_scale = floor_scale / 2.0        
        floor_ceil_c = torch.linalg.norm(floor_xy, ord=2, dim=-1)
        floor_ceil_v = top_corners[:, :, 1] * (-0.5 * np.pi)
        floor_ceil_z = (floor_ceil_c * torch.tan(floor_ceil_v)).mean(dim=1, keepdim=True)
        floor_ceil_z = floor_ceil_z.unsqueeze(1).expand(b, 4, 1).contiguous()
        # floor_centroid = floor_xy.mean(dim=1)
        # floor_xy = floor_xy - floor_centroid.unsqueeze(1)

        ceil_u_t = top_corners[:, :, 0] * np.pi
        ceil_v_t = top_corners[:, :, 1] * (-0.5 * np.pi)
        ceil_c = ceil_z / torch.tan(ceil_v_t)
        ceil_x = ceil_c * torch.sin(ceil_u_t)
        ceil_y = -ceil_c * torch.cos(ceil_u_t)
        ceil_xy = torch.stack([ceil_x, ceil_y], dim=-1)
        ceil_floor_c = torch.linalg.norm(ceil_xy, ord=2, dim=-1)
        ceil_v_b = bottom_corners[:, :, 1] * (-0.5 * np.pi)
        ceil_floor_z = (ceil_floor_c * torch.tan(ceil_v_b)).mean(dim=1, keepdim=True)
        fix_ceil = -ceil_z / ceil_floor_z
        ceil_z_fix = ceil_z * fix_ceil
        ceil_z_fix = ceil_z_fix.unsqueeze(1).expand(b, 4, 1).contiguous()
        ceil_floor_fixed_c = ceil_z_fix.squeeze(-1) / torch.tan(ceil_v_t)
        ceil_x = ceil_floor_fixed_c * torch.sin(ceil_u_t)
        ceil_y = -ceil_floor_fixed_c * torch.cos(ceil_u_t)
        ceil_xy = torch.stack([ceil_x, ceil_y], dim=-1)
        ceil_scale = get_scale_all(ceil_xy)
        ceil_scale = ceil_scale / 2.0
        # ceil_centroid = ceil_xy.mean(dim=1)
        # ceil_xy = ceil_xy - ceil_centroid.unsqueeze(1)

        joint_xy = 0.5 * (floor_xy + ceil_xy)
        joint_scale = 0.5 * (floor_scale + ceil_scale)        
        joint_centroid = joint_xy.mean(dim=1)
        joint_xy = joint_xy - joint_centroid.unsqueeze(1)

        inds = torch.sort(torch.atan2(joint_xy[..., 0], joint_xy[..., 1] + 1e-12))[1]
        # axes = self.cuboid_axes[:, inds.squeeze(), :]
        # homography = kornia.get_perspective_transform(joint_xy, axes)    
        homography = kornia.get_perspective_transform(joint_xy, dst.expand(b, N, 2)) 
        homogeneous = torch.cat([joint_xy, torch.ones_like(joint_xy[..., -1:])], dim=2)
        xformed = (homography @ homogeneous.transpose(1, 2)).transpose(1, 2)
        xformed = xformed[:, :, :2] / xformed[:, :, 2].unsqueeze(-1)
        rect_joint_xy = xformed * joint_scale.unsqueeze(1) + joint_centroid.unsqueeze(1)
        original_xy = joint_xy + joint_centroid.unsqueeze(1)
        R, t, s = svd(rect_joint_xy, original_xy[:, inds.squeeze(), :])
        rect_joint_xy = transform_points(rect_joint_xy, R, t, s)
        bottom_points = torch.cat([rect_joint_xy, floor_z * torch.ones_like(floor_c.unsqueeze(-1))], dim=-1)
        top_points = torch.cat([rect_joint_xy, ceil_z_fix], dim=-1)
        return top_points, bottom_points

def homography_ceil_based(
    top_corners: torch.Tensor, # in [-1, 1]
    bottom_corners: torch.Tensor, # in [-1, 1]
    ceil_z: float=1.6,
):
    b, N, _ = top_corners.size()
    u_t = top_corners[:, :, 0] * np.pi
    v_t = top_corners[:, :, 1] * -0.5 * np.pi
    c = ceil_z / torch.tan(v_t)
    x = c * torch.sin(u_t)
    y = -c * torch.cos(u_t)
    ceil_xy = torch.stack([x, y], dim=-1)
    scale = get_scale_all(ceil_xy)
    c = torch.linalg.norm(ceil_xy, ord=2, dim=-1)
    v_b = bottom_corners[:, :, 1] * -0.5 * np.pi
    floor_z = (c * torch.tan(v_b)).mean(dim=1, keepdim=True)
    fix_ceil = -ceil_z / floor_z
    floor_z = -1.6
    ceil_z_fix = ceil_z * fix_ceil
    ceil_z_fix = ceil_z_fix.unsqueeze(1).expand(b, 4, 1).contiguous()
    c = ceil_z_fix.squeeze(-1) / torch.tan(v_t)
    x = c * torch.sin(u_t)
    y = -c * torch.cos(u_t)
    ceil_xy = torch.stack([x, y], dim=-1)
    homography = kornia.get_perspective_transform(ceil_xy, dst.expand(b, N, 2))
    homogeneous = torch.cat([ceil_xy, torch.ones(b, N, 1)], dim=2)
    xformed = (homography @ homogeneous.transpose(1, 2)).transpose(1, 2)
    xformed = xformed[:, :, :2] / xformed[:, :, 2].unsqueeze(-1)
    # circ_shifted = torch.cat([ceil_xy[:, 1:, :], ceil_xy[:, :1, :]], dim=1)
    # scales = torch.pairwise_distance(ceil_xy.transpose(-1, -2), circ_shifted.transpose(-1, -2))
    # scales = torch.pairwise_distance(ceil_xy, circ_shifted)
    # x_scale = torch.mean(scales[:, ::2], dim=1)
    # y_scale = torch.mean(scales[:, 1::2], dim=1)
    # scale = torch.stack([x_scale, y_scale], dim=1)
    # scale = torch.pairwise_distance(ceil_xy[:, :2, :], ceil_xy[:, 1:3, :])
    # scale = get_scale_123(ceil_xy)
    
    # scale = get_scale_all(ceil_xy) #NOTE: previous SCALE (after floor rescaling)

    # scale = scale.flip(-1) / 2.0
    scale = scale / 2.0
    centroid = ceil_xy.mean(dim=1)
    rect_ceil_xy = xformed * scale.unsqueeze(1) + centroid.unsqueeze(1)
    bottom_points = torch.cat([rect_ceil_xy, floor_z * torch.ones_like(c.unsqueeze(-1))], dim=-1)
    top_points = torch.cat([rect_ceil_xy, ceil_z_fix], dim=-1)
    return top_points, bottom_points

def homography_joint(
    top_corners: torch.Tensor, # in [-1, 1]
    bottom_corners: torch.Tensor, # in [-1, 1]
):
    top_ceil, bottom_ceil = homography_ceil_based(top_corners, bottom_corners)
    top_floor, bottom_floor = homography_floor_based(top_corners, bottom_corners)
    return (top_ceil + top_floor) * 0.5, (bottom_ceil + bottom_floor) * 0.5
    
__POINT_FUNCS__ = {
    "floor": floor_xyz,
    "ceil": ceil_xyz,
    "homo_floor": homography_floor_based,
    "homo_ceil": homography_ceil_based,
    "homo_joint": homography_joint,
    "floor_aligned": homography_floor_aligned,
    "ceil_svd": homography_ceil_svd,
    "floor_svd": homography_floor_svd,
    "joint_svd": homography_joint_svd,
}

def xyz_2_coorxy(xs, ys, zs, H, W):
    us = np.arctan2(xs, -ys)
    vs = -np.arctan(zs / np.sqrt(xs**2 + ys**2))
    coorx = (us / (2 * np.pi) + 0.5) * W
    coory = (vs / np.pi + 0.5) * H
    return coorx, coory

def write_points(
    top_pts3d: torch.Tensor,
    bottom_pts3d: torch.Tensor,
    path: str,
    key: str,
    color: int,
):
    b, N, _ = top_pts3d.size()
    for i in range(b):
        points = torch.cat([top_pts3d[i, ...], bottom_pts3d[i, ...]], dim=0)            
        wf_lines = [[i, (i+1)%N] for i in range(N)] +\
                [[i+N, (i+1)%N+N] for i in range(N)] +\
                [[i, i+N] for i in range(N)]
        wf_colors = [[c / 255.0 for c in colors[color]] for i in range(len(wf_lines))]
        wf_line_set = open3d.geometry.LineSet()
        wf_line_set.points = open3d.utility.Vector3dVector(points.detach().numpy())
        wf_line_set.lines = open3d.utility.Vector2iVector(wf_lines)
        wf_line_set.colors = open3d.utility.Vector3dVector(wf_colors)
        open3d.io.write_line_set(os.path.join(
                path, f"test_{key}.ply"
            ),
            wf_line_set, write_ascii=True,
        )

def project_points(
    points3d: torch.Tensor,
    height: int,
    width: int,
):
    phi = torch.atan2(points3d[:, :, 0], -1.0 * points3d[:, :, 1])
    xy_dist = torch.linalg.norm(points3d[:, :, :2], ord=2, dim=-1)
    theta = -1.0 * torch.atan2(points3d[:, :, 2], xy_dist)
    u = (phi / (2.0 * np.pi) + 0.5) * width
    v = (theta / np.pi + 0.5) * height
    return torch.stack([u, v], dim=-1)

def draw_points(
    points2d: torch.Tensor, # [B, N, 2]
    images: np.array, # [B, H, W, 3]
    color: np.array
):
    b = images.shape[0]
    N = points2d.shape[1]
    for i in range(b):
        img = images[i]
        pts2d = points2d[i]        
        for k in range(N):
            pt2d = pts2d[k]
            x = int(pt2d[0])
            y = int(pt2d[1])
            cv2.drawMarker(img, (x, y),color, cv2.MARKER_SQUARE, 15, 2)

def _xyz2uv(x: float, y: float, z: float, eps: float=1e-17) -> typing.Tuple[float, float]:
    normXZ = math.sqrt( math.pow(x, 2) + math.pow(z, 2) )
    normXZ = eps if normXZ < eps else normXZ
    normXYZ = math.sqrt(
        math.pow(x, 2) + 
        math.pow(y, 2) + 
        math.pow(z, 2) 
    )
    normXYZ = eps if normXYZ < eps else normXYZ
    v = math.asin(y / normXYZ)
    u = math.asin(x / normXZ)
    if z > 0 and u > 0:
        u = math.pi - u
    elif z > 0 and u < 0:
        u = -math.pi - u
    if x == 0 and z > 0:
        u = math.pi
    return (u, v)

def _coords2uv(coords, width, height) -> np.array:
    middleX = width / 2 + 0.5
    middleY = height / 2 + 0.5
    uv = np.hstack([
        (coords[:, [0]] - middleX) / width * 2 * np.pi,
        -(coords[:, [1]] - middleY) / height * np.pi
    ])
    return uv

def _xyz2uvN(xyz: np.array, planeID: int=1) -> np.array:
    ID1 = (int(planeID) - 1 + 0) % 3
    ID2 = (int(planeID) - 1 + 1) % 3
    ID3 = (int(planeID) - 1 + 2) % 3
    normXY = np.sqrt(xyz[:, [ID1]] ** 2 + xyz[:, [ID2]] ** 2)
    normXY[normXY < 0.000001] = 0.000001
    normXYZ = np.sqrt(xyz[:, [ID1]] ** 2 + xyz[:, [ID2]] ** 2 + xyz[:, [ID3]] ** 2)
    v = np.arcsin(xyz[:, [ID3]] / normXYZ)
    u = np.arcsin(xyz[:, [ID1]] / normXY)
    valid = (xyz[:, [ID2]] < 0) & (u >= 0)
    u[valid] = np.pi - u[valid]
    valid = (xyz[:, [ID2]] < 0) & (u <= 0)
    u[valid] = -np.pi - u[valid]
    uv = np.hstack([u, v])
    uv[np.isnan(uv[:, 0]), 0] = 0
    return uv

def _computeUVN(n: np.array, in_: float, planeID: int) -> np.array:
    """
    compute v given u and normal.
    """
    if planeID == 2:
        n = np.array([n[1], n[2], n[0]])
    elif planeID == 3:
        n = np.array([n[2], n[0], n[1]])
    bc = n[0] * np.sin(in_) + n[1] * np.cos(in_)
    bs = n[2]
    out = np.arctan(-bc / (bs + 1e-9))
    return out

def _line_from_two_points(pt1: np.array, pt2: np.array) -> np.array:
    """
    Generate line segment based on two points on panorama
    pt1, pt2: two points on panorama
    line:
        1~3-th dim: normal of the line
        4-th dim: the projection dimension ID
        5~6-th dim: the u of line segment endpoints in projection plane
    """
    numLine = pt1.shape[0]
    lines = np.zeros((numLine, 6))
    n = np.cross(pt1, pt2)
    n = n / (matlib.repmat(np.sqrt(np.sum(n ** 2, 1, keepdims=True)), 1, 3) + 1e-9)
    lines[:, 0:3] = n

    areaXY = np.abs(np.sum(n * matlib.repmat([0, 0, 1], numLine, 1), 1, keepdims=True))
    areaYZ = np.abs(np.sum(n * matlib.repmat([1, 0, 0], numLine, 1), 1, keepdims=True))
    areaZX = np.abs(np.sum(n * matlib.repmat([0, 1, 0], numLine, 1), 1, keepdims=True))
    planeIDs = np.argmax(np.hstack([areaXY, areaYZ, areaZX]), axis=1) + 1
    lines[:, 3] = planeIDs

    for i in range(numLine):
        uv = _xyz2uvN(np.vstack([pt1[i, :], pt2[i, :]]), lines[i, 3])
        umax = uv[:, 0].max() + np.pi
        umin = uv[:, 0].min() + np.pi
        if umax - umin > np.pi:
            lines[i, 4:6] = np.array([umax, umin]) / 2 / np.pi
        else:
            lines[i, 4:6] = np.array([umin, umax]) / 2 / np.pi

    return lines

def _uv2xyzN(uv: np.array, planeID: int=1) -> np.array:
    ID1 = (int(planeID) - 1 + 0) % 3
    ID2 = (int(planeID) - 1 + 1) % 3
    ID3 = (int(planeID) - 1 + 2) % 3
    xyz = np.zeros((uv.shape[0], 3))
    xyz[:, ID1] = np.cos(uv[:, 1]) * np.sin(uv[:, 0])
    xyz[:, ID2] = np.cos(uv[:, 1]) * np.cos(uv[:, 0])
    xyz[:, ID3] = np.sin(uv[:, 1])
    return xyz

def _line_index_from_coords(cor_all: np.array, im_w: int, im_h: int):
    assert len(cor_all) % 2 == 0
    uv = _coords2uv(cor_all, im_w, im_h)
    xyz = _uv2xyzN(uv)
    lines = _line_from_two_points(xyz[0::2], xyz[1::2])
    num_sample = max(im_h, im_w)

    cs, rs = [], []
    for i in range(lines.shape[0]):
        n = lines[i, 0:3]
        sid = lines[i, 4] * 2 * np.pi
        eid = lines[i, 5] * 2 * np.pi
        if eid < sid:
            x = np.linspace(sid, eid + 2 * np.pi, num_sample)
            x = x % (2 * np.pi)
        else:
            x = np.linspace(sid, eid, num_sample)

        u = -np.pi + x.reshape(-1, 1)
        v = _computeUVN(n, u, lines[i, 3])
        xyz = _uv2xyzN(np.hstack([u, v]), lines[i, 3])
        uv = _xyz2uvN(xyz, 1)

        r = np.minimum(np.floor((uv[:, 0] + np.pi) / (2 * np.pi) * im_w) + 1,
                       im_w).astype(np.int32)
        c = np.minimum(np.floor((np.pi / 2 - uv[:, 1]) / np.pi * im_h) + 1,
                       im_h).astype(np.int32)
        cs.extend(r - 1)
        rs.extend(c - 1)
    return rs, cs

def draw_boundary_from_corners(
    img: np.array,
    corners: np.array,
    color: typing.List[float],
):
    im_h, im_w = img.shape[:2]
    cor_all = [corners]
    for i in range(len(corners)):
        cor_all.append(corners[i, :])
        cor_all.append(corners[(i+2) % len(corners), :])
    cor_all = np.vstack(cor_all)

    rs, cs = _line_index_from_coords(cor_all, im_w, im_h)
    rs = np.array(rs)
    cs = np.array(cs)

    for dx, dy in [[-1, 0], [1, 0], [0, 0], [0, 1], [0, -1]]:
        img[np.clip(rs + dx, 0, im_h - 1), np.clip(cs + dy, 0, im_w - 1), 0] = color[0]
        img[np.clip(rs + dx, 0, im_h - 1), np.clip(cs + dy, 0, im_w - 1), 1] = color[1]
        img[np.clip(rs + dx, 0, im_h - 1), np.clip(cs + dy, 0, im_w - 1), 2] = color[2]


if __name__ == '__main__':

    import argparse
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('--img_glob', required=True,
                        help='Image texture in equirectangular format')
    parser.add_argument('--layout_glob', required=True,
                        help='Txt file containing layout corners (cor_id)')
    parser.add_argument('--output_path', required=True,
                        help='Folder where the output files will be saved at')
    parser.add_argument('--scale', default=1, type=float,
                        help='Scale texture for visualization')
    parser.add_argument('--merge_wireframes', action='store_true',
                    help='Merge all wireframes')                  
    args, _ = parser.parse_known_args()

    def test11():
        aligned = torch.Tensor([
            [ 0.9846, -0.5946],
            [ 0.0052, -0.2709],
            [ 0.2570, -0.1933],
            [ 0.5987, -0.2493],
            [ 0.9846,  0.5342],
            [ 0.0052,  0.2275],
            [ 0.2570,  0.1608],
            [ 0.5987,  0.2087]
        ])
        # print(f"Aligned: {aligned}")
        unaligned = torch.Tensor([
            [-0.9846, -0.5418],
            [ 0.0205, -0.2194],
            [ 0.2252, -0.2062],
            [ 0.5488, -0.2653],
            [-0.9846,  0.4780],
            [ 0.0205,  0.2186],
            [ 0.2252,  0.1963],
            [ 0.5488,  0.1743]
        ])
        return torch.chunk(unaligned.unsqueeze(0), 2, dim=-2)

    def test12():
        aligned = torch.Tensor([
            [-0.5102, -0.0457],
            [-0.3152, -0.0382],
            [ 0.2878, -0.0431],
            [ 0.5122, -0.0547],
            [-0.5102,  0.0476],
            [-0.3152,  0.0399],
            [ 0.2878,  0.0449],
            [ 0.5122,  0.0571]
        ])
        # print(f"Aligned: {aligned}")
        unaligned = torch.Tensor([
            [-0.9819, -0.0570],
            [-0.8568, -0.0825],
            [-0.0103, -0.0342],
            [ 0.0277, -0.0350],
            [-0.9819,  0.1501],
            [-0.8568,  0.0491],
            [-0.0103,  0.0539],
            [ 0.0277,  0.0307]
        ])
        return torch.chunk(unaligned.unsqueeze(0), 2, dim=-2)

    def test14():
        aligned = torch.Tensor([
            [-0.6743, -0.2757],
            [-0.5121, -0.3158],
            [ 0.5236, -0.5184],
            [ 0.7783, -0.3808],
            [-0.6743,  0.2207],
            [-0.5121,  0.2547],
            [ 0.5236,  0.4402],
            [ 0.7783,  0.3115]
        ])
        # print(f"Aligned: {aligned}")
        unaligned = torch.Tensor([
            [-0.9225, -0.2832],
            [-0.7755, -0.2755],
            [-0.1321, -0.4738],
            [ 0.2905, -0.6044],
            [-0.9225,  0.2201],
            [-0.7755,  0.2172],
            [-0.1321,  0.4091],
            [ 0.2905,  0.5314]
        ])
        return torch.chunk(unaligned.unsqueeze(0), 2, dim=-2)

    def test15():
        unaligned = torch.Tensor([
            [-0.7225, -0.2832],
            [-0.3755, -0.2755],
            [-0.4321, -0.4738],
            [ 0.8905, -0.6044],
            [-0.7225, 0.2832],
            [-0.3755, 0.2755],
            [-0.4321, 0.4738],
            [ 0.8905, 0.6044],
        ])
        return torch.chunk(unaligned.unsqueeze(0), 2, dim=-2)

    # top_corners, bottom_corners = test11()
    # top_corners, bottom_corners = test12()
    # top_corners, bottom_corners = test14()
    top_corners, bottom_corners = test15()

    H = 256
    W = 512
    images = torch.zeros(1, H, W, 3).numpy().astype(np.uint8)
    # subset = ['floor', 'homo_joint']
    subset = [
        'floor', # 'ceil', 
        # 'homo_floor', 'homo_joint', 'homo_ceil',
        # 'floor_aligned', 
        'floor_svd', # 'ceil_svd', 'joint_svd'
    ]
    for idx, (k, f) in enumerate(toolz.keyfilter(
        lambda k: k in subset, __POINT_FUNCS__).items()
    ):
        top_pts3d, bottom_pts3d = f(top_corners, bottom_corners)
        write_points(top_pts3d, bottom_pts3d, path=args.output_path,
            color=idx, key=k,
        )
        top_pts2d = project_points(top_pts3d, H, W)
        bottom_pts2d = project_points(bottom_pts3d, H, W)
        draw_points(top_pts2d, images, colors[idx])
        draw_points(bottom_pts2d, images, colors[idx])
        draw_boundary_from_corners(images[0], top_pts2d[0], colors[idx])
    for i, img in enumerate(images):
        cv2.imwrite(os.path.join(args.output_path, f"test.jpg"), img)