import torch
import numpy as np
import math
import cv2

world_mat_object = torch.tensor([
    [0.5085, 0.3226, 0.7984, 0.0000],
    [-0.3479, 0.9251, -0.1522, 0.0000],
    [-0.7877, -0.2003, 0.5826, 0.3384],
    [0.0000, 0.0000, 0.0000, 1.0000]
])

world_mat_inv = torch.tensor([
    [0.4019, 0.9157, 0.0000, 0.3359],
    [-0.1932, 0.0848, 0.9775, -1.0227],
    [0.8951, -0.3928, 0.2110, -7.0748],
    [-0.0000, 0.0000, -0.0000, 1.0000]
])


proj = torch.tensor([
    [2.1875, 0.0000, 0.0000, 0.0000],
    [0.0000, 3.8889, 0.0000, 0.0000],
    [0.0000, 0.0000, -1.0020, -0.2002],
    [0.0000, 0.0000, -1.0000, 0.0000]
])


RANGES = [[0, 540], [100, 960]]

TARGET = [500, -1]


def resize(img):
    if TARGET[1] == -1:
        r = img.shape[0] / img.shape[1]
        img = cv2.resize(img, (TARGET[0], int(r * TARGET[0])))
    else:
        img = cv2.resize(img, (TARGET[0], TARGET[1]))

    return img


def scatter(u_pix, v_pix, distances, res, radius=5):
    distances -= 6
    img = np.zeros(res)
    for (u, v, d) in zip(u_pix, v_pix, distances):
        v, u = int(v), int(u)
        f = np.exp(-d / 0.7)
        if radius == 0:
            img[v, u] = max(img[v, u], f)
        else:
            for t1 in range(-radius, radius):
                for t2 in range(-radius, radius):
                    ty, tx = v - t1, u - t2
                    ty, tx = max(0, ty), max(0, tx)
                    ty, tx = min(res[0] - 1, ty), min(res[1] - 1, tx)
                    img[ty, tx] = max(img[ty, tx], f)

    return img


def generate_roation(phi_x, phi_y, phi_z):
    def Rx(theta):
        return torch.tensor([[1, 0, 0],
                          [0, math.cos(theta), -math.sin(theta)],
                          [0, math.sin(theta), math.cos(theta)]])

    def Ry(theta):
        return torch.tensor([[math.cos(theta), 0, math.sin(theta)],
                          [0, 1, 0],
                          [-math.sin(theta), 0, math.cos(theta)]])

    def Rz(theta):
        return torch.tensor([[math.cos(theta), -math.sin(theta), 0],
                          [math.sin(theta), math.cos(theta), 0],
                          [0, 0, 1]])

    return Rz(phi_z) @ Ry(phi_y) @ Rx(phi_x)


def rotate_pc(pc, rx, ry, rz):
    rotation = generate_roation(rx, ry, rz)
    rotated = pc.clone()
    rotated[:, :3] = rotated[:, :3] @ rotation.T
    if rotated.shape[-1] == 6:
        rotated[:, 3:] = rotated[:, 3:] @ rotation.T
    return rotated


def draw_pc(pc: torch.Tensor, res=(540, 960), radius=5, timer=None, dy=0, scale=1):
    xyz = pc[:, :3]
    xyz -= xyz.mean(dim=0)
    t_scale = xyz.norm(dim=-1).max()
    xyz /= t_scale
    xyz *= scale

    xyz[:, -1] += xyz[:, -1].min()

    n, _ = xyz.shape

    if timer is not None:
        with timer('project'):
            xyz_pad = torch.cat([xyz, torch.ones_like(pc[:, :1])], dim=-1)
            xyz_local = xyz_pad @ world_mat_inv.T
            distances = -xyz_local[:, 2]

            projected = xyz_local @ proj.T
            projected = projected / projected[:, 3:4]
            projected = projected[:, :3]

            u_pix = ((projected[0] + 1) / 2) * res[1]
            v_pix = ((projected[1] + 1) / 2) * res[0] + dy

        with timer('z-buffer'):
            z_buffer = scatter(u_pix, v_pix, distances, res, radius=radius)[:, :]
    else:
        xyz_pad = torch.cat([xyz, torch.ones_like(pc[:, :1])], dim=-1)
        xyz_local = xyz_pad @ world_mat_inv.T
        distances = -xyz_local[:, 2]

        projected = xyz_local @ proj.T
        projected = projected / projected[:, 3:4]
        projected = projected[:, :3]

        u_pix = ((projected[:, 0] + 1) / 2) * res[1]
        v_pix = ((projected[:, 1] + 1) / 2) * res[0] + dy

        z_buffer = scatter(u_pix, v_pix, distances, res, radius=radius)[:, :]

    z_buffer = z_buffer[RANGES[0][0]: RANGES[0][1], :]
    z_buffer = z_buffer[:, RANGES[1][0]:RANGES[1][1]]
    z_buffer = resize(z_buffer)
    return z_buffer