### Competition submission entry point.
### Loads checkpoint, runs inference, writes submission.json.

import sys
import os
import subprocess
sys.path.insert(0, os.path.dirname(__file__))

# ── Install deps ───────────────────────────────────────────────────────
def _pip(*pkgs):
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q"] + list(pkgs))

for pkg, mod in [
    ("torch==2.4.0", "torch"),
    ("numpy<2", "numpy"),
    ("scipy", "scipy"),
    ("plyfile", "plyfile"),
    ("pycolmap", "pycolmap"),
    ("tqdm", "tqdm"),
    ("joblib", "joblib"),
    ("huggingface-hub", "huggingface_hub"),
    ("datasets==3.6.0", "datasets"),
]:
    try:
        __import__(mod.split("=")[0])
    except ImportError:
        _pip(pkg)

import json
from pathlib import Path

import numpy as np
import torch
from datasets import load_dataset
from joblib import Parallel, delayed
from tqdm import tqdm

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

CHECKPOINT   = Path(__file__).parent / "wireframe_detr_cdn_multiscale_384d_128q.pth"
TTA_ROTATIONS = 1       # TTA off — use 4 for +HSS at cost of 4x inference time
NUM_POINTS   = 7168
SCORE_THRESH = 0.9
MERGE_DIST   = 0.5


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


def predict_wireframe_safely(sample, model):
    try:
        scene = Scene(sample)
        verts, edges = predict_wireframe_v2(
            scene, model, "cuda",
            pt_type="colmap",
            num_points=NUM_POINTS,
            score_threshold=SCORE_THRESH,
            merge_distance_threshold=MERGE_DIST,
        )
        if len(edges) == 0:
            verts, edges, _ = empty_solution(sample)
    except Exception as e:
        print(f"Failed ({sample.get('order_id', '?')}): {e} — empty solution")
        verts, edges, _ = empty_solution(sample)
    edges = [(int(a), int(b)) for a, b in edges]
    return verts, edges, sample["order_id"]


if __name__ == "__main__":
    print("------------ Loading dataset ------------")
    with open("params.json") as f:
        params = json.load(f)
    print(params)

    data_path = Path("/tmp/data")
    if not data_path.exists():
        from huggingface_hub import snapshot_download
        snapshot_download(repo_id=params["dataset"], local_dir=str(data_path), 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")],
    }
    print(data_files)
    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("------------ Loading model ------------")
    set_random_seed(0)
    checkpoint = torch.load(CHECKPOINT, map_location="cuda")
    model = get_model(checkpoint, num_classes=1)
    load_checkpoint_compat(model, checkpoint)
    model.to("cuda").eval()
    print(f"Model loaded from {CHECKPOINT}")
    print(f"Inference: TTA={TTA_ROTATIONS}, {NUM_POINTS} pts, threshold={SCORE_THRESH}")

    print("------------ Predicting ------------")
    solution = []
    for subset_name in dataset:
        print(f"Predicting on {subset_name}")
        preds = Parallel(n_jobs=1, prefer="processes")(
            delayed(predict_wireframe_safely)(sample, model)
            for sample in tqdm(dataset[subset_name])
        )
        for verts, edges, order_id in preds:
            print(f"{order_id}: {len(verts)} verts, {len(edges)} edges")
            solution.append({
                "order_id":    order_id,
                "wf_vertices": verts.tolist(),
                "wf_edges":    edges,
            })

    print("------------ Saving ------------")
    output_path = Path(params.get("output_path", "."))
    output_path.mkdir(parents=True, exist_ok=True)
    for save_path in [Path("submission.json"), output_path / "submission.json"]:
        save_path.parent.mkdir(parents=True, exist_ok=True)
        with open(save_path, "w") as f:
            json.dump(solution, f)
        print(f"Saved → {save_path}")

    print("------------ Done ------------")