import os
import subprocess
import sys

# Allocator config for transient memory spikes in the folding pipeline.
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
# Persist checkpoints / common runtime files between calls.
OPENDDE_ROOT_DIR = os.environ.setdefault("OPENDDE_ROOT_DIR", "/tmp/opendde_data")
os.environ.setdefault("LAYERNORM_TYPE", "torch")

# Install the OpenDDE package WITHOUT its dependencies, so its pinned
# torch==2.7.1 does not clobber the ZeroGPU-managed torch. This makes the
# `opendde` and `runner` modules importable. Done before `import spaces`
# because pip itself does not initialize CUDA.
try:
    import runner.batch_inference  # noqa: F401
except Exception:
    subprocess.run(
        [
            sys.executable, "-m", "pip", "install", "--no-deps",
            "git+https://github.com/aurekaresearch/OpenDDE.git",
        ],
        check=True,
    )

import spaces  # noqa: E402  (must precede torch / CUDA-touching imports)

import glob  # noqa: E402
import json  # noqa: E402
import tempfile  # noqa: E402
import time  # noqa: E402
import uuid  # noqa: E402

import gradio as gr  # noqa: E402

# ---------------------------------------------------------------------------
# Model / runtime setup
# ---------------------------------------------------------------------------
CHECKPOINT_DIR = os.path.join(OPENDDE_ROOT_DIR, "checkpoint")
COMMON_DIR = os.path.join(OPENDDE_ROOT_DIR, "common")
os.makedirs(CHECKPOINT_DIR, exist_ok=True)
os.makedirs(COMMON_DIR, exist_ok=True)

MODEL_REPO = "aurekaresearch/OpenDDE"


def _prefetch_runtime_assets():
    """Download checkpoint + common runtime files once at startup so the GPU
    call only performs inference (no download inside the time-limited window)."""
    from huggingface_hub import hf_hub_download

    targets = [
        ("opendde.pt", os.path.join(CHECKPOINT_DIR, "opendde.pt")),
        ("common/components.cif", os.path.join(COMMON_DIR, "components.cif")),
        (
            "common/components.cif.rdkit_mol.pkl",
            os.path.join(COMMON_DIR, "components.cif.rdkit_mol.pkl"),
        ),
    ]
    for repo_file, dest in targets:
        if os.path.exists(dest):
            continue
        cached = hf_hub_download(MODEL_REPO, repo_file, repo_type="model")
        os.makedirs(os.path.dirname(dest), exist_ok=True)
        if not os.path.exists(dest):
            try:
                os.symlink(cached, dest)
            except OSError:
                import shutil

                shutil.copy(cached, dest)


try:
    _prefetch_runtime_assets()
    print("OpenDDE runtime assets ready.", flush=True)
except Exception as e:  # non-fatal: CLI can still download lazily
    print(f"Asset prefetch failed ({e!r}); CLI will download on demand.", flush=True)

# Amino-acid alphabet used for validation (20 standard + X).
AA_ALPHABET = set("ACDEFGHIKLMNPQRSTVWYX")

MAX_RESIDUES = 400  # keep single-request runtime within the ZeroGPU window

DESCRIPTION = """
# 🧬 OpenDDE — Biomolecular Structure Prediction

Predict the all-atom 3D structure of a **protein sequence** with
[**OpenDDE**](https://huggingface.co/aurekaresearch/OpenDDE), an open-source
AlphaFold3-family biomolecular foundation model from Aureka Research.

Paste a single-letter amino-acid sequence, hit **Predict structure**, and get an
interactive 3D structure plus a downloadable `.cif` file and confidence metrics
(pLDDT / pTM). This demo runs **single-sequence** (no MSA / templates) so it
stays fast on ZeroGPU.
""".strip()


def _pymol_style_html(cif_text: str) -> str:
    """Build a self-contained py3Dmol viewer for a CIF structure.

    The viewer is rendered inside an ``<iframe srcdoc=...>`` (the canonical
    Gradio + 3Dmol.js embedding pattern). Injecting the WebGL canvas straight
    into the Gradio DOM is unreliable — the ``<script>`` may not re-run on
    output updates and the viewer often initializes before the div has layout,
    producing a 0-sized canvas that paints only its dark background (i.e. a
    solid black box). The iframe isolates the viewer, guarantees the script
    runs on every update, gives the container an explicit size, and lets us
    call ``viewer.resize()`` once layout settles.
    """
    # HTML-escape the CIF so it can be safely embedded inside the iframe's
    # single-quoted ``srcdoc`` attribute, and inside a JS template literal.
    safe = (
        cif_text.replace("&", "&amp;")
        .replace("'", "&#39;")
        .replace("\\", "\\\\")
        .replace("`", "\\`")
        .replace("</", "<\\/")
    )
    srcdoc = f"""<!DOCTYPE html>
<html>
<head>
  <meta charset="utf-8" />
  <style>
    html, body {{ margin: 0; padding: 0; height: 100%; background: #0d1117; }}
    #mol {{ width: 100%; height: 100%; position: relative; }}
  </style>
  <script src="https://cdnjs.cloudflare.com/ajax/libs/jquery/3.6.3/jquery.min.js"></script>
  <script src="https://cdnjs.cloudflare.com/ajax/libs/3Dmol/2.4.0/3Dmol-min.js"></script>
</head>
<body>
  <div id="mol"></div>
  <script>
    var CIF = `{safe}`;
    function boot() {{
      if (typeof $3Dmol === 'undefined') {{ setTimeout(boot, 100); return; }}
      var el = document.getElementById('mol');
      var viewer = $3Dmol.createViewer(el, {{ backgroundColor: '#0d1117' }});
      viewer.addModel(CIF, 'cif');
      // Colour cartoon by pLDDT stored in the B-factor column.
      viewer.setStyle({{}}, {{ cartoon: {{ colorscheme: {{
          prop: 'b', gradient: 'roygb', min: 50, max: 90
      }} }} }});
      viewer.addStyle({{ hetflag: true }}, {{ stick: {{}} }});
      viewer.zoomTo();
      viewer.resize();
      viewer.render();
      viewer.zoom(1.15, 800);
      // Re-fit after the iframe has fully laid out.
      setTimeout(function () {{ viewer.resize(); viewer.render(); }}, 300);
      window.addEventListener('resize', function () {{
        viewer.resize(); viewer.render();
      }});
    }}
    boot();
  </script>
</body>
</html>"""
    # Embed via srcdoc; escape double quotes so the attribute stays intact.
    srcdoc_attr = srcdoc.replace('"', "&quot;")
    return (
        f'<iframe srcdoc="{srcdoc_attr}" '
        'style="width:100%;height:520px;border:none;border-radius:12px;'
        'overflow:hidden;background:#0d1117;" '
        'sandbox="allow-scripts allow-same-origin"></iframe>'
    )


def _validate_sequence(sequence: str) -> str:
    seq = "".join(sequence.split()).upper()
    if not seq:
        raise gr.Error("Please enter a protein amino-acid sequence.")
    bad = sorted(set(seq) - AA_ALPHABET)
    if bad:
        raise gr.Error(
            f"Invalid amino-acid letter(s): {', '.join(bad)}. "
            "Use the 20 standard amino acids (and X)."
        )
    if len(seq) > MAX_RESIDUES:
        raise gr.Error(
            f"Sequence is {len(seq)} residues; this demo caps at {MAX_RESIDUES} "
            "residues to fit the ZeroGPU time window."
        )
    return seq


def _estimate_duration(sequence, *args, **kwargs):
    seq = "".join(str(sequence).split())
    n = max(1, len(seq))
    # fp32 torch-kernel path scales roughly with length; be generous.
    return int(min(600, 120 + n * 0.9))


@spaces.GPU(duration=_estimate_duration)
def _run_opendde(sequence: str, steps: int, cycles: int, seed: int) -> dict:
    """Run OpenDDE single-sequence prediction inside the GPU worker."""
    work = tempfile.mkdtemp(prefix="opendde_run_")
    out_dir = os.path.join(work, "output")
    os.makedirs(out_dir, exist_ok=True)

    job_name = "prediction"
    job = [
        {
            "name": job_name,
            "modelSeeds": [int(seed)],
            "sequences": [
                {"proteinChain": {"sequence": sequence, "count": 1}}
            ],
        }
    ]
    input_json = os.path.join(work, "input.json")
    with open(input_json, "w") as f:
        json.dump(job, f)

    cmd = [
        sys.executable,
        "-m",
        "runner.batch_inference",
        "pred",
        "-i", input_json,
        "-o", out_dir,
        "-n", "opendde_v1",
        "--use_msa", "false",
        "--use_template", "false",
        "--use_rna_msa", "false",
        "--triatt_kernel", "torch",
        "--trimul_kernel", "torch",
        "--dtype", "fp32",
        "--sample", "1",
        "--step", str(int(steps)),
        "--cycle", str(int(cycles)),
    ]
    env = dict(os.environ)
    env["OPENDDE_ROOT_DIR"] = OPENDDE_ROOT_DIR
    env["LAYERNORM_TYPE"] = "torch"

    proc = subprocess.run(cmd, env=env, capture_output=True, text=True)
    tail = (proc.stdout or "")[-2000:] + "\n" + (proc.stderr or "")[-2000:]

    pred_dir = os.path.join(out_dir, job_name, f"seed_{int(seed)}", "predictions")
    cifs = sorted(glob.glob(os.path.join(pred_dir, "*.cif")))
    if not cifs:
        # surface an error directory message if present
        err_dir = os.path.join(out_dir, "ERR")
        err_msg = ""
        if os.path.isdir(err_dir):
            for fn in glob.glob(os.path.join(err_dir, "*.txt")):
                err_msg += open(fn).read()
        raise RuntimeError(
            f"OpenDDE produced no structure.\n{err_msg}\n--- log ---\n{tail}"
        )

    cif_path = cifs[0]
    with open(cif_path) as f:
        cif_text = f.read()

    confidence = {}
    conf_files = sorted(
        glob.glob(os.path.join(pred_dir, "*summary_confidence*.json"))
    )
    if conf_files:
        try:
            confidence = json.load(open(conf_files[0]))
        except Exception:
            confidence = {}

    return {"cif_text": cif_text, "confidence": confidence, "log": tail}


def predict(sequence, steps=100, cycles=4, seed=101, randomize_seed=False):
    """Predict a 3D protein structure from an amino-acid sequence.

    Args:
        sequence: Single-letter amino-acid sequence (20 standard AAs + X).
        steps: Number of diffusion sampling steps.
        cycles: Number of Pairformer recycling iterations.
        seed: Random seed for the diffusion sampler.
        randomize_seed: If True, pick a fresh random seed for this run.
    """
    import random

    seq = _validate_sequence(sequence)
    if randomize_seed:
        seed = random.randint(1, 65535)
    seed = int(seed)

    t0 = time.perf_counter()
    result = _run_opendde(seq, int(steps), int(cycles), seed)
    elapsed = time.perf_counter() - t0

    cif_text = result["cif_text"]
    conf = result["confidence"] or {}

    # Save CIF to a temp file for download.
    tmp_cif = tempfile.NamedTemporaryFile(
        delete=False, suffix=".cif", prefix="opendde_structure_"
    )
    tmp_cif.write(cif_text.encode())
    tmp_cif.close()

    def _fmt(v):
        try:
            return round(float(v), 4)
        except Exception:
            return v

    metrics = {
        "residues": len(seq),
        "seed": seed,
        "runtime_seconds": round(elapsed, 1),
    }
    for key in ("plddt", "ptm", "iptm", "gpde", "ranking_score", "has_clash"):
        if key in conf:
            metrics[key] = _fmt(conf[key])

    viewer_html = _pymol_style_html(cif_text)
    return viewer_html, tmp_cif.name, metrics, gr.update(value=seed)


CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

EXAMPLE_SEQS = [
    # Villin headpiece subdomain (HP36) — small fast-folding protein.
    ["MLSDEDFKAVFGMTRSAFANLPLWKQQNLKKEKGLF"],
    # Trp-cage (TC5b) miniprotein.
    ["NLYIQWLKDGGPSSGRPPPS"],
    # Chignolin-like / short helical peptide.
    ["ACDEFGHIKLMNPQRSTVWY"],
]

with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(DESCRIPTION)
        with gr.Row():
            sequence = gr.Textbox(
                label="Protein sequence",
                placeholder="Paste a single-letter amino-acid sequence, e.g. MLSDEDFK...",
                lines=3,
                scale=4,
            )
        run = gr.Button("Predict structure", variant="primary")

        with gr.Row():
            with gr.Column(scale=3):
                viewer = gr.HTML(label="Predicted 3D structure")
            with gr.Column(scale=1):
                metrics_out = gr.JSON(label="Confidence metrics")
                cif_file = gr.File(label="Download structure (.cif)")

        with gr.Accordion("Advanced settings", open=False):
            steps = gr.Slider(
                20, 200, value=100, step=10, label="Diffusion steps"
            )
            cycles = gr.Slider(
                1, 10, value=4, step=1, label="Recycling cycles"
            )
            with gr.Row():
                seed = gr.Number(label="Seed", value=101, precision=0)
                randomize_seed = gr.Checkbox(label="Randomize seed", value=False)

        gr.Examples(
            examples=EXAMPLE_SEQS,
            inputs=[sequence],
            outputs=[viewer, cif_file, metrics_out, seed],
            fn=predict,
            cache_examples=False,
            run_on_click=True,
        )

    run.click(
        predict,
        inputs=[sequence, steps, cycles, seed, randomize_seed],
        outputs=[viewer, cif_file, metrics_out, seed],
        api_name="predict",
    )

if __name__ == "__main__":
    demo.launch(mcp_server=True)