"""AlphaDynamics interactive demo on Hugging Face Spaces.

User pastes a peptide sequence, gets:
- interactive Ramachandran density plot (Plotly)
- per-residue panels for short peptides
- basin populations table (alpha-R, beta, PPII, alpha-L)
- downloadable .npz with the raw trajectory tensor

Free CPU tier. Limits: max 50 residues, max 16 trajectories, max 1000 steps.
"""
from __future__ import annotations

import os
import tempfile
from pathlib import Path

import gradio as gr
import numpy as np
import plotly.graph_objects as go
from plotly.subplots import make_subplots

# alphadynamics is in requirements.txt — installed at Space build time
from alphadynamics import predict_torsion_ensemble


# ----------------------------------------------------------------------------
# Constants — protect free CPU tier
# ----------------------------------------------------------------------------
MAX_RESIDUES = 50
MAX_ENSEMBLE = 16
MAX_STEPS = 1000

EXAMPLES = [
    ["AAAY", 8, 500],
    ["KLVFFAE", 8, 500],          # amyloid-β fragment
    ["GNNQQNY", 8, 500],          # Sup35 prion peptide
    ["AAYAA", 4, 300],            # quick test
    ["MAEHLLLY", 8, 500],         # random short
    ["FVNQHLCGSHLVEALYL", 4, 300],  # insulin B chain N-terminus
]


# ----------------------------------------------------------------------------
# Plotting
# ----------------------------------------------------------------------------
def make_ramachandran_figure(traj: np.ndarray, sequence: str) -> go.Figure:
    n_ens, n_t, n_res, _ = traj.shape
    has_per_res = 1 <= n_res <= 12

    if has_per_res:
        ncols = min(4, n_res)
        nrows_per_res = (n_res + ncols - 1) // ncols
        specs = [[{"colspan": ncols, "rowspan": 1}] + [None] * (ncols - 1)]
        for _ in range(nrows_per_res):
            specs.append([{} for _ in range(ncols)])
        titles = [f"<b>{sequence}</b> — aggregate density"]
        for i in range(n_res):
            aa = sequence[i] if i < len(sequence) else "?"
            titles.append(f"{aa}{i+1}")
        # pad
        total_cells = ncols + nrows_per_res * ncols
        titles += [""] * max(0, total_cells - len(titles))
        fig = make_subplots(
            rows=1 + nrows_per_res, cols=ncols, specs=specs,
            subplot_titles=titles,
            horizontal_spacing=0.04, vertical_spacing=0.10,
        )
    else:
        fig = make_subplots(
            rows=1, cols=1,
            subplot_titles=[
                f"<b>{sequence}</b> — Ramachandran ({n_res} residues, {n_ens}×{n_t} samples)"
            ],
        )
        ncols = 1
        nrows_per_res = 0

    phi_all = np.degrees(traj[..., 0].flatten())
    psi_all = np.degrees(traj[..., 1].flatten())

    fig.add_trace(
        go.Histogram2dContour(
            x=phi_all, y=psi_all,
            colorscale="Viridis",
            ncontours=20,
            contours=dict(coloring="fill", showlines=False),
            line=dict(width=0),
            showscale=True,
            colorbar=dict(title="density",
                          len=0.45 if has_per_res else 0.85,
                          y=0.78 if has_per_res else 0.5,
                          x=1.02),
            hovertemplate="φ=%{x:.0f}°<br>ψ=%{y:.0f}°<br>density=%{z:.4f}<extra></extra>",
        ),
        row=1, col=1,
    )
    for px, py, label, color in [
        (-60, -45, "α-R",   "white"),
        (-120, 130, "β",    "white"),
        (-60, 140, "PPII",  "white"),
        (60, 50, "α-L (forbidden)", "orange"),
    ]:
        fig.add_annotation(
            x=px, y=py, text=f"<b>{label}</b>",
            showarrow=False, font=dict(color=color, size=12),
            xref="x1", yref="y1",
        )

    if has_per_res:
        for i in range(n_res):
            row = 2 + (i // ncols)
            col = 1 + (i % ncols)
            phi_i = np.degrees(traj[:, :, i, 0].flatten())
            psi_i = np.degrees(traj[:, :, i, 1].flatten())
            fig.add_trace(
                go.Histogram2dContour(
                    x=phi_i, y=psi_i,
                    colorscale="Viridis",
                    ncontours=15,
                    contours=dict(coloring="fill", showlines=False),
                    line=dict(width=0),
                    showscale=False,
                    hovertemplate="φ=%{x:.0f}°<br>ψ=%{y:.0f}°<extra></extra>",
                ),
                row=row, col=col,
            )

    for axis in fig.layout:
        if axis.startswith("xaxis"):
            fig.layout[axis].update(
                range=[-180, 180], tickvals=[-180, -90, 0, 90, 180],
                ticksuffix="°", gridcolor="rgba(255,255,255,0.15)",
                zerolinecolor="rgba(255,255,255,0.4)", zerolinewidth=1,
            )
        elif axis.startswith("yaxis"):
            suffix = axis[len("yaxis"):]
            fig.layout[axis].update(
                range=[-180, 180], tickvals=[-180, -90, 0, 90, 180],
                ticksuffix="°", gridcolor="rgba(255,255,255,0.15)",
                zerolinecolor="rgba(255,255,255,0.4)", zerolinewidth=1,
                scaleanchor="x" + suffix, scaleratio=1,
            )

    fig.update_layout(
        template="plotly_dark",
        height=350 + (250 * nrows_per_res if has_per_res else 0),
        margin=dict(l=70, r=120, t=80, b=60),
        font=dict(family="Inter, system-ui, -apple-system, sans-serif", size=12),
        hovermode="closest",
    )
    return fig


def basin_table_md(traj: np.ndarray) -> str:
    phi = np.degrees(traj[..., 0].flatten())
    psi = np.degrees(traj[..., 1].flatten())

    def b(plo, phi_, slo, shi):
        return ((phi >= plo) & (phi <= phi_) & (psi >= slo) & (psi <= shi)).mean() * 100

    return (
        "| Basin | Population |\n"
        "|---|---:|\n"
        f"| α-R helix     (φ ≈ -60, ψ ≈ -45)  | **{b(-130,-30,-90,30):.1f}%** |\n"
        f"| β-sheet       (φ ≈ -120, ψ ≈ 120) | **{b(-180,-90,70,180):.1f}%** |\n"
        f"| PPII extended (φ ≈ -60, ψ ≈ 140)  | **{b(-90,-30,100,180):.1f}%** |\n"
        f"| α-L (sterically forbidden region) | **{b(30,100,-10,90):.1f}%** ← should be near 0 |\n"
    )


# ----------------------------------------------------------------------------
# Inference
# ----------------------------------------------------------------------------
_AA_VOCAB = set("ACDEFGHIKLMNPQRSTVWYX")


def predict(sequence: str, n_ensemble: int, rollout_steps: int):
    sequence = (sequence or "").strip().upper()
    if not sequence:
        raise gr.Error("Please enter a peptide sequence (1-letter amino-acid code).")
    if len(sequence) > MAX_RESIDUES:
        raise gr.Error(
            f"Free demo limited to {MAX_RESIDUES} residues. "
            f"For longer peptides, install locally: pip install alphadynamics"
        )
    bad = sorted(set(sequence) - _AA_VOCAB)
    if bad:
        raise gr.Error(
            f"Sequence contains non-standard residues: {bad}. "
            f"Use only one-letter codes from {sorted(_AA_VOCAB)}."
        )

    n_ensemble = max(1, min(int(n_ensemble), MAX_ENSEMBLE))
    rollout_steps = max(50, min(int(rollout_steps), MAX_STEPS))

    traj = predict_torsion_ensemble(
        sequence,
        n_ensemble=n_ensemble,
        rollout_steps=rollout_steps,
        seed=42,
        device="cpu",
        show_progress=False,
    )

    # save NPZ for download
    out_path = Path(tempfile.gettempdir()) / f"alphadynamics_{sequence}_torsions.npz"
    np.savez_compressed(
        out_path,
        sequence=sequence,
        torsions=traj,
        torsion_units="radians",
        torsion_axes="(ensemble, time, residues, [phi, psi])",
        n_ensemble=n_ensemble,
        rollout_steps=rollout_steps,
        model_name="ad_transfer_v2_clean",
    )

    # Generate backbone PDB (NEW v0.4.0)
    try:
        from alphadynamics import trajectory_to_pdb, trajectory_diagnostics
        pdb_path = Path(tempfile.gettempdir()) / f"alphadynamics_{sequence}_backbone.pdb"
        # Use ensemble member 0, subsample to 50 frames for fast download
        member = traj[0]
        if len(member) > 50:
            idx = np.linspace(0, len(member) - 1, 50).astype(int)
            member = member[idx]
        trajectory_to_pdb(member, sequence, str(pdb_path))
        diag = trajectory_diagnostics(member)
        diag_md = (
            f"\n\n**3D backbone diagnostics** (Cα-only, 50 frames):\n"
            f"- Radius of gyration: {diag['rg_mean']:.2f} ± {diag['rg_std']:.2f} Å\n"
            f"- End-to-end distance: {diag['end_to_end_mean']:.2f} ± {diag['end_to_end_std']:.2f} Å"
        )
    except Exception as e:
        pdb_path = None
        diag_md = f"\n\n*(PDB rebuild unavailable: {e})*"

    fig = make_ramachandran_figure(traj, sequence)
    info = (
        f"### `{sequence}` — {len(sequence)} residue{'s' if len(sequence) > 1 else ''}\n\n"
        f"**{n_ensemble}** trajectories × **{rollout_steps}** steps "
        f"= **{n_ensemble * rollout_steps * len(sequence):,}** torsion samples\n\n"
        + basin_table_md(traj)
        + diag_md
    )
    return fig, info, str(out_path), str(pdb_path) if pdb_path else None


# ----------------------------------------------------------------------------
# Gradio UI
# ----------------------------------------------------------------------------
DESCRIPTION = """
# 🧬 AlphaDynamics — Protein Torsion Dynamics from Sequence

A tiny (~123K param) neural propagator that predicts the **Ramachandran
density** of a peptide's backbone (φ, ψ) angles **from sequence alone**.

On the canonical 4AA benchmark: **2.39× lower JSD** than Microsoft Timewarp at
**3000× fewer parameters**. Cross-validated against Top8000 PDB statistics.

📦 `pip install alphadynamics` · 🐙 [GitHub](https://github.com/krisss0mecom/AlphaDynamics)
· 🤗 [Model card](https://huggingface.co/krissss0/alphadynamics)

**This demo is CPU-only** — limited to 50 residues / 16 trajectories / 1000
steps. For longer peptides or larger ensembles, install locally.
"""

NOTES = """
### What this tool does
Predicts an ensemble of (φ, ψ) torsion-angle trajectories for any peptide
sequence (4–100 residues recommended, capped at 50 on this demo). Useful for:

- **Quick conformational triage** before launching expensive MD simulations
- **Comparing sequence variants / mutants** side by side
- **Estimating α-helix / β-sheet / PPII basin populations**
- **Teaching biochemistry** with live, interactive Ramachandran plots
- **AI-for-biology baselines and benchmarks**

### Honest limits
- Density only, **not** kinetics (no transition rates / dwell times)
- Backbone only, **no** side-chain rotamers (χ angles)
- Monomer only, **no** multimer / aggregation
- Best for 4–100 residue peptides; reliability degrades outside

### How to read the Ramachandran plot
- **α-R region** (φ ≈ -60, ψ ≈ -45) → right-handed alpha-helix
- **β-sheet region** (φ ≈ -120, ψ ≈ 120) → extended / beta-strand conformations
- **PPII region** (φ ≈ -60, ψ ≈ 140) → polyproline-II extended (very common in
  short peptides in solution)
- **α-L region** (φ ≈ 60, ψ ≈ 50) → left-handed helix, sterically forbidden
  for almost all amino acids; should be close to 0% if the model honors physics

### Architecture (one paragraph)
A residue's (φ, ψ) is treated as a phase pair on a torus. An MLP emits per-residue
oscillator parameters from sequence + position + current angles. A phase-flow ODE
integrates 64 coupled phase oscillators with RK4. The result is decoded into a
mixture of axis-independent von Mises distributions, sampled, and rolled out
autoregressively.

This is the protein-dynamics application of a multi-year line of work on phase
oscillators (REZON hardware, phase-entanglement-rc, theta-gamma neural coupling).
"""

with gr.Blocks(title="AlphaDynamics") as demo:
    gr.Markdown(DESCRIPTION)
    with gr.Row():
        with gr.Column(scale=1):
            seq_input = gr.Textbox(
                label="Peptide sequence",
                placeholder="e.g. KLVFFAE",
                value="AAAY",
                max_lines=1,
            )
            with gr.Row():
                n_ens_input = gr.Slider(
                    minimum=1, maximum=MAX_ENSEMBLE, step=1, value=8,
                    label="Trajectories (more = smoother density)",
                )
                rs_input = gr.Slider(
                    minimum=50, maximum=MAX_STEPS, step=50, value=500,
                    label="Steps per trajectory",
                )
            predict_btn = gr.Button("Predict torsion dynamics 🚀", variant="primary")
            info_md = gr.Markdown(label="Basin populations")
            file_out = gr.File(label="Download trajectory (.npz)")
            pdb_out = gr.File(label="Download backbone PDB (NEW v0.4.0)")
            gr.Markdown(
                "💡 **3D viewer:** open the downloaded `.pdb` in "
                "[PyMOL](https://pymol.org/) / VMD / ChimeraX, or browse online at "
                "[krisss0mecom.github.io/AlphaDynamics/examples/3d_movie_demo](https://krisss0mecom.github.io/AlphaDynamics/examples/3d_movie_demo/viewer.html)"
            )
        with gr.Column(scale=2):
            plot_out = gr.Plot(label="Ramachandran density")

    predict_btn.click(
        fn=predict,
        inputs=[seq_input, n_ens_input, rs_input],
        outputs=[plot_out, info_md, file_out, pdb_out],
    )

    gr.Examples(
        examples=EXAMPLES,
        inputs=[seq_input, n_ens_input, rs_input],
        label="Try a known peptide:",
    )

    with gr.Accordion("📖 About this tool / how to read the plot", open=False):
        gr.Markdown(NOTES)

    gr.Markdown(
        "---\n"
        "Created by **Krzysztof Gwozdz** 🇵🇱 · Apache 2.0 · "
        "[Cite](https://huggingface.co/krissss0/alphadynamics) · "
        "[GitHub Issues](https://github.com/krisss0mecom/AlphaDynamics/issues) for feedback"
    )


if __name__ == "__main__":
    demo.queue(max_size=10).launch(server_name="0.0.0.0")