#!/usr/bin/env python
"""PALIMPSESTE — Train the 2GB mega model with creative generation support."""
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 palimseste.creative import CreativeGenerator
from build_mega_corpus import get_mega_corpus

# Delete old model to make disk space
if os.path.exists('palimpseste-max/palimpseste_memory.bin'):
    os.remove('palimpseste-max/palimpseste_memory.bin')
    print('Deleted old model', flush=True)

# Load ALL data
pairs = get_mega_corpus()
print(f'Total pairs: {len(pairs)}', flush=True)

# Config
cfg = PalimpsesteConfig(D=20_000, context_window=256, kernel_radius=400, temperature=0.0)
lm = PalimpsesteForCausalLM(config=cfg)

# BPE tokenizer
print('Training BPE tokenizer...', flush=True)
full_text = ' '.join(q + ' ' + a for q, a in pairs)
bpe = BPETokenizer(encoder=lm.encoder, vocab_size=3000)
bpe.train(full_text, verbose=False)
lm.attach_tokenizer(bpe)
print(f'BPE vocab: {bpe.vocab_size_actual}', flush=True)

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

# Standard generation tests
print('\n=== STANDARD GENERATION ===', flush=True)
tests = ['who are you', 'what is python', 'write a poem about the sea', 'what is gravity']
for q in tests:
    r = lm.respond(q, max_new_tokens=200)
    print(f'  {q} -> {r[:60]}', flush=True)

# Creative generation tests
print('\n=== CREATIVE GENERATION (mixture logits) ===', flush=True)
gen = CreativeGenerator(lm=lm, top_k=5, diversity=0.3)
creative_tests = [
    ('what is ruby', 0.0),
    ('write a poem about wind', 0.0),
    ('explain love', 0.3),
    ('what is consciousness', 0.0),
]
for q, temp in creative_tests:
    result = gen.generate(q, max_new_tokens=60, temperature=temp, seed=42)
    print(f'  {q} (temp={temp}) -> {result.text[:60]}', flush=True)
    print(f'    novel_tokens={result.novel_tokens}, diversity={result.mixture_diversity:.2f}', flush=True)

# Save
print('\nSaving model...', flush=True)
hf = HFPalimpsesteLM(config=lm.config)
hf.mem = lm.mem; hf.phi = lm.phi; hf.encoder = lm.encoder
hf.tokenizer = lm.tokenizer; hf._self_hv = lm._self_hv; hf._token_bits_cache = None
hf.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)

# Tune LSH
print('Tuning LSH...', flush=True)
t0 = time.perf_counter()
lm.tune_lsh()
print(f'LSH tuned in {time.perf_counter()-t0:.0f}s', flush=True)

# Re-save with tuned LSH
hf.mem = lm.mem
hf.save_pretrained('./palimpseste-max', tune_lsh=False)
print(f'Final size: {os.path.getsize("./palimpseste-max/palimpseste_memory.bin")/1024/1024:.0f} MB', flush=True)

# Speed test
print('\n=== SPEED TEST ===', flush=True)
for q in ['who are you', 'hello', 'what is python']:
    t0 = time.perf_counter()
    r = lm.respond(q, max_new_tokens=200)
    print(f'  [{time.perf_counter()-t0:.1f}s] {q} -> {r[:50]}', flush=True)

print('\nDONE', flush=True)