#!/usr/bin/env python3
"""Run reproducible CupOFMoCA PDC experiments.

This wrapper keeps the submitted generation programs intact while providing one
portable interface, bundled checkpoint defaults, named targets, and collision-
safe outputs.
"""

from __future__ import annotations

import argparse
import json
import shlex
import subprocess
import sys
from pathlib import Path


ROOT = Path(__file__).resolve().parents[1]
CONFIG = json.loads((ROOT / "configs" / "paper.json").read_text())

CHECKPOINTS = {
    "smiles": ROOT / "checkpoints" / "SMILESReDi_base.ckpt",
    "aa": ROOT / "checkpoints" / "PepReDi_base.pt",
    "uaa": ROOT / "checkpoints" / "UAAReDi_base.pt",
}


def read_fasta(path: Path) -> dict[str, str]:
    records: dict[str, str] = {}
    name = None
    chunks: list[str] = []
    for raw in path.read_text().splitlines():
        line = raw.strip()
        if not line:
            continue
        if line.startswith(">"):
            if name is not None:
                records[name] = "".join(chunks)
            name, chunks = line[1:].split()[0], []
        elif name is None:
            raise ValueError(f"Sequence before FASTA header in {path}")
        else:
            chunks.append(line)
    if name is not None:
        records[name] = "".join(chunks)
    return records


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--representation", choices=("smiles", "aa", "uaa"), required=True)
    parser.add_argument("--variant", choices=("baseline", "cupofmoca"), default="cupofmoca")
    parser.add_argument("--method", choices=("guided", "staged"), default="guided")
    parser.add_argument("--target", choices=tuple(CONFIG["targets"]), required=True)
    parser.add_argument(
        "--seeds", type=int, nargs="+", default=CONFIG["default_seeds"],
        help="Random seeds for generation.",
    )
    parser.add_argument("--num-candidates", type=int, default=CONFIG["num_candidates_per_seed"])
    parser.add_argument("--gen-len", type=int, default=None)
    parser.add_argument("--checkpoint", type=Path, default=None)
    parser.add_argument("--output-root", type=Path, default=ROOT / "outputs")
    parser.add_argument("--python", default=sys.executable)
    parser.add_argument("--dry-run", action="store_true")
    parser.add_argument("--force", action="store_true", help="Allow an existing output to be replaced.")
    parser.add_argument(
        "extra_args", nargs=argparse.REMAINDER,
        help="Arguments after '--' are appended to the underlying generator command.",
    )
    args = parser.parse_args()
    if args.num_candidates < 1:
        parser.error("--num-candidates must be positive")
    return args


def build_command(args: argparse.Namespace, seed: int, target_seq: str, output: Path) -> list[str]:
    suffix = "" if args.representation == "smiles" else f"_{args.representation}"
    prefix = "cupofmoca_" if args.variant == "cupofmoca" else ""
    program = f"smiles.generation.{prefix}{args.method}{suffix}"
    checkpoint = (args.checkpoint or CHECKPOINTS[args.representation]).resolve()
    target_cfg = CONFIG["targets"][args.target]
    gen_len = args.gen_len or (
        target_cfg["smiles_length"]
        if args.representation == "smiles"
        else CONFIG["representation_lengths"][args.representation]
    )

    cmd = [
        args.python, "-u", "-m", program,
        "--checkpoint", str(checkpoint),
        "--target", target_seq,
        "--num_samples", str(CONFIG["num_samples_per_batch"]),
        "--num_batches", str(args.num_candidates),
        "--gen_len", str(gen_len),
        "--optimization_steps", str(CONFIG["optimization_steps"]),
        "--weights", *map(str, CONFIG["weights"]),
        "--top_p", str(CONFIG["top_p"]),
        "--seed", str(seed),
        "--output_file", str(output),
    ]

    if args.representation in {"aa", "uaa"}:
        cmd += ["--tokenizer_type", "amino"]
    if args.representation == "aa":
        # This is the scorer used by the PDC paper runs. The optional AA-input
        # scorer is not needed and its multi-GB checkpoint is intentionally omitted.
        cmd += ["--affinity_type", "smiles"]

    if args.method == "guided":
        cmd += [
            "--builder_type", target_cfg["builder"],
            "--require_valid_x0",
            "--max_x0_retries", str(CONFIG["max_x0_retries"]),
            "--output_dir", str(output.parent / "logs"),
        ]

    if args.variant == "cupofmoca":
        cmd += [
            "--use_acceptance",
            "--beta_min", str(CONFIG["beta_min"]),
            "--beta_max", str(CONFIG["beta_max"]),
        ]
        if args.method == "guided":
            cmd += ["--guidance_method", "exp_tilt"]
            if args.representation in {"aa", "uaa"}:
                cmd += ["--cascade_k", str(CONFIG["cascade_k"])]
        else:
            cmd += ["--output_dir", str(output.parent / "logs")]

    extra = list(args.extra_args)
    if extra and extra[0] == "--":
        extra = extra[1:]
    return cmd + extra


def main() -> int:
    args = parse_args()
    targets = read_fasta(ROOT / "configs" / "targets.fasta")
    target_seq = targets[args.target]
    failures = 0

    for seed in args.seeds:
        run_dir = (
            args.output_root.resolve()
            / args.target
            / args.representation
            / f"{args.variant}_{args.method}"
            / f"seed_{seed}"
        )
        output = run_dir / "molecules.txt"
        if output.exists() and not args.force and not args.dry_run:
            print(f"Refusing to overwrite {output}; pass --force to replace it.", file=sys.stderr)
            failures += 1
            continue
        if output.exists() and args.force:
            output.unlink()
        command = build_command(args, seed, target_seq, output)
        display = " ".join(
            f"<{args.target} sequence, {len(target_seq)} aa>" if item == target_seq else shlex.quote(item)
            for item in command
        )
        print(display)
        if args.dry_run:
            continue
        run_dir.mkdir(parents=True, exist_ok=True)
        completed = subprocess.run(command, cwd=ROOT, check=False)
        if completed.returncode:
            failures += 1
            print(f"Seed {seed} failed with exit code {completed.returncode}.", file=sys.stderr)

    return 1 if failures else 0


if __name__ == "__main__":
    raise SystemExit(main())