"""Regenerate knapsack_cache/ and knapsack_cache.zip against the currently
installed instanovo version.

When InstaNovo bumps its Knapsack on-disk format (e.g. v1.1 -> v1.2.2 added a
field to parameters.pkl, giving the "not enough values to unpack" error on
load), the bundled cache becomes unreadable and the runtime falls back to
multi-minute regeneration on every container cold-start. This script does
that regeneration once on a workstation so the committed zip stays in sync.

Run with:
    .venv/bin/python scripts/regenerate_knapsack.py

Expect a few minutes on CPU. Output:
    ./knapsack_cache/{parameters.pkl,masses.npy,chart.npy}
    ./knapsack_cache.zip   (the three files above, repacked for HF Spaces)
"""

from __future__ import annotations

import shutil
import time
import zipfile
from pathlib import Path

from instanovo.constants import MASS_SCALE, MAX_MASS
from instanovo.inference import Knapsack
from instanovo.transformer.model import InstaNovo

TRANSFORMER_MODEL_ID = "instanovo-v1.2.0"
KNAPSACK_DIR = Path(__file__).resolve().parent.parent / "knapsack_cache"
KNAPSACK_ZIP = KNAPSACK_DIR.with_suffix(".zip")


def main() -> None:
    print(f"Loading {TRANSFORMER_MODEL_ID} for residue set...")
    model, _ = InstaNovo.from_pretrained(TRANSFORMER_MODEL_ID)
    residue_set = model.residue_set

    residue_masses = dict(residue_set.residue_masses)
    excluded = set(residue_set.special_tokens) | {
        k for k, v in residue_masses.items() if v <= 0
    }
    for res in excluded:
        residue_masses.pop(res, None)
    if not residue_masses:
        raise SystemExit("No valid residues with positive mass; cannot regenerate.")
    print(f"Excluding {sorted(excluded)} from knapsack; {len(residue_masses)} residues remain.")

    if KNAPSACK_DIR.exists():
        print(f"Clearing existing {KNAPSACK_DIR}...")
        shutil.rmtree(KNAPSACK_DIR)
    KNAPSACK_DIR.parent.mkdir(parents=True, exist_ok=True)
    # Knapsack.save() creates KNAPSACK_DIR itself and refuses if it already exists.

    start = time.time()
    print("Generating knapsack (this can take several minutes)...")
    knapsack = Knapsack.construct_knapsack(
        residue_masses=residue_masses,
        residue_indices=residue_set.residue_to_index,
        max_mass=MAX_MASS,
        mass_scale=MASS_SCALE,
    )
    knapsack.save(str(KNAPSACK_DIR))
    print(f"Saved knapsack to {KNAPSACK_DIR} ({time.time() - start:.1f}s).")

    print(f"Packing {KNAPSACK_ZIP}...")
    if KNAPSACK_ZIP.exists():
        KNAPSACK_ZIP.unlink()
    with zipfile.ZipFile(KNAPSACK_ZIP, "w", zipfile.ZIP_DEFLATED) as zf:
        for f in sorted(KNAPSACK_DIR.iterdir()):
            zf.write(f, arcname=f.name)
    print(f"Wrote {KNAPSACK_ZIP} ({KNAPSACK_ZIP.stat().st_size / 1e6:.1f} MB).")

    print("Done. Commit the new zip with git (LFS-tracked) and push to the Space.")


if __name__ == "__main__":
    main()