import argparse
from pathlib import Path

import hoho2025.metric_helper
import numpy as np
import torch
from datasets import load_dataset
from tqdm import tqdm

from .model import get_model, load_checkpoint_compat
from .inference import predict_wireframe_v2, predict_wireframe_v2_tta
from .scene import Scene
from .utils import set_random_seed, start_debug

# Scenes with known pose/annotation misalignment, grouped by severity.
# Used to produce a clean validation2 split from the public validation set.
_MISALIGNED_FILE = Path(__file__).resolve().parent.parent / "assets" / "misaligned.txt"
_SEVERITY_ORDER = ("extreme", "severe", "medium")


def load_validation2(stream: bool = False, misalignment_severity: str = "medium") -> object:
    """Load the validation split filtered to validation2 (misalignment-cleaned subset)."""
    dataset = load_dataset("usm3d/hoho22k_2026_trainval", trust_remote_code=True, streaming=stream)
    return _drop_misaligned(dataset["validation"], misalignment_severity)


def _drop_misaligned(data, severity: str = "medium"):
    """Drop scenes with pose/annotation misalignment up to the given severity level."""
    min_idx = _SEVERITY_ORDER.index(severity)
    drop: set[str] = set()
    section = None
    for line in _MISALIGNED_FILE.read_text().splitlines():
        line = line.strip()
        if not line or line.startswith("#"):
            section = line.lstrip("#").strip().lower() if line.startswith("#") else section
            continue
        if section in _SEVERITY_ORDER and _SEVERITY_ORDER.index(section) <= min_idx:
            drop.add(line)
    print(f"Dropping {len(drop)} misaligned scenes (severity <= {severity})")
    return data.filter(lambda s: s["order_id"] not in drop)


def run_eval(model, dataset, device, args) -> list[tuple[str, object]]:
    results = []
    for data in tqdm(dataset, desc="Evaluating"):
        scene = Scene(data)
        if args.tta_rotations > 1:
            verts, edges = predict_wireframe_v2_tta(
                scene, model, device,
                pt_type=args.pt_type,
                num_points=args.num_points,
                score_threshold=args.threshold,
                merge_distance_threshold=args.merge_distance_threshold,

                n_rotations=args.tta_rotations,
            )
        else:
            verts, edges = predict_wireframe_v2(
                scene, model, device,
                pt_type=args.pt_type,
                num_points=args.num_points,
                score_threshold=args.threshold,

                merge_distance_threshold=args.merge_distance_threshold,
                merge_method=args.merge_method,
                feature_alpha=args.feature_alpha,
            )
        result = hoho2025.metric_helper.hss(verts, edges, scene.verts, scene.edges)
        results.append((scene.order_id, result))
        print(f"{scene.order_id}: HSS={result.hss:.4f} F1={result.f1:.4f} IoU={result.iou:.4f}")
    return results


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Evaluate WireframeDETR on a dataset split.")
    parser.add_argument("--split", required=True, help="Dataset split (train / validation / validation2).")
    parser.add_argument("--checkpoint", type=Path, required=True, help="Trained model checkpoint.")
    parser.add_argument("--device", default="cuda")
    parser.add_argument("--seed", type=int, help="Shuffle seed.")
    parser.add_argument("--stream", action="store_true", help="Stream dataset instead of downloading.")
    parser.add_argument("-d", "--debug", action="store_true")
    parser.add_argument("--num_points", "-np", type=int, default=7168)
    parser.add_argument("--threshold", "-t", type=float, default=0.9, help="Edge confidence threshold.")
    parser.add_argument("--merge_distance_threshold", type=float, default=0.5)
    parser.add_argument("--pt_type", type=str, default="colmap", choices=["monodepth", "colmap"])
    parser.add_argument("--merge_method", default="distance", choices=["distance", "feature", "auto"])
    parser.add_argument("--feature_alpha", type=float, default=0.95)
    parser.add_argument("--tta_rotations", type=int, default=1,
                        help="Y-axis TTA rotations (1=off, 4=0/90/180/270deg).")
    parser.add_argument("--misalignment_severity", default="medium", choices=list(_SEVERITY_ORDER),
                        help="Max misalignment severity to drop when split=validation2.")
    parser.add_argument("--results", type=Path, help="CSV path for per-scene scores.")
    args = parser.parse_args()

    if args.debug:
        start_debug()

    if args.split == "validation2":
        dataset = load_validation2(stream=args.stream, misalignment_severity=args.misalignment_severity)
    else:
        dataset = load_dataset("usm3d/hoho22k_2026_trainval", trust_remote_code=True, streaming=args.stream)[args.split]

    set_random_seed(args.seed or 0)
    if args.seed is not None:
        dataset = dataset.shuffle(seed=args.seed)

    checkpoint = torch.load(args.checkpoint, map_location=args.device)
    model = get_model(checkpoint, num_classes=1)
    load_checkpoint_compat(model, checkpoint)
    model.to(args.device).eval()

    results = run_eval(model, dataset, args.device, args)

    mean_hss = np.mean([r.hss for _, r in results])
    mean_f1 = np.mean([r.f1 for _, r in results])
    mean_iou = np.mean([r.iou for _, r in results])
    print(f"\nHSS={mean_hss:.4f}  F1={mean_f1:.4f}  IoU={mean_iou:.4f}")

    if args.results:
        with open(args.results, "w") as f:
            f.write("order_id,hss,f1,iou\n")
            for order_id, result in results:
                f.write(f"{order_id},{result.hss:.4f},{result.f1:.4f},{result.iou:.4f}\n")
        print(f"Saved to {args.results}")