import os

# Environment preamble before any heavy import: these gate AtomWorks/numba/biotite startup and
# keep the ZeroGPU parent process free of CUDA initialization.
os.environ.setdefault("NUMBA_DISABLE_CUDA", "1")
os.environ.setdefault("ALLOW_BIOTITE_CCD", "true")
os.environ.setdefault("DISABLE_CUEQUIVARIANCE", "true")
os.environ.setdefault("WANDB_MODE", "disabled")
os.environ.setdefault("GRADIO_ANALYTICS_ENABLED", "false")
os.environ.setdefault("CLUSTER", "local")
os.environ.setdefault("OMP_NUM_THREADS", "1")
os.environ.setdefault("OPENBLAS_NUM_THREADS", "1")

import spaces  # noqa: E402  # MUST precede torch / diffusers / atomworks

import time
from pathlib import Path

import gradio as gr
import torch

MODEL_ID = "nvidia/RFD4-Proteina-Diffusers"
EXAMPLES_DIR = Path(__file__).parent / "examples"

from diffusers import RFD4ProteinaPipeline  # noqa: E402
from rfd4_proteina_prep import prepare_cif, write_cif  # noqa: E402

# The feature tensors the pipeline consumes as direct keyword arguments; everything else the
# preprocessor emits rides inside `conditioning_features`.
REQUIRED_KEYS = (
    "padding_mask",
    "is_in_context_conditioning",
    "supports_local_latent",
    "token_type",
    "atom37_mask",
)

MODES = (
    "Full design",
    "Inverse folding",
    "Partial diffusion",
)

MODE_INFO = {
    "Full design": (
        "Design every free token in the input: sequence and structure are generated from noise, "
        "conditioned on whatever the annotated CIF holds fixed."
    ),
    "Inverse folding": (
        "The backbone of every free token is held fixed; RFD4 predicts only the sequence. Use it on "
        "a structure whose fold you want to keep."
    ),
    "Partial diffusion": (
        "Start from the input structure at a partly-noised state: the beginning of the trajectory "
        "is skipped, so designs stay close to the input fold."
    ),
}

# The released recipe fixes the number of diffusion steps (len(config.recycle_depths) = 60).
N_STEPS = 60

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

_PLACEHOLDER_HTML = (
    '<div style="width:100%;height:520px;display:flex;align-items:center;justify-content:center;'
    "color:var(--body-text-color);border:1px dashed var(--border-color-primary);border-radius:8px;"
    'font-size:14px;">Press “Design” to see the generated structure here</div>'
)

# The pipeline is loaded once at module scope, eagerly on the GPU: no offloading, no device_map.
pipe = RFD4ProteinaPipeline.from_pretrained(MODEL_ID, dtype=torch.bfloat16).to("cuda")


def _describe(mode: str) -> str:
    """One-line explanation of a design mode, shown under the dropdown."""
    return MODE_INFO.get(mode, "")


def _viewer_for(cif_path: str | Path, height: int = 520) -> str:
    """Render a designed CIF as an interactive 3Dmol.js cartoon view."""
    import py3Dmol

    with open(cif_path) as handle:
        data = handle.read()
    view = py3Dmol.view(width="100%", height=height)
    view.addModel(data, "cif")
    view.setStyle({"cartoon": {"colorscheme": "chain"}})
    view.zoomTo()
    return view._make_html()


@spaces.GPU(duration=120)
def design_protein(
    cif_file,
    mode="Full design",
    seed=42,
    num_designs=1,
    remaining_fraction=0.4,
    progress=gr.Progress(track_tqdm=True),
):
    """Run RFD4-Proteina on an annotated CIF and return the designed structure(s).

    Args:
        cif_file: Path to the uploaded annotated CIF (AtomWorks Studio format).
        mode: "Full design", "Inverse folding" or "Partial diffusion".
        seed: Random seed; the same input and seed reproduce the same design.
        num_designs: How many independent designs to draw in one call (1-8).
        remaining_fraction: "Partial diffusion" only — the fraction of the trajectory left to
            run. The released example uses 0.4.
        progress: Gradio progress bar.

    Returns:
        The design CIF path (or list of paths), an interactive 3D view of the first design, and a
        short text summary.
    """
    started = time.perf_counter()
    if cif_file is None:
        raise gr.Error("Upload an annotated CIF first, or pick one of the examples below.")
    num_designs = int(num_designs)
    if not 1 <= num_designs <= 8:
        raise gr.Error("Number of designs must be between 1 and 8.")

    # Feature preparation is CPU-side AtomWorks parsing; the pipeline moves every tensor to its
    # execution device itself, so CPU tensors are fine here.
    features, atoms = prepare_cif(str(cif_file), seed=int(seed), device="cpu")
    tokens = int(features["padding_mask"].shape[1])

    required = {name: features[name] for name in REQUIRED_KEYS}
    conditioning = {name: value for name, value in features.items() if name not in REQUIRED_KEYS}

    if mode == "Inverse folding":
        kwargs = {
            "inverse_folding": True,
            "partial_t_bb_ca": 0.0,
            "partial_t_local_latents": 0.0,
        }
        required["coords_nm"] = features["coords_nm"]
    elif mode == "Partial diffusion":
        # Reproduce the released partial-diffusion entry point: start each modality at the time
        # its schedule reaches after skipping (1 - remaining_fraction) of the trajectory.
        fraction = float(min(max(remaining_fraction, 0.05), 1.0))
        index = round((1 - fraction) * N_STEPS)
        pipe.bb_ca_scheduler.set_timesteps(N_STEPS)
        pipe.local_latents_scheduler.set_timesteps(N_STEPS)
        kwargs = {
            "partial_t_bb_ca": float(pipe.bb_ca_scheduler.timesteps[index]),
            "partial_t_local_latents": float(pipe.local_latents_scheduler.timesteps[index]),
        }
        required["coords_nm"] = features["coords_nm"]
    else:
        kwargs = {}
        conditioning["coords_nm"] = features["coords_nm"]

    if num_designs == 1:
        generator = torch.Generator("cuda").manual_seed(int(seed))
    else:
        generator = [
            torch.Generator("cuda").manual_seed(int(seed) + offset)
            for offset in range(num_designs)
        ]

    with torch.no_grad():
        output = pipe(
            padding_mask=required["padding_mask"],
            is_in_context_conditioning=required["is_in_context_conditioning"],
            supports_local_latent=required["supports_local_latent"],
            token_type=required["token_type"],
            atom37_mask=required["atom37_mask"],
            coords_nm=required.get("coords_nm"),
            conditioning_features=conditioning,
            num_designs_per_prompt=num_designs,
            generator=generator,
            **kwargs,
        )

    # `write_cif` reads one feature row per output row: repeat the input features per design.
    if num_designs > 1:
        features = {
            name: value.repeat_interleave(num_designs, dim=0) for name, value in features.items()
        }

    out_paths = []
    for row in range(num_designs):
        out_paths.append(write_cif(output, features, atoms, _unique_path(row, seed), row=row))

    elapsed = time.perf_counter() - started
    summary = (
        f"**{num_designs}** design{'s' if num_designs > 1 else ''} · **{tokens}** tokens · "
        f"**{mode.lower()}** · seed **{int(seed)}** · **{elapsed:.1f} s** on this GPU"
    )
    viewer_html = _viewer_for(out_paths[0])

    if num_designs == 1:
        return out_paths[0], viewer_html, summary
    return out_paths, viewer_html, summary


def _unique_path(row: int, seed) -> Path:
    """A fresh, non-existent output path in this worker's temp directory."""
    base = Path(os.environ.get("GRADIO_TEMP_DIR") or "/tmp/gradio") / f"rfd4_design_{int(seed)}_{row}.cif"
    base.parent.mkdir(parents=True, exist_ok=True)
    candidate = base
    suffix = 0
    while candidate.exists():
        suffix += 1
        candidate = base.with_name(f"{base.stem}_{suffix}{base.suffix}")
    return candidate


with gr.Blocks(title="RFD4-Proteina · Diffusers") as demo:
    gr.HTML(
        """
    <h1 style="margin: 0; font-size: 28px;">RFD4-Proteina · Diffusers</h1>
    <p style="margin-top: 4px;">
      Joint sequence–structure generation for biomolecular design, conditioned on an
      <b>annotated CIF</b> that says what to hold fixed and what to design. This Space runs
      <code>RFD4ProteinaPipeline</code> from 🤗 Diffusers on ZeroGPU.
    </p>
    """
    )
    with gr.Column(elem_id="col-container"):
        with gr.Row():
            with gr.Column(scale=1):
                cif_input = gr.File(
                    label="Annotated CIF (input)",
                    file_types=[".cif"],
                    type="filepath",
                    file_count="single",
                )
                mode = gr.Dropdown(
                    choices=list(MODES),
                    value="Full design",
                    label="Design mode",
                    info="What RFD4 generates versus keeps from the input.",
                )
                mode_info = gr.Markdown(_describe("Full design"))
                with gr.Accordion("Advanced options", open=False):
                    seed = gr.Slider(0, 2**31 - 1, value=42, step=1, label="Seed")
                    num_designs = gr.Slider(1, 8, value=1, step=1, label="Number of designs")
                    remaining_fraction = gr.Slider(
                        0.05,
                        1.0,
                        value=0.4,
                        step=0.05,
                        label="Remaining schedule (partial diffusion only)",
                        info="The released partial-diffusion example uses 0.4.",
                    )
                go = gr.Button("Design", variant="primary")
            with gr.Column(scale=2):
                viewer = gr.HTML(value=_PLACEHOLDER_HTML, show_label=False)
                out_files = gr.File(
                    label="Designed structure(s) · annotated CIF",
                    file_count="single",
                    interactive=False,
                )
                summary = gr.Markdown()
        mode.change(_describe, inputs=mode, outputs=mode_info)
        go.click(
            design_protein,
            inputs=[cif_input, mode, seed, num_designs, remaining_fraction],
            outputs=[out_files, viewer, summary],
            api_name="design",
        )

        gr.Examples(
            examples=[
                [str(EXAMPLES_DIR / "motif_scaffolding.cif"), "Full design"],
                [str(EXAMPLES_DIR / "inverse_folding.cif"), "Inverse folding"],
                [str(EXAMPLES_DIR / "inverse_folding.cif"), "Partial diffusion"],
                [str(EXAMPLES_DIR / "protein_binder.cif"), "Full design"],
                [str(EXAMPLES_DIR / "enzyme_design.cif"), "Full design"],
                [str(EXAMPLES_DIR / "small_molecule_binder.cif"), "Full design"],
                [str(EXAMPLES_DIR / "unconditional.cif"), "Full design"],
            ],
            inputs=[cif_input, mode],
            outputs=[out_files, viewer, summary],
            fn=design_protein,
            cache_examples=True,
            cache_mode="lazy",
            example_labels=[
                "Motif scaffolding (1QJG)",
                "Inverse folding (2KL8)",
                "Partial diffusion (2KL8)",
                "Protein binder (PD-L1)",
                "Enzyme (catalytic motif)",
                "Small-molecule binder (SAM)",
                "Unconditional protein",
            ],
            label="Author examples (from the model card, BSD-3-Clause)",
        )

demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)