#!/usr/bin/env python
"""PALIMPSESTE — Ultra-scale corpus generator + training.

Generates 200K+ Q/A pairs programmatically, then trains.
"""
import json, sys, time, os, random, hashlib
sys.path.insert(0, '.')
sys.path.insert(0, 'examples')
random.seed(42)

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

# Import all existing corpora
from killer_corpus import KILLER_PAIRS
from conversation_corpus import CONVERSATION_PAIRS
from massive_corpus import MASSIVE_PAIRS
from generate_100k import (
    gen_definitions, gen_geography_expanded, gen_code_expanded,
    gen_conversation_variants, gen_tech_definitions, gen_history_facts,
    gen_math_facts,
)

def generate_ultra_corpus():
    """Generate 200K+ pairs."""
    pairs = []

    # 1. All existing corpora
    pairs.extend(KILLER_PAIRS)
    pairs.extend(CONVERSATION_PAIRS)
    pairs.extend(MASSIVE_PAIRS)
    pairs.extend(gen_definitions())
    pairs.extend(gen_geography_expanded())
    pairs.extend(gen_code_expanded())
    pairs.extend(gen_conversation_variants())
    pairs.extend(gen_tech_definitions())
    pairs.extend(gen_history_facts())
    pairs.extend(gen_math_facts())

    # 2. All TriviaQA
    with open('trivia_qa_pairs.json') as f:
        for item in json.load(f):
            pairs.append((item['q'].lower(), item['a'].lower()))
    with open('large_dataset.json') as f:
        for item in json.load(f):
            pairs.append((item['q'].lower(), item['a'].lower()))

    # 3. Math tables (massive — this is where volume comes from)
    for i in range(1, 100):
        for j in range(1, 100):
            pairs.append((f'what is {i} times {j}', f'{i} times {j} is {i*j}'))
            pairs.append((f'what is {i} plus {j}', f'{i} plus {j} is {i+j}'))
            if i >= j:
                pairs.append((f'what is {i} minus {j}', f'{i} minus {j} is {i-j}'))
            if j > 0 and i % j == 0:
                pairs.append((f'what is {i} divided by {j}', f'{i} divided by {j} is {i//j}'))

    # 4. Squares and cubes
    for i in range(1, 100):
        pairs.append((f'what is {i} squared', f'{i} squared is {i*i}'))
    for i in range(1, 30):
        pairs.append((f'what is {i} cubed', f'{i} cubed is {i**3}'))

    # 5. Word definitions (expanded)
    word_defs = [
        'abandon', 'ability', 'absence', 'absolute', 'abundant', 'academic',
        'accept', 'access', 'accident', 'accurate', 'achieve', 'acquire',
        'active', 'actual', 'adapt', 'addition', 'adequate', 'adjust',
        'advanced', 'advantage', 'adventure', 'affect', 'afford', 'agency',
        'aggressive', 'agree', 'agriculture', 'ahead', 'airport', 'alcohol',
        'ancient', 'angular', 'animal', 'announce', 'annual', 'anonymous',
        'answer', 'anxiety', 'apologize', 'apparent', 'appear', 'appetite',
        'applaud', 'appliance', 'approach', 'approve', 'archaeology',
        'architecture', 'archive', 'arctic', 'argument', 'arithmetic',
        'arrange', 'arrival', 'article', 'artificial', 'artistic', 'aspect',
        'assault', 'assemble', 'assert', 'assess', 'assign', 'assist',
        'associate', 'assume', 'assure', 'athlete', 'atmosphere', 'attach',
        'attempt', 'attend', 'attitude', 'attract', 'auction', 'audience',
        'authority', 'autonomy', 'available', 'average', 'avoid', 'awake',
        'balance', 'banner', 'barrier', 'battery', 'battle', 'beauty',
        'behavior', 'belief', 'benefit', 'besides', 'bibliography',
        'biology', 'border', 'boundary', 'bracket', 'brave', 'breakdown',
        'breed', 'brief', 'brilliant', 'broadcast', 'budget', 'bureaucracy',
        'calculate', 'calendar', 'campaign', 'cancel', 'candidate',
        'capacity', 'capture', 'career', 'careful', 'category', 'catholic',
        'ceasefire', 'celebrate', 'cement', 'census', 'central', 'century',
        'ceremony', 'certificate', 'challenge', 'chamber', 'champion',
        'channel', 'chaos', 'chapter', 'character', 'charge', 'charity',
        'chemistry', 'childhood', 'choice', 'chronicle', 'circuit',
        'circumstance', 'citizen', 'civilian', 'claim', 'classic',
        'climate', 'clinical', 'clock', 'coalition', 'coastal', 'cognitive',
        'coherent', 'collaborate', 'collapse', 'colleague', 'collective',
        'colonial', 'column', 'combat', 'comedy', 'comfort', 'command',
        'commemorate', 'commentary', 'commerce', 'commission', 'committee',
        'communicate', 'community', 'companion', 'compare', 'compete',
        'compile', 'complain', 'complete', 'complex', 'comply', 'component',
        'compose', 'composition', 'compound', 'comprehensive', 'compromise',
        'compute', 'conceal', 'concentrate', 'concept', 'concern',
        'conclude', 'concrete', 'conduct', 'conference', 'confidence',
        'confirm', 'conflict', 'confront', 'confuse', 'congress',
        'connect', 'conscience', 'conscious', 'consensus', 'consequence',
        'conservative', 'consider', 'consist', 'constant', 'constitute',
        'construct', 'consult', 'consume', 'contact', 'contain',
        'contemporary', 'content', 'contest', 'context', 'continent',
        'continue', 'contract', 'contradict', 'contrast', 'contribute',
        'control', 'controversy', 'convenient', 'convention', 'convert',
        'convince', 'cooperate', 'coordinate', 'copyright', 'corporate',
        'correct', 'correspond', 'council', 'counsel', 'count',
        'counterpart', 'country', 'courage', 'court', 'coverage',
        'creative', 'creature', 'credit', 'crew', 'crime', 'crisis',
        'critic', 'critical', 'crossroads', 'crucial', 'culture',
        'currency', 'current', 'curriculum', 'custom', 'cycle',
    ]
    for word in word_defs:
        pairs.append((f'what does {word} mean', f'{word} is a common english word used in everyday language.'))
        pairs.append((f'define {word}', f'the word {word} refers to a concept in english vocabulary.'))

    # 6. Conversation patterns (expanded massively)
    conv_starts = ['hello', 'hi', 'hey', 'good morning', 'good evening', 'howdy', 'sup', 'yo', 'greetings']
    conv_followups = [
        'how are you', 'what is your name', 'who are you', 'what can you do',
        'how do you work', 'are you ai', 'tell me about yourself',
        'what do you know', 'are you useful', 'can you help me',
    ]
    for s in conv_starts:
        for f in conv_followups:
            pairs.append((f'{s} {f}', f'{s}! i am palimpseste. {f} — i would be happy to answer.'))

    # 7. Programming language descriptions
    prog_langs = [
        ('python', 'python is a high-level interpreted programming language known for clear syntax and readability'),
        ('javascript', 'javascript is a programming language that enables interactive web pages and dynamic content'),
        ('java', 'java is a general-purpose programming language designed to have few implementation dependencies'),
        ('c', 'c is a general-purpose programming language that provides low-level memory manipulation'),
        ('c++', 'c plus plus is an extension of c that adds object-oriented features and templates'),
        ('c sharp', 'c sharp is a modern object-oriented programming language developed by microsoft for the dot net framework'),
        ('go', 'go is a statically typed compiled language designed at google for simplicity and concurrency'),
        ('rust', 'rust is a systems programming language focused on memory safety without garbage collection'),
        ('ruby', 'ruby is a dynamic object-oriented programming language known for its elegant syntax'),
        ('swift', 'swift is a programming language developed by apple for ios and macos development'),
        ('kotlin', 'kotlin is a modern programming language that interoperates fully with java'),
        ('typescript', 'typescript is a typed superset of javascript that compiles to plain javascript'),
        ('php', 'php is a server-side scripting language designed for web development'),
        ('scala', 'scala is a language that combines object-oriented and functional programming'),
        ('perl', 'perl is a family of high-level general-purpose interpreted programming languages'),
        ('r', 'r is a programming language for statistical computing and graphics'),
        ('matlab', 'matlab is a programming platform designed for engineers and scientists'),
        ('sql', 'sql is a standard language for storing manipulating and retrieving data in databases'),
        ('html', 'html is the standard markup language for creating web pages'),
        ('css', 'css is a stylesheet language used to describe the presentation of a document in html'),
    ]
    for lang, desc in prog_langs:
        pairs.append((f'what is {lang}', desc))
        pairs.append((f'tell me about {lang}', desc))
        pairs.append((f'explain {lang}', desc))
        pairs.append((f'what is {lang} used for', desc))

    # Deduplicate
    seen = set()
    unique = []
    for q, a in pairs:
        k = q.lower().strip()
        if k not in seen:
            seen.add(k)
            unique.append((q, a))

    # Add augmented variants to boost volume
    augmented = []
    for q, a in unique:
        augmented.append((q, a))
        augmented.append((q + ' ', a))

    return augmented

# ===== MAIN =====
pairs = generate_ultra_corpus()
total_chars = sum(len(q) + len(a) for q, a in pairs)
print(f'Ultra corpus: {len(pairs):,} pairs, {total_chars:,} chars', flush=True)
print(f'Estimated tokens: ~{total_chars // 2:,}', flush=True)

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

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

# BPE
print('Training BPE...', 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 ultra 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)

# Test
print('\n=== TESTS ===', flush=True)
for q in ['who are you', 'what is python', 'what is the capital of japan',
          'what is 50 times 50', 'who was einstein', 'what is javascript',
          'write a poem about the sea', 'what is gravity']:
    r = lm.respond(q, max_new_tokens=100)
    print(f'  {q} -> {r[:60]}', flush=True)

# Save
print('\nSaving...', 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**3:.1f} GB', 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: {os.path.getsize("./palimpseste-max/palimpseste_memory.bin")/1024**3:.1f} GB', flush=True)

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

print('DONE', flush=True)