"""Read actual LTF MIL attention and render its original-image coordinates."""

from __future__ import annotations

import json
import math
from pathlib import Path

import numpy as np
from PIL import Image, ImageDraw, ImageFont


def capture_prediction(detector, path, branch, threshold):
    """Observe one ordinary prediction; temporary hooks are always removed.

    Call under the app's model lock. return_viz only appends existing model
    tensors to the output; it does not change logits or rerun the backbone.
    """
    captured = {}

    def enable_viz(module, args, kwargs):
        return args, dict(kwargs, return_viz=True)

    def collect(module, args, output):
        viz = output["viz"]
        for key in ("tile_attn", "level1_patch_attn", "level2_attn", "gate"):
            captured[key] = viz[key].detach().float().cpu().numpy()

    def collect_heads(module, args, output):
        captured["head_attn"] = output[1][0, :, 0, :].detach().float().cpu().numpy()

    handles = []
    try:
        handles.append(
            detector.model.register_forward_pre_hook(enable_viz, with_kwargs=True)
        )
        handles.append(detector.model.register_forward_hook(collect))
        handles.append(
            detector.model.core.level1_aggregate_attn.register_forward_hook(
                collect_heads
            )
        )
        result = detector.predict(path, branch, threshold)
    finally:
        for handle in handles:
            handle.remove()
    nt = result["views"] - 1
    patches = captured["level1_patch_attn"].reshape(nt, -1)
    side = math.isqrt(patches.shape[1])
    if side * side != patches.shape[1]:
        raise ValueError("Patch attention does not form a square grid.")
    attention = dict(
        l1=captured["tile_attn"].reshape(nt),
        l1_heads=captured["head_attn"],
        l2=captured["level2_attn"].reshape(nt + 1),
        patch=patches.reshape(nt, side, side),
        gate=float(captured["gate"].reshape(-1)[0]),
    )
    for key in ("l1", "l2", "patch", "l1_heads"):
        weights = attention[key]
        if not np.isfinite(weights).all() or (weights < 0).any():
            raise ValueError("Attention contains invalid values.")
    return result, attention


def tile_boxes(width, height, tile_size=224):
    # Match inference: non-overlapping full blocks, then right/bottom boundary blocks.
    from utils.tile_dataset_srm import get_start_positions

    return [
        (x, y, x + tile_size, y + tile_size)
        for y in get_start_positions(height, tile_size)
        for x in get_start_positions(width, tile_size)
    ]


def relative_scale(values):
    values = np.asarray(values, dtype=np.float32)
    low, high = float(values.min()), float(values.max())
    if high - low < 1e-12:
        return np.full_like(values, 0.5)
    return (values - low) / (high - low)


def colors(values):
    """Fixed blue → neutral → orange-red scale, centered on uniform attention."""
    x = np.clip(np.asarray(values), 0, 1)
    stops = np.array(
        [[35, 57, 93], [104, 134, 171], [180, 191, 203], [247, 163, 66], [238, 80, 60]],
        dtype=np.float32,
    )
    first = np.minimum((x * 4).astype(int), 3)
    fraction = (x * 4 - first)[..., None]
    return (stops[first] * (1 - fraction) + stops[first + 1] * fraction).astype(
        np.uint8
    )


def attention_levels(values, reference_count):
    """A common reference: 1/N is neutral; clip only outside [0.25×, 4×]."""
    ratios = np.maximum(np.asarray(values, dtype=np.float32) * reference_count, 1e-12)
    return np.clip((np.log2(ratios) + 2) / 4, 0, 1)


def preview_image(path):
    with Image.open(path) as source:
        image = source.convert("RGB")
    image.thumbnail((1200, 1200))
    return image


def draw_grid(path, boxes):
    image = preview_image(path)
    with Image.open(path) as original:
        sx, sy = image.width / original.width, image.height / original.height
    draw = ImageDraw.Draw(image)
    font = ImageFont.load_default(size=18)
    labels = []
    for index, (x0, y0, x1, y1) in enumerate(boxes, 1):
        rect = (
            int(x0 * sx),
            int(y0 * sy),
            min(image.width - 1, int(x1 * sx) - 1),
            min(image.height - 1, int(y1 * sy) - 1),
        )
        draw.rectangle(rect, outline=(11, 35, 59), width=4)
        draw.rectangle(rect, outline=(234, 252, 255), width=2)
        if len(boxes) > 80 and (index - 1) % math.ceil(len(boxes) / 80):
            continue
        label = f"T{index}"
        left, top = rect[0] + 3, rect[1] + 3
        # Edge alignment can place two rows just one pixel apart. Keep their
        # exact crop borders while moving text labels to avoid hiding earlier IDs.
        for candidate_top in range(top, max(top + 1, rect[3] - 21), 24):
            extent = draw.textbbox((left, candidate_top), label, font=font)
            if not any(
                left < old[2] + 3
                and extent[2] + 3 > old[0]
                and candidate_top < old[3] + 3
                and extent[3] + 3 > old[1]
                for old in labels
            ):
                top = candidate_top
                break
        extent = draw.textbbox((left, top), label, font=font)
        labels.append((left, top, extent[2], extent[3]))
        draw.rectangle(
            (left - 2, top - 2, extent[2] + 4, extent[3] + 4), fill=(11, 35, 59)
        )
        draw.text((left, top), label, fill="white", font=font)
    return image


def tile_overlay(path, boxes, weights, reference_count=None):
    image = preview_image(path)
    with Image.open(path) as original:
        sx, sy = image.width / original.width, image.height / original.height
    summed = np.zeros((image.height, image.width), dtype=np.float32)
    coverage = np.zeros_like(summed)
    for box, weight in zip(boxes, weights, strict=True):
        x0, y0, x1, y1 = box
        x0, y0 = int(x0 * sx), int(y0 * sy)
        x1, y1 = (
            min(image.width, math.ceil(x1 * sx)),
            min(image.height, math.ceil(y1 * sy)),
        )
        summed[y0:y1, x0:x1] += float(weight)
        coverage[y0:y1, x0:x1] += 1
    averaged = summed / np.maximum(coverage, 1)
    count = reference_count or len(weights)
    ratio = np.maximum(averaged * count, 1e-12)
    heat = colors(attention_levels(averaged, count)).astype(np.float32)
    alpha = np.clip(0.15 + np.abs(np.log2(ratio)) * 0.3, 0.15, 0.75)[..., None]
    pixels = np.asarray(image, dtype=np.float32) * (1 - alpha) + heat * alpha
    overlay = Image.fromarray(np.clip(pixels, 0, 255).astype(np.uint8))
    # Rank outlines are useful only when attention differences are meaningful.
    if float(np.ptp(weights)) * count >= 0.10:
        draw = ImageDraw.Draw(overlay)
        for rank, index in enumerate(
            np.argsort(-np.asarray(weights), kind="stable")[:3], 1
        ):
            x0, y0, x1, y1 = boxes[index]
            rect = (
                int(x0 * sx),
                int(y0 * sy),
                min(image.width - 1, int(x1 * sx) - 1),
                min(image.height - 1, int(y1 * sy) - 1),
            )
            draw.rectangle(rect, outline=(12, 26, 42), width=7)
            draw.rectangle(rect, outline=(255, 198, 65), width=4)
    return overlay


def patch_view(state, tile_number):
    if not state:
        return None, None, "Run analysis first."
    index = int(tile_number) - 1
    if index < 0 or index >= len(state["boxes"]):
        raise ValueError("Tile number is out of range.")
    box = state["boxes"][index]
    with Image.open(state["path"]) as source:
        tile = source.convert("RGB").crop(box)
    weights = state["attention"]["patch"][index]
    heat = Image.fromarray(colors(attention_levels(weights, weights.size))).resize(
        tile.size, Image.Resampling.NEAREST
    )
    overlay = Image.blend(tile, heat, 0.65)
    description = (
        f"T{index + 1} · original box {list(box)} · patch grid "
        f"{weights.shape[0]} × {weights.shape[1]} · "
        f"patch weights sum = {float(weights.sum()):.6f}. "
        f" Uniform attention = {100/weights.size:.4f}%. "
        "All tiles share a scale from 0.25× to 4× uniform attention. Small images retain the black padding used by the model."
    )
    return tile, overlay, description


def top_tiles(state, k=4):
    order = np.argsort(-state["attention"]["l1"], kind="stable")[:k]
    gallery = []
    for index in order:
        tile, _, _ = patch_view(state, int(index) + 1)
        weight = float(state["attention"]["l1"][index])
        gallery.append(
            (tile, f"#{len(gallery)+1} · T{index + 1} · weight {weight*100:.2f}%")
        )
    return gallery


def head_view(state, selection):
    if not state:
        return None, "Run analysis first."
    if selection == "mean":
        weights = state["attention"]["l1"]
        label = "Mean across L1 aggregation heads"
    else:
        index = int(selection.split("_")[-1]) - 1
        weights = state["attention"]["l1_heads"][index]
        label = f"L1 aggregation head {index + 1}"
    return tile_overlay(state["path"], state["boxes"], weights), (
        f"{label} · uniform reference = {100/len(weights):.2f}%. "
        "Blue is below the reference, gray is near it, and orange-red is above it. Overlapping regions use mean tile attention."
        "Tile attention is not a pixel-level forgery probability."
    )


def parameter_rows(detector, branch, threshold, views):
    core, visual = detector.model.core, detector.model.visual
    mil1, mil2 = core.mil_level1.head.head, core.mil_level2.head.head
    rows = [
        ["Inference", "Branch / threshold", f"{branch} / {threshold:.2f}"],
        [
            "Inference",
            "Precision / device",
            f"{visual.conv1.weight.dtype} / {detector.device}",
        ],
        ["Inference", "CLIP chunk size", str(detector.model.eval_chunk_size)],
        ["Inference", "Views: global + local", f"1 + {views - 1} = {views}"],
        ["Inference", "Tiling policy", "Non-overlapping full blocks; overlap only at right/bottom boundaries"],
        ["Architecture", "Backbone", "Frozen CLIP ViT-L/14"],
        [
            "Architecture",
            "Local tile size",
            f"{int(detector.model.input_resolution)} × {int(detector.model.input_resolution)}",
        ],
        ["Architecture", "CLIP patch size", str(visual.conv1.kernel_size[0])],
        ["Architecture", "Feature dimension", str(detector.model.embed_dim)],
        [
            "Architecture",
            "Feature extraction layer",
            str(
                detector.model.clip_mil_stop_after_block_1based
                if detector.model.use_clip_intermediate_for_mil
                else len(visual.transformer.resblocks)
            ),
        ],
        [
            "Architecture",
            "L1 across-tile aggregation heads",
            str(core.level1_aggregate_attn.num_heads),
        ],
        [
            "Architecture",
            "L1 MIL pooling / hidden",
            f"{mil1.pooling} / {mil1.attn_V.out_features} ({mil1.attn_bag_mode})",
        ],
        [
            "Architecture",
            "L2 MIL pooling / hidden",
            f"{mil2.pooling} / {mil2.attn_V.out_features} ({mil2.attn_bag_mode})",
        ],
        [
            "Architecture",
            "Fusion / gate hidden",
            f"pool-logit gated / {core.pool_gate[0].out_features}",
        ],
        ["Architecture", "Fusion MLP hidden", str(core.fusion_head[0].out_features)],
        ["Inference", "Cross-shuffle / dropout", "Disabled in evaluation mode"],
    ]
    reference = Path(__file__).with_name("training_reference.json")
    if reference.is_file():
        training = json.loads(reference.read_text())
        rows.extend(
            ["Training only", key, json.dumps(value, ensure_ascii=False)]
            for key, value in training.items()
        )
    return rows


def public_attention(state):
    weights = state["attention"]
    return dict(
        meaning="Relative model attention, not a forgery probability or causal explanation.",
        display="Fixed log2 scale relative to uniform 1/N attention: 0.25x to 4x; overlap uses mean tile weight.",
        local_tile_boxes=[list(box) for box in state["boxes"]],
        l1_tile_weights=weights["l1"].tolist(),
        l1_aggregation_head_weights=weights["l1_heads"].tolist(),
        l2_view_weights=weights["l2"].tolist(),
        l2_view_order=["global"] + [f"T{i + 1}" for i in range(len(state["boxes"]))],
        patches_per_tile=list(weights["patch"].shape[1:]),
        fusion_embedding_gate_l1=weights["gate"],
        fusion_embedding_gate_l2=1 - weights["gate"],
    )