"""Test-time augmentation via multi-seed priority sampling.

Runs the model N times with different priority-sample seeds, concatenates the
world-space segment predictions from each pass, and lets the standard
merge_vertices_iterative do the deduplication. Because the iterative merge
takes union-find clusters and uses each cluster's centroid as the merged
position, this is effectively Hungarian-averaging for matched segments —
without needing to solve a real assignment problem.

The previous TTA attempt (commit 857514e, reverted) failed because it picked
ONE pass's output. This implementation aggregates ALL passes' segments and
lets the established merge logic combine them.

Why it should work:
  - Stochastic variation comes from priority-sample seed (the model itself is
    deterministic). Different seeds give different points → slightly different
    model predictions for the same scene.
  - Matched segments (true edges) appear in all passes near each other → they
    cluster in the merge and get averaged toward consensus.
  - Spurious segments (hallucinations) appear in only 1 pass → they survive
    individually but are typically not high-confidence enough to win.
  - The iterative merge thresholds 0.15→0.6 m are appropriate for the typical
    inter-pass jitter of correctly-predicted segments.
"""

from __future__ import annotations

import numpy as np
import torch

import script
from s23dr_2026_example.segment_postprocess import merge_vertices_iterative
from s23dr_2026_example.varifold import segments_to_vertices_edges
from s23dr_2026_example.postprocess_v2 import snap_to_point_cloud, snap_horizontal


def _model_to_world_segments(sample_dict, model, device):
    """Run the model on a single fused sample, return (N, 2, 3) world segments.

    Returns None if no segments pass the confidence threshold.
    """
    tokens, masks = script.build_tokens_single(sample_dict, model, device)
    scale = float(sample_dict["scale"])
    center = sample_dict["center"]
    with torch.no_grad(), torch.autocast(
        device_type='cuda', dtype=torch.float16,
        enabled=(device.type == 'cuda'),
    ):
        out = model.forward_tokens(tokens, masks)
    segs = out["segments"][0].float().cpu()
    conf = (
        torch.sigmoid(out["conf"][0].float()).cpu().numpy()
        if "conf" in out else None
    )
    if conf is not None:
        segs = segs[conf > script.CONF_THRESH]
    if len(segs) < 1:
        return None
    return segs.numpy() * scale + center  # (N, 2, 3)


def _match_segments(anchor_segs, other_segs, max_endpoint_dist=0.4):
    """Hungarian-match segments between two passes (flip-invariant distance).

    Returns list of (i_anchor, j_other, was_flipped) for accepted matches.
    Matches with cost > 2*max_endpoint_dist are rejected.
    """
    if len(anchor_segs) == 0 or len(other_segs) == 0:
        return []
    from scipy.optimize import linear_sum_assignment

    N, M = len(anchor_segs), len(other_segs)
    # vectorized cost computation
    a0 = anchor_segs[:, 0][:, None, :]  # (N, 1, 3)
    a1 = anchor_segs[:, 1][:, None, :]
    b0 = other_segs[None, :, 0]          # (1, M, 3)
    b1 = other_segs[None, :, 1]
    d_same = (np.linalg.norm(a0 - b0, axis=-1) +
              np.linalg.norm(a1 - b1, axis=-1))
    d_flip = (np.linalg.norm(a0 - b1, axis=-1) +
              np.linalg.norm(a1 - b0, axis=-1))
    cost = np.minimum(d_same, d_flip)
    flipped = d_flip < d_same

    row, col = linear_sum_assignment(cost)
    threshold = 2.0 * max_endpoint_dist
    return [(int(i), int(j), bool(flipped[i, j]))
            for i, j in zip(row, col) if cost[i, j] <= threshold]


def predict_sample_tta_hungarian(sample, cfg, model, device,
                                  seeds=(2718, 31415, 42),
                                  match_dist: float = 0.4,
                                  min_passes_for_keep: int = 1,
                                  snap_target_classes=(0, 1, 2)):
    """Hungarian-averaged TTA. Aggregates segments via flip-invariant matching.

    Pass 0 is the anchor. Pass 1+ segments are matched to anchor via Hungarian
    on endpoint distance. Each anchor segment gets averaged with its matches
    (orientation-aligned). Anchor segments with < min_passes_for_keep matches
    are kept only if min_passes_for_keep == 0; otherwise dropped.

    Args:
        min_passes_for_keep: 0 = keep all anchor segments, 1 = require matching
            in at least 1 other pass.
    """
    sample_dicts = []
    all_segs = []
    for seed in seeds:
        rng = np.random.RandomState(int(seed))
        sd = script.fuse_and_sample(sample, cfg, rng)
        if sd is None:
            continue
        sw = _model_to_world_segments(sd, model, device)
        if sw is None or len(sw) == 0:
            continue
        sample_dicts.append(sd)
        all_segs.append(sw)

    if not all_segs:
        return script.empty_solution()

    if len(all_segs) == 1:
        # No TTA gain possible; just run the normal post-process.
        return _post_segments_to_wireframe(
            all_segs[0], sample_dicts[0], snap_target_classes)

    anchor = all_segs[0]
    matches_per_anchor = [[] for _ in range(len(anchor))]  # list of (other_seg, flipped)
    for p_idx in range(1, len(all_segs)):
        for i_a, j_o, flipped in _match_segments(anchor, all_segs[p_idx],
                                                   max_endpoint_dist=match_dist):
            matches_per_anchor[i_a].append((all_segs[p_idx][j_o], flipped))

    averaged = []
    for i, matches in enumerate(matches_per_anchor):
        if len(matches) < min_passes_for_keep:
            continue
        seg = anchor[i]
        if not matches:
            averaged.append(seg)
            continue
        # Align orientations to anchor and average
        aligned = [seg]
        for other_seg, flipped in matches:
            aligned.append(other_seg[::-1] if flipped else other_seg)
        averaged.append(np.mean(aligned, axis=0))

    if not averaged:
        # Defensive: fall back to concat-and-merge if matching dropped everything
        return _post_segments_to_wireframe(
            np.concatenate(all_segs, axis=0), sample_dicts[0],
            snap_target_classes)

    return _post_segments_to_wireframe(
        np.asarray(averaged), sample_dicts[0], snap_target_classes)


def _post_segments_to_wireframe(segments, sd0, snap_target_classes):
    """Standard post-process: segments -> vertices/edges -> merge -> snap."""
    pv, pe = segments_to_vertices_edges(torch.tensor(segments))
    pv, pe = pv.numpy(), np.array(pe, dtype=np.int32)
    pv, pe = merge_vertices_iterative(pv, pe)

    xyz_norm = sd0["xyz_norm"]
    mask = sd0["mask"]
    cid = sd0["class_id"]
    xyz_world = xyz_norm[mask] * float(sd0["scale"]) + sd0["center"]
    cid_valid = cid[mask]
    pv = snap_to_point_cloud(
        pv, xyz_world, cid_valid,
        snap_radius=script.SNAP_RADIUS,
        target_classes=list(snap_target_classes),
    )
    pv = snap_horizontal(pv, pe)

    if len(pv) < 2 or len(pe) < 1:
        return script.empty_solution()
    return pv, [(int(a), int(b)) for a, b in pe]


def predict_sample_tta(sample, cfg, model, device,
                        seeds=(2718, 31415, 42),
                        snap_target_classes=(0, 1, 2)):
    """Multi-seed TTA prediction. Returns (vertices, edges) in world space.

    Args:
        sample: raw dataset entry.
        cfg: FuserConfig.
        model: loaded model.
        device: torch device.
        seeds: tuple of priority-sample seeds. Length = # TTA passes.
        snap_target_classes: classes for snap_to_point_cloud (default
            [apex, eave_end_point, flashing_end_point] = [0, 1, 2]).
    """
    # 1) Run fuse + model for each seed, collect world-space segments.
    sample_dicts = []
    all_segs = []
    for seed in seeds:
        rng = np.random.RandomState(int(seed))
        sd = script.fuse_and_sample(sample, cfg, rng)
        if sd is None:
            continue
        sw = _model_to_world_segments(sd, model, device)
        if sw is None or len(sw) == 0:
            continue
        sample_dicts.append(sd)
        all_segs.append(sw)

    if not all_segs:
        return script.empty_solution()

    # 2) Concatenate segments across passes. The downstream merge will cluster
    #    near-duplicate vertices (matched across passes) and take centroids,
    #    yielding the Hungarian-average behavior.
    combined = np.concatenate(all_segs, axis=0)

    # 3) Standard post-process: segments -> vertices/edges, iterative merge.
    pv, pe = segments_to_vertices_edges(torch.tensor(combined))
    pv, pe = pv.numpy(), np.array(pe, dtype=np.int32)
    pv, pe = merge_vertices_iterative(pv, pe)

    # 4) Snap to point cloud. Use the FIRST sample_dict's context (the
    #    fused points are roughly similar across seeds so this is a reasonable
    #    proxy; using merged xyz would require re-fusing).
    sd0 = sample_dicts[0]
    xyz_norm = sd0["xyz_norm"]
    mask = sd0["mask"]
    cid = sd0["class_id"]
    scale0 = float(sd0["scale"])
    center0 = sd0["center"]
    xyz_world = xyz_norm[mask] * scale0 + center0
    cid_valid = cid[mask]
    pv = snap_to_point_cloud(
        pv, xyz_world, cid_valid,
        snap_radius=script.SNAP_RADIUS,
        target_classes=list(snap_target_classes),
    )
    pv = snap_horizontal(pv, pe)

    if len(pv) < 2 or len(pe) < 1:
        return script.empty_solution()

    edges = [(int(a), int(b)) for a, b in pe]
    return pv, edges