#!/usr/bin/env python3
"""Evaluate blind baseline/experience STEP reconstructions against held-out STEP."""
from __future__ import annotations

import argparse
import csv
import importlib.util
import itertools
import json
import math
from pathlib import Path
from typing import Any

import numpy as np
from build123d import import_step
from scipy.spatial import cKDTree

HERE = Path(__file__).resolve().parent
EXTRACTOR = (
    HERE.parents[1]
    / "skills/cad-experience-builder/scripts/step_to_case.py"
)


def load_extractor():
    spec = importlib.util.spec_from_file_location("ab50_step_to_case", EXTRACTOR)
    module = importlib.util.module_from_spec(spec)
    assert spec and spec.loader
    spec.loader.exec_module(module)
    return module


EXTRACT = load_extractor()


def rel_error(actual: float, expected: float) -> float:
    return abs(actual - expected) / max(abs(expected), 1e-9)


def feature_names(payload: dict[str, Any]) -> set[str]:
    result: set[str] = set()
    for item in payload.get("design_ir", {}).get("features", []):
        if isinstance(item, str):
            result.add(item)
        elif isinstance(item, dict):
            name = item.get("kind") or item.get("type") or item.get("name")
            if name:
                result.add(str(name))
    return result


def f1_score(left: set[str], right: set[str]) -> float:
    if not left and not right:
        return 1.0
    if not left or not right:
        return 0.0
    overlap = len(left & right)
    precision = overlap / len(right)
    recall = overlap / len(left)
    return 2 * precision * recall / max(precision + recall, 1e-9)


def mesh(shape, tolerance: float = 0.35) -> tuple[np.ndarray, np.ndarray]:
    vertices, triangles = shape.tessellate(tolerance)
    points = np.asarray([[v.X, v.Y, v.Z] for v in vertices], dtype=float)
    faces = np.asarray(triangles, dtype=int)
    return points, faces


def sample_surface(
    vertices: np.ndarray, faces: np.ndarray, count: int, seed: int
) -> np.ndarray:
    tri = vertices[faces]
    cross = np.cross(tri[:, 1] - tri[:, 0], tri[:, 2] - tri[:, 0])
    areas = np.linalg.norm(cross, axis=1) / 2.0
    total = float(areas.sum())
    if total <= 1e-12:
        return vertices.copy()
    rng = np.random.default_rng(seed)
    chosen = rng.choice(len(faces), size=count, p=areas / total)
    selected = tri[chosen]
    u = rng.random(count)
    v = rng.random(count)
    reflected = u + v > 1.0
    u[reflected] = 1.0 - u[reflected]
    v[reflected] = 1.0 - v[reflected]
    return (
        selected[:, 0]
        + u[:, None] * (selected[:, 1] - selected[:, 0])
        + v[:, None] * (selected[:, 2] - selected[:, 0])
    )


def cube_transforms() -> list[np.ndarray]:
    transforms = []
    for permutation in itertools.permutations(range(3)):
        base = np.eye(3)[:, permutation]
        for signs in itertools.product((-1.0, 1.0), repeat=3):
            transforms.append(base @ np.diag(signs))
    return transforms


TRANSFORMS = cube_transforms()


def surface_distance(
    reference, generated, seed: int, samples: int
) -> tuple[float, float]:
    ref_v, ref_f = mesh(reference)
    gen_v, gen_f = mesh(generated)
    ref = sample_surface(ref_v, ref_f, samples, seed)
    gen = sample_surface(gen_v, gen_f, samples, seed + 1)
    ref -= (ref.min(axis=0) + ref.max(axis=0)) / 2.0
    gen -= (gen.min(axis=0) + gen.max(axis=0)) / 2.0
    ref_tree = cKDTree(ref)
    best_chamfer = math.inf
    best_hausdorff = math.inf
    for transform in TRANSFORMS:
        candidate = gen @ transform
        gen_tree = cKDTree(candidate)
        d_gr = ref_tree.query(candidate, workers=-1)[0]
        d_rg = gen_tree.query(ref, workers=-1)[0]
        chamfer = float((d_gr.mean() + d_rg.mean()) / 2.0)
        hausdorff = float(max(np.quantile(d_gr, 0.95), np.quantile(d_rg, 0.95)))
        if chamfer < best_chamfer:
            best_chamfer = chamfer
            best_hausdorff = hausdorff
    diagonal = max(float(np.linalg.norm(np.ptp(ref, axis=0))), 1e-9)
    return best_chamfer / diagonal, best_hausdorff / diagonal


def task_experience(step_path: Path) -> tuple[int, int]:
    task = step_path.parent / "cad-task.json"
    if not task.is_file():
        return 0, 0
    payload = json.loads(task.read_text())
    exp = payload.get("experience", {})
    returned = exp.get("methods_returned")
    if returned is None:
        returned = len(exp.get("methods", exp.get("methods_applied", [])))
    if isinstance(returned, list):
        returned = len(returned)
    return int(returned or 0), len(exp.get("methods_applied", []))


def evaluate_one(
    case: dict[str, Any], variant: str, samples: int
) -> dict[str, Any]:
    generated_path = Path(case[f"{variant}_step"])
    row: dict[str, Any] = {
        "id": case["id"],
        "family": case["family"],
        "variant": variant,
        "generated": generated_path.is_file(),
    }
    if not generated_path.is_file():
        row["error"] = "missing STEP"
        return row
    try:
        reference = import_step(case["reference_step"])
        generated = import_step(str(generated_path))
        ref_case = json.loads(Path(case["case_json"]).read_text())
        gen_case = EXTRACT.extract_step_case(generated_path)
        ref_bbox = sorted(float(v) for v in reference.bounding_box().size)
        gen_bbox = sorted(float(v) for v in generated.bounding_box().size)
        bbox_error = float(
            np.mean([rel_error(a, b) for a, b in zip(gen_bbox, ref_bbox)])
        )
        volume_error = rel_error(float(generated.volume), float(reference.volume))
        area_error = rel_error(float(generated.area), float(reference.area))
        face_error = rel_error(len(generated.faces()), len(reference.faces()))
        edge_error = rel_error(len(generated.edges()), len(reference.edges()))
        chamfer, hausdorff95 = surface_distance(
            reference, generated, int(case["id"]), samples
        )
        feature_f1 = f1_score(feature_names(ref_case), feature_names(gen_case))
        family_match = float(
            ref_case["design_ir"]["part_family"]
            == gen_case["design_ir"]["part_family"]
        )
        returned, applied = task_experience(generated_path)

        # Transparent 0–100 similarity index. It is a benchmark summary, not a
        # claim of manufacturing interchangeability or exact B-rep identity.
        score = 100.0 * (
            0.18 * math.exp(-volume_error)
            + 0.12 * math.exp(-area_error)
            + 0.16 * math.exp(-bbox_error)
            + 0.24 * math.exp(-8.0 * chamfer)
            + 0.10 * math.exp(-(face_error + edge_error) / 2.0)
            + 0.15 * feature_f1
            + 0.05 * family_match
        )
        row.update(
            {
                "valid": bool(generated.is_valid),
                "solid_count": len(generated.solids()),
                "reference_volume": float(reference.volume),
                "generated_volume": float(generated.volume),
                "volume_rel_error": volume_error,
                "surface_area_rel_error": area_error,
                "bbox_mape": bbox_error,
                "face_rel_error": face_error,
                "edge_rel_error": edge_error,
                "surface_chamfer_norm": chamfer,
                "surface_hausdorff95_norm": hausdorff95,
                "feature_f1": feature_f1,
                "family_match": family_match,
                "experience_methods_returned": returned,
                "experience_methods_applied": applied,
                "similarity_score": score,
            }
        )
    except Exception as exc:
        row["error"] = f"{type(exc).__name__}: {exc}"
    return row


def mean(rows: list[dict[str, Any]], key: str) -> float:
    values = [float(row[key]) for row in rows if key in row]
    return sum(values) / len(values) if values else math.nan


def write_report(rows: list[dict[str, Any]]) -> None:
    complete = [row for row in rows if "similarity_score" in row]
    baseline = [row for row in complete if row["variant"] == "baseline"]
    experience = [row for row in complete if row["variant"] == "experience"]
    by_id = {}
    for row in complete:
        by_id.setdefault(row["id"], {})[row["variant"]] = row
    pairs = [
        item for item in by_id.values() if {"baseline", "experience"} <= item.keys()
    ]
    paired_baseline = [item["baseline"] for item in pairs]
    paired_experience = [item["experience"] for item in pairs]
    deltas = [
        item["experience"]["similarity_score"]
        - item["baseline"]["similarity_score"]
        for item in pairs
    ]
    wins = sum(delta > 0.25 for delta in deltas)
    ties = sum(abs(delta) <= 0.25 for delta in deltas)
    losses = sum(delta < -0.25 for delta in deltas)
    methods_returned = sum(
        row.get("experience_methods_returned", 0) for row in experience
    )
    methods_applied = sum(
        row.get("experience_methods_applied", 0) for row in experience
    )
    report = f"""# CAD 经验库 50 件 A/B 盲重建报告

## 口径

- 测试集：固定分层抽取 50 件；这 50 件不进入经验库。
- 训练经验：其余 950 件生成的通用经验库。
- 两组输入完全相同：仅多视图图片和文字描述；生成过程禁止读取参考 STEP 和解析 JSON。
- Baseline：空经验库；Experience：950 件经验库。
- 参考 STEP 只在生成完成后用于评分。
- “相似度”是体积、面积、包围盒、表面距离、拓扑和语义特征的综合诊断分，不代表制造级 1:1 或可互换性。

## 当前完成度

- 有效 Baseline：{len(baseline)}/50
- 有效 Experience：{len(experience)}/50
- 完整配对：{len(pairs)}/50
- 经验方法返回/实际采用：{methods_returned}/{methods_applied}

## 汇总

| 指标 | 无经验 | 有经验 | 差值 |
|---|---:|---:|---:|
| 综合相似度（0–100，越高越好） | {mean(paired_baseline, "similarity_score"):.2f} | {mean(paired_experience, "similarity_score"):.2f} | {mean(paired_experience, "similarity_score") - mean(paired_baseline, "similarity_score"):+.2f} |
| 体积相对误差（越低越好） | {mean(paired_baseline, "volume_rel_error"):.3f} | {mean(paired_experience, "volume_rel_error"):.3f} | {mean(paired_experience, "volume_rel_error") - mean(paired_baseline, "volume_rel_error"):+.3f} |
| 表面积相对误差 | {mean(paired_baseline, "surface_area_rel_error"):.3f} | {mean(paired_experience, "surface_area_rel_error"):.3f} | {mean(paired_experience, "surface_area_rel_error") - mean(paired_baseline, "surface_area_rel_error"):+.3f} |
| 包围盒尺寸 MAPE | {mean(paired_baseline, "bbox_mape"):.3f} | {mean(paired_experience, "bbox_mape"):.3f} | {mean(paired_experience, "bbox_mape") - mean(paired_baseline, "bbox_mape"):+.3f} |
| 归一化表面 Chamfer | {mean(paired_baseline, "surface_chamfer_norm"):.4f} | {mean(paired_experience, "surface_chamfer_norm"):.4f} | {mean(paired_experience, "surface_chamfer_norm") - mean(paired_baseline, "surface_chamfer_norm"):+.4f} |
| 特征 F1 | {mean(paired_baseline, "feature_f1"):.3f} | {mean(paired_experience, "feature_f1"):.3f} | {mean(paired_experience, "feature_f1") - mean(paired_baseline, "feature_f1"):+.3f} |

经验组胜/平/负（±0.25 分视为平）：**{wins}/{ties}/{losses}**。
"""
    (HERE / "report.md").write_text(report)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--samples", type=int, default=1600)
    args = parser.parse_args()
    manifest = json.loads((HERE / "manifest.json").read_text())
    rows = [
        evaluate_one(case, variant, args.samples)
        for case in manifest["cases"]
        for variant in ("baseline", "experience")
    ]
    (HERE / "metrics.json").write_text(
        json.dumps(rows, indent=2, ensure_ascii=False) + "\n"
    )
    keys = sorted({key for row in rows for key in row})
    with (HERE / "metrics.csv").open("w", newline="") as stream:
        writer = csv.DictWriter(stream, fieldnames=keys)
        writer.writeheader()
        writer.writerows(rows)
    write_report(rows)
    completed = sum("similarity_score" in row for row in rows)
    print(json.dumps({"evaluated": completed, "expected": 100}))


if __name__ == "__main__":
    main()