#!/usr/bin/env python
"""PALIMPSESTE — Train the killer model with BPE tokenizer.

BPE reduces tokens 5x → 5x fewer prediction errors → much better generation.
Combined with the killer knowledge corpus (quality over quantity).
"""
import json, sys, time, os
sys.path.insert(0, '.')
sys.path.insert(0, 'examples')

from palimseste.lm import PalimpsesteForCausalLM, PalimpsesteConfig
from palimseste.hf import HFPalimpsesteLM
from palimseste.bpe import BPETokenizer
from killer_corpus import get_killer_corpus

# Load data
pairs = list(get_killer_corpus())

# Also load trivia data for breadth
try:
    with open('trivia_qa_pairs.json') as f:
        trivia = json.load(f)
    pairs.extend([(item['q'].lower(), item['a'].lower()) for item in trivia[:3000]])
    print(f'Loaded {len(trivia[:3000])} trivia pairs', flush=True)
except FileNotFoundError:
    pass

print(f'Total pairs: {len(pairs)}', flush=True)

# Config: ctx=256 for complete answers, temp=0.0 for greedy
cfg = PalimpsesteConfig(
    D=20_000,
    context_window=256,
    kernel_radius=400,
    kernel_min_weight=1e-6,
    temperature=0.0,
)
print(f'Config: D={cfg.D} ctx={cfg.context_window} temp={cfg.temperature}', flush=True)

# Build model
lm = PalimpsesteForCausalLM(config=cfg)

# ================================================================
# BPE TOKENIZER — train once on the full corpus
# ================================================================
print('Training BPE tokenizer...', flush=True)
full_text = ' '.join(q + ' ' + a for q, a in pairs)
bpe = BPETokenizer(encoder=lm.encoder, vocab_size=2000)
bpe.train(full_text, verbose=False)
lm.attach_tokenizer(bpe)
print(f'BPE vocab: {bpe.vocab_size_actual} tokens', flush=True)

# Compare BPE vs char-level token count
sample = 'who are you'
char_ids = [ord(c) for c in sample]  # rough char count
bpe_ids = bpe.encode(sample, add_bos=True, add_eos=True)
print(f'BPE reduction: "{sample}" = {len(sample)+2} chars → {len(bpe_ids)} BPE tokens '
      f'({len(sample)/max(len(bpe_ids),1):.1f}x)', flush=True)

# ================================================================
# TRAIN with BPE
# ================================================================
print('Training with BPE...', flush=True)
t0 = time.perf_counter()
n_tokens = lm.train_on_qa_pairs(pairs, verbose=False)
dt = time.perf_counter() - t0
print(f'Trained: {n_tokens:,} tokens in {dt:.0f}s ({n_tokens/dt:.0f} tok/s)', flush=True)
print(f'|M| = {len(lm.mem):,} traces', flush=True)

# ================================================================
# TEST
# ================================================================
print('\n=== KILLER MODEL TESTS ===', flush=True)
tests = [
    'who are you', 'hello', 'what is python', 'how do you learn',
    'do you use a gpu', 'what is the capital of france',
    'what can you do', 'are you conscious',
    'what is a black hole', 'what is recursion',
    'what is gravity', 'what is dna',
    'how are you different from chatgpt', 'tell a joke',
]
for q in tests:
    t0 = time.perf_counter()
    r = lm.respond(q, max_new_tokens=200)
    dt = time.perf_counter() - t0
    print(f'  [{dt:.1f}s] {q}', flush=True)
    print(f'        -> {r}', flush=True)

# ================================================================
# SAVE
# ================================================================
print('\nSaving killer model...', flush=True)
hf_lm = HFPalimpsesteLM(config=lm.config)
hf_lm.mem = lm.mem
hf_lm.phi = lm.phi
hf_lm.encoder = lm.encoder
hf_lm.tokenizer = lm.tokenizer
hf_lm._self_hv = lm._self_hv
hf_lm._token_bits_cache = None
hf_lm.save_pretrained('./palimpseste-max', tune_lsh=False)
sz = os.path.getsize('./palimpseste-max/palimpseste_memory.bin')
print(f'Saved: {sz/1024/1024:.0f} MB', flush=True)
print('DONE', flush=True)