#!/usr/bin/env python3
import argparse
import json
import sys
import time
from pathlib import Path

import numpy as np
from scipy.spatial import cKDTree

sys.path.append(str(Path(__file__).resolve().parent))
import eval_tnt_wrapper as W


def sigmoid(x):
    return 1.0 / (1.0 + np.exp(-x))


def query_tree(tree, points, k=1, batch_size=200000):
    all_d = []
    all_i = []
    for s in range(0, len(points), batch_size):
        e = min(s + batch_size, len(points))
        try:
            d, i = tree.query(points[s:e], k=k, workers=-1)
        except TypeError:
            d, i = tree.query(points[s:e], k=k)
        all_d.append(d)
        all_i.append(i)
    return np.concatenate(all_d, axis=0), np.concatenate(all_i, axis=0)


def gaussian_normals_from_vertex(v, names, indices, transform_mat):
    required = ["scale_0", "scale_1", "scale_2", "rot_0", "rot_1", "rot_2", "rot_3"]
    missing = [k for k in required if k not in names]
    if missing:
        raise ValueError(f"reconstruction PLY missing Gaussian fields: {missing}")

    idx = indices

    scales = np.stack(
        [v["scale_0"][idx], v["scale_1"][idx], v["scale_2"][idx]],
        axis=1,
    ).astype(np.float64)

    min_axis = np.argmin(scales, axis=1)

    q = np.stack(
        [v["rot_0"][idx], v["rot_1"][idx], v["rot_2"][idx], v["rot_3"][idx]],
        axis=1,
    ).astype(np.float64)

    q = q / (np.linalg.norm(q, axis=1, keepdims=True) + 1e-12)

    # 3DGS convention is usually [w, x, y, z].
    w, x, y, z = q[:, 0], q[:, 1], q[:, 2], q[:, 3]

    r00 = 1 - 2 * (y * y + z * z)
    r01 = 2 * (x * y - w * z)
    r02 = 2 * (x * z + w * y)

    r10 = 2 * (x * y + w * z)
    r11 = 1 - 2 * (x * x + z * z)
    r12 = 2 * (y * z - w * x)

    r20 = 2 * (x * z - w * y)
    r21 = 2 * (y * z + w * x)
    r22 = 1 - 2 * (x * x + y * y)

    normals = np.empty((len(idx), 3), dtype=np.float64)

    m0 = min_axis == 0
    normals[m0, 0] = r00[m0]
    normals[m0, 1] = r10[m0]
    normals[m0, 2] = r20[m0]

    m1 = min_axis == 1
    normals[m1, 0] = r01[m1]
    normals[m1, 1] = r11[m1]
    normals[m1, 2] = r21[m1]

    m2 = min_axis == 2
    normals[m2, 0] = r02[m2]
    normals[m2, 1] = r12[m2]
    normals[m2, 2] = r22[m2]

    # T&T trans is similarity-like. Transform directions by linear part.
    A = transform_mat[:3, :3].astype(np.float64)
    normals = normals @ A.T
    normals = normals / (np.linalg.norm(normals, axis=1, keepdims=True) + 1e-12)

    return normals.astype(np.float32), scales[min_axis, np.arange(len(idx))] if False else scales


def estimate_gt_normals_pca(gt_points, gt_tree, gt_query_points, normal_k, batch_size):
    """
    For each query point on/near GT, find K GT neighbors and estimate PCA normal.
    """
    _, nn = query_tree(gt_tree, gt_query_points, k=normal_k, batch_size=batch_size)

    if nn.ndim == 1:
        raise ValueError("normal_k must be > 1")

    normals = np.empty((len(gt_query_points), 3), dtype=np.float32)

    for s in range(0, len(gt_query_points), batch_size):
        e = min(s + batch_size, len(gt_query_points))
        neigh = gt_points[nn[s:e]]  # [B, K, 3]
        centered = neigh - neigh.mean(axis=1, keepdims=True)
        cov = np.einsum("bki,bkj->bij", centered, centered) / max(normal_k - 1, 1)
        vals, vecs = np.linalg.eigh(cov)
        n = vecs[:, :, 0]
        n = n / (np.linalg.norm(n, axis=1, keepdims=True) + 1e-12)
        normals[s:e] = n.astype(np.float32)

    return normals


def summarize_angles(angles):
    return {
        "normal_angular_error_mean_deg": float(np.mean(angles)),
        "normal_angular_error_median_deg": float(np.median(angles)),
        "normal_angular_error_q10_deg": float(np.quantile(angles, 0.10)),
        "normal_angular_error_q25_deg": float(np.quantile(angles, 0.25)),
        "normal_angular_error_q75_deg": float(np.quantile(angles, 0.75)),
        "normal_angular_error_q90_deg": float(np.quantile(angles, 0.90)),
        "normal_angular_error_iqr_deg": float(np.quantile(angles, 0.75) - np.quantile(angles, 0.25)),
    }


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--method", required=True)
    ap.add_argument("--scene", required=True, choices=sorted(W.SCENE_MAP.keys()))
    ap.add_argument("--project-root", default="/root/autodl-tmp/SplatAtlas")
    ap.add_argument("--outputs-root", default=None)
    ap.add_argument("--tnt-eval-root", default=None)
    ap.add_argument("--iteration", type=int, default=None)

    ap.add_argument("--mode", choices=["all", "subsample"], default="subsample")
    ap.add_argument("--n-sample", type=int, default=200000)
    ap.add_argument("--seed", type=int, default=0)

    ap.add_argument("--distance-multiplier", type=float, default=2.0)
    ap.add_argument("--normal-k", type=int, default=30)
    ap.add_argument("--max-normal-points", type=int, default=50000)
    ap.add_argument("--batch-size", type=int, default=50000)
    ap.add_argument("--verbose", action="store_true")
    args = ap.parse_args()

    t0 = time.time()

    project_root = Path(args.project_root)
    outputs_root = Path(args.outputs_root) if args.outputs_root else project_root / "outputs"
    tnt_eval_root = Path(args.tnt_eval_root) if args.tnt_eval_root else project_root / "data" / "tnt_eval"

    scene_lower = args.scene.lower()
    official_scene = W.SCENE_MAP[scene_lower]
    tau = W.TAU_DICT[scene_lower]
    near_threshold = args.distance_multiplier * tau

    ply_path = W.locate_recon_ply(outputs_root, args.method, scene_lower, args.iteration)

    scene_eval_dir = tnt_eval_root / official_scene
    gt_ply_path = scene_eval_dir / f"{official_scene}.ply"
    crop_path = scene_eval_dir / f"{official_scene}.json"
    trans_path = scene_eval_dir / f"{official_scene}_trans.txt"

    if args.verbose:
        print("=" * 80)
        print("Normal angular error evaluation")
        print("method:", args.method)
        print("scene:", scene_lower, "->", official_scene)
        print("tau:", tau)
        print("near threshold:", near_threshold)
        print("recon:", ply_path)
        print("gt:", gt_ply_path)

    trans = W.read_transform(trans_path)
    crop = W.load_crop(crop_path)

    # Load recon.
    recon_raw, recon_vertex, recon_names = W.load_vertex_data(ply_path)
    recon_aligned = W.apply_transform(recon_raw, trans)
    recon_crop_mask = W.crop_mask_tnt(recon_aligned, crop)
    recon_crop = recon_aligned[recon_crop_mask]
    recon_crop_raw_idx = np.where(recon_crop_mask)[0].astype(np.int64)

    if len(recon_crop) == 0:
        raise RuntimeError("Reconstruction crop is empty.")

    # Subsample reconstruction before NN and normal calculation.
    eval_idx_in_crop = W.choose_eval_indices(
        len(recon_crop), args.mode, args.n_sample, args.seed
    )
    recon_eval = recon_crop[eval_idx_in_crop]
    recon_eval_raw_idx = recon_crop_raw_idx[eval_idx_in_crop]

    # Load GT.
    gt_raw, _, _ = W.load_vertex_data(gt_ply_path)
    gt_crop_mask = W.crop_mask_tnt(gt_raw, crop)
    gt_crop = gt_raw[gt_crop_mask]

    if len(gt_crop) == 0:
        raise RuntimeError("GT crop is empty.")

    if args.verbose:
        print("n_gaussians_recon:", len(recon_raw))
        print("n_recon_after_crop:", len(recon_crop))
        print("n_recon_eval:", len(recon_eval))
        print("n_gt_after_crop:", len(gt_crop))

    # Nearest GT distance for each evaluated Gaussian center.
    gt_tree = cKDTree(gt_crop)
    d_r2g, nn_gt_idx = query_tree(
        gt_tree, recon_eval, k=1, batch_size=max(args.batch_size, 1)
    )

    surface_mask = d_r2g < near_threshold
    surface_indices_eval = np.where(surface_mask)[0].astype(np.int64)

    if args.verbose:
        print("surface-near filter:", f"d < {near_threshold}")
        print("n_surface_near:", len(surface_indices_eval), "/", len(recon_eval))
        if len(d_r2g):
            print("distance median:", float(np.median(d_r2g)))
            print("distance q10:", float(np.quantile(d_r2g, 0.10)))
            print("distance q90:", float(np.quantile(d_r2g, 0.90)))

    if len(surface_indices_eval) == 0:
        raise RuntimeError("No surface-near Gaussians found. Increase --distance-multiplier.")

    # Cap normal points for speed and comparable statistics.
    rng = np.random.default_rng(args.seed)
    if len(surface_indices_eval) > args.max_normal_points:
        chosen = np.sort(
            rng.choice(surface_indices_eval, size=args.max_normal_points, replace=False)
        ).astype(np.int64)
    else:
        chosen = surface_indices_eval

    recon_normal_points = recon_eval[chosen]
    recon_raw_indices_for_normals = recon_eval_raw_idx[chosen]
    nearest_gt_points = gt_crop[nn_gt_idx[chosen]]

    # Gaussian normals.
    gauss_normals, gauss_scales = gaussian_normals_from_vertex(
        recon_vertex,
        recon_names,
        recon_raw_indices_for_normals,
        trans,
    )

    # GT normals by local PCA.
    gt_normals = estimate_gt_normals_pca(
        gt_points=gt_crop,
        gt_tree=gt_tree,
        gt_query_points=nearest_gt_points,
        normal_k=args.normal_k,
        batch_size=args.batch_size,
    )

    dots = np.sum(gauss_normals * gt_normals, axis=1)
    dots = np.clip(np.abs(dots), 0.0, 1.0)
    angles = np.degrees(np.arccos(dots))

    scale_min = np.min(gauss_scales, axis=1)
    scale_max = np.max(gauss_scales, axis=1)
    anisotropy = np.exp(scale_max - scale_min)

    result = {
        "method": args.method,
        "scene": scene_lower,
        "official_scene": official_scene,
        "eval_protocol": "auxiliary_gaussian_minor_axis_vs_tnt_gt_pca_normal",
        "is_official_tnt_metric": False,
        "ply_path": str(ply_path),
        "gt_ply_path": str(gt_ply_path),
        "crop_path": str(crop_path),
        "trans_path": str(trans_path),

        "tau": float(tau),
        "surface_near_threshold": float(near_threshold),
        "surface_near_rule": f"d_recon_to_gt < {args.distance_multiplier} * tau",

        "mode": args.mode,
        "n_sample_requested": int(args.n_sample),
        "seed": int(args.seed),
        "normal_k": int(args.normal_k),
        "max_normal_points": int(args.max_normal_points),

        "n_gaussians_recon": int(len(recon_raw)),
        "n_recon_after_crop": int(len(recon_crop)),
        "n_recon_eval": int(len(recon_eval)),
        "n_gt_after_crop": int(len(gt_crop)),
        "n_surface_near": int(len(surface_indices_eval)),
        "surface_near_ratio": float(len(surface_indices_eval) / max(len(recon_eval), 1)),
        "n_normal_eval": int(len(chosen)),

        "distance_recon_to_gt_mean": float(np.mean(d_r2g)),
        "distance_recon_to_gt_median": float(np.median(d_r2g)),
        "distance_recon_to_gt_q10": float(np.quantile(d_r2g, 0.10)),
        "distance_recon_to_gt_q90": float(np.quantile(d_r2g, 0.90)),

        **summarize_angles(angles),

        "gaussian_log_scale_min_median": float(np.median(scale_min)),
        "gaussian_log_scale_max_median": float(np.median(scale_max)),
        "gaussian_anisotropy_exp_scale_range_median": float(np.median(anisotropy)),

        "wall_time_seconds": float(time.time() - t0),
    }

    if "opacity" in recon_names:
        opacity_raw = recon_vertex["opacity"][recon_raw_indices_for_normals].astype(np.float64)
        opacity_sigmoid = sigmoid(opacity_raw)
        result.update({
            "opacity_raw_median": float(np.median(opacity_raw)),
            "opacity_sigmoid_median": float(np.median(opacity_sigmoid)),
            "opacity_sigmoid_q25": float(np.quantile(opacity_sigmoid, 0.25)),
            "opacity_sigmoid_q75": float(np.quantile(opacity_sigmoid, 0.75)),
        })

    out_dir = outputs_root / "tnt_eval_normals" / f"{args.method}_{scene_lower}"
    out_dir.mkdir(parents=True, exist_ok=True)
    out_path = out_dir / "normal_eval.json"
    with open(out_path, "w") as f:
        json.dump(result, f, indent=2, sort_keys=True)

    print(json.dumps(result, indent=2, sort_keys=True))
    print(f"\n[WROTE] {out_path}")


if __name__ == "__main__":
    main()