#!/usr/bin/env python3
"""Live ball-occupancy + synthetic k* sweep.

Synthetic (no weights)::

    python examples/occupancy.py --sweep

Live store (HF weights or a local Memory)::

    python examples/occupancy.py --live
    python examples/occupancy.py --live --limit 306 --model thefinalboss/palimpseste-max

The first query whose Hamming ball exceeds k*_exact ≈ 0.02 D is the first
predictable exact-recall failure. Soft cleanup holds longer (~0.03 D).
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from palimseste.capacity import (  # noqa: E402
    DEFAULT_D,
    DEFAULT_M,
    format_report,
    format_sweep,
    kstar_exact,
    kstar_soft,
    kstar_theory,
    occupancy,
    occupancy_from_addresses,
    sweep_exact_soft,
)


def cmd_sweep(args: argparse.Namespace) -> int:
    rows = sweep_exact_soft(D=args.D, M=args.M, reps=args.reps)
    print(format_sweep(rows, args.D, args.M))
    if args.json:
        print(json.dumps(rows, indent=2))
    return 0


def cmd_live(args: argparse.Namespace) -> int:
    try:
        from palimseste.hf import HFPalimpsesteLM
        from palimseste.phi import Phi
    except ImportError as e:
        print(
            "live path needs the full palimseste package "
            "(phi.py / memory.py / hf.py).\n"
            f"Import failed: {e}\n"
            "Run --sweep until you drop capacity.py into the HF tree.",
            file=sys.stderr,
        )
        return 2

    model_id = args.model
    print(f"loading {model_id} …", file=sys.stderr)
    lm = HFPalimpsesteLM.from_pretrained(model_id)
    mem = lm.mem
    phi = getattr(lm, "phi", None) or Phi()
    D = int(getattr(next(iter(mem.traces)).address, "D", DEFAULT_D))
    print(
        f"|M|={len(mem.traces)}  D={D}  radius={phi.config.radius}  "
        f"k*_exact={kstar_exact(D)}  k*_soft={kstar_soft(D)}  "
        f"k*_theory={kstar_theory(D)}",
        file=sys.stderr,
    )

    if args.queries:
        queries = []
        tags = []
        # encode each line as the model would — fall back to learner if present
        encode = getattr(lm, "encode", None) or getattr(lm, "embed", None)
        if encode is None:
            print("model has no encode(); probing stored addresses instead", file=sys.stderr)
            rep = occupancy_from_addresses(mem, phi, limit=args.limit)
        else:
            for line in Path(args.queries).read_text(encoding="utf-8").splitlines():
                line = line.strip()
                if not line:
                    continue
                queries.append(encode(line))
                tags.append(line[:80])
                if args.limit and len(queries) >= args.limit:
                    break
            rep = occupancy(mem, phi, queries, tags=tags)
    else:
        rep = occupancy_from_addresses(mem, phi, limit=args.limit)

    print(format_report(rep))
    if args.json:
        print(json.dumps(rep.as_dict(), indent=2))
    return 1 if rep.n_over_exact else 0


def main() -> int:
    p = argparse.ArgumentParser(description="PALIMPSESTE capacity / occupancy")
    sub = p.add_subparsers(dest="cmd")

    ps = sub.add_parser("sweep", help="synthetic exact/soft codebook sweep")
    ps.add_argument("--D", type=int, default=8000)
    ps.add_argument("--M", type=int, default=DEFAULT_M)
    ps.add_argument("--reps", type=int, default=60)
    ps.add_argument("--json", action="store_true")

    pl = sub.add_parser("live", help="histogram |N_r(q)| on a loaded store")
    pl.add_argument("--model", default="thefinalboss/palimpseste-max")
    pl.add_argument("--limit", type=int, default=306)
    pl.add_argument("--queries", default=None, help="text file, one query per line")
    pl.add_argument("--json", action="store_true")

    # also allow --sweep flag without subcommand
    p.add_argument("--sweep", action="store_true")
    p.add_argument("--live", action="store_true")
    p.add_argument("--D", type=int, default=8000)
    p.add_argument("--M", type=int, default=DEFAULT_M)
    p.add_argument("--reps", type=int, default=60)
    p.add_argument("--model", default="thefinalboss/palimpseste-max")
    p.add_argument("--limit", type=int, default=306)
    p.add_argument("--queries", default=None)
    p.add_argument("--json", action="store_true")

    args = p.parse_args()
    if args.cmd == "sweep" or (args.sweep and args.cmd is None):
        return cmd_sweep(args)
    if args.cmd == "live" or args.live:
        return cmd_live(args)
    # default: sweep, cheap
    return cmd_sweep(args)


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