#!/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_json(path: Path) -> dict[str, Any]:
    return json.loads(path.read_text(encoding="utf-8"))


def crop_sprite(atlas: Image.Image, meta: dict[str, Any]) -> Image.Image:
    crop = atlas.crop((
        int(meta["x"]), int(meta["y"]),
        int(meta["x"]) + int(meta["width"]),
        int(meta["y"]) + int(meta["height"]),
    )).convert("RGBA")
    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))
    return crop


def resize_center(crop: Image.Image, size: int = 64, pad: int = 5) -> Image.Image:
    w, h = crop.size
    target = size - pad * 2
    scale = min(target / max(w, 1), target / max(h, 1))
    nw, nh = max(1, round(w * scale)), max(1, round(h * scale))
    crop = crop.resize((nw, nh), Image.Resampling.LANCZOS)
    canvas = Image.new("RGBA", (size, size), (255, 255, 255, 0))
    canvas.alpha_composite(crop, ((size - nw) // 2, (size - nh) // 2))
    return canvas


def dct_matrix(n: int) -> np.ndarray:
    mat = np.zeros((n, n), dtype=np.float32)
    factor = math.pi / n
    for k in range(n):
        scale = math.sqrt(1 / n) if k == 0 else math.sqrt(2 / n)
        for i in range(n):
            mat[k, i] = scale * math.cos((i + 0.5) * k * factor)
    return mat


DCT32 = dct_matrix(32)


def phash(gray32: np.ndarray) -> np.ndarray:
    coeff = DCT32 @ gray32 @ DCT32.T
    low = coeff[:8, :8].copy()
    vals = low.ravel()[1:]
    med = float(np.median(vals))
    return (low.ravel()[1:] > med).astype(np.uint8)


def color_hist(rgb: np.ndarray, alpha: np.ndarray, bins: int = 4) -> np.ndarray:
    mask = alpha > 0.08
    if not np.any(mask):
        return np.zeros(bins ** 3, dtype=np.float32)
    pix = np.clip(rgb[mask], 0, 0.999999)
    q = (pix * bins).astype(np.int32)
    idx = q[:, 0] * bins * bins + q[:, 1] * bins + q[:, 2]
    hist = np.bincount(idx, minlength=bins ** 3).astype(np.float32)
    s = hist.sum()
    return hist / s if s else hist


def descriptor(crop: Image.Image) -> dict[str, Any]:
    original_w, original_h = crop.size
    canvas = resize_center(crop, 64, 5)
    arr = np.asarray(canvas, dtype=np.float32) / 255.0
    alpha = arr[:, :, 3]
    rgb = arr[:, :, :3]
    composited = rgb * alpha[:, :, None] + (1.0 - alpha[:, :, None])
    gray_img = Image.fromarray(np.uint8(np.clip(
        (0.299 * composited[:, :, 0] + 0.587 * composited[:, :, 1] + 0.114 * composited[:, :, 2]) * 255,
        0, 255
    )))
    gray32 = np.asarray(gray_img.resize((32, 32), Image.Resampling.LANCZOS), dtype=np.float32) / 255.0
    edge = np.asarray(gray_img.filter(ImageFilter.FIND_EDGES), dtype=np.float32) / 255.0

    mask = alpha > 0.08
    xproj = mask.sum(axis=0).astype(np.float32)
    yproj = mask.sum(axis=1).astype(np.float32)
    if xproj.sum():
        xproj /= xproj.sum()
    if yproj.sum():
        yproj /= yproj.sum()

    return {
        "alpha": alpha,
        "gray": gray32,
        "edge": edge,
        "phash": phash(gray32),
        "hist": color_hist(rgb, alpha),
        "xproj": xproj,
        "yproj": yproj,
        "aspect": original_w / max(1, original_h),
        "occupancy": float(mask.mean()),
    }


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


def similarity(a: dict[str, Any], b: dict[str, Any]) -> float:
    alpha = cosine(a["alpha"], b["alpha"])
    edge = cosine(a["edge"], b["edge"])
    gray = 1.0 - float(np.mean((a["gray"] - b["gray"]) ** 2))
    hamming = float(np.mean(a["phash"] != b["phash"]))
    ph = 1.0 - hamming
    hist = cosine(a["hist"], b["hist"])
    proj = 0.5 * (cosine(a["xproj"], b["xproj"]) + cosine(a["yproj"], b["yproj"]))
    ratio = min(a["aspect"], b["aspect"]) / max(a["aspect"], b["aspect"])
    occ = 1.0 - min(1.0, abs(a["occupancy"] - b["occupancy"]) / max(a["occupancy"], b["occupancy"], 0.02))
    return (
        0.24 * alpha +
        0.17 * edge +
        0.16 * gray +
        0.14 * ph +
        0.12 * hist +
        0.08 * proj +
        0.05 * ratio +
        0.04 * occ
    )


def resource_records(catalog: dict[str, Any], standard: bool) -> list[dict[str, Any]]:
    seen: set[tuple[Any, ...]] = set()
    out = []
    for obj in catalog.get("visual_objects", []):
        if not obj.get("resource_key") or not obj.get("resource_meta"):
            continue
        layer = obj.get("standard_layer_name") if standard else None
        key = (
            obj.get("component_role"),
            obj.get("resource_key"),
            layer,
        )
        if key in seen:
            continue
        seen.add(key)
        out.append({
            "resource_key": obj["resource_key"],
            "resource_meta": obj["resource_meta"],
            "component_role": obj.get("component_role"),
            "standard_layer_name": layer,
            "visual_object_id": obj.get("visual_object_id"),
        })
    return out


def build_desc(records: list[dict[str, Any]], atlas: Image.Image) -> list[dict[str, Any]]:
    return [descriptor(crop_sprite(atlas, r["resource_meta"])) for r in records]


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--standard", required=True)
    ap.add_argument("--cm95", required=True)
    ap.add_argument("--output", required=True)
    ap.add_argument("--min-score", type=float, default=0.84)
    ap.add_argument("--min-margin", type=float, default=0.025)
    args = ap.parse_args()

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

    std_cat = load_json(std_path)
    cm_cat = load_json(cm_path)
    std_dir = std_path.parent.parent / "navsea-anonymous"
    cm_dir = cm_path.parent
    std_atlas = Image.open(std_dir / "_sprite" / "sprite.png").convert("RGBA")
    cm_atlas = Image.open(cm_dir / "_sprite" / "sprite.png").convert("RGBA")

    std = resource_records(std_cat, True)
    cm = resource_records(cm_cat, False)
    std_d = build_desc(std, std_atlas)
    cm_d = build_desc(cm, cm_atlas)

    std_by_role: dict[str, list[int]] = defaultdict(list)
    cm_by_role: dict[str, list[int]] = defaultdict(list)
    for i, r in enumerate(std):
        std_by_role[r["component_role"]].append(i)
    for i, r in enumerate(cm):
        cm_by_role[r["component_role"]].append(i)

    cm_rankings: dict[int, list[tuple[float, int]]] = {}
    std_best: dict[int, tuple[float, int]] = {}

    for role, cm_indices in cm_by_role.items():
        std_indices = std_by_role.get(role, [])
        if not std_indices:
            continue
        for ci in cm_indices:
            vals = sorted(
                ((similarity(cm_d[ci], std_d[si]), si) for si in std_indices),
                reverse=True,
            )
            cm_rankings[ci] = vals[:5]
        for si in std_indices:
            vals = max(
                ((similarity(std_d[si], cm_d[ci]), ci) for ci in cm_indices),
                default=(0.0, -1),
            )
            std_best[si] = vals

    anchors = []
    candidates = []
    for ci, ranked in cm_rankings.items():
        if not ranked:
            continue
        best_score, si = ranked[0]
        second_score = ranked[1][0] if len(ranked) > 1 else 0.0
        mutual = std_best.get(si, (0.0, -1))[1] == ci
        margin = best_score - second_score
        row = {
            "cm95_resource_key": cm[ci]["resource_key"],
            "component_role": cm[ci]["component_role"],
            "best_standard_resource_key": std[si]["resource_key"],
            "best_standard_layer": std[si]["standard_layer_name"],
            "score": best_score,
            "margin": margin,
            "mutual_nearest": mutual,
            "standard_visual_object_id": std[si]["visual_object_id"],
            "alternatives": [
                {
                    "standard_resource_key": std[sj]["resource_key"],
                    "standard_layer": std[sj]["standard_layer_name"],
                    "score": score,
                }
                for score, sj in ranked[1:4]
            ],
        }
        candidates.append(row)
        if mutual and best_score >= args.min_score and margin >= args.min_margin:
            row = dict(row)
            row["status"] = "MATCH"
            anchors.append(row)

    # Map accepted visual anchors back to CM95 anonymous Visual Object occurrences.
    anchor_by_key_role = {
        (x["cm95_resource_key"], x["component_role"]): x for x in anchors
    }
    occurrences = []
    for obj in cm_cat.get("visual_objects", []):
        k = (obj.get("resource_key"), obj.get("component_role"))
        anchor = anchor_by_key_role.get(k)
        if not anchor:
            continue
        occurrences.append({
            "visual_object_id": obj["visual_object_id"],
            "group_id": obj["group_id"],
            "cm95_resource_key": obj["resource_key"],
            "assigned_standard_layer": anchor["best_standard_layer"],
            "score": anchor["score"],
            "basis": "VISUAL_ANCHOR_MUTUAL_NEAREST",
        })

    summary = {
        "standard_visual_resources": len(std),
        "cm95_visual_resources": len(cm),
        "candidate_pairs": len(candidates),
        "accepted_visual_anchors": len(anchors),
        "cm95_visual_object_occurrences_named": len(occurrences),
        "standard_layers_reached": sorted(set(x["best_standard_layer"] for x in anchors if x.get("best_standard_layer"))),
    }
    result = {
        "schema": "style-visual-anchor-match/v1",
        "rules": {
            "cm95_source_layer_used": False,
            "provider_name_used": False,
            "resource_key_text_semantics_used": False,
            "standard_layer_is_output_only": True,
            "min_score": args.min_score,
            "min_margin": args.min_margin,
            "requires_mutual_nearest": True,
        },
        "summary": summary,
        "anchors": sorted(anchors, key=lambda x: x["score"], reverse=True),
        "cm95_named_occurrences": occurrences,
        "top_candidates": sorted(candidates, key=lambda x: x["score"], reverse=True)[:200],
    }
    (out / "visual-anchor-matches.json").write_text(
        json.dumps(result, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
    )
    print(json.dumps(summary, ensure_ascii=False, indent=2))
    for x in result["anchors"][:80]:
        print(
            f'{x["score"]:.4f} margin={x["margin"]:.4f} '
            f'{x["cm95_resource_key"]} -> {x["best_standard_resource_key"]} '
            f'[{x["best_standard_layer"]}]'
        )


if __name__ == "__main__":
    main()