"""
ForensicAI Morphic Simulation Engine
Gray-Scott reaction-diffusion parameterized by forensic case data.
Palette: UV lab blue, evidence amber, biological green, digital cyan, crime red.
"""

from __future__ import annotations

import io
from typing import Optional

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

# ── Forensic color palette ────────────────────────────────────────────────────

_BLACK   = np.array([5,   5,   10],  dtype=np.float32)   # near black
_NAVY    = np.array([10,  15,  40],  dtype=np.float32)
_BLUE    = np.array([30,  80,  200], dtype=np.float32)
_CYAN    = np.array([0,   180, 220], dtype=np.float32)   # UV cyan
_GREEN   = np.array([20,  200, 100], dtype=np.float32)   # bio green
_AMBER   = np.array([245, 160, 10],  dtype=np.float32)   # evidence amber
_RED     = np.array([210, 40,  40],  dtype=np.float32)   # crime red
_PURPLE  = np.array([120, 40,  200], dtype=np.float32)   # forensic purple
_WHITE   = np.array([230, 240, 255], dtype=np.float32)   # cold white

# Case type → color identity
CASE_PALETTES: dict[str, tuple[np.ndarray, np.ndarray, np.ndarray]] = {
    "homicide":       (_RED,    _PURPLE, _WHITE),
    "sexual_assault": (_PURPLE, _CYAN,   _WHITE),
    "digital_crime":  (_CYAN,   _BLUE,   _WHITE),
    "drug_offense":   (_AMBER,  _GREEN,  _WHITE),
    "arson":          (_AMBER,  _RED,    _WHITE),
    "burglary":       (_BLUE,   _AMBER,  _WHITE),
    "fraud":          (_GREEN,  _CYAN,   _WHITE),
    "hit_and_run":    (_AMBER,  _RED,    _CYAN),
    "default":        (_CYAN,   _BLUE,   _WHITE),
}

# Case type → Gray-Scott parameters
# Homicide/serious → complex labyrinthine (low f, mid k)
# Digital → grid-like / spotted (mid f, high k)
# Bio/sexual → worm-like tendrils (mid f, mid k)
CASE_PARAMS: dict[str, tuple[float, float]] = {
    "homicide":       (0.025, 0.055),
    "sexual_assault": (0.037, 0.061),
    "digital_crime":  (0.029, 0.057),
    "drug_offense":   (0.040, 0.060),
    "arson":          (0.033, 0.058),
    "burglary":       (0.038, 0.062),
    "fraud":          (0.026, 0.053),
    "hit_and_run":    (0.035, 0.059),
    "default":        (0.037, 0.060),
}

# ── Color LUT builder ─────────────────────────────────────────────────────────

def build_lut(c1: np.ndarray, c2: np.ndarray, c3: np.ndarray) -> np.ndarray:
    """Build 256×3 uint8 LUT: black → navy → c1 → c2 → c3 → cold white."""
    stops = [
        (0,   _BLACK),
        (45,  _NAVY),
        (90,  c1 * 0.5),
        (130, c1),
        (170, c2),
        (210, c3 * 0.9),
        (245, c3),
        (255, _WHITE),
    ]
    lut = np.zeros((256, 3), dtype=np.float32)
    for i in range(len(stops) - 1):
        i0, a = stops[i]
        i1, b = stops[i + 1]
        span = i1 - i0
        for j in range(i0, i1):
            t = (j - i0) / span
            lut[j] = a * (1 - t) + b * t
    lut[255] = _WHITE
    return np.clip(lut, 0, 255).astype(np.uint8)

# ── Gray-Scott engine ─────────────────────────────────────────────────────────

def _laplacian(z: np.ndarray) -> np.ndarray:
    return (
        np.roll(z,  1, 0) + np.roll(z, -1, 0) +
        np.roll(z,  1, 1) + np.roll(z, -1, 1) - 4.0 * z
    )


def run_gray_scott(
    w: int, h: int, iters: int,
    feed: float, kill: float,
    seed: int, n_patches: int,
) -> np.ndarray:
    rng = np.random.default_rng(seed)
    u = np.ones((h, w),  dtype=np.float32)
    v = np.zeros((h, w), dtype=np.float32)
    r = max(5, min(w, h) // 16)
    for _ in range(n_patches):
        cx = int(rng.integers(r, w - r))
        cy = int(rng.integers(r, h - r))
        noise = 0.05 * rng.random((2*r, 2*r)).astype(np.float32)
        u[cy-r:cy+r, cx-r:cx+r] = 0.50 + noise
        v[cy-r:cy+r, cx-r:cx+r] = 0.25 + noise
    Du, Dv = 0.16, 0.08
    for _ in range(iters):
        uvv = u * v * v
        u  += Du * _laplacian(u) - uvv + feed * (1.0 - u)
        v  += Dv * _laplacian(v) + uvv - (feed + kill) * v
        np.clip(u, 0.0, 1.0, out=u)
        np.clip(v, 0.0, 1.0, out=v)
    return v

# ── Renderer ──────────────────────────────────────────────────────────────────

_FONT = None
def _font(size=11):
    global _FONT
    if _FONT is None:
        try:
            _FONT = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSansMono.ttf", size)
        except Exception:
            _FONT = ImageFont.load_default()
    return _FONT


def _render(v: np.ndarray, lut: np.ndarray, size: int) -> Image.Image:
    vn = (v - v.min()) / (v.max() - v.min() + 1e-9)
    vs = np.clip(vn * 1.9 - 0.1, 0.0, 1.0)
    idx = (vs * 255).astype(np.uint8)
    rgb = lut[idx]
    img = Image.fromarray(rgb, "RGB").resize((size, size), Image.BILINEAR)
    glow = img.filter(ImageFilter.GaussianBlur(3))
    return Image.blend(img, glow, 0.3)


def _hud(img: Image.Image, case_type: str, generation: int,
         label: str, feed: float, kill: float) -> Image.Image:
    w, h = img.size
    draw = ImageDraw.Draw(img, "RGBA")
    font = _font(11)
    bracket = (0, 180, 220, 160)        # UV cyan
    blen = 22
    for bx, by in [(6, 6), (w-6, 6), (6, h-6), (w-6, h-6)]:
        sx = 1 if bx < w//2 else -1
        sy = 1 if by < h//2 else -1
        draw.line([(bx, by), (bx + sx*blen, by)], fill=bracket, width=2)
        draw.line([(bx, by), (bx, by + sy*blen)], fill=bracket, width=2)

    # Top bar
    draw.rectangle([(0, 0), (w, 26)], fill=(5, 5, 15, 210))
    draw.text((10, 6), "FORENSICAI · 2ND BRAIN", fill=(0, 180, 220, 220), font=font)
    draw.text((w - 72, 6), f"GEN {generation:03d}", fill=(200, 100, 240, 200), font=font)

    # Bottom bar
    draw.rectangle([(0, h-26), (w, h)], fill=(5, 5, 15, 210))
    draw.text((10, h-19), label, fill=(180, 200, 230, 200), font=font)
    draw.text((w-130, h-19), f"f={feed:.4f}  k={kill:.4f}", fill=(80, 120, 180, 160), font=font)

    # Case type badge
    case_colors = {
        "homicide": (210, 40, 40), "digital_crime": (0, 180, 220),
        "drug_offense": (245, 160, 10), "sexual_assault": (160, 60, 220),
        "arson": (245, 100, 20), "default": (100, 140, 255),
    }
    cc = case_colors.get(case_type, (100, 140, 255))
    draw.rectangle([(w-95, 28), (w-6, 48)], fill=(*cc, 180))
    draw.text((w-90, 32), case_type[:12].upper(), fill=(255, 255, 255, 230), font=font)
    return img


def generate_forensic_image(
    dominant_case: str = "default",
    evidence_confidence: float = 0.70,
    session_count: int = 1,
    pattern_generation: int = 1,
    case_complexity: float = 0.5,
    out_size: int = 512,
    grid: int = 220,
    iterations: int = 900,
) -> bytes:
    base_f, base_k = CASE_PARAMS.get(dominant_case, CASE_PARAMS["default"])

    # Confidence shifts kill (higher confidence → more defined patterns)
    conf_norm = max(0.0, min(1.0, evidence_confidence))
    cmplx_norm = max(0.0, min(1.0, case_complexity))
    feed = round(base_f + cmplx_norm * 0.008, 5)
    kill = round(base_k + conf_norm  * 0.006, 5)

    seed     = (session_count * 17 + pattern_generation * 11 + 7) % (2**31)
    n_patches = min(2 + session_count // 6, 14)

    v = run_gray_scott(grid, grid, iterations, feed, kill, seed, n_patches)

    c1, c2, c3 = CASE_PALETTES.get(dominant_case, CASE_PALETTES["default"])
    lut = build_lut(c1, c2, c3)
    img = _render(v, lut, out_size)

    label = f"{dominant_case}  ·  conf {evidence_confidence:.2f}  ·  complexity {case_complexity:.2f}"
    img = _hud(img, dominant_case, pattern_generation, label, feed, kill)

    buf = io.BytesIO()
    img.save(buf, "PNG", optimize=True)
    return buf.getvalue()


def generate_case_heatmap(
    case_history: list[str],
    confidence_history: list[float],
    case_frequency: dict[str, int],
    out_size: int = 512,
) -> bytes:
    img = Image.new("RGB", (out_size, out_size), (5, 5, 15))
    draw = ImageDraw.Draw(img)
    font = _font(11)
    w, h = out_size, out_size

    # Title
    draw.rectangle([(0, 0), (w, 30)], fill=(10, 20, 50))
    draw.text((10, 8), "FORENSICAI 2ND BRAIN  —  CASE INTELLIGENCE MAP", fill=(0, 180, 220), font=font)

    # Confidence history strip (left half)
    pw = w // 2 - 10
    strip_top, strip_bot = 50, h - 60
    strip_h = strip_bot - strip_top
    draw.text((10, 36), "EVIDENCE CONFIDENCE HISTORY", fill=(80, 140, 220), font=font)
    hist = confidence_history[-pw:] if len(confidence_history) > pw else confidence_history
    for i, cf in enumerate(hist):
        t = max(0.0, min(1.0, cf))
        r = int(30  + (210 - 30)  * t)
        g = int(80  + (180 - 80)  * t)
        b = int(200 + (220 - 200) * t)
        bh = int(strip_h * t)
        draw.line([(10+i, strip_bot), (10+i, strip_bot - bh)], fill=(r, g, b))
    draw.rectangle([(8, strip_top), (10+pw, strip_bot)], outline=(30, 50, 100), width=1)

    # Case frequency bars (right half)
    rx = w // 2 + 5
    draw.text((rx, 36), "CASE TYPE FREQUENCY", fill=(80, 140, 220), font=font)
    case_colors_map = {
        "homicide":       (210, 40,  40),
        "digital_crime":  (0,   180, 220),
        "drug_offense":   (245, 160, 10),
        "sexual_assault": (160, 60,  220),
        "arson":          (245, 100, 20),
        "burglary":       (80,  140, 255),
        "fraud":          (20,  200, 100),
        "hit_and_run":    (245, 130, 50),
    }
    mx = max(case_frequency.values()) if case_frequency else 1
    bw = (w - rx - 15) // max(len(case_frequency), 1)
    bx = rx + 5
    for case, n in case_frequency.items():
        ratio = n / max(mx, 1)
        bh    = int((strip_h - 20) * ratio)
        col   = case_colors_map.get(case, (100, 140, 220))
        draw.rectangle([(bx, strip_bot-bh), (bx+bw-5, strip_bot)], fill=col)
        draw.text((bx, strip_bot+4), case[:6], fill=(160, 190, 230), font=font)
        draw.text((bx, strip_bot-bh-14), str(n), fill=(210, 220, 240), font=font)
        bx += bw

    # Confidence trend line (bottom)
    bt = h - 56
    draw.text((10, bt-16), "CONFIDENCE TREND", fill=(80, 140, 220), font=font)
    scores = confidence_history[-(w-20):] if len(confidence_history) > w-20 else confidence_history
    if len(scores) > 1:
        pts = [(10 + int(i*(w-20)/len(scores)), bt + int((1-s)*40)) for i, s in enumerate(scores)]
        for i in range(len(pts)-1):
            draw.line([pts[i], pts[i+1]], fill=(0, 180, 220), width=2)
    draw.rectangle([(8, bt), (w-8, h-12)], outline=(30, 50, 100), width=1)

    buf = io.BytesIO()
    img.save(buf, "PNG", optimize=True)
    return buf.getvalue()