#!/usr/bin/env python3
"""Run AF3-NA+ from Hugging Face files and emit one K=4 artifact."""

from __future__ import annotations

import argparse
from dataclasses import asdict, dataclass
import json
import os
import pathlib
import platform
import tempfile
import time
from typing import Any, Mapping, Sequence

import numpy as np

from hub_loader import LoadedRelease, load_release
from evotemplate_na.inference import (
    load_artifact,
    materialize_torch_tile_artifact_arrays,
    write_artifact,
)
from evotemplate_na.model import NucleicAcidChain, NucleicAcidTokenizer


class HubPredictionError(ValueError):
    """Raised before publishing an invalid or ambiguous artifact."""


@dataclass(frozen=True)
class PredictionSummary:
    artifact_directory: str
    arrays_file: str
    manifest_file: str
    artifact_id: str
    length: int
    num_chains: int
    device: str
    dtype: str
    wall_seconds: float
    gpu_seconds: float | None
    peak_memory_mb: float | None


def _load_json(path: pathlib.Path) -> Any:
    try:
        return json.loads(path.read_text(encoding="utf-8"))
    except OSError as exc:
        raise HubPredictionError(f"cannot read input JSON: {exc}") from exc
    except (UnicodeDecodeError, json.JSONDecodeError) as exc:
        raise HubPredictionError(f"invalid input JSON: {exc}") from exc


def _canonical_chain(value: Mapping[str, Any], index: int) -> dict[str, str]:
    if not isinstance(value, Mapping):
        raise HubPredictionError(f"chains[{index}] must be an object")
    if set(value) != {"chain_id", "polymer_type", "sequence"}:
        raise HubPredictionError(
            f"chains[{index}] requires exactly chain_id/polymer_type/sequence"
        )
    chain_id = value["chain_id"]
    polymer = value["polymer_type"]
    sequence = value["sequence"]
    if not isinstance(chain_id, str) or not chain_id.strip():
        raise HubPredictionError("chain_id must be a non-empty string")
    if not isinstance(polymer, str) or polymer.strip().upper() not in ("RNA", "DNA"):
        raise HubPredictionError("polymer_type must be RNA or DNA")
    if not isinstance(sequence, str):
        raise HubPredictionError("sequence must be a string")
    canonical = "".join(sequence.split()).upper()
    normalized_polymer = polymer.strip().upper()
    alphabet = set("ACGU" if normalized_polymer == "RNA" else "ACGT")
    if not canonical or not set(canonical) <= alphabet:
        raise HubPredictionError(
            f"chain {chain_id!r} is not canonical {normalized_polymer}"
        )
    return {
        "chain_id": chain_id.strip(),
        "polymer_type": normalized_polymer,
        "sequence": canonical,
    }


def _af3_chains(payload: Mapping[str, Any]) -> list[dict[str, str]]:
    if payload.get("dialect") != "alphafold3" or not isinstance(
        payload.get("sequences"), list
    ):
        raise HubPredictionError("input is not an AlphaFold 3 JSON object")
    raw: list[dict[str, Any]] = []
    for entry in payload["sequences"]:
        if not isinstance(entry, Mapping) or len(entry) != 1:
            raise HubPredictionError("malformed AlphaFold 3 sequence entry")
        kind = next(iter(entry))
        if kind not in ("rna", "dna"):
            continue
        molecule = entry[kind]
        if not isinstance(molecule, Mapping):
            raise HubPredictionError("malformed AlphaFold 3 NA molecule")
        raw_ids = molecule.get("id")
        ids = raw_ids if isinstance(raw_ids, list) else [raw_ids]
        for chain_id in ids:
            raw.append(
                {
                    "chain_id": chain_id,
                    "polymer_type": kind.upper(),
                    "sequence": molecule.get("sequence"),
                }
            )
    if not raw:
        raise HubPredictionError("AlphaFold 3 input contains no RNA/DNA chains")
    return [_canonical_chain(chain, index) for index, chain in enumerate(raw)]


def normalize_input(
    payload: Any,
    *,
    fallback_artifact_id: str,
    artifact_id_override: str | None = None,
) -> tuple[str, tuple[dict[str, str], ...], tuple[tuple[str, str], ...], int]:
    """Normalize either AF3 JSON or ``{name, chains}`` release input."""

    if not isinstance(payload, Mapping):
        raise HubPredictionError("input JSON root must be an object")
    if payload.get("dialect") == "alphafold3":
        chains = _af3_chains(payload)
        explicit_pairs = None
        declared_max_length = None
    else:
        raw_chains = payload.get("chains")
        if not isinstance(raw_chains, list) or not raw_chains:
            raise HubPredictionError("simple input requires a non-empty chains list")
        chains = [
            _canonical_chain(chain, index) for index, chain in enumerate(raw_chains)
        ]
        explicit_pairs = payload.get("authorized_chain_pairs")
        declared_max_length = payload.get("max_length")
    ids = [chain["chain_id"] for chain in chains]
    if len(set(ids)) != len(ids):
        raise HubPredictionError("chain_id values must be unique")
    if explicit_pairs is None:
        pairs = [
            (left, right)
            for index, left in enumerate(ids)
            for right in ids[index:]
        ]
    else:
        if not isinstance(explicit_pairs, list) or not explicit_pairs:
            raise HubPredictionError("authorized_chain_pairs must be a non-empty list")
        pairs = []
        seen: set[tuple[str, str]] = set()
        for pair in explicit_pairs:
            if (
                not isinstance(pair, list)
                or len(pair) != 2
                or pair[0] not in ids
                or pair[1] not in ids
            ):
                raise HubPredictionError("authorized_chain_pairs contains an invalid pair")
            canonical = tuple(sorted((str(pair[0]), str(pair[1]))))
            if canonical in seen:
                raise HubPredictionError("authorized_chain_pairs contains duplicates")
            seen.add(canonical)
            pairs.append((str(pair[0]), str(pair[1])))
    length = sum(len(chain["sequence"]) for chain in chains)
    max_length = length if declared_max_length is None else declared_max_length
    if (
        isinstance(max_length, bool)
        or not isinstance(max_length, int)
        or max_length < length
    ):
        raise HubPredictionError("max_length must be an integer covering all residues")
    raw_name = artifact_id_override or payload.get("name") or fallback_artifact_id
    if not isinstance(raw_name, str) or not raw_name.strip():
        raise HubPredictionError("artifact_id/name must be a non-empty string")
    return raw_name.strip(), tuple(chains), tuple(pairs), max_length


def authorized_pair_mask(
    chains: Sequence[Mapping[str, str]], pairs: Sequence[Sequence[str]]
) -> np.ndarray:
    offsets: dict[str, slice] = {}
    cursor = 0
    for chain in chains:
        stop = cursor + len(chain["sequence"])
        offsets[chain["chain_id"]] = slice(cursor, stop)
        cursor = stop
    result = np.zeros((cursor, cursor), dtype=np.bool_)
    for left, right in pairs:
        left_slice, right_slice = offsets[left], offsets[right]
        result[left_slice, right_slice] = True
        result[right_slice, left_slice] = True
    return result


def predict(
    *,
    model_dir: str | pathlib.Path,
    input_path: str | pathlib.Path,
    output_directory: str | pathlib.Path,
    device: str = "cpu",
    dtype: str = "auto",
    block_size: int | None = None,
    artifact_id: str | None = None,
    overwrite: bool = False,
) -> PredictionSummary:
    """Load one release and publish one AF3-compatible K=4 artifact."""

    try:
        import torch
    except ImportError as exc:
        raise RuntimeError("AF3-NA+ requires torch") from exc
    input_source = pathlib.Path(input_path).expanduser().resolve(strict=True)
    payload = _load_json(input_source)
    name, chains, pairs, max_length = normalize_input(
        payload,
        fallback_artifact_id=input_source.stem,
        artifact_id_override=artifact_id,
    )
    if block_size is not None and (
        isinstance(block_size, bool) or not isinstance(block_size, int) or block_size <= 0
    ):
        raise HubPredictionError("block_size must be a positive integer")
    if device == "auto":
        device = "cuda" if torch.cuda.is_available() else "cpu"
    loaded: LoadedRelease = load_release(model_dir, device=device, dtype=dtype)
    if len(chains) > loaded.spec.model_config.student.max_chains:
        raise HubPredictionError("input chain count exceeds the release max_chains")
    destination = pathlib.Path(output_directory).expanduser().resolve()
    if not overwrite and any(
        (destination / name).exists() for name in ("manifest.json", "arrays.npz")
    ):
        raise HubPredictionError("artifact exists; pass --overwrite to replace it")
    destination.mkdir(parents=True, exist_ok=True)
    tokenizer = NucleicAcidTokenizer(
        max_chains=loaded.spec.model_config.student.max_chains
    )
    model_chains = [
        NucleicAcidChain(
            chain_id=chain["chain_id"],
            polymer_type=chain["polymer_type"],
            sequence=chain["sequence"],
        )
        for chain in chains
    ]
    batch = tokenizer.batch_encode([model_chains]).to(loaded.device)
    length = sum(len(chain["sequence"]) for chain in chains)
    residue_mask = np.ones(length, dtype=np.bool_)
    pair_mask = authorized_pair_mask(chains, pairs)
    tile_size = block_size or loaded.spec.model_config.pair_decoder.block_size
    with tempfile.TemporaryDirectory(prefix=".af3-na-plus-", dir=destination) as scratch:
        start = time.perf_counter()
        gpu_seconds: float | None = None
        peak_memory_mb: float | None = None
        start_event = end_event = None
        if loaded.device.type == "cuda":
            if loaded.device.index is not None:
                torch.cuda.set_device(loaded.device.index)
            torch.cuda.reset_peak_memory_stats(loaded.device)
            start_event = torch.cuda.Event(enable_timing=True)
            end_event = torch.cuda.Event(enable_timing=True)
            start_event.record()
        with torch.inference_mode():
            student = loaded.model.encode_full_chain(batch)
            context = loaded.model.prepare_context(batch, student_output=student)
            arrays = materialize_torch_tile_artifact_arrays(
                tile_forward=loaded.model.forward_tile,
                context=context,
                slot_logits=context.decoder_context.slot_logits,
                length=length,
                active_pair_mask=pair_mask,
                residue_mask=residue_mask,
                scratch_directory=pathlib.Path(scratch),
                block_size=tile_size,
            )
        if start_event is not None and end_event is not None:
            end_event.record()
            torch.cuda.synchronize(loaded.device)
            gpu_seconds = float(start_event.elapsed_time(end_event)) / 1000.0
            peak_memory_mb = float(torch.cuda.max_memory_allocated(loaded.device)) / (1024.0**2)
        wall_seconds = time.perf_counter() - start
        precision = "bf16" if loaded.dtype == torch.bfloat16 else "fp32"
        write_artifact(
            destination,
            arrays=arrays,
            artifact_id=name,
            chains=chains,
            authorized_chain_pairs=pairs,
            max_length=max_length,
            view_mode="forward+reverse_complement",
            runtime={
                "python": platform.python_version(),
                "torch": torch.__version__,
                "pid": os.getpid(),
                "wall_seconds": wall_seconds,
                "gpu_seconds": gpu_seconds,
                "peak_memory_mb": peak_memory_mb,
            },
            device=str(loaded.device),
            precision=precision,
            quantization=None,
            slot_weight_is_normalized=True,
        )
    # Re-open through the existing scientific schema validator.  This checks
    # shapes, masks, and probabilities.
    artifact = load_artifact(destination)
    return PredictionSummary(
        artifact_directory=str(destination),
        arrays_file=str(destination / "arrays.npz"),
        manifest_file=str(destination / "manifest.json"),
        artifact_id=name,
        length=artifact.length,
        num_chains=len(chains),
        device=str(loaded.device),
        dtype=str(loaded.dtype).removeprefix("torch."),
        wall_seconds=wall_seconds,
        gpu_seconds=gpu_seconds,
        peak_memory_mb=peak_memory_mb,
    )


def main(argv: Sequence[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model-dir", type=pathlib.Path, default=pathlib.Path("."))
    parser.add_argument("--input", type=pathlib.Path, required=True)
    parser.add_argument("--output", type=pathlib.Path, required=True)
    parser.add_argument("--device", default="auto")
    parser.add_argument(
        "--dtype", choices=("auto", "bfloat16", "float32"), default="auto"
    )
    parser.add_argument("--block-size", type=int)
    parser.add_argument("--artifact-id")
    parser.add_argument("--overwrite", action="store_true")
    args = parser.parse_args(argv)
    summary = predict(
        model_dir=args.model_dir,
        input_path=args.input,
        output_directory=args.output,
        device=args.device,
        dtype=args.dtype,
        block_size=args.block_size,
        artifact_id=args.artifact_id,
        overwrite=args.overwrite,
    )
    print(json.dumps(asdict(summary), sort_keys=True))
    return 0


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


__all__ = [
    "HubPredictionError",
    "PredictionSummary",
    "authorized_pair_mask",
    "main",
    "normalize_input",
    "predict",
]