#!/usr/bin/env python3
"""Audit fused and per-view binary PLY point clouds without loading them fully."""

from __future__ import annotations

import argparse
import json
from dataclasses import dataclass
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
from scipy.spatial import cKDTree


PLY_DTYPE = np.dtype([
    ("x", "<f4"), ("y", "<f4"), ("z", "<f4"),
    ("red", "u1"), ("green", "u1"), ("blue", "u1"),
])


@dataclass
class PlyInfo:
    path: Path
    count: int
    header_bytes: int


def ply_info(path: Path) -> PlyInfo:
    with path.open("rb") as fh:
        header = b""
        while b"end_header\n" not in header:
            block = fh.read(4096)
            if not block:
                raise ValueError(f"Missing PLY end_header: {path}")
            header += block
            if len(header) > 65536:
                raise ValueError(f"Unexpectedly large PLY header: {path}")
    end = header.index(b"end_header\n") + len(b"end_header\n")
    text = header[:end].decode("ascii")
    if "format binary_little_endian 1.0" not in text:
        raise ValueError(f"Only binary little-endian PLY is supported: {path}")
    line = next(x for x in text.splitlines() if x.startswith("element vertex "))
    count = int(line.rsplit(" ", 1)[1])
    expected = end + count * PLY_DTYPE.itemsize
    if path.stat().st_size != expected:
        raise ValueError(f"PLY size/header mismatch: {path} ({path.stat().st_size} != {expected})")
    return PlyInfo(path, count, end)


def mmap_ply(info: PlyInfo) -> np.memmap:
    return np.memmap(info.path, dtype=PLY_DTYPE, mode="r", offset=info.header_bytes,
                     shape=(info.count,))


def sample_ply(info: PlyInfo, limit: int) -> tuple[np.ndarray, np.ndarray]:
    data = mmap_ply(info)
    n = min(info.count, limit)
    # Midpoints of equal-width bins avoid first/last-frame and random-seed bias.
    idx = ((np.arange(n, dtype=np.int64) * info.count + info.count // 2) // n)
    rows = data[idx]
    xyz = np.column_stack((rows["x"], rows["y"], rows["z"])).astype(np.float64)
    rgb = np.column_stack((rows["red"], rows["green"], rows["blue"])).astype(np.uint8)
    keep = np.isfinite(xyz).all(axis=1)
    return xyz[keep], rgb[keep]


def stats(info: PlyInfo, xyz: np.ndarray, rgb: np.ndarray) -> dict:
    q = np.percentile(xyz, [0, .1, 1, 50, 99, 99.9, 100], axis=0)
    centered = xyz - xyz.mean(axis=0)
    cov = centered.T @ centered / max(1, len(xyz) - 1)
    eigval, eigvec = np.linalg.eigh(cov)
    order = np.argsort(eigval)[::-1]
    eigval, eigvec = eigval[order], eigvec[:, order]

    # Work in a central trimmed box for a stable nearest-neighbour spacing estimate.
    lo, hi = np.percentile(xyz, [1, 99], axis=0)
    central = xyz[np.all((xyz >= lo) & (xyz <= hi), axis=1)]
    if len(central) > 50000:
        central = central[np.linspace(0, len(central) - 1, 50000).astype(int)]
    tree = cKDTree(central)
    nn = tree.query(central, k=2, workers=-1)[0][:, 1]

    voxel = {}
    for size in (0.1, 0.5, 1.0, 2.0, 5.0):
        keys = np.floor(xyz / size).astype(np.int64)
        unique = len(np.unique(keys, axis=0))
        voxel[str(size)] = {
            "unique_sample_voxels": unique,
            "sample_occupancy_ratio": unique / len(xyz),
        }
    return {
        "path": str(info.path),
        "file_bytes": info.path.stat().st_size,
        "vertex_count": info.count,
        "sample_count": len(xyz),
        "sample_quantiles_xyz": {str(p): q[i].tolist() for i, p in enumerate((0, .1, 1, 50, 99, 99.9, 100))},
        "sample_extent_xyz": (q[-1] - q[0]).tolist(),
        "sample_centroid_xyz": xyz.mean(axis=0).tolist(),
        "pca_variance": eigval.tolist(),
        "pca_axes_columns": eigvec.tolist(),
        "nearest_neighbor_sample": {
            "p50": float(np.percentile(nn, 50)),
            "p90": float(np.percentile(nn, 90)),
            "p99": float(np.percentile(nn, 99)),
            "zero_ratio": float(np.mean(nn == 0)),
        },
        "rgb_min": rgb.min(axis=0).tolist(),
        "rgb_max": rgb.max(axis=0).tolist(),
        "voxel_sample": voxel,
    }


def full_scan_check(info: PlyInfo, chunk_size: int = 5_000_000) -> dict:
    """Read every fused vertex to validate coordinates and exact bounds."""
    data = mmap_ply(info)
    lower = np.full(3, np.inf)
    upper = np.full(3, -np.inf)
    nonfinite = 0
    adjacent_duplicates = 0
    previous = None
    for start in range(0, info.count, chunk_size):
        rows = data[start:min(start + chunk_size, info.count)]
        xyz = np.column_stack((rows["x"], rows["y"], rows["z"]))
        finite = np.isfinite(xyz).all(axis=1)
        nonfinite += int((~finite).sum())
        xyz = xyz[finite]
        if not len(xyz):
            continue
        lower = np.minimum(lower, xyz.min(axis=0))
        upper = np.maximum(upper, xyz.max(axis=0))
        if previous is not None:
            adjacent_duplicates += int(np.all(previous == xyz[0]))
        adjacent_duplicates += int(np.all(xyz[1:] == xyz[:-1], axis=1).sum())
        previous = xyz[-1]
    return {
        "vertex_count": info.count,
        "min_xyz": lower.tolist(),
        "max_xyz": upper.tolist(),
        "extent_xyz": (upper - lower).tolist(),
        "nonfinite_vertices": nonfinite,
        "adjacent_exact_duplicates": adjacent_duplicates,
    }


def sampled_distance_to_fused(fused: np.ndarray, component: np.ndarray) -> dict:
    """Diagnostic distance only; samples are not known point correspondences."""
    lower, upper = np.percentile(fused, [.2, 99.8], axis=0)
    fused = fused[np.all((fused >= lower) & (fused <= upper), axis=1)]
    component = component[np.all((component >= lower - 5) & (component <= upper + 5), axis=1)]
    distance = cKDTree(fused).query(component, k=1, workers=-1)[0]
    return {
        "component_sample_kept": len(component),
        "p50_mm": float(np.percentile(distance, 50)),
        "p90_mm": float(np.percentile(distance, 90)),
        "p95_mm": float(np.percentile(distance, 95)),
        "within_2mm": float(np.mean(distance <= 2)),
        "within_5mm": float(np.mean(distance <= 5)),
        "interpretation": "sample-to-sample diagnostic; not a calibrated correspondence residual",
    }


def overlap_score(a: np.ndarray, b: np.ndarray, voxel_size: float) -> dict:
    origin = np.minimum(np.percentile(a, 1, axis=0), np.percentile(b, 1, axis=0))
    qa = np.floor((a - origin) / voxel_size).astype(np.int32)
    qb = np.floor((b - origin) / voxel_size).astype(np.int32)
    # Structured 3-int view makes set operations substantially cheaper.
    row = np.dtype((np.void, qa.dtype.itemsize * qa.shape[1]))
    ua = np.unique(np.ascontiguousarray(qa).view(row).ravel())
    ub = np.unique(np.ascontiguousarray(qb).view(row).ravel())
    inter = np.intersect1d(ua, ub, assume_unique=True).size
    union = len(ua) + len(ub) - inter
    return {
        "voxel_size": voxel_size,
        "intersection": int(inter),
        "union": int(union),
        "jaccard": inter / union if union else 1.0,
        "a_covered": inter / len(ua) if len(ua) else 1.0,
        "b_covered": inter / len(ub) if len(ub) else 1.0,
    }


def robust_limits(groups: list[np.ndarray], axes=(0, 1)):
    cat = np.concatenate([g[:, axes] for g in groups], axis=0)
    lo = np.percentile(cat, .2, axis=0)
    hi = np.percentile(cat, 99.8, axis=0)
    pad = np.maximum((hi - lo) * .03, 1e-6)
    return lo - pad, hi + pad


def scatter_projection(ax, xyz, axes, color, title, size=.12, alpha=.45):
    ax.scatter(xyz[:, axes[0]], xyz[:, axes[1]], s=size, c=color, alpha=alpha,
               linewidths=0, rasterized=True)
    ax.set_title(title, fontsize=10, weight="bold")
    ax.set_aspect("equal", adjustable="box")
    ax.set_xlabel("XYZ"[axes[0]])
    ax.set_ylabel("XYZ"[axes[1]])
    ax.grid(alpha=.12, linewidth=.5)


def find_runs(root: Path) -> list[dict]:
    fused = sorted(root.glob("**/all_merge_pointcloud_full.ply"))
    # Viewer copies and the root-level 18h duplicate do not need separate analysis.
    fused = [p for p in fused if "viewer" not in p.parts and p != root / "datas/all_merge_pointcloud_full.ply"]
    runs = []
    for p in fused:
        comps = sorted(p.parent.glob("merged_pointcloud_*.ply"))
        if "ac28_ply" in str(p.parent):
            label = "ac28_ply"
        elif "18" in p.parent.name:
            label = "18h01"
        elif "16" in p.parent.name:
            label = "16h"
        else:
            label = p.parent.name
        runs.append({"label": label, "fused": p, "components": comps})
    return runs


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--root", type=Path, default=Path(__file__).resolve().parents[1])
    ap.add_argument("--out", type=Path, default=None)
    ap.add_argument("--sample", type=int, default=180000)
    args = ap.parse_args()
    root = args.root.resolve()
    out = (args.out or root / "results/ply_audit").resolve()
    out.mkdir(parents=True, exist_ok=True)

    runs = find_runs(root)
    report = {"method": {"sampling": "deterministic equal-width bin midpoints", "sample_limit": args.sample}, "runs": []}
    full_checks = {}
    sensor_distances = {}
    sampled: dict[str, tuple[np.ndarray, np.ndarray]] = {}

    for run in runs:
        files = [run["fused"], *run["components"]]
        row = {"label": run["label"], "fused": str(run["fused"].relative_to(root)), "components": [], "files": {}}
        for path in files:
            info = ply_info(path)
            xyz, rgb = sample_ply(info, args.sample)
            key = str(path)
            sampled[key] = (xyz, rgb)
            row["files"][path.name] = stats(info, xyz, rgb)
            if path != run["fused"]:
                row["components"].append(path.name)
        fxyz = sampled[str(run["fused"])][0]
        full_checks[run["label"]] = full_scan_check(ply_info(run["fused"]))
        row["component_vs_fused_overlap"] = {}
        sensor_distances[run["label"]] = {}
        for path in run["components"]:
            cxyz = sampled[str(path)][0]
            row["component_vs_fused_overlap"][path.stem] = overlap_score(fxyz, cxyz, 5.0)
            sensor = path.stem.replace("merged_pointcloud_", "")
            sensor_distances[run["label"]][sensor] = sampled_distance_to_fused(fxyz, cxyz)
        report["runs"].append(row)

    # Overview: each fused cloud in its three principal Cartesian projections.
    fig, axs = plt.subplots(len(runs), 3, figsize=(15, 4.4 * len(runs)), constrained_layout=True)
    if len(runs) == 1:
        axs = axs[None, :]
    for i, run in enumerate(runs):
        xyz, rgb = sampled[str(run["fused"])]
        colors = np.clip(rgb.astype(float) / 255.0, .04, 1)
        for j, axes in enumerate(((0, 1), (0, 2), (1, 2))):
            scatter_projection(axs[i, j], xyz, axes, colors, f"{run['label']} fused · {'XYZ'[axes[0]]}{'XYZ'[axes[1]]}")
            lo, hi = robust_limits([xyz], axes)
            axs[i, j].set_xlim(lo[0], hi[0]); axs[i, j].set_ylim(lo[1], hi[1])
    fig.suptitle("all_merge_pointcloud_full.ply — fused cloud audit", fontsize=15, weight="bold")
    fig.savefig(out / "fused_overview.png", dpi=180)
    plt.close(fig)

    # Per-run overlays use sensor identity as colour to expose misregistration/seams.
    palette = ["#2E86DE", "#E74C3C", "#27AE60", "#8E44AD", "#F39C12"]
    for run in runs:
        if not run["components"]:
            continue
        groups = [sampled[str(p)][0] for p in run["components"]]
        # Cap only for rendering; all sampled points remain in JSON statistics.
        groups = [g[::max(1, len(g)//60000)] for g in groups]
        fig, axs = plt.subplots(1, 3, figsize=(16, 5), constrained_layout=True)
        for ax, axes in zip(axs, ((0, 1), (0, 2), (1, 2))):
            for g, path, color in zip(groups, run["components"], palette):
                ax.scatter(g[:, axes[0]], g[:, axes[1]], s=.15, color=color, alpha=.25,
                           linewidths=0, rasterized=True, label=path.stem.replace("merged_pointcloud_", ""))
            ax.set_title(f"{run['label']} sensors · {'XYZ'[axes[0]]}{'XYZ'[axes[1]]}", weight="bold")
            ax.set_aspect("equal", adjustable="box")
            ax.set_xlabel("XYZ"[axes[0]]); ax.set_ylabel("XYZ"[axes[1]])
            lo, hi = robust_limits(groups, axes)
            ax.set_xlim(lo[0], hi[0]); ax.set_ylim(lo[1], hi[1]); ax.grid(alpha=.12)
        axs[-1].legend(markerscale=20, fontsize=8, loc="best")
        fig.savefig(out / f"{run['label']}_sensor_overlay.png", dpi=180)
        plt.close(fig)

    with (out / "report.json").open("w", encoding="utf-8") as fh:
        json.dump(report, fh, ensure_ascii=False, indent=2)
    with (out / "full_scan_checks.json").open("w", encoding="utf-8") as fh:
        json.dump(full_checks, fh, ensure_ascii=False, indent=2)
    with (out / "sensor_to_fused_sample_distances.json").open("w", encoding="utf-8") as fh:
        json.dump(sensor_distances, fh, ensure_ascii=False, indent=2)
    print(out / "report.json")


if __name__ == "__main__":
    main()