from __future__ import annotations

from dataclasses import dataclass
from typing import Any

import numpy as np
import open3d as o3d
import torch
import trimesh
from scipy.spatial import KDTree


@dataclass
class DistanceStats:
    mean: float
    max: float
    rmse: float
    min: float
    std: float


class Evaluator:
    """Evaluation metrics for point-cloud workflows."""

    def chamfer_distance(
        self,
        pointcloud_a: o3d.geometry.PointCloud,
        pointcloud_b: o3d.geometry.PointCloud,
        sqrt: bool = True,
    ) -> tuple[float, float, float]:
        points_a = np.asarray(pointcloud_a.points)
        points_b = np.asarray(pointcloud_b.points)
        if len(points_a) == 0 or len(points_b) == 0:
            return float("inf"), 0.0, 0.0

        tree_a = KDTree(points_a)
        tree_b = KDTree(points_b)
        dist_a, _ = tree_b.query(points_a, k=1)
        dist_b, _ = tree_a.query(points_b, k=1)
        cd_a = float(np.mean(dist_a ** 2))
        cd_b = float(np.mean(dist_b ** 2))
        cd = cd_a + cd_b
        return (float(np.sqrt(cd)), cd_a, cd_b) if sqrt else (float(cd), cd_a, cd_b)

    def point_to_mesh_distance_trimesh(self, points: np.ndarray, mesh_path: str) -> tuple[np.ndarray, DistanceStats]:
        mesh = trimesh.load(mesh_path, force="mesh")
        _, distances, _ = mesh.nearest.on_surface(points)
        distances = np.asarray(distances, dtype=np.float64)
        return distances, self.distance_stats(distances)

    def point_to_mesh_distance_torch(
        self,
        points: np.ndarray,
        mesh: trimesh.Trimesh,
        chunk_size: int = 2048,
        triangle_chunk_size: int | None = 4096,
    ) -> tuple[np.ndarray, DistanceStats]:
        vertices, faces = self._mesh_arrays(mesh)
        return self.point_to_triangles_distance_torch(
            points,
            vertices,
            faces,
            point_chunk_size=chunk_size,
            triangle_chunk_size=triangle_chunk_size,
        )

    def point_to_triangles_distance_torch(
        self,
        points: np.ndarray,
        vertices: np.ndarray,
        faces: np.ndarray,
        point_chunk_size: int = 2048,
        triangle_chunk_size: int | None = 4096,
    ) -> tuple[np.ndarray, DistanceStats]:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        dtype = torch.float64
        vertices_t = torch.as_tensor(np.asarray(vertices, dtype=np.float64), dtype=dtype, device=device)
        faces_t = torch.as_tensor(np.asarray(faces, dtype=np.int64), dtype=torch.long, device=device)
        triangles = vertices_t[faces_t]
        points_t = torch.as_tensor(points, dtype=dtype, device=device)

        distances = []
        with torch.no_grad():
            for start in range(0, points_t.shape[0], point_chunk_size):
                chunk = points_t[start:start + point_chunk_size]
                if triangle_chunk_size is None or triangle_chunk_size <= 0:
                    dist2 = self._point_to_triangle_distance_squared(chunk, triangles)
                    min_dist2 = dist2.min(dim=1).values
                else:
                    min_dist2 = torch.full(
                        (chunk.shape[0],),
                        torch.inf,
                        dtype=points_t.dtype,
                        device=device,
                    )
                    for triangle_start in range(0, triangles.shape[0], triangle_chunk_size):
                        triangle_chunk = triangles[triangle_start:triangle_start + triangle_chunk_size]
                        dist2 = self._point_to_triangle_distance_squared(chunk, triangle_chunk)
                        min_dist2 = torch.minimum(min_dist2, dist2.min(dim=1).values)

                distances.append(torch.sqrt(min_dist2).cpu().numpy())

        all_distances = np.concatenate(distances) if distances else np.array([], dtype=np.float64)
        return all_distances, self.distance_stats(all_distances)

    def colorize_distances(self, points: np.ndarray, distances: np.ndarray) -> o3d.geometry.PointCloud:
        if len(points) != len(distances):
            raise ValueError("points and distances must have the same length.")
        pointcloud = o3d.geometry.PointCloud()
        pointcloud.points = o3d.utility.Vector3dVector(np.asarray(points, dtype=np.float64))

        if len(distances) == 0:
            pointcloud.colors = o3d.utility.Vector3dVector(np.empty((0, 3), dtype=np.float64))
            return pointcloud

        lower = float(np.min(distances))
        upper = float(np.max(distances))
        normalized = (distances - lower) / (upper - lower + 1e-12)
        colors = self._blue_to_red_colormap(normalized)
        pointcloud.colors = o3d.utility.Vector3dVector(colors)
        return pointcloud

    def distance_stats(self, distances: np.ndarray) -> DistanceStats:
        if len(distances) == 0:
            return DistanceStats(0.0, 0.0, 0.0, 0.0, 0.0)
        return DistanceStats(
            mean=float(np.mean(distances)),
            max=float(np.max(distances)),
            rmse=float(np.sqrt(np.mean(distances ** 2))),
            min=float(np.min(distances)),
            std=float(np.std(distances)),
        )

    def registration_summary(self, transformation: np.ndarray, fitness: float, rmse: float) -> dict[str, Any]:
        return {
            "fitness": fitness,
            "rmse": rmse,
            "transformation": transformation,
        }

    def _mesh_arrays(self, mesh: trimesh.Trimesh) -> tuple[np.ndarray, np.ndarray]:
        vertices = np.asarray(mesh.vertices, dtype=np.float64)
        faces = np.asarray(mesh.faces, dtype=np.int64)
        if len(vertices) == 0 or len(faces) == 0:
            raise ValueError("Mesh must contain vertices and faces.")
        return vertices, faces

    def _blue_to_red_colormap(self, values: np.ndarray) -> np.ndarray:
        values = np.clip(np.asarray(values, dtype=np.float64), 0.0, 1.0)
        colors = np.zeros((len(values), 3), dtype=np.float64)
        colors[:, 0] = np.clip(1.8 * values - 0.2, 0.0, 1.0)
        colors[:, 1] = np.clip(1.0 - np.abs(values - 0.5) * 2.0, 0.0, 1.0)
        colors[:, 2] = np.clip(1.2 - 1.8 * values, 0.0, 1.0)
        return colors

    def _point_to_triangle_distance_squared(self, points: torch.Tensor, triangles: torch.Tensor) -> torch.Tensor:
        p = points[:, None, :]
        a = triangles[None, :, 0, :]
        b = triangles[None, :, 1, :]
        c = triangles[None, :, 2, :]

        ab = b - a
        ac = c - a
        ap = p - a

        d1 = (ab * ap).sum(dim=-1)
        d2 = (ac * ap).sum(dim=-1)
        inf = torch.full_like(d1, torch.inf)
        dist2 = torch.where((d1 <= 0.0) & (d2 <= 0.0), (ap * ap).sum(dim=-1), inf)

        bp = p - b
        d3 = (ab * bp).sum(dim=-1)
        d4 = (ac * bp).sum(dim=-1)
        dist2 = torch.minimum(dist2, torch.where((d3 >= 0.0) & (d4 <= d3), (bp * bp).sum(dim=-1), inf))

        vc = d1 * d4 - d3 * d2
        mask = (vc <= 0.0) & (d1 >= 0.0) & (d3 <= 0.0)
        v = d1 / (d1 - d3).clamp_min(torch.finfo(points.dtype).eps)
        closest = a + v[..., None] * ab
        dist2 = torch.minimum(dist2, torch.where(mask, ((p - closest) ** 2).sum(dim=-1), inf))

        cp = p - c
        d5 = (ab * cp).sum(dim=-1)
        d6 = (ac * cp).sum(dim=-1)
        dist2 = torch.minimum(dist2, torch.where((d6 >= 0.0) & (d5 <= d6), (cp * cp).sum(dim=-1), inf))

        vb = d5 * d2 - d1 * d6
        mask = (vb <= 0.0) & (d2 >= 0.0) & (d6 <= 0.0)
        w = d2 / (d2 - d6).clamp_min(torch.finfo(points.dtype).eps)
        closest = a + w[..., None] * ac
        dist2 = torch.minimum(dist2, torch.where(mask, ((p - closest) ** 2).sum(dim=-1), inf))

        va = d3 * d6 - d5 * d4
        mask = (va <= 0.0) & ((d4 - d3) >= 0.0) & ((d5 - d6) >= 0.0)
        w = (d4 - d3) / ((d4 - d3) + (d5 - d6)).clamp_min(torch.finfo(points.dtype).eps)
        closest = b + w[..., None] * (c - b)
        dist2 = torch.minimum(dist2, torch.where(mask, ((p - closest) ** 2).sum(dim=-1), inf))

        denom = (va + vb + vc).clamp_min(torch.finfo(points.dtype).eps)
        v = vb / denom
        w = vc / denom
        closest = a + ab * v[..., None] + ac * w[..., None]
        return torch.minimum(dist2, ((p - closest) ** 2).sum(dim=-1))