from __future__ import annotations

import cv2
import numpy as np
import torch
import torch.nn.functional as F
import nvdiffrast.torch as dr

from camera import Camera
from .mesh import compute_face_normal, vis_shading_mesh


iris_vertex_array = np.array([
    3931,3932,3933,3935,3936,3937,3939,3940,3941,3943,3944,3945,3947,3948,
    3949,3951,3952,3953,3955,3956,3957,3959,3960,3961,3963,3964,3965,3967,
    3968,3969,3971,3972,3973,3975,3976,3977,3979,3980,3981,3983,3984,3985,
    3987,3988,3989,3991,3992,3993,3995,3996,3997,3999,4000,4001,4003,4004,
    4005,4007,4008,4009,4011,4012,4013,4015,4016,4017,4019,4020,4021,4023,
    4024,4025,4027,4028,4029,4031,4032,4033,4035,4036,4037,4039,4040,4041,
    4043,4044,4045,4047,4048,4049,4051,4052,4053,4054,4056,4057,4058,4477,
    4478,4479,4481,4482,4483,4485,4486,4487,4489,4490,4491,4493,4494,4495,
    4497,4498,4499,4501,4502,4503,4505,4506,4507,4509,4510,4511,4513,4514,
    4515,4517,4518,4519,4521,4522,4523,4525,4526,4527,4529,4530,4531,4533,
    4534,4535,4537,4538,4539,4541,4542,4543,4545,4546,4547,4549,4550,4551,
    4553,4554,4555,4557,4558,4559,4561,4562,4563,4565,4566,4567,4569,4570,
    4571,4573,4574,4575,4577,4578,4579,4581,4582,4583,4585,4586,4587,4589,
    4590,4591,4593,4594,4595,4597,4598,4599,4600,4602,4603,4604
], dtype=np.int32)


class Visualizer:
    def __init__(
        self,
        flame_assets_path: str = "weights/flame/flame_assets.npz",
        device: str = "cuda",
    ) -> None:
        self.device = torch.device(device)
        self.glctx = dr.RasterizeCudaContext()

        flame_assets = np.load(flame_assets_path)
        self.vis_edges = np.asarray(flame_assets["vis_edges"], dtype=np.int32)
        faces_t = torch.from_numpy(np.asarray(flame_assets["faces"]))
        self.faces_int32 = faces_t.to(dtype=torch.int32, device=self.device)
        self.faces_int64 = faces_t.to(dtype=torch.int64, device=self.device)

        num_vertices = int(self.faces_int64.max().item()) + 1
        base_vertex_colors = torch.ones((num_vertices, 3), dtype=torch.float32, device=self.device)
        iris_idx = torch.from_numpy(iris_vertex_array).to(dtype=torch.long, device=self.device)
        base_vertex_colors[iris_idx] = torch.tensor([0.0, 1.0, 0.0], dtype=torch.float32, device=self.device)
        self.base_vertex_colors = base_vertex_colors.contiguous()

    def vis_align(self, image: np.ndarray, vertex_coord: np.ndarray, line_width: int = 1) -> np.ndarray:
        image_vis = image.copy()
        vertex_xy = np.asarray(vertex_coord, dtype=np.int32)
        for edge in self.vis_edges:
            x0, y0, x1, y1 = vertex_xy[edge].reshape(-1).tolist()
            cv2.line(image_vis, (x0, y0), (x1, y1), (0, 127, 255), line_width, lineType=cv2.LINE_AA)
        return image_vis

    def vis_contour(self, image: np.ndarray, vertex_coord: np.ndarray, line_width: int = 1) -> np.ndarray:
        image_vis = image.copy()
        h, w = image_vis.shape[:2]

        vertices = torch.from_numpy(np.asarray(vertex_coord, dtype=np.float32)).to(self.device)
        vertices_ndc = torch.zeros((vertices.shape[0], 4), dtype=torch.float32, device=self.device)
        vertices_ndc[:, 0] = vertices[:, 0] / float(w) * 2.0 - 1.0
        vertices_ndc[:, 1] = vertices[:, 1] / float(h) * 2.0 - 1.0
        vertices_ndc[:, 2] = 0.5
        vertices_ndc[:, 3] = 1.0

        rast_out = dr.rasterize(
            self.glctx,
            vertices_ndc.unsqueeze(0),
            self.faces_int32,
            resolution=[h, w],
        )[0].squeeze(0)
        mask = (rast_out[..., 3:] > 0).to(dtype=torch.uint8)
        mask = (mask * 255).cpu().numpy()

        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        if contours:
            cnt = max(contours, key=cv2.contourArea)
            cv2.drawContours(image_vis, [cnt], -1, (0, 255, 0), line_width)
        return image_vis

    def vis_flame_normal_overlay(
        self,
        camera: Camera,
        image: np.ndarray,
        vertices: np.ndarray | torch.Tensor,
        alpha: float = 0.5,
        return_normal_map: bool = False,
    ) -> np.ndarray | tuple[np.ndarray, np.ndarray]:
        image_tf = torch.from_numpy(np.asarray(image, dtype=np.uint8)).to(dtype=torch.float32, device=self.device) / 255.0
        if isinstance(vertices, np.ndarray):
            vertices_tf = torch.from_numpy(vertices).to(dtype=torch.float32, device=self.device)
        else:
            vertices_tf = vertices.to(dtype=torch.float32, device=self.device)

        normal_map, mask = vis_shading_mesh(self.glctx, camera, vertices_tf, self.faces_int64)
        overlay = (image_tf * (1.0 - alpha) + normal_map * alpha) * mask + image_tf * (~mask)
        normal_map = (normal_map * 255.0).to(dtype=torch.uint8, device="cpu").numpy()
        overlay = (overlay * 255.0).to(dtype=torch.uint8, device="cpu").numpy()
        overlay = np.ascontiguousarray(overlay)
        normal_map = np.ascontiguousarray(normal_map)
        if return_normal_map:
            return overlay, normal_map
        return overlay

    def vis_flame_shading_overlay(
        self,
        camera: Camera,
        image: np.ndarray,
        vertices: np.ndarray | torch.Tensor,
        alpha: float = 0.5,
        return_shading: bool = False,
        shading_background: str = "black",
    ) -> np.ndarray | tuple[np.ndarray, np.ndarray]:
        image_tf = torch.from_numpy(np.asarray(image, dtype=np.uint8)).to(dtype=torch.float32, device=self.device) / 255.0
        if isinstance(vertices, np.ndarray):
            vertices_tf = torch.from_numpy(vertices).to(dtype=torch.float32, device=self.device)
        else:
            vertices_tf = vertices.to(dtype=torch.float32, device=self.device)

        vertices_pad = F.pad(vertices_tf, [0, 1], value=1.0)
        viewproj = camera.get_full_proj.to(device=self.device, dtype=torch.float32)
        vertices_clip = torch.matmul(vertices_pad, viewproj.transpose(0, 1))
        rast_out, _ = dr.rasterize(
            self.glctx,
            vertices_clip.unsqueeze(0),
            self.faces_int32,
            resolution=(camera.height, camera.width),
        )
        rast_squeezed = rast_out.squeeze(0)

        face_id = rast_squeezed[..., 3:]
        mask = face_id > 0
        face_id = torch.clamp_min(face_id - 1, 0).long()

        face_normals = compute_face_normal(vertices_tf, self.faces_int64)
        normal_map_world = face_normals[face_id].squeeze(-2)
        view_rot = camera.get_w2v[:3, :3].to(device=self.device, dtype=torch.float32)
        normal_map_view = torch.matmul(normal_map_world, view_rot.transpose(0, 1))
        normal_map_view = F.normalize(normal_map_view, dim=-1) * -1.0

        light_dir = torch.tensor([0.0, 0.25, 1.0], dtype=torch.float32, device=self.device)
        light_dir = F.normalize(light_dir, dim=0).view(1, 1, 3)
        diffuse = torch.clamp_min(torch.sum(normal_map_view * light_dir, dim=-1, keepdim=True), 0.0)
        shading_scalar = torch.clamp(0.15 + 0.85 * diffuse, 0.0, 1.0)

        color_map, _ = dr.interpolate(self.base_vertex_colors, rast_out, self.faces_int32)
        color_map = color_map.squeeze(0)
        shading = color_map * shading_scalar

        overlay = (image_tf * (1.0 - alpha) + shading * alpha) * mask + image_tf * (~mask)
        if shading_background == "black":
            shading = shading * mask
        elif shading_background == "white":
            shading = shading * mask + (~mask).to(dtype=shading.dtype)
        else:
            raise ValueError("shading_background must be 'black' or 'white'")
        shading = (shading * 255.0).to(dtype=torch.uint8, device="cpu").numpy()
        overlay = (overlay * 255.0).to(dtype=torch.uint8, device="cpu").numpy()
        overlay = np.ascontiguousarray(overlay)
        shading = np.ascontiguousarray(shading)
        if return_shading:
            return overlay, shading
        return overlay