"""Run IgFold structure prediction from antibody sequences or a FASTA file."""

import argparse
import json
from pathlib import Path

import torch

from igfold import IgFoldRunner


DEFAULT_HEAVY = (
    "EVQLVQSGPEVKKPGTSVKVSCKASGFTFMSSAVQWVRQARGQRLEWIGWIVIGSGNTNYAQKF"
    "QERVTITRDMSTSTAYMELSSLRSEDTAVYYCAAPYCSSISCNDGFDIWGQGTMVTVS"
)
DEFAULT_LIGHT = (
    "DVVMTQTPFSLPVSLGDQASISCRSSQSLVHSNGNTYLHWYLQKPGQSPKLLIYKVSNRFSGV"
    "PDRFSGSGSGTDFTLKISRVEAEDLGVYFCSQSTHVPYTFGGGTKLEIK"
)


def parse_args():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--fasta", type=Path, help="Input FASTA containing H/L chains.")
    parser.add_argument("--heavy", help="Heavy-chain amino-acid sequence.")
    parser.add_argument("--light", help="Light-chain amino-acid sequence.")
    parser.add_argument("--output", type=Path, default=Path("output/inference/antibody.pdb"))
    parser.add_argument("--num-models", type=int, choices=range(1, 5), default=4)
    parser.add_argument("--refine", action="store_true", help="Enable structure refinement.")
    parser.add_argument("--openmm", action="store_true", help="Use OpenMM for refinement.")
    parser.add_argument("--renum", action="store_true", help="Apply Chothia numbering.")
    parser.add_argument("--cpu", action="store_true", help="Force CPU inference.")
    return parser.parse_args()


def main():
    args = parse_args()
    root = Path(__file__).resolve().parents[1]
    checkpoints = sorted((root / "weight" / "IgFold").glob("*.ckpt"))[: args.num_models]
    if len(checkpoints) != args.num_models:
        raise FileNotFoundError(
            f"Expected {args.num_models} checkpoints in {root / 'weight' / 'IgFold'}, "
            f"found {len(checkpoints)}."
        )

    if args.fasta is not None:
        fasta_file = str(args.fasta.resolve())
        sequences = None
    else:
        fasta_file = None
        sequences = {"H": args.heavy or DEFAULT_HEAVY}
        light = args.light if args.light is not None else DEFAULT_LIGHT
        if light:
            sequences["L"] = light

    if args.refine and not args.openmm:
        from igfold.refine.pyrosetta_ref import init_pyrosetta

        init_pyrosetta()

    args.output.parent.mkdir(parents=True, exist_ok=True)
    runner = IgFoldRunner(
        model_ckpts=[str(path) for path in checkpoints],
        try_gpu=not args.cpu,
    )
    output = runner.fold(
        str(args.output),
        fasta_file=fasta_file,
        sequences=sequences,
        do_refine=args.refine,
        use_openmm=args.openmm,
        do_renum=args.renum,
    )

    result = {
        "status": "PASS",
        "output_pdb": str(args.output.resolve()),
        "models": len(runner.models),
        "device": str(next(runner.models[0].parameters()).device),
        "torch": torch.__version__,
        "coords_shape": list(output.coords.shape),
        "prmsd_shape": list(output.prmsd.shape),
    }
    print("IGFOLD_INFERENCE_RESULT=" + json.dumps(result, sort_keys=True))


if __name__ == "__main__":
    main()