#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import math
from collections import defaultdict
from pathlib import Path
from typing import Any

import numpy as np
from PIL import Image, ImageFilter


def load_catalog(path: Path) -> dict[str, Any]:
    return json.loads(path.read_text(encoding="utf-8"))


def load_atlas(catalog_dir: Path) -> Image.Image:
    return Image.open(catalog_dir / "_sprite" / "sprite.png").convert("RGBA")


def visual_feature(atlas: Image.Image, meta: dict[str, Any]) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    crop = atlas.crop((
        int(meta["x"]), int(meta["y"]),
        int(meta["x"]) + int(meta["width"]),
        int(meta["y"]) + int(meta["height"]),
    ))
    alpha = np.asarray(crop.getchannel("A"))
    ys, xs = np.where(alpha > 8)
    if len(xs):
        crop = crop.crop((int(xs.min()), int(ys.min()), int(xs.max()) + 1, int(ys.max()) + 1))

    w, h = crop.size
    scale = min(54.0 / max(w, 1), 54.0 / max(h, 1))
    nw = max(1, round(w * scale))
    nh = max(1, round(h * scale))
    crop = crop.resize((nw, nh), Image.Resampling.LANCZOS)

    canvas = Image.new("RGBA", (64, 64), (255, 255, 255, 0))
    canvas.alpha_composite(crop, ((64 - nw) // 2, (64 - nh) // 2))
    arr = np.asarray(canvas, dtype=np.float32) / 255.0
    a = arr[:, :, 3]
    rgb = arr[:, :, :3] * a[:, :, None] + (1.0 - a[:, :, None])
    edge = np.asarray(canvas.convert("L").filter(ImageFilter.FIND_EDGES), dtype=np.float32) / 255.0
    return rgb, a, edge


def cosine(a: np.ndarray, b: np.ndarray) -> float:
    aa = a.ravel()
    bb = b.ravel()
    denom = float(np.linalg.norm(aa) * np.linalg.norm(bb))
    return float(np.dot(aa, bb) / denom) if denom else 0.0


def image_similarity(a: tuple[np.ndarray, np.ndarray, np.ndarray], b: tuple[np.ndarray, np.ndarray, np.ndarray]) -> float:
    rgb1, alpha1, edge1 = a
    rgb2, alpha2, edge2 = b
    rgb_score = 1.0 - float(np.mean((rgb1 - rgb2) ** 2))
    alpha_score = cosine(alpha1, alpha2)
    edge_score = cosine(edge1, edge2)
    return 0.45 * rgb_score + 0.35 * alpha_score + 0.20 * edge_score


def unique_sprite_objects(catalog: dict[str, Any]) -> list[dict[str, Any]]:
    seen: set[tuple[str, str, str | None]] = set()
    out: list[dict[str, Any]] = []
    for obj in catalog.get("visual_objects", []):
        if not obj.get("resource_key") or not obj.get("resource_meta"):
            continue
        key = (
            obj.get("component_role") or "",
            obj.get("visual_hash") or "",
            obj.get("standard_layer_name"),
            obj.get("group_id"),
        )
        if key in seen:
            continue
        seen.add(key)
        out.append(obj)
    return out


def build_features(catalog: dict[str, Any], catalog_dir: Path) -> tuple[list[dict[str, Any]], list[tuple[np.ndarray, np.ndarray, np.ndarray]]]:
    atlas = load_atlas(catalog_dir)
    objects = unique_sprite_objects(catalog)
    features = [visual_feature(atlas, obj["resource_meta"]) for obj in objects]
    return objects, features


def compatible(cm: dict[str, Any], std: dict[str, Any]) -> bool:
    if cm.get("component_role") != std.get("component_role"):
        return False
    cg = cm.get("geometry_role")
    sg = std.get("geometry_role")
    if cg and sg and cg != sg:
        return False
    return True


def classify_object(
    cm_obj: dict[str, Any],
    cm_feat: tuple[np.ndarray, np.ndarray, np.ndarray],
    std_objects: list[dict[str, Any]],
    std_features: list[tuple[np.ndarray, np.ndarray, np.ndarray]],
) -> dict[str, Any]:
    layer_best: dict[str, tuple[float, str]] = {}
    for std, feat in zip(std_objects, std_features):
        layer = std.get("standard_layer_name")
        if not layer or not compatible(cm_obj, std):
            continue
        score = image_similarity(cm_feat, feat)
        prev = layer_best.get(layer)
        if prev is None or score > prev[0]:
            layer_best[layer] = (score, std["visual_object_id"])

    ranked = sorted(
        (
            {"layer": layer, "score": score, "standard_visual_object_id": vo_id}
            for layer, (score, vo_id) in layer_best.items()
        ),
        key=lambda x: x["score"],
        reverse=True,
    )
    top = ranked[0] if ranked else None
    second = ranked[1] if len(ranked) > 1 else None
    if not top:
        status = "UNKNOWN"
    else:
        margin = top["score"] - (second["score"] if second else 0.0)
        if top["score"] >= 0.90 and margin >= 0.025:
            status = "MATCH"
        elif top["score"] >= 0.82:
            status = "PARTIAL"
        else:
            status = "UNKNOWN"
    return {
        "visual_object_id": cm_obj["visual_object_id"],
        "group_id": cm_obj["group_id"],
        "component_role": cm_obj.get("component_role"),
        "geometry_role": cm_obj.get("geometry_role"),
        "resource_key": cm_obj.get("resource_key"),
        "status": status,
        "best": top,
        "alternatives": ranked[1:4],
    }


def group_scores(object_matches: list[dict[str, Any]]) -> dict[str, dict[str, Any]]:
    grouped: dict[str, list[dict[str, Any]]] = defaultdict(list)
    for item in object_matches:
        grouped[item["group_id"]].append(item)

    results: dict[str, dict[str, Any]] = {}
    for group_id, items in grouped.items():
        evidence: dict[str, list[float]] = defaultdict(list)
        for item in items:
            for cand in ([item["best"]] if item.get("best") else []) + item.get("alternatives", []):
                if cand and cand["score"] >= 0.78:
                    evidence[cand["layer"]].append(cand["score"])
        ranking = []
        for layer, scores in evidence.items():
            scores = sorted(scores, reverse=True)
            strongest = scores[: min(12, len(scores))]
            coverage = min(1.0, len(scores) / max(3.0, len(items) * 0.25))
            score = (sum(strongest) / len(strongest)) * (0.65 + 0.35 * coverage)
            ranking.append({
                "layer": layer,
                "score": score,
                "evidence_count": len(scores),
                "best_evidence": scores[:5],
            })
        ranking.sort(key=lambda x: x["score"], reverse=True)
        top = ranking[0] if ranking else None
        second = ranking[1] if len(ranking) > 1 else None
        if top:
            margin = top["score"] - (second["score"] if second else 0.0)
            dominance = top["evidence_count"] / max(1, sum(x["evidence_count"] for x in ranking))
            if top["score"] >= 0.86 and margin >= 0.035 and dominance >= 0.40:
                status = "MATCH"
            elif top["score"] >= 0.80:
                status = "PARTIAL"
            else:
                status = "UNKNOWN"
        else:
            status = "UNKNOWN"
        results[group_id] = {
            "group_id": group_id,
            "status": status,
            "best": top,
            "alternatives": ranking[1:4],
            "object_count": len(items),
        }
    return results


def apply_group_result(object_matches: list[dict[str, Any]], groups: dict[str, dict[str, Any]]) -> None:
    for item in object_matches:
        g = groups.get(item["group_id"])
        if g and g["status"] == "MATCH":
            item["final_layer"] = g["best"]["layer"]
            item["final_basis"] = "GROUP_VISUAL_CONSENSUS"
            item["final_status"] = "MATCH"
        elif item["status"] == "MATCH":
            item["final_layer"] = item["best"]["layer"]
            item["final_basis"] = "OBJECT_VISUAL_MATCH"
            item["final_status"] = "MATCH"
        elif item["status"] == "PARTIAL":
            item["final_layer"] = item["best"]["layer"] if item.get("best") else None
            item["final_basis"] = "OBJECT_VISUAL_PARTIAL"
            item["final_status"] = "PARTIAL"
        else:
            item["final_layer"] = None
            item["final_basis"] = "INSUFFICIENT_VISUAL_EVIDENCE"
            item["final_status"] = "UNKNOWN"


def blind_validation(
    cm_catalog: dict[str, Any],
    group_results: dict[str, dict[str, Any]],
) -> dict[str, Any]:
    # This is deliberately post-match audit only. Raw source-layer names are read
    # after predictions have already been produced.
    provenance = {g["group_id"]: g.get("provenance", {}) for g in cm_catalog.get("groups", [])}
    rows = []
    comparable = correct = 0
    for group_id, pred in sorted(group_results.items()):
        raw_name = provenance.get(group_id, {}).get("source_layer")
        predicted = pred.get("best", {}).get("layer") if pred.get("best") else None
        row = {
            "group_id": group_id,
            "raw_source_layer_for_audit_only": raw_name,
            "predicted_standard_layer": predicted,
            "status": pred.get("status"),
            "same_name_after_blind_match": bool(raw_name and predicted and raw_name == predicted),
        }
        if raw_name and predicted and not str(raw_name).startswith("s52_"):
            comparable += 1
            if raw_name == predicted:
                correct += 1
        rows.append(row)
    return {
        "note": "Raw CM95 source-layer is used only after matching for blind validation; it is not matcher input.",
        "comparable_group_count": comparable,
        "same_name_count": correct,
        "same_name_rate": (correct / comparable) if comparable else None,
        "rows": rows,
    }


def html_report(result: dict[str, Any]) -> str:
    stats = result["summary"]
    val = result["blind_validation"]
    rows = []
    for g in result["group_matches"]:
        best = g.get("best") or {}
        alts = ", ".join(f'{x["layer"]} {x["score"]:.3f}' for x in g.get("alternatives", []))
        rows.append(
            "<tr>"
            f"<td><code>{g['group_id']}</code></td>"
            f"<td>{g['status']}</td>"
            f"<td>{best.get('layer','')}</td>"
            f"<td>{best.get('score',0):.3f}</td>"
            f"<td>{g['object_count']}</td>"
            f"<td>{alts}</td>"
            "</tr>"
        )
    return f"""<!doctype html><meta charset='utf-8'><title>CM95 Visual Object Layer Match</title>
<style>
body{{font-family:Arial,sans-serif;margin:24px;color:#17202a}} table{{border-collapse:collapse;width:100%}}
th,td{{border:1px solid #d8dee6;padding:7px;text-align:left;font-size:13px}} th{{background:#eef2f6}}
code{{font-size:11px}} .box{{display:inline-block;margin:0 12px 12px 0;padding:10px 14px;border:1px solid #d8dee6;border-radius:8px}}
</style>
<h1>CM95 Visual Object → NavSea Standard Layer</h1>
<p>匹配程序不读取 CM95 source-layer 名称；Layer 名仅来自 NavSea 已冻结 Standard Visual Object Catalog。</p>
<div class='box'>CM95 Visual Objects <b>{stats['cm95_sprite_visual_objects']}</b></div>
<div class='box'>MATCH <b>{stats['object_match']}</b></div>
<div class='box'>PARTIAL <b>{stats['object_partial']}</b></div>
<div class='box'>UNKNOWN <b>{stats['object_unknown']}</b></div>
<div class='box'>Blind audit {val['same_name_count']}/{val['comparable_group_count']}</div>
<h2>匿名 Group 的视觉匹配</h2>
<table><thead><tr><th>匿名 Group</th><th>状态</th><th>预测标准 Layer</th><th>分数</th><th>VO 数</th><th>候选</th></tr></thead>
<tbody>{''.join(rows)}</tbody></table>
<p>详细到每个 Visual Object 的匹配、候选和 blind validation 见 <code>match-results.json</code>。</p>
"""


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--standard", required=True)
    ap.add_argument("--cm95", required=True)
    ap.add_argument("--output", required=True)
    args = ap.parse_args()

    std_path = Path(args.standard)
    cm_path = Path(args.cm95)
    output = Path(args.output)
    output.mkdir(parents=True, exist_ok=True)

    std_catalog = load_catalog(std_path)
    cm_catalog = load_catalog(cm_path)
    std_objects, std_features = build_features(std_catalog, std_path.parent.parent / "navsea-anonymous")
    cm_objects, cm_features = build_features(cm_catalog, cm_path.parent)

    # Sanitized match inputs deliberately omit CM95 provenance/source-layer.
    sanitized = [
        {
            "visual_object_id": o["visual_object_id"],
            "group_id": o["group_id"],
            "component_role": o.get("component_role"),
            "geometry_role": o.get("geometry_role"),
            "resource_key": o.get("resource_key"),
        }
        for o in cm_objects
    ]
    (output / "cm95-match-input-sanitized.json").write_text(
        json.dumps(sanitized, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
    )

    object_matches = [
        classify_object(obj, feat, std_objects, std_features)
        for obj, feat in zip(cm_objects, cm_features)
    ]
    groups = group_scores(object_matches)
    apply_group_result(object_matches, groups)
    validation = blind_validation(cm_catalog, groups)

    result = {
        "schema": "style-visual-object-match/v1",
        "method": {
            "cm95_source_layer_used_for_matching": False,
            "provider_name_used_for_matching": False,
            "resource_key_text_semantics_used_for_matching": False,
            "visual_inputs": ["sprite pixels", "component role", "geometry role", "anonymous group membership"],
            "standard_labels_from": "frozen NavSea Standard Visual Object Catalog",
        },
        "summary": {
            "standard_sprite_visual_objects": len(std_objects),
            "cm95_sprite_visual_objects": len(cm_objects),
            "object_match": sum(x["final_status"] == "MATCH" for x in object_matches),
            "object_partial": sum(x["final_status"] == "PARTIAL" for x in object_matches),
            "object_unknown": sum(x["final_status"] == "UNKNOWN" for x in object_matches),
            "group_match": sum(x["status"] == "MATCH" for x in groups.values()),
            "group_partial": sum(x["status"] == "PARTIAL" for x in groups.values()),
            "group_unknown": sum(x["status"] == "UNKNOWN" for x in groups.values()),
        },
        "group_matches": sorted(groups.values(), key=lambda x: x["group_id"]),
        "visual_object_matches": object_matches,
        "blind_validation": validation,
    }
    (output / "match-results.json").write_text(json.dumps(result, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    (output / "report.html").write_text(html_report(result), encoding="utf-8")
    print(json.dumps(result["summary"], ensure_ascii=False, indent=2))
    print(json.dumps({k: validation[k] for k in ("comparable_group_count", "same_name_count", "same_name_rate")}, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()