"""
Character-Level Next-Letter Predictor
On-the-fly generation during training (no precomputed dataset)
ODE-CCT framework: character sequences as trajectories
"""

import os
import re
import random
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from collections import defaultdict, deque
from pathlib import Path

# ═══════════════════════════════════════════════════════════════════
# CONFIGURATION
# ═══════════════════════════════════════════════════════════════════

class Config:
    book_folder = "./books"
    
    # Model
    char_embed_dim = 64
    hidden_dim = 256
    num_layers = 2
    
    # Sequence
    seq_len = 50  # Context length for prediction
    
    # Training
    batch_size = 64
    epochs = 20
    lr = 0.001
    
    # ODE-CCT
    entropy_threshold = 0.5
    
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

cfg = Config()

# ═══════════════════════════════════════════════════════════════════
# TEXT LOADING
# ═══════════════════════════════════════════════════════════════════

def load_text_files(folder_path):
    """Load raw text from markdown files."""
    text = ""
    folder = Path(folder_path)
    
    if not folder.exists():
        print(f"⚠️ Folder '{folder_path}' not found. Using sample text.")
        return sample_text()
    
    for md_file in folder.glob("*.md"):
        print(f"📖 Loading: {md_file.name}")
        with open(md_file, 'r', encoding='utf-8', errors='ignore') as f:
            content = f.read()
            # Remove markdown syntax but keep structure
            content = re.sub(r'#+', '', content)  # Remove headers
            content = re.sub(r'\[([^\]]+)\]\([^\)]+\)', r'\1', content)  # Links to text
            content = re.sub(r'[*_`~]', '', content)  # Remove formatting
            text += content + "\n"
    
    if len(text) == 0:
        return sample_text()
    
    print(f"✅ Loaded {len(text)} characters")
    return text

def sample_text():
    return """
    In the beginning God created the heavens and the earth. Now the earth was 
    formless and empty, darkness was over the surface of the deep, and the 
    Spirit of God was hovering over the waters. And God said, Let there be 
    light, and there was light. God saw that the light was good, and he 
    separated the light from the darkness. God called the light day, and 
    the darkness he called night. And there was evening, and there was morning
    """

# ═══════════════════════════════════════════════════════════════════
# CHARACTER VOCABULARY
# ═══════════════════════════════════════════════════════════════════

def build_char_vocab(text):
    """Build character-to-index mapping."""
    # Filter to printable characters
    chars = sorted(set(text))
    
    # Special tokens
    char2idx = {
        '<PAD>': 0,
        '<UNK>': 1,
        '<BOS>': 2,
        '<EOS>': 3
    }
    
    for i, char in enumerate(chars):
        char2idx[char] = i + 4
    
    idx2char = {v: k for k, v in char2idx.items()}
    
    print(f"📚 Character vocabulary size: {len(char2idx)}")
    return char2idx, idx2char, chars

# ═══════════════════════════════════════════════════════════════════
# ON-THE-FLY DATA GENERATOR
# ═══════════════════════════════════════════════════════════════════

class StreamingTextGenerator:
    """
    Generates training samples on-the-fly from raw text.
    No precomputed dataset - yields batches during training.
    """
    
    def __init__(self, text, char2idx, seq_len=50, stride=3):
        self.text = text
        self.char2idx = char2idx
        self.seq_len = seq_len
        self.stride = stride  # Step size between sequences
        self.vocab_size = len(char2idx)
        
        # Pre-compute valid indices
        self.valid_indices = [
            i for i in range(len(text) - seq_len)
            if all(c in char2idx for c in text[i:i+seq_len+1])
        ]
        
        if len(self.valid_indices) == 0:
            # Fallback: use all indices
            self.valid_indices = list(range(0, len(text) - seq_len, stride))
        
        print(f"🖼️ Streaming generator: {len(self.valid_indices)} valid sequences")
    
    def __iter__(self):
        return self
    
    def __next__(self):
        return self.get_batch(cfg.batch_size)
    
    def get_batch(self, batch_size):
        """Get a random batch of sequences."""
        # Randomly sample starting positions
        indices = random.sample(self.valid_indices, min(batch_size, len(self.valid_indices)))
        
        # Extract sequences
        input_seqs = []
        target_seqs = []
        
        for start_idx in indices:
            # Input: seq_len characters
            input_chars = self.text[start_idx:start_idx + self.seq_len]
            # Target: next character
            target_char = self.text[start_idx + self.seq_len]
            
            # Convert to indices
            input_indices = [self.char2idx.get(c, 1) for c in input_chars]
            target_index = self.char2idx.get(target_char, 1)
            
            input_seqs.append(input_indices)
            target_seqs.append(target_index)
        
        # Convert to tensors
        input_tensor = torch.LongTensor(input_seqs)
        target_tensor = torch.LongTensor(target_seqs)
        
        return input_tensor, target_tensor
    
    def get_epoch_batches(self, batch_size, num_batches=None):
        """Yield all batches for one epoch."""
        if num_batches is None:
            num_batches = len(self.valid_indices) // batch_size
        
        for _ in range(num_batches):
            yield self.get_batch(batch_size)

# ═══════════════════════════════════════════════════════════════════
# MODEL: Character-Level RNN with ODE-CCT
# ═══════════════════════════════════════════════════════════════════

class CharPredictor(nn.Module):
    """
    Character-level next-letter predictor.
    Uses LSTM to encode character sequence trajectory.
    """
    
    def __init__(self, vocab_size, embed_dim=64, hidden_dim=256, num_layers=2):
        super().__init__()
        
        self.vocab_size = vocab_size
        self.hidden_dim = hidden_dim
        self.num_layers = num_layers
        
        # Character embedding
        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
        
        # LSTM for sequence encoding
        self.lstm = nn.LSTM(
            input_size=embed_dim,
            hidden_size=hidden_dim,
            num_layers=num_layers,
            batch_first=True,
            dropout=0.2 if num_layers > 1 else 0
        )
        
        # Prediction head
        self.predictor = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(hidden_dim, vocab_size)
        )
        
        # Entropy tracking for ODE-CCT
        self.register_buffer('entropy_buffer', torch.zeros(100))
        self.entropy_idx = 0
    
    def forward(self, x, return_entropy=False):
        """
        x: [batch, seq_len] - character indices
        Returns: logits [batch, vocab_size], entropy (optional)
        """
        # Embed characters
        embedded = self.embedding(x)  # [batch, seq_len, embed_dim]
        
        # LSTM encoding
        lstm_out, (hidden, cell) = self.lstm(embedded)
        
        # Use last hidden state as trajectory representation
        last_hidden = hidden[-1]  # [batch, hidden_dim]
        
        # Predict next character
        logits = self.predictor(last_hidden)  # [batch, vocab_size]
        
        if return_entropy:
            probs = F.softmax(logits, dim=1)
            entropy = -(probs * torch.log(probs + 1e-8)).sum(dim=1).mean()
            
            # Update entropy buffer
            idx = self.entropy_idx % 100
            self.entropy_buffer[idx] = entropy
            self.entropy_idx += 1
            
            return logits, entropy.item()
        
        return logits
    
    def predict_next(self, x):
        """Predict next character with confidence."""
        logits = self.forward(x)
        probs = F.softmax(logits, dim=1)
        
        top_prob, top_idx = probs.max(dim=1)
        
        return top_idx, top_prob, logits
    
    def get_trajectory_state(self, x):
        """Get the LSTM hidden state as trajectory representation."""
        embedded = self.embedding(x)
        _, (hidden, cell) = self.lstm(embedded)
        return hidden, cell


class ConditionalCollapseCharPredictor(CharPredictor):
    """
    Enhanced model with CCT (Conditional Collapse Theory):
    - Detects periodic character patterns
    - Skips computation for common patterns
    - Tracks entropy for adaptive compute
    """
    
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        
        # Pattern cache (n-gram -> next char)
        self.ngram_cache = {}
        self.ngram_order = 3  # Use 3-grams for caching
        
        # Compute budget tracking
        self.compute_budget = 1.0
        self.compute_used = 0.0
        
    def check_cache(self, seq):
        """Check if sequence matches cached pattern."""
        seq_tuple = tuple(seq.tolist()[-self.ngram_order:])
        return self.ngram_cache.get(seq_tuple)
    
    def update_cache(self, seq, next_char):
        """Update pattern cache."""
        seq_tuple = tuple(seq.tolist()[-self.ngram_order:])
        self.ngram_cache[seq_tuple] = next_char
    
    def forward_with_cct(self, x):
        """
        Forward with Conditional Collapse:
        - Check cache first
        - If hit: low compute (cache lookup)
        - If miss: full LSTM forward
        """
        batch_size = x.size(0)
        results = []
        
        for i in range(batch_size):
            seq = x[i]
            
            # Check cache
            cached = self.check_cache(seq)
            
            if cached is not None and random.random() < 0.5:
                # Use cached prediction
                self.compute_used += 0.1  # Cache hit is cheap
                results.append(cached)
            else:
                # Full forward pass
                self.compute_used += 1.0
                logits = self.forward(seq.unsqueeze(0))
                pred = logits.argmax(dim=1).item()
                
                # Update cache
                self.update_cache(seq, pred)
                results.append(pred)
        
        return torch.LongTensor(results).to(x.device)

# ═══════════════════════════════════════════════════════════════════
# TRAINING (ON-THE-FLY)
# ═══════════════════════════════════════════════════════════════════

def train_on_the_fly(model, text, char2idx, epochs=20, lr=0.001):
    """
    Train model with on-the-fly data generation.
    No precomputed dataset - generates batches during training.
    """
    
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    criterion = nn.CrossEntropyLoss()
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=5, factor=0.5)
    
    model.to(cfg.device)
    
    # Create streaming generator
    generator = StreamingTextGenerator(text, char2idx, seq_len=cfg.seq_len, stride=3)
    
    for epoch in range(epochs):
        model.train()
        
        epoch_loss = 0.0
        epoch_correct = 0
        epoch_total = 0
        
        entropy_sum = 0.0
        batches_done = 0
        
        # Generate batches on-the-fly
        num_batches = len(generator.valid_indices) // cfg.batch_size
        
        for batch_idx in range(num_batches):
            # Generate batch on-the-fly
            input_seq, target_char = generator.get_batch(cfg.batch_size)
            
            # Move to device
            input_seq = input_seq.to(cfg.device)
            target_char = target_char.to(cfg.device)
            
            # Forward pass
            optimizer.zero_grad()
            logits, entropy = model(input_seq, return_entropy=True)
            
            # Loss
            loss = criterion(logits, target_char)
            
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            optimizer.step()
            
            # Metrics
            epoch_loss += loss.item()
            preds = logits.argmax(dim=1)
            epoch_correct += (preds == target_char).sum().item()
            epoch_total += target_char.size(0)
            
            entropy_sum += entropy
            batches_done += 1
            
            # Progress update every 500 batches
            if (batch_idx + 1) % 500 == 0:
                print(f"  Batch {batch_idx+1}/{num_batches} | "
                      f"Loss: {epoch_loss/batches_done:.4f} | "
                      f"Acc: {100*epoch_correct/epoch_total:.2f}%")
        
        # Epoch metrics
        train_acc = 100 * epoch_correct / epoch_total
        avg_loss = epoch_loss / batches_done
        avg_entropy = entropy_sum / batches_done
        
        # Calculate perplexity
        perplexity = np.exp(avg_entropy)
        
        scheduler.step(avg_loss)
        
        print(f"Epoch {epoch+1}/{epochs} | "
              f"Loss: {avg_loss:.4f} | Acc: {train_acc:.2f}% | "
              f"Entropy: {avg_entropy:.4f} | PPL: {perplexity:.2f}")
    
    return model

# ═══════════════════════════════════════════════════════════════════
# TESTING WITH ODE-CCT ANALYSIS
# ═══════════════════════════════════════════════════════════════════

def test_model(model, text, char2idx, idx2char, num_samples=1000):
    """Test model on held-out text."""
    
    model.eval()
    criterion = nn.CrossEntropyLoss()
    
    # Use last 10% of text for testing
    test_text = text[int(len(text) * 0.9):]
    
    correct = 0
    total = 0
    entropies = []
    
    # Create test generator
    test_gen = StreamingTextGenerator(test_text, char2idx, seq_len=cfg.seq_len, stride=5)
    
    with torch.no_grad():
        for batch_idx, (input_seq, target_char) in enumerate(test_gen.get_epoch_batches(cfg.batch_size, num_batches=num_samples // cfg.batch_size)):
            input_seq = input_seq.to(cfg.device)
            target_char = target_char.to(cfg.device)
            
            logits, entropy = model(input_seq, return_entropy=True)
            
            loss = criterion(logits, target_char)
            
            preds = logits.argmax(dim=1)
            correct += (preds == target_char).sum().item()
            total += target_char.size(0)
            
            probs = F.softmax(logits, dim=1)
            max_probs = probs.max(dim=1)[0]
            entropies.extend(max_probs.tolist())
    
    accuracy = 100 * correct / total
    avg_conf = np.mean(entropies)
    perplexity = np.exp(entropy)
    
    print("\n" + "="*60)
    print("📊 CHARACTER PREDICTION TEST RESULTS")
    print("="*60)
    print(f"Test Accuracy: {accuracy:.2f}% ({correct}/{total})")
    print(f"Average Confidence: {avg_conf:.4f}")
    print(f"Perplexity: {perplexity:.2f}")
    
    # Character-specific accuracy
    print("\n🔤 Character Accuracy Analysis:")
    char_correct = defaultdict(int)
    char_total = defaultdict(int)
    
    # Re-run for detailed analysis
    with torch.no_grad():
        for batch_idx, (input_seq, target_char) in enumerate(test_gen.get_epoch_batches(cfg.batch_size, num_batches=50)):
            input_seq = input_seq.to(cfg.device)
            target_char = target_char.to(cfg.device)
            
            logits, _ = model(input_seq, return_entropy=True)
            preds = logits.argmax(dim=1)
            
            for pred, true in zip(preds.cpu(), target_char.cpu()):
                true_char = idx2char.get(true.item(), '?')
                char_total[true_char] += 1
                if pred == true:
                    char_correct[true_char] += 1
    
    # Most/least predicted characters
    char_acc = {c: char_correct.get(c, 0) / max(char_total.get(c, 1), 1) for c in char_total}
    sorted_chars = sorted(char_acc.items(), key=lambda x: -x[1])
    
    print(f"  Best predicted ({len(sorted_chars)} chars):")
    for char, acc in sorted_chars[:5]:
        c = char if char not in '<PAD><UNK><BOS><EOS>' else f"'{char}'"
        print(f"    '{c}': {acc*100:.1f}%")
    
    print(f"  Worst predicted:")
    for char, acc in sorted_chars[-5:]:
        c = char if char not in '<PAD><UNK><BOS><EOS>' else f"'{char}'"
        print(f"    '{c}': {acc*100:.1f}%")
    
    print("="*60)
    
    return accuracy, perplexity

# ═══════════════════════════════════════════════════════════════════
# SAMPLE GENERATION
# ═══════════════════════════════════════════════════════════════════

def generate_text(model, seed_text, char2idx, idx2char, length=200, temperature=1.0):
    """Generate text given a seed."""
    
    model.eval()
    
    # Initialize with seed
    current_seq = [char2idx.get(c, 1) for c in seed_text[-cfg.seq_len:]]
    
    if len(current_seq) < cfg.seq_len:
        current_seq = [0] * (cfg.seq_len - len(current_seq)) + current_seq
    
    generated = seed_text
    
    with torch.no_grad():
        for _ in range(length):
            input_tensor = torch.LongTensor([current_seq[-cfg.seq_len:]]).to(cfg.device)
            
            logits, _ = model(input_tensor, return_entropy=False)
            
            # Temperature sampling
            probs = F.softmax(logits / temperature, dim=1)
            pred_idx = torch.multinomial(probs, 1).item()
            
            pred_char = idx2char.get(pred_idx, '')
            
            # Stop at EOS
            if pred_idx == 3:  # <EOS>
                break
            
            generated += pred_char
            current_seq.append(pred_idx)
    
    return generated

# ═══════════════════════════════════════════════════════════════════
# MAIN
# ═══════════════════════════════════════════════════════════════════

def main():
    print("🧠 Character-Level Next-Letter Predictor (On-The-Fly Training)")
    print("="*60)
    
    # Step 1: Load text
    print("\n📂 Step 1: Loading text files...")
    text = load_text_files(cfg.book_folder)
    
    # Step 2: Build character vocabulary
    print("\n📚 Step 2: Building character vocabulary...")
    char2idx, idx2char, chars = build_char_vocab(text)
    vocab_size = len(char2idx)
    
    # Step 3: Build model
    print("\n🏗️ Step 3: Building model...")
    model = CharPredictor(
        vocab_size=vocab_size,
        embed_dim=cfg.char_embed_dim,
        hidden_dim=cfg.hidden_dim,
        num_layers=cfg.num_layers
    )
    print(model)
    
    # Step 4: Train (on-the-fly)
    print("\n🚀 Step 4: Training (on-the-fly data generation)...")
    print(f"   Generating batches from {len(text)} characters")
    model = train_on_the_fly(model, text, char2idx, epochs=cfg.epochs, lr=cfg.lr)
    
    # Step 5: Test
    print("\n🧪 Step 5: Testing...")
    accuracy, perplexity = test_model(model, text, char2idx, idx2char)
    
    # Step 6: Generate samples
    print("\n🎨 Text Generation Samples:")
    seed_texts = [
        "In the",
        "God said",
        "the light",
        "and the"
    ]
    
    for seed in seed_texts:
        generated = generate_text(model, seed, char2idx, idx2char, length=100)
        print(f"\n  Seed: '{seed}'")
        print(f"  Generated: '{generated}'")
    
    # Save model
    torch.save({
        'model_state_dict': model.state_dict(),
        'char2idx': char2idx,
        'idx2char': idx2char,
        'config': {
            'vocab_size': vocab_size,
            'embed_dim': cfg.char_embed_dim,
            'hidden_dim': cfg.hidden_dim,
            'num_layers': cfg.num_layers,
            'seq_len': cfg.seq_len
        }
    }, 'char_predictor.pt')
    
    print("\n💾 Model saved to 'char_predictor.pt'")
    print(f"Final Test Accuracy: {accuracy:.2f}%")

if __name__ == "__main__":
    main()