#!/usr/bin/env python
"""PALIMPSESTE — Train a conversational chat model.

Trains on Q/A pairs (from corpus_chat.py or a JSON file) so the model learns
to respond to user messages. Each pair is encoded as a self-contained episode:
BOS question EOS BOS answer EOS.

Usage:
  # train on the built-in corpus
  python examples/train_chat.py --preset small --output ./chat_model

  # train on custom Q/A pairs from a JSON file: [{"q": "...", "a": "..."}, ...]
  python examples/train_chat.py --qa-file my_pairs.json --output ./chat_model

  # train at 1B-scale capacity
  python examples/train_chat.py --preset 1b --output ./chat_model_1b

  # train and push to HF Hub
  python examples/train_chat.py --preset small --output ./chat_model --push-to-hub user/palimpseste-chat

  # then chat interactively:
  python examples/chat.py --model ./chat_model
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parent))
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

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


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description="Train a PALIMPSESTE conversational chat model.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    p.add_argument("--output", "-o", type=str, required=True,
                   help="output directory for the saved model")
    p.add_argument("--preset", type=str, default="small",
                   choices=list(PRESETS.keys()),
                   help="capacity preset (default: small)")
    p.add_argument("--qa-file", type=str, default=None,
                   help="JSON file with [{q, a}, ...]; default: built-in corpus")
    p.add_argument("--D", type=int, default=None, help="override D")
    p.add_argument("--context-window", type=int, default=None,
                   help="override context window (default: 64 for chat)")
    p.add_argument("--kernel-radius", type=int, default=None,
                   help="override kernel radius")
    p.add_argument("--temperature", type=float, default=0.3,
                   help="sampling temperature (default: 0.3 for focused answers)")
    p.add_argument("--push-to-hub", type=str, default=None)
    p.add_argument("--hub-token", type=str, default=None)
    p.add_argument("--seed", type=int, default=42)
    p.add_argument("--epochs", type=int, default=1,
                   help="number of passes over the corpus (default: 1; "
                        "PALIMPSESTE is append-only so >1 reinforces)")
    return p.parse_args()


def load_pairs(args: argparse.Namespace) -> list[tuple[str, str]]:
    if args.qa_file:
        with open(args.qa_file, "r", encoding="utf-8") as f:
            data = json.load(f)
        return [(d["q"], d["a"]) for d in data]
    return get_corpus()


def main() -> None:
    args = parse_args()
    pairs = load_pairs(args)

    # build config — chat benefits from a larger context window
    config = PRESETS[args.preset]
    if args.D is not None:
        config = PalimpsesteConfig(D=args.D, context_window=config.context_window,
                                   kernel_radius=config.kernel_radius)
    if args.context_window is not None:
        config.context_window = args.context_window
    elif args.context_window is None and config.context_window < 48:
        config.context_window = 64  # chat needs room for Q+A
    if args.kernel_radius is not None:
        config.kernel_radius = args.kernel_radius
    config.temperature = args.temperature

    print("=" * 70)
    print("PALIMPSESTE — conversational chat training")
    print("=" * 70)
    print(f"config: D={config.D:,}  context_window={config.context_window}  "
          f"kernel_radius={config.kernel_radius}  temperature={config.temperature}")
    print(f"Q/A pairs: {len(pairs)}")

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

    # build vocab from all questions + answers
    full_text = "".join(q + a for q, a in pairs)
    lm.build_tokenizer(full_text)
    print(f"vocab: {lm.tokenizer.vocab_size} chars")

    # train (O(1) per token, no gradient)
    import time
    total = 0
    for epoch in range(args.epochs):
        t0 = time.perf_counter()
        n = lm.train_on_qa_pairs(pairs, verbose=(args.epochs == 1))
        dt = time.perf_counter() - t0
        total += n
        print(f"  epoch {epoch+1}/{args.epochs}: {n} tokens in {dt:.2f}s "
              f"({n/max(dt,1e-9):,.0f} tok/s), |M|={len(lm.mem)}")

    print(f"\ntotal: {total} tokens, |M| = {len(lm.mem):,}")

    # test responses on a few sample questions using the Conversation layer
    print("\n--- sample responses ---")
    from palimseste.chat import Conversation
    conv = Conversation(model=lm, learn_live=True)
    conv.register_questions(pairs)
    test_inputs = [
        "bonjour",
        "qui es-tu",
        "comment tu apprends",
        "tu utilises un gpu",
        "tu utilise un gpu",  # typo
        "who are you",
        "how do you learn",
        "c quoi palimpseste",
        "merci",
    ]
    for q in test_inputs:
        resp = conv.respond(q, temperature=0.0, seed=0)
        print(f"  Q: {q}")
        print(f"  A: {resp}")
        print()
    conv.reset()  # don't save conversation history into the model

    # save
    out_dir = Path(args.output)
    print(f"saving to {out_dir} ...")
    lm.save_pretrained(out_dir)

    # 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. Chat with:")
    print(f"  python examples/chat.py --model {out_dir}")
    print("=" * 70)


if __name__ == "__main__":
    main()