#!/usr/bin/env python3
"""Validate the selected size contract and list every generated NPZ artifact."""
from __future__ import annotations

import html
import json
import os
import sys
from pathlib import Path

import numpy as np

HERE = Path(__file__).resolve().parent
PARTS = ("AC28", "ME1a_DT")
sys.path.insert(0, str(HERE.parents[1]))
from models.architectures import CNN2DUNet, MeshGraphProcessor, PointNetSeg, UNet3D
import torch


def directory_page(folder: Path, title: str) -> None:
    page = folder / "index.html"
    links = []
    for path in sorted(folder.rglob("*.npz")):
        relative = os.path.relpath(path, page.parent).replace(os.sep, "/")
        label = f"{path.relative_to(folder)} ({path.stat().st_size/1e6:.2f} MB)"
        links.append(f'<li><a href="{html.escape(relative, quote=True)}">{html.escape(label)}</a></li>')
    page.write_text(f'''<!doctype html><html lang="en"><meta charset="utf-8"><meta name="viewport" content="width=device-width,initial-scale=1">
<title>{html.escape(title)}</title><style>body{{font:16px/1.5 system-ui,sans-serif;background:#f4f7fb;color:#183146;max-width:1050px;margin:auto;padding:30px}}li{{margin:8px}}a{{color:#126193}}</style>
<p><a href="../index.html">Size decision report</a></p><h1>{html.escape(title)}</h1><ul>{''.join(links)}</ul></html>''', encoding="utf-8")


def main() -> None:
    contract = json.loads((HERE / "standard_sizes.json").read_text())
    image_size = int(contract["image"]["height"])
    checks = []
    signatures = {}
    for part in PARTS:
        root = HERE / "standardized" / part
        with np.load(root / f"image_{image_size}.npz", allow_pickle=False) as d:
            assert d["x"].shape == (6, 17, image_size, image_size)
            assert np.allclose(d["field_of_view_mm"], 980)
            assert np.isfinite(d["x"][d["feature_valid"]]).all()
            assert not d["training_valid"].any()
            signatures.setdefault("image", []).append(tuple(d["feature_names"].astype(str)))
        checks.append(f"{part}: six {image_size}² masked image views, common 980 mm field of view")
        with np.load(root / "voxel_sparse_64.npz", allow_pickle=False) as sparse, np.load(
                root / "voxel_dense_64.npz", allow_pickle=False) as dense:
            assert tuple(sparse["grid_dims"]) == (64, 8, 40)
            assert dense["x"].shape == (12, 64, 8, 40)
            assert np.isclose(float(sparse["voxel_pitch_mm"]), 15.0)
            assert np.isclose(float(dense["voxel_pitch_mm"]), 15.0)
            coord = sparse["voxel_coords"]
            assert np.all((coord >= 0) & (coord < np.asarray([64, 8, 40])))
            assert np.array_equal(dense["x"][:, coord[:, 0], coord[:, 1], coord[:, 2]].T, sparse["x"])
            assert not sparse["training_valid"].any()
            signatures.setdefault("voxel", []).append(tuple(sparse["feature_names"].astype(str)))
        checks.append(f"{part}: sparse/dense voxel match at 15 mm and 64×8×40")
        with np.load(root / "graph_input_5000.npz", allow_pickle=False) as d:
            assert len(d["x"]) == contract["graph_context"]["actual_nodes"][part]
            assert d["edge_index"].shape[0] == 2
            assert not d["training_valid"].any()
            signatures.setdefault("graph_node", []).append(tuple(d["feature_names"].astype(str)))
            signatures.setdefault("graph_edge", []).append(tuple(d["edge_feature_names"].astype(str)))
        checks.append(f"{part}: 5k requested graph with measured actual nodes and common schema")
        with np.load(root / "graph_reference_12000.npz", allow_pickle=False) as d:
            assert len(d["x"]) > len(np.load(root / "graph_input_5000.npz")["x"])
        with np.load(root / "points_8192.npz", allow_pickle=False) as d:
            assert d["x"].shape[0] == 8192
            assert not d["training_valid"].any()
            signatures.setdefault("point", []).append(tuple(d["feature_names"].astype(str)))
        checks.append(f"{part}: 8,192 CAD surface samples")
    for family, values in signatures.items():
        assert values[0] == values[1], f"cross-part {family} features differ"
        checks.append(f"cross-part {family} feature schema matches")
    forward_checks = model_forward_checks()
    checks.append("same model weights forward on both parts for image, voxel, graph and point inputs")
    report = {"status": "pass", "checks": checks, "contract": "standard_sizes.json",
              "model_forward_checks": forward_checks,
              "interpretation": "shape/finite execution only; this does not measure predictive accuracy"}
    (HERE / "validation_report.json").write_text(json.dumps(report, indent=2), encoding="utf-8")
    if not (HERE / "candidates/index.html").exists():
        directory_page(HERE / "candidates", "All tested preprocessing sizes")
    directory_page(HERE / "standardized", "Selected standardized inputs")
    print(f"{len(checks)} checks passed")


def model_forward_checks() -> list[dict]:
    torch.set_num_threads(4)
    torch.manual_seed(29)
    models = {"image": CNN2DUNet(34).eval(), "voxel": UNet3D(24).eval(),
              "graph": MeshGraphProcessor(32, 8, 1, steps=2).eval(),
              "point": PointNetSeg(26).eval()}
    image_size = int(json.loads((HERE / "standard_sizes.json").read_text())["image"]["height"])
    results = []
    with torch.no_grad():
        for part in PARTS:
            root = HERE / "standardized" / part
            with np.load(root / f"image_{image_size}.npz", allow_pickle=False) as d:
                x = np.concatenate([d["x"][0:1], d["feature_valid"][0:1].astype(np.float32)], axis=1)
                output = models["image"](torch.from_numpy(x))
                assert output.shape == (1, 1, image_size, image_size) and torch.isfinite(output).all()
                results.append({"part": part, "modality": "image", "output_shape": list(output.shape)})
            with np.load(root / "voxel_dense_64.npz", allow_pickle=False) as d:
                x = np.concatenate([d["x"], d["feature_valid"].astype(np.float32)], axis=0)[None]
                output = models["voxel"](torch.from_numpy(x))
                assert output.shape == (1, 1, 64, 8, 40) and torch.isfinite(output).all()
                results.append({"part": part, "modality": "voxel", "output_shape": list(output.shape)})
            with np.load(root / "graph_input_5000.npz", allow_pickle=False) as d:
                x = np.concatenate([d["x"], d["feature_valid"].astype(np.float32)], axis=1)
                output = models["graph"](torch.from_numpy(x),
                                          torch.from_numpy(np.asarray(d["edge_index"], np.int64)),
                                          torch.from_numpy(np.asarray(d["edge_attr"], np.float32)),
                                          torch.zeros((1, 1)), torch.zeros(len(x), dtype=torch.long))
                assert output.shape == (len(x), 1) and torch.isfinite(output).all()
                results.append({"part": part, "modality": "graph", "output_shape": list(output.shape)})
            with np.load(root / "points_8192.npz", allow_pickle=False) as d:
                x = np.concatenate([d["x"], d["feature_valid"].astype(np.float32)], axis=1)[None]
                output = models["point"](torch.from_numpy(x))
                assert output.shape == (1, 8192, 1) and torch.isfinite(output).all()
                results.append({"part": part, "modality": "point", "output_shape": list(output.shape)})
    return results


if __name__ == "__main__":
    main()