"""Dataset-free sampling from a downloaded maltefranke/morph bundle."""

from __future__ import annotations

import argparse
import json
import os
from pathlib import Path
import sys


def load_bundle(root, variant, epoch, device="cpu", steps=100):
    """Load the exact bundled runtime, statistics, current weights and EMA."""
    root = Path(root).resolve()
    bundle = root / "chembl3d" / variant
    runtime = bundle / "runtime"
    if "morph" in sys.modules:
        raise RuntimeError("Load each runtime in a fresh Python process.")
    sys.path.insert(0, str(runtime / "src"))
    os.environ.setdefault("MORPH_COMPILE", "0")
    os.environ.setdefault("MORPH_ATTENTION_BACKEND", "sdpa")
    os.environ.setdefault("MORPH_FLASH_ATTN", "0")
    os.environ.setdefault("MORPH_MUON_EXECUTION", "native")

    import torch
    from omegaconf import OmegaConf
    import morph.config  # registers the training config's resolvers
    from morph.inference.generate import _build_module, load_checkpoint

    if device == "auto":
        device = "cuda" if torch.cuda.is_available() else "cpu"
    cfg = OmegaConf.load(bundle / "config.yaml")
    cfg.data.preprocessing.root = str(bundle)
    for field, filename in {
        "atom_tokens_path": "atom_tokens.txt",
        "edge_tokens_path": "edge_tokens.txt",
        "charge_tokens_path": "charge_tokens.txt",
        "distributions_path": "distributions.pt",
    }.items():
        path = bundle / filename
        if not path.is_file():
            raise FileNotFoundError(path)
        cfg.data.preprocessing[field] = str(path)
    cfg.integrator.integrator.num_integration_steps = steps
    module, vocab, distributions = _build_module(cfg)
    checkpoint = bundle / f"epoch_{epoch:03d}.ckpt"
    load_checkpoint(module, checkpoint)
    module.eval().to(device)
    return module, vocab, distributions, cfg


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--root", type=Path, default=Path(__file__).resolve().parent)
    parser.add_argument("--variant", choices=["explicit_h", "implicit_h"], required=True)
    parser.add_argument("--epoch", type=int, required=True)
    parser.add_argument("--num-mols", type=int, default=16)
    parser.add_argument("--batch-size", type=int, default=16)
    parser.add_argument("--steps", type=int, default=100)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--device", choices=["auto", "cpu", "cuda"], default="auto")
    parser.add_argument("--output", type=Path, default=Path("molecules.sdf"))
    args = parser.parse_args()
    if min(args.num_mols, args.batch_size, args.steps) < 1:
        parser.error("num-mols, batch-size, and steps must be positive")

    import torch

    torch.manual_seed(args.seed)
    module, vocab, _, cfg = load_bundle(
        args.root, args.variant, args.epoch, args.device, args.steps
    )
    from morph.inference.generate import _prior_loader, final_mol
    from rdkit import Chem

    loader = _prior_loader(
        module.distributions, args.num_mols, args.batch_size,
        uniform_n_atoms_padding=int(cfg.data.uniform_n_atoms_padding),
        n_atoms_strategy=str(cfg.data.n_atoms_strategy),
    )
    records = []
    args.output.parent.mkdir(parents=True, exist_ok=True)
    with Chem.SDWriter(str(args.output)) as writer, torch.inference_mode():
        for batch_index, batch in enumerate(loader):
            batch = tuple(graph.to(module.device) for graph in batch)
            result = module.predict_step(batch, batch_index)
            for trajectory, valid in zip(result["all_mols"], result["is_valid"], strict=True):
                frame = trajectory[-1]
                if not torch.isfinite(frame.x).all():
                    raise RuntimeError("Non-finite generated coordinates")
                mol = final_mol(trajectory, vocab)
                index = len(records)
                records.append({
                    "index": index, "valid": bool(valid),
                    "smiles": Chem.MolToSmiles(mol) if mol is not None else None,
                    "represented_atoms": int(frame.x.shape[0]),
                })
                if valid and mol is not None:
                    mol.SetProp("_Name", f"sample_{index}")
                    writer.write(mol)
    assert len(records) == args.num_mols
    report = {
        "variant": args.variant, "completed_epochs": args.epoch,
        "seed": args.seed, "integration_steps": args.steps,
        "weights": "EMA", "coordinate_units": "angstrom",
        "num_attempted": len(records),
        "num_valid": sum(row["valid"] for row in records),
        "samples": records,
    }
    args.output.with_suffix(".json").write_text(json.dumps(report, indent=2) + "\n")
    print(json.dumps({k: v for k, v in report.items() if k != "samples"}), flush=True)


if __name__ == "__main__":
    main()