#!/usr/bin/env python
"""Smoke-test the packaged WiSER example scene.

This script intentionally has two levels:
1. Always validate that the example manifests and assets are readable.
2. Optionally load a checkpoint and report its structure.

Full neural inference needs the CUDA sparse backend and the paper checkpoint;
for release QA, asset loading and checkpoint readability catch most packaging
errors without requiring a large GPU run.
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path
from typing import Any

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

import torch

from wiser.data.csi_path_targets import MergedPathTargetConfig
from wiser.data.dataset import MultiSceneTripleDataset
from wiser.data.radiomap_dataset import RadiomapDataset


def _load_json(path: Path) -> dict[str, Any]:
    return json.loads(path.read_text())


def _resolve_record_paths(records: list[dict[str, Any]], manifest_path: Path) -> list[dict[str, Any]]:
    out: list[dict[str, Any]] = []
    base = manifest_path.parent
    for rec in records:
        r = dict(rec)
        tx_path = r.get("tx_path")
        if tx_path and not Path(tx_path).is_absolute():
            r["tx_path"] = str((base / tx_path).resolve())
        out.append(r)
    return out


def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--example-root", default="example")
    p.add_argument("--checkpoint", default=None)
    p.add_argument("--out-json", default="outputs/example_summary.json")
    p.add_argument("--channels", type=int, default=512)
    p.add_argument("--max-cir-samples", type=int, default=16)
    args = p.parse_args()

    example_root = Path(args.example_root).resolve()
    manifests = example_root / "manifests"
    radiomap_manifest_path = manifests / "radiomap_example.json"
    cir_manifest_path = manifests / "cir_example.json"
    scene3d_root = example_root / "data" / "scene3d"
    wireless_root = example_root / "data" / "wireless"

    summary: dict[str, Any] = {
        "example_root": str(example_root),
        "radiomap_manifest": str(radiomap_manifest_path),
        "cir_manifest": str(cir_manifest_path),
    }

    rm_manifest = _load_json(radiomap_manifest_path)
    rm_records = _resolve_record_paths(rm_manifest["val_heldout"], radiomap_manifest_path)
    rm_ds = RadiomapDataset(
        rm_records,
        channels=args.channels,
        grid_h=int(rm_manifest.get("grid_h", 36)),
        grid_w=int(rm_manifest.get("grid_w", 36)),
        db_floor=-300.0,
        scene3d_root=scene3d_root,
        dataset_kind=rm_manifest.get("dataset_kind", "sionna_radiomap"),
    )
    rm_sample = rm_ds[0]
    summary["radiomap"] = {
        "num_records": len(rm_ds),
        "scene_id": rm_sample["scene_id"],
        "grid_shape": list(rm_sample["gt_radiomap_db"].shape),
        "valid_cells": int(rm_sample["extent_mask"].sum().item()),
        "voxel_count": int(rm_sample["voxel_feats"].shape[0]),
    }

    cir_manifest = _load_json(cir_manifest_path)
    cir_records = _resolve_record_paths(cir_manifest["triples"], cir_manifest_path)
    cir_manifest_for_loader = dict(cir_manifest)
    cir_manifest_for_loader["triples"] = cir_records[: args.max_cir_samples]
    cir_ds = MultiSceneTripleDataset(
        cir_manifest_for_loader,
        target_config=MergedPathTargetConfig(),
        voxel_channels=args.channels,
        wireless_root=wireless_root,
        scene3d_root=scene3d_root,
        dataset_tag=cir_manifest.get("dataset_tag", "voxel_original_csi_path_10cm_1e6"),
        precompute_stats=False,
    )
    cir_sample = cir_ds[0]
    summary["cir"] = {
        "num_records_checked": len(cir_ds),
        "scene_id": cir_sample["scene_id"],
        "tx_id": int(cir_sample["tx_id"]),
        "rx_id": int(cir_sample["rx_id"]),
        "num_paths": int(cir_sample["csi_path_count"]),
        "voxel_count": int(cir_sample["voxel_level"]["feats"].shape[1]),
    }

    if args.checkpoint:
        ckpt_path = Path(args.checkpoint).resolve()
        ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
        state = ckpt.get("model_state_dict", ckpt.get("state_dict", ckpt))
        summary["checkpoint"] = {
            "path": str(ckpt_path),
            "top_level_keys": sorted(list(ckpt.keys())) if isinstance(ckpt, dict) else [],
            "num_state_tensors": len(state) if isinstance(state, dict) else None,
            "phase_name": ckpt.get("phase_name") if isinstance(ckpt, dict) else None,
            "epoch": ckpt.get("epoch") if isinstance(ckpt, dict) else None,
        }

    out_json = Path(args.out_json)
    out_json.parent.mkdir(parents=True, exist_ok=True)
    out_json.write_text(json.dumps(summary, indent=2))
    print(json.dumps(summary, indent=2))


if __name__ == "__main__":
    main()