"""
Morphic DNA simulation engine — Gray-Scott reaction-diffusion parameterized by genomic data.

Color palette mirrors the AdaptiveDNA desktop app (#0f172a bg, DNA green/blue/purple).
"""

from __future__ import annotations

import io
from typing import Optional

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

# ── AdaptiveDNA color palette ─────────────────────────────────────────────────

_NAVY   = np.array([15,  23,  42],  dtype=np.float32)
_GREEN  = np.array([20,  180, 118], dtype=np.float32)
_BLUE   = np.array([59,  130, 246], dtype=np.float32)
_PURPLE = np.array([168, 85,  247], dtype=np.float32)
_YELLOW = np.array([245, 158, 11],  dtype=np.float32)
_CYAN   = np.array([6,   182, 212], dtype=np.float32)
_CREAM  = np.array([255, 248, 220], dtype=np.float32)
_RED    = np.array([239, 68,  68],  dtype=np.float32)
_AMBER  = np.array([251, 191, 36],  dtype=np.float32)

CROP_TINTS: dict[str, np.ndarray] = {
    "Rice":    _GREEN,
    "Maize":   _YELLOW,
    "Wheat":   _AMBER,
    "Tomato":  _RED,
    "Soybean": _PURPLE,
    "default": _CYAN,
}

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

def build_lut(tint: np.ndarray) -> np.ndarray:
    """Build a 256×3 uint8 LUT: navy → tint → blue → purple → cyan → cream."""
    stops = [
        (0,   _NAVY),
        (60,  tint * 0.35),
        (110, tint * 0.75),
        (150, tint),
        (185, _BLUE),
        (215, _PURPLE),
        (238, _CYAN),
        (255, _CREAM),
    ]
    lut = np.zeros((256, 3), dtype=np.float32)
    for i in range(len(stops) - 1):
        i0, c0 = stops[i]
        i1, c1 = stops[i + 1]
        span = i1 - i0
        for j in range(i0, i1):
            t = (j - i0) / span
            lut[j] = c0 * (1 - t) + c1 * t
    lut[255] = _CREAM
    return np.clip(lut, 0, 255).astype(np.uint8)


# ── Gray-Scott reaction-diffusion ─────────────────────────────────────────────

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


def run_gray_scott(
    width: int,
    height: int,
    iterations: int,
    feed: float,
    kill: float,
    seed: int,
    n_patches: int,
) -> np.ndarray:
    """Return the v-field (activator) after `iterations` steps."""
    rng = np.random.default_rng(seed)
    u = np.ones((height, width), dtype=np.float32)
    v = np.zeros((height, width), dtype=np.float32)

    r = max(6, min(width, height) // 14)
    for _ in range(n_patches):
        cx = int(rng.integers(r, width  - r))
        cy = int(rng.integers(r, height - r))
        noise = 0.04 * 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(iterations):
        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


# ── Image renderer ────────────────────────────────────────────────────────────

_FONT_MONO = None  # loaded lazily; gracefully degrades to default

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


def _render_v_field(v: np.ndarray, lut: np.ndarray, out_size: int) -> Image.Image:
    v_min, v_max = v.min(), v.max()
    v_norm = (v - v_min) / (v_max - v_min + 1e-9)
    # Amplify contrast: sigmoid-like stretch
    v_stretched = np.clip(v_norm * 1.8 - 0.1, 0, 1)
    idx = (v_stretched * 255).astype(np.uint8)
    rgb = lut[idx]
    img = Image.fromarray(rgb, mode="RGB")
    img = img.resize((out_size, out_size), Image.BILINEAR)
    # Subtle glow blend
    glow = img.filter(ImageFilter.GaussianBlur(radius=4))
    return Image.blend(img, glow, 0.25)


def _overlay_hud(img: Image.Image, label: str, generation: int, feed: float, kill: float) -> Image.Image:
    """Draw a minimal sci-fi HUD onto the image."""
    w, h = img.size
    draw = ImageDraw.Draw(img, "RGBA")
    font = _get_font(12)

    # Corner brackets
    bracket_c = (20, 180, 118, 180)   # DNA green
    blen = 24
    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_c, width=2)
        draw.line([(bx, by), (bx, by + sy * blen)], fill=bracket_c, width=2)

    # Top banner
    draw.rectangle([(0, 0), (w, 26)], fill=(15, 23, 42, 200))
    draw.text((10, 6), "ADAPTIVEDNA · 2ND BRAIN", fill=(20, 180, 118, 230), font=font)
    gen_txt = f"GEN {generation:03d}"
    draw.text((w - 70, 6), gen_txt, fill=(168, 85, 247, 200), font=font)

    # Bottom banner
    draw.rectangle([(0, h - 26), (w, h)], fill=(15, 23, 42, 200))
    draw.text((10, h - 19), label, fill=(200, 220, 255, 200), font=font)
    param_txt = f"f={feed:.4f}  k={kill:.4f}"
    draw.text((w - 130, h - 19), param_txt, fill=(100, 150, 200, 180), font=font)

    return img


def generate_morphic_image(
    gc_content: float = 50.0,
    avg_score: float = 0.70,
    dominant_crop: str = "default",
    session_count: int = 1,
    pattern_generation: int = 1,
    out_size: int = 512,
    grid: int = 220,
    iterations: int = 900,
) -> bytes:
    """
    Generate a morphic DNA simulation PNG.

    Parameters are mapped to Gray-Scott feed/kill rates so the visual evolves
    as genomic data accumulates in the 2nd Brain.
    """
    gc_norm = max(0.0, min(1.0, gc_content / 100.0))
    score_norm = max(0.0, min(1.0, avg_score))

    # Parameter mapping:
    #  High GC → more worm-like patterns (higher feed)
    #  High guide score → tighter, more structured patterns (higher kill)
    feed = 0.029 + gc_norm  * 0.022      # 0.029 – 0.051
    kill = 0.053 + score_norm * 0.012    # 0.053 – 0.065

    seed     = (session_count * 13 + pattern_generation * 7 + 42) % (2 ** 31)
    n_patches = min(3 + session_count // 8, 12)

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

    tint = CROP_TINTS.get(dominant_crop, CROP_TINTS["default"])
    lut  = build_lut(tint)
    img  = _render_v_field(v, lut, out_size)

    label = f"{dominant_crop}  ·  GC {gc_content:.1f}%  ·  Score {avg_score:.2f}"
    img   = _overlay_hud(img, label, pattern_generation, feed, kill)

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


def generate_brain_heatmap(
    gc_history: list[float],
    score_history: list[float],
    crop_frequency: dict[str, int],
    out_size: int = 512,
) -> bytes:
    """
    Generate a 2-panel heatmap visualization of the brain's accumulated state:
    left = GC history strip, right = crop frequency bars.
    """
    img = Image.new("RGB", (out_size, out_size), color=(15, 23, 42))
    draw = ImageDraw.Draw(img)
    font = _get_font(12)
    w, h = out_size, out_size

    # ── Title ────────────────────────────────────────────────────────────────
    draw.rectangle([(0, 0), (w, 30)], fill=(20, 30, 60))
    draw.text((10, 8), "ADAPTIVEDNA 2ND BRAIN  —  STATE SNAPSHOT", fill=(20, 180, 118), font=font)

    # ── GC history strip (left half) ─────────────────────────────────────────
    panel_w = w // 2 - 10
    strip_top, strip_bot = 50, h - 60
    strip_h = strip_bot - strip_top

    draw.text((10, 36), "GC CONTENT HISTORY", fill=(100, 150, 255), font=font)

    hist = gc_history[-panel_w:] if len(gc_history) > panel_w else gc_history
    if hist:
        for i, gc in enumerate(hist):
            t  = max(0.0, min(1.0, gc / 100.0))
            r  = int(59  + (168 - 59)  * t)
            g  = int(130 + (85  - 130) * t)
            b  = int(246 + (247 - 246) * t)
            bar_h = int(strip_h * t)
            x = 10 + i
            draw.line([(x, strip_bot), (x, strip_bot - bar_h)], fill=(r, g, b), width=1)

    draw.rectangle([(8, strip_top), (10 + panel_w, strip_bot)], outline=(40, 60, 100), width=1)

    # ── Crop frequency bars (right half) ─────────────────────────────────────
    rx_start = w // 2 + 5
    draw.text((rx_start, 36), "CROP FREQUENCY", fill=(100, 150, 255), font=font)

    crop_colors = {
        "Rice":    (20,  180, 118),
        "Maize":   (245, 158, 11),
        "Wheat":   (251, 191, 36),
        "Tomato":  (239, 68,  68),
        "Soybean": (168, 85,  247),
    }
    max_freq = max(crop_frequency.values()) if crop_frequency else 1
    bar_w    = (w - rx_start - 15) // max(len(crop_frequency), 1)
    bar_x    = rx_start + 5

    for crop, freq in crop_frequency.items():
        ratio  = freq / max(max_freq, 1)
        bar_h  = int((strip_h - 20) * ratio)
        color  = crop_colors.get(crop, (100, 100, 200))
        draw.rectangle([(bar_x, strip_bot - bar_h), (bar_x + bar_w - 6, strip_bot)], fill=color)
        draw.text((bar_x, strip_bot + 4), crop[:4], fill=(180, 200, 255), font=font)
        draw.text((bar_x, strip_bot - bar_h - 14), str(freq), fill=(220, 220, 220), font=font)
        bar_x += bar_w

    # ── Score history line (bottom strip) ─────────────────────────────────────
    bot_top = h - 56
    draw.text((10, bot_top - 16), "GUIDE SCORE HISTORY", fill=(100, 150, 255), font=font)
    scores = score_history[-(w - 20):] if len(score_history) > w - 20 else score_history
    if len(scores) > 1:
        pts = [
            (10 + int(i * (w - 20) / len(scores)),
             bot_top + 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=(6, 182, 212), width=2)

    draw.rectangle([(8, bot_top), (w - 8, h - 12)], outline=(40, 60, 100), width=1)

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