"""
Failure analysis: turn the per-metric spatial maps into an interpretable,
clinician-facing overlay that answers "*where* is this image bad, and *why*?".

Several metrics stash a spatial `_map` (local focus energy, illumination shadow,
artifact outliers, vessel signal). Here we:
  * build a per-axis heatmap the clinician can toggle,
  * build a single composite "problem map" that highlights the regions the
    dominant failing axis complains about,
  * annotate the composite with the driving reason(s).
"""

from __future__ import annotations
import numpy as np
import cv2
from .qc_metrics import inner_disc_mask


def _norm(x):
    x = x.astype(np.float32)
    lo, hi = np.percentile(x, 1), np.percentile(x, 99)
    if hi <= lo:
        return np.zeros_like(x)
    return np.clip((x - lo) / (hi - lo), 0, 1)


def _heat_overlay(rgb, heat, mask, color=(230, 60, 60), alpha=0.55):
    """Blend a heat map (0-1) onto the RGB image within the mask."""
    heat = heat.copy()
    heat[~mask] = 0
    heat = cv2.GaussianBlur(heat, (0, 0), 3)
    heat = _norm(heat)
    layer = np.zeros_like(rgb, np.float32)
    for c in range(3):
        layer[..., c] = color[c]
    out = rgb.astype(np.float32) * (1 - alpha * heat[..., None]) + layer * (alpha * heat[..., None])
    return np.clip(out, 0, 255).astype(np.uint8)


def per_axis_heatmap(rgb, fov, metric):
    """Return an overlay image for a single metric that carries a `_map`."""
    mask = inner_disc_mask(rgb.shape, fov, 0.98)
    m = metric.get("_map")
    if m is None:
        return rgb.copy()
    name = metric["name"]
    if name == "Focus / Defocus":
        # invert: low local sharpness == problem
        loc = cv2.GaussianBlur(m.astype(np.float32), (0, 0), 9)
        problem = 1 - _norm(loc)
        return _heat_overlay(rgb, problem, mask, color=(60, 120, 240))
    if name == "Illumination Uniformity":
        return _heat_overlay(rgb, _norm(m), mask, color=(40, 40, 90), alpha=0.6)
    if name == "Artifact Burden":
        heat = cv2.dilate(m.astype(np.float32), np.ones((5, 5), np.float32))
        return _heat_overlay(rgb, heat, mask, color=(255, 210, 40), alpha=0.8)
    if name == "Vessel Visibility":
        return _heat_overlay(rgb, _norm(m), mask, color=(60, 220, 120), alpha=0.7)
    return _heat_overlay(rgb, _norm(m), mask)


def composite_problem_map(rgb, fov, metrics, score_summary):
    """Highlight the regions responsible for the worst failing axes, with a
    legend of the driving reasons. Returns (overlay_rgb, caption)."""
    mask = inner_disc_mask(rgb.shape, fov, 0.98)
    failing = [m for m in metrics if m["score"] < 0.5 and m.get("_map") is not None]
    failing = sorted(failing, key=lambda m: m["score"])[:3]

    if not failing:
        cap = "No focal quality defect localised - any limitation is global."
        return rgb.copy(), cap

    accum = np.zeros(rgb.shape[:2], np.float32)
    palette = {
        "Focus / Defocus": (60, 120, 240),
        "Illumination Uniformity": (40, 40, 90),
        "Artifact Burden": (255, 210, 40),
        "Vessel Visibility": (60, 220, 120),
    }
    out = rgb.astype(np.float32).copy()
    legend = []
    for metric in failing:
        m = metric["_map"].astype(np.float32)
        name = metric["name"]
        if name == "Focus / Defocus":
            loc = cv2.GaussianBlur(m, (0, 0), 9)
            heat = 1 - _norm(loc)
        else:
            heat = _norm(m)
        heat[~mask] = 0
        heat = cv2.GaussianBlur(heat, (0, 0), 3)
        heat = _norm(heat)
        weight = (1 - metric["score"])
        color = palette.get(name, (230, 60, 60))
        alpha = 0.55 * weight
        for c in range(3):
            out[..., c] = out[..., c] * (1 - alpha * heat) + color[c] * (alpha * heat)
        accum = np.maximum(accum, heat * weight)
        legend.append(f"{name} ({metric['status']})")

    out = np.clip(out, 0, 255).astype(np.uint8)
    cap = "Highlighted regions drive the verdict - " + "; ".join(legend)
    return out, cap


def fov_overlay(rgb, fov):
    """Draw the detected retinal field boundary + inner analysis disc."""
    out = rgb.copy()
    cx, cy, r = int(fov["cx"]), int(fov["cy"]), int(fov["radius"])
    cv2.circle(out, (cx, cy), r, (0, 200, 255), 2)
    cv2.circle(out, (cx, cy), int(r * 0.90), (0, 255, 180), 1, cv2.LINE_AA)
    cv2.drawMarker(out, (cx, cy), (0, 255, 180), cv2.MARKER_CROSS, 14, 2)
    return out