#!/usr/bin/env python
"""PALIMPSESTE — Train a causal LM and save it (HF-compatible format).

Usage:
  # train on a text file with the 'small' preset (D=10000)
  python examples/train_lm.py --input corpus.txt --preset small --output ./my_model

  # train on inline text
  python examples/train_lm.py --text "your training text here" --preset tiny --output ./my_model

  # train at 1B-scale capacity (D=100000) — large memory, slow but no GPU
  python examples/train_lm.py --input corpus.txt --preset 1b --output ./palimpseste-1b

  # train and push to Hugging Face Hub
  python examples/train_lm.py --input corpus.txt --preset small --output ./my_model --push-to-hub user/palimpseste-small

  # custom config
  python examples/train_lm.py --input corpus.txt --output ./my_model --D 20000 --context-window 48 --kernel-radius 300

Training is O(1) per token — no epochs, no gradient, no GPU. The "model" is an
append-only memory of (context_hv -> token_hv) associations. The millionth
token costs exactly as much as the first.
"""

from __future__ import annotations

import argparse
import sys
import time
from pathlib import Path

import numpy as np

# allow running from the repo root without installing
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from palimseste.lm import PalimpsesteConfig, PRESETS
from palimseste.hf import HFPalimpsesteLM


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description="Train a PALIMPSESTE causal LM (no GPU, no gradient).",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    g = p.add_mutually_exclusive_group(required=True)
    g.add_argument("--input", type=str, help="path to a UTF-8 text file to train on")
    g.add_argument("--text", type=str, help="inline training text")

    p.add_argument("--output", "-o", type=str, required=True,
                   help="output directory for the saved model")
    p.add_argument("--preset", type=str, default=None,
                   choices=list(PRESETS.keys()),
                   help="use a capacity preset (tiny/small/medium/large/1b)")

    p.add_argument("--D", type=int, default=None, help="hypervector dimensionality")
    p.add_argument("--context-window", type=int, default=None)
    p.add_argument("--kernel-radius", type=int, default=None)
    p.add_argument("--temperature", type=float, default=0.5)

    p.add_argument("--push-to-hub", type=str, default=None,
                   help="push to HF Hub repo (e.g. user/palimpseste-small)")
    p.add_argument("--hub-token", type=str, default=None)
    p.add_argument("--seed", type=int, default=42)
    return p.parse_args()


def build_config(args: argparse.Namespace) -> PalimpsesteConfig:
    if args.preset:
        cfg = PRESETS[args.preset]
        # allow overrides
        if args.D is not None:
            cfg = PalimpsesteConfig(D=args.D, context_window=cfg.context_window,
                                    kernel_radius=cfg.kernel_radius)
        if args.context_window is not None:
            cfg.context_window = args.context_window
        if args.kernel_radius is not None:
            cfg.kernel_radius = args.kernel_radius
        cfg.temperature = args.temperature
        return cfg
    # fully custom
    D = args.D or 10_000
    return PalimpsesteConfig(
        D=D,
        context_window=args.context_window or 32,
        kernel_radius=args.kernel_radius or 200,
        temperature=args.temperature,
    )


def load_text(args: argparse.Namespace) -> str:
    if args.text is not None:
        return args.text
    p = Path(args.input)
    if not p.exists():
        sys.exit(f"error: input file not found: {p}")
    return p.read_text(encoding="utf-8")


def main() -> None:
    args = parse_args()
    config = build_config(args)
    text = load_text(args)

    print("=" * 70)
    print("PALIMPSESTE — causal LM training")
    print("=" * 70)
    print(f"config: D={config.D:,}  context_window={config.context_window}  "
          f"kernel_radius={config.kernel_radius}  temperature={config.temperature}")
    cap_log2 = float(np.log2(0.14 * config.D) + config.D / 4.0)
    print(f"theoretical capacity: 2^{cap_log2:.1f} "
          f"≈ 10^{cap_log2 * 0.30103:.1f} associations")
    print(f"corpus: {len(text):,} chars")
    print()

    rng = np.random.default_rng(args.seed)
    lm = HFPalimpsesteLM(config=config, rng=rng)

    # build tokenizer
    t0 = time.perf_counter()
    lm.build_tokenizer(text)
    print(f"vocab built: {lm.tokenizer.vocab_size} chars "
          f"({time.perf_counter() - t0:.2f}s)")

    # train (O(1) per token)
    t0 = time.perf_counter()
    n_tokens = lm.train_on_text(text, verbose=True)
    dt = time.perf_counter() - t0
    print(f"\ntrained: {n_tokens:,} tokens in {dt:.2f}s "
          f"({n_tokens / max(dt, 1e-9):,.0f} tok/s)")
    print(f"|M| = {len(lm.mem):,} traces")

    # quick self-eval
    ev = lm.evaluate(text)
    print(f"\nself-eval: next_token_accuracy={ev['next_token_accuracy']:.3f}  "
          f"mean_surprise={ev['mean_surprise']:.3f}")

    # sample generation
    sample_prompts = ["", "the ", "a "]
    print("\n--- sample generation ---")
    for prompt in sample_prompts:
        out = lm.generate(prompt, max_new_tokens=40, temperature=0.5, seed=args.seed)
        print(f"  {prompt!r:8s} -> {out.text!r}")

    # save
    out_dir = Path(args.output)
    print(f"\nsaving to {out_dir} ...")
    lm.save_pretrained(out_dir)
    print(f"saved. files: {sorted(p.name for p in out_dir.iterdir())}")

    # push to hub
    if args.push_to_hub:
        print(f"\npushing to HF Hub: {args.push_to_hub} ...")
        url = lm.push_to_hub(args.push_to_hub, token=args.hub_token)
        print(f"pushed: {url}")

    print("\n" + "=" * 70)
    print("done. Load with:")
    print(f"  from palimseste.hf import HFPalimpsesteLM")
    print(f"  lm = HFPalimpsesteLM.from_pretrained('{out_dir}')")
    print("=" * 70)


if __name__ == "__main__":
    main()