"""Run `RFD4ProteinaPipeline` on the AtomWorks Studio example inputs and write annotated CIFs plus a JSON report.

```bash
python run_examples.py --checkpoint nvidia/RFD4-Proteina-Diffusers --inputs-dir inputs --out outputs
```
"""

import argparse
import json
import time
from pathlib import Path

import torch
from diffusers import RFD4ProteinaPipeline
from rfd4_proteina_prep import prepare_cif, write_cif


# Example id -> (input file, pipeline arguments), for the inputs this repository ships under
# `examples/inputs/`. `examples/manifest.json` is the catalog of record; `check_against_manifest()`
# below fails the run if the two disagree, because they have drifted apart before -- this table
# still listed two antibody tasks whose inputs the repository does not ship.
TASKS = {
    "protein_binder": ("protein_binder.cif", {}),
    "enzyme_design": ("enzyme_design.cif", {}),
    "motif_scaffolding": ("motif_scaffolding.cif", {}),
    "small_molecule_binder": ("small_molecule_binder.cif", {}),
    "unconditional": ("unconditional.cif", {}),
    "inverse_folding": (
        "inverse_folding.cif",
        {"inverse_folding": True, "partial_t_bb_ca": 0.0, "partial_t_local_latents": 0.0},
    ),
    "partial_diffusion": ("partial_diffusion.cif", {"partial_fraction": 0.4}),
}

# The tensors the pipeline takes by name; everything else prepare_cif returns is conditioning.
# `coords_nm` belongs here: partial diffusion and inverse folding read it as a named argument, and
# a full run ignores it, so passing it always is correct for every task.
REQUIRED = [
    "padding_mask",
    "is_in_context_conditioning",
    "supports_local_latent",
    "token_type",
    "atom37_mask",
    "coords_nm",
]


def check_against_manifest():
    """Fail loudly if TASKS and the shipped manifest disagree about which inputs exist.

    Ids and files only, deliberately. The manifest's `pipeline_kwargs` are written for the pinned
    Diffusers revision that `run_hub_examples.py` enforces, which still takes one `partial_t` dict;
    this script runs against whatever Diffusers is installed, where the start times are separate
    arguments. Comparing the options would fail on that difference alone.
    """
    manifest_path = Path(__file__).resolve().parent / "manifest.json"
    if not manifest_path.is_file():
        return
    catalog = {entry["id"]: entry for entry in json.loads(manifest_path.read_text())["examples"]}
    for name, (file, options) in TASKS.items():
        entry = catalog.get(name)
        if entry is None:
            raise SystemExit(f"{name} is not in {manifest_path}; it may have been excluded upstream.")
        if Path(entry["file"]).name != file:
            raise SystemExit(f"{name} reads {file}, but the manifest names {entry['file']}.")
        fraction = entry.get("remaining_schedule_fraction")
        if (fraction is not None) != ("partial_fraction" in options):
            raise SystemExit(f"{name} and the manifest disagree about whether it starts part-way through.")
        if fraction is not None and fraction != options["partial_fraction"]:
            raise SystemExit(f"{name} starts at {options['partial_fraction']}, the manifest says {fraction}.")


def main():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--checkpoint", default="nvidia/RFD4-Proteina-Diffusers")
    parser.add_argument("--inputs-dir", type=Path, required=True)
    parser.add_argument("--out", type=Path, required=True)
    parser.add_argument("--tasks", nargs="*", default=list(TASKS))
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--dtype", choices=["bfloat16", "float32"], default="bfloat16")
    args = parser.parse_args()
    check_against_manifest()
    unknown = set(args.tasks) - TASKS.keys()
    if unknown:
        parser.error(f"Unknown tasks: {sorted(unknown)}. Choose from {list(TASKS)}.")
    args.out.mkdir(parents=True, exist_ok=True)

    device = torch.device("cuda:0")
    pipe = RFD4ProteinaPipeline.from_pretrained(args.checkpoint, dtype=getattr(torch, args.dtype)).to(device)
    steps = pipe.scheduler.num_inference_steps
    report = {}
    for name in args.tasks:
        file, options = TASKS[name]
        options = dict(options)
        fraction = options.pop("partial_fraction", None)
        if fraction is not None:
            # Start from the noised input at the point of the schedule where `fraction` of the trajectory remains.
            pipe.scheduler.set_timesteps(steps)
            index = round((1 - fraction) * steps)
            # One start time per modality, each its own argument: `partial_t_bb_ca`, `partial_t_local_latents`.
            for modality in pipe.scheduler.modalities:
                options[f"partial_t_{modality}"] = float(pipe.scheduler.modality_timestep(modality, index))
        features, atoms = prepare_cif(args.inputs_dir / file, seed=args.seed, device=device)
        generator = torch.Generator(device).manual_seed(args.seed)
        torch.cuda.synchronize()
        start = time.time()
        # The pipeline takes the masks and molecule tensors by name; the annotations go in one
        # mapping. Passing `features` positionally binds the whole dict to `padding_mask`.
        inputs = {name: features[name] for name in REQUIRED}
        conditioning = {name: value for name, value in features.items() if name not in REQUIRED}
        output = pipe(**inputs, conditioning_features=conditioning, generator=generator, **options)
        torch.cuda.synchronize()
        seconds = time.time() - start
        destination = args.out / f"{name}.cif"
        destination.unlink(missing_ok=True)
        write_cif(output, features, atoms, destination)
        report[name] = {
            "input": file,
            "tokens": int(features["padding_mask"].shape[1]),
            "designed_tokens": int(output.emit_mask.sum()),
            "seconds": round(seconds, 1),
            "options": {k: v for k, v in options.items()},
        }
        print(name, report[name], flush=True)
    (args.out / "report.json").write_text(json.dumps(report, indent=2) + "\n")


if __name__ == "__main__":
    main()