"""
Morphic simulation engine.
Generates Gray-Scott reaction-diffusion patterns that self-assemble
based on protein properties — each protein produces a unique visual fingerprint.
"""
from __future__ import annotations
import hashlib, io, math
import numpy as np
from PIL import Image, ImageDraw, ImageFont
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap

# ── Protein-property → Gray-Scott parameter mapping ───────────────────────────

FUNCTION_PALETTES = {
    "enzyme":        [(0.02, 0.055), "#00f5ff", "#0a0a2e"],
    "receptor":      [(0.035, 0.065), "#c084fc", "#0d0520"],
    "transporter":   [(0.025, 0.06),  "#10b981", "#031a0e"],
    "structural":    [(0.04, 0.063),  "#fbbf24", "#1a1000"],
    "transcription": [(0.03, 0.057),  "#f97316", "#1a0a00"],
    "signaling":     [(0.022, 0.051), "#3b82f6", "#00020f"],
    "unknown":       [(0.028, 0.058), "#94a3b8", "#0c0c14"],
}

def _protein_to_gs_params(
    uniprot_id: str,
    plddt: float,
    length: int,
    function: str,
) -> dict:
    """Map protein properties to Gray-Scott f/k parameters and palette."""
    palette_key = function.lower() if function.lower() in FUNCTION_PALETTES else "unknown"
    (f_base, k_base), fg_hex, bg_hex = FUNCTION_PALETTES[palette_key]

    # pLDDT (0–100) shifts f: high confidence → more stable patterns
    f = f_base + (plddt / 100 - 0.5) * 0.008
    # Protein length shifts k: longer → finer pattern
    k = k_base + min(length / 10000, 0.006)
    # Unique seed from UniProt ID
    seed = int(hashlib.md5(uniprot_id.encode()).hexdigest()[:8], 16) % 10000

    return {"f": f, "k": k, "fg": fg_hex, "bg": bg_hex, "seed": seed}


def _run_gs(f: float, k: float, seed: int, steps: int = 3000, N: int = 128) -> np.ndarray:
    """Run Gray-Scott PDE on an N×N grid, return V concentration."""
    rng = np.random.default_rng(seed)
    U = np.ones((N, N), dtype=np.float32)
    V = np.zeros((N, N), dtype=np.float32)

    # Seed perturbation
    cx, cy = N // 2, N // 2
    r = max(4, N // 16)
    U[cx - r:cx + r, cy - r:cy + r] = 0.5
    V[cx - r:cx + r, cy - r:cy + r] = 0.25
    V += rng.uniform(0, 0.01, (N, N)).astype(np.float32)

    Du, Dv, dt = 0.16, 0.08, 1.0

    for _ in range(steps):
        Lu = (np.roll(U, 1, 0) + np.roll(U, -1, 0) +
              np.roll(U, 1, 1) + np.roll(U, -1, 1) - 4 * U)
        Lv = (np.roll(V, 1, 0) + np.roll(V, -1, 0) +
              np.roll(V, 1, 1) + np.roll(V, -1, 1) - 4 * V)
        uvv = U * V * V
        U += dt * (Du * Lu - uvv + f * (1 - U))
        V += dt * (Dv * Lv + uvv - (f + k) * V)
        np.clip(U, 0, 1, out=U)
        np.clip(V, 0, 1, out=V)

    return V


def _hex_to_rgb(h: str) -> tuple:
    h = h.lstrip("#")
    return tuple(int(h[i:i+2], 16) / 255.0 for i in (0, 2, 4))


def generate_morphic_image(
    uniprot_id: str,
    plddt: float = 75.0,
    length: int = 300,
    function: str = "unknown",
    protein_name: str = "",
    width: int = 800,
    height: int = 800,
) -> bytes:
    """
    Return PNG bytes of a morphic simulation unique to this protein.
    The pattern self-assembles based on protein properties.
    """
    params = _protein_to_gs_params(uniprot_id, plddt, length, function)
    V = _run_gs(params["f"], params["k"], params["seed"])

    # Build custom colormap from protein palette
    bg = _hex_to_rgb(params["bg"])
    fg = _hex_to_rgb(params["fg"])
    mid = tuple(min(1.0, c * 1.6) for c in fg)
    cmap = LinearSegmentedColormap.from_list(
        "protein", [bg, params["bg"], fg, mid, "#ffffff"], N=512
    )

    fig, ax = plt.subplots(figsize=(width / 100, height / 100), dpi=100)
    fig.patch.set_facecolor(params["bg"])
    ax.set_facecolor(params["bg"])
    ax.imshow(V, cmap=cmap, interpolation="bilinear", aspect="auto", vmin=0, vmax=V.max())
    ax.axis("off")

    # Overlay protein ID and pLDDT badge
    label = f"{uniprot_id}  ·  pLDDT {plddt:.1f}  ·  {length} aa"
    ax.text(0.5, 0.03, label, transform=ax.transAxes, color=params["fg"],
            fontsize=10, ha="center", va="bottom",
            fontfamily="monospace", alpha=0.85)
    if protein_name:
        ax.text(0.5, 0.96, protein_name, transform=ax.transAxes, color="#ffffff",
                fontsize=13, ha="center", va="top", fontweight="bold", alpha=0.9)

    buf = io.BytesIO()
    plt.savefig(buf, format="png", bbox_inches="tight", pad_inches=0,
                facecolor=params["bg"])
    plt.close(fig)
    buf.seek(0)
    return buf.read()


def generate_confidence_heatmap(
    residue_scores: list[float],
    uniprot_id: str,
    protein_name: str = "",
    width: int = 900,
    height: int = 280,
) -> bytes:
    """Per-residue pLDDT heatmap with adaptive colour coding."""
    arr = np.array(residue_scores, dtype=np.float32).reshape(1, -1)

    cmap = LinearSegmentedColormap.from_list(
        "plddt", ["#f97316", "#eab308", "#22c55e", "#1d4ed8"], N=256
    )

    fig, ax = plt.subplots(figsize=(width / 100, height / 100), dpi=100)
    fig.patch.set_facecolor("#020817")
    ax.set_facecolor("#020817")
    im = ax.imshow(arr, cmap=cmap, aspect="auto", vmin=0, vmax=100)
    ax.set_yticks([])
    ax.tick_params(colors="#64748b", labelsize=8)
    ax.set_xlabel("Residue index", color="#64748b", fontsize=9)
    title = f"{protein_name} ({uniprot_id}) — per-residue pLDDT" if protein_name else f"{uniprot_id} — per-residue pLDDT"
    ax.set_title(title, color="#94a3b8", fontsize=10, pad=8)
    plt.colorbar(im, ax=ax, orientation="horizontal", fraction=0.04, pad=0.18,
                 label="pLDDT confidence").ax.xaxis.label.set_color("#64748b")

    buf = io.BytesIO()
    plt.savefig(buf, format="png", bbox_inches="tight", facecolor="#020817")
    plt.close(fig)
    buf.seek(0)
    return buf.read()