"""S23DR 2026 submission: rayv9_learnt_baseline_snap

Pipeline per scene:
  1. v1d heatmap -> spatial NMS -> rays
  2. rays + SfM -> voxel volume -> RayVoxelTransformer -> hybrid NMS -> vertices
  3. fuse_and_sample -> learned baseline -> wireframe
  4. midpoint-snap baseline verts to v9, append unmatched v9 verts

Val HSS (100 scenes): baseline=0.352  snap=0.411  (+0.059, wins=85/100)
Vertex F1@0.5=0.494  F1@1.0=0.685

Usage:
  # local smoke-test on training split (n scenes)
  python script.py --mode local --n_scenes 10

  # competition submission (reads params.json, writes submission.json)
  python script.py
"""
import argparse
import json
import time
from pathlib import Path

import numpy as np
import torch

import v9_inference as v9
import baseline_inference as bl
from snap import snap_midpoint_plus_unmatched
from s23dr_2026_example.point_fusion import FuserConfig

SCRIPT_DIR    = Path(__file__).resolve().parent
V1D_CKPT      = SCRIPT_DIR / "v1d_checkpoint.pt"
V9_CKPT       = SCRIPT_DIR / "v9_checkpoint.pt"
BASELINE_CKPT = SCRIPT_DIR / "baseline_checkpoint.pt"


def empty_solution():
    return np.zeros((2, 3)), [(0, 1)]


def run(dataset, v1d_model, v9_model, v9_img_size, bl_model, device, n_scenes=None):
    cfg = FuserConfig()
    rng = np.random.RandomState(2718)
    solution = []
    processed = 0
    t0 = time.time()

    for subset_name in dataset:
        print(f"\nProcessing {subset_name}...", flush=True)
        for sample in dataset[subset_name]:
            if n_scenes is not None and processed >= n_scenes:
                break
            order_id = sample["order_id"]

            v9_verts = np.zeros((0, 3))
            try:
                v9_verts = v9.predict_vertices(
                    sample, v1d_model, v9_model, v9_img_size, device)
            except Exception as e:
                print(f"  v9 failed {order_id}: {e}", flush=True)

            bl_result = None
            try:
                fused = bl.fuse_and_sample(sample, cfg, rng)
                if fused is not None:
                    bl_result = bl.predict(fused, bl_model, device)
            except Exception as e:
                print(f"  baseline failed {order_id}: {e}", flush=True)

            if bl_result is None:
                pred_v, pred_e = empty_solution()
            else:
                pred_v, pred_e = snap_midpoint_plus_unmatched(bl_result[0], bl_result[1], v9_verts)

            solution.append({
                "order_id":    order_id,
                "wf_vertices": pred_v.tolist(),
                "wf_edges":    [(int(a), int(b)) for a, b in pred_e],
            })
            processed += 1
            elapsed = time.time() - t0
            print(f"  [{processed}] {order_id}  "
                  f"v9={len(v9_verts)}  bl={'ok' if bl_result else 'fail'}  "
                  f"{elapsed:.0f}s elapsed", flush=True)

    return solution


def load_models(device):
    print("Loading v1d...", flush=True)
    v1d_model = v9.load_v1d(V1D_CKPT, device)
    print("Loading v9...", flush=True)
    v9_model, v9_img_size = v9.load_v9(V9_CKPT, device)
    print("Loading baseline...", flush=True)
    bl_model = bl.load_model(BASELINE_CKPT, device)
    return v1d_model, v9_model, v9_img_size, bl_model


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--mode", choices=["local", "submit"], default="submit",
                        help="local: stream from hoho22k_2026_trainval training split; "
                             "submit: read params.json and use test data")
    parser.add_argument("--n_scenes", type=int, default=None,
                        help="cap number of scenes (local mode)")
    args = parser.parse_args()

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Device: {device}", flush=True)

    v1d_model, v9_model, v9_img_size, bl_model = load_models(device)

    from datasets import load_dataset

    if args.mode == "local":
        print("Mode: local (hoho22k_2026_trainval / train split)", flush=True)
        dataset = load_dataset("usm3d/hoho22k_2026_trainval", split="train",
                               streaming=True, trust_remote_code=True)
        # Wrap in dict to match the multi-subset loop
        dataset = {"train": dataset}
    else:
        print("Mode: submit", flush=True)
        with open("params.json") as f:
            params = json.load(f)
        print(f"Competition: {params.get('competition_id', '?')}", flush=True)
        data_path = Path("/tmp/data")
        if not data_path.exists():
            from huggingface_hub import snapshot_download
            snapshot_download(repo_id=params["dataset"], local_dir="/tmp/data",
                              repo_type="dataset")
        data_files = {
            "validation": [str(p) for p in data_path.rglob("*public*/**/*.tar")],
            "test":       [str(p) for p in data_path.rglob("*private*/**/*.tar")],
        }
        dataset = load_dataset(
            str(data_path / "hoho22k_2026_test_x_anon.py"),
            data_files=data_files,
            trust_remote_code=True,
            writer_batch_size=100,
        )
        print(f"Loaded: {dataset}", flush=True)

    solution = run(dataset, v1d_model, v9_model, v9_img_size, bl_model,
                   device, n_scenes=args.n_scenes)

    with open("submission.json", "w") as f:
        json.dump(solution, f)
    print(f"\nSaved submission.json ({len(solution)} entries)", flush=True)