"""
Word-MLP Sequence Model: Conv2D + LSTM for Language Prediction
ODE-CCT framework: word sequences as trajectories, next word as trajectory prediction
"""

import os
import re
import random
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from torch.nn.utils.rnn import pad_sequence, pack_padded_sequence, pad_packed_sequence
from PIL import Image, ImageDraw, ImageFont
from collections import defaultdict
from pathlib import Path

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

class Config:
    book_folder = "./books"
    
    # Image
    img_size = 28
    font_size = 20
    
    # Model
    embed_dim = 128
    hidden_dim = 256
    num_classes = None
    
    # Sequence
    seq_len = 8  # Input sequence length (predict next word)
    
    # Training
    batch_size = 64
    epochs = 20
    lr = 0.001
    
    # ODE-CCT parameters
    entropy_threshold = 0.3  # Collapse threshold for prediction confidence
    periodicity_window = 3   # Detect repeating patterns of this length
    
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

cfg = Config()

# ═══════════════════════════════════════════════════════════════════
# DATA LOADING
# ═══════════════════════════════════════════════════════════════════

def load_markdown_files(folder_path):
    """Load all .md files from folder and extract words."""
    words = []
    folder = Path(folder_path)
    
    if not folder.exists():
        print(f"⚠️ Folder '{folder_path}' not found. Using sample text.")
        return sample_words()
    
    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()
        
        extracted = re.findall(r'[a-zA-Z]{2,}', content)
        words.extend([w.lower() for w in extracted])
    
    if len(words) == 0:
        return sample_words()
    
    print(f"✅ Extracted {len(words)} words from markdown files.")
    return words

def sample_words():
    sample_text = """
    intelligence artificial machine learning neural network deep 
    consciousness awareness thinking reasoning planning algorithm 
    """
    return re.findall(r'[a-zA-Z]{2,}', sample_text.lower())

def build_vocabulary(words, max_vocab_size=5000):
    """Build word-to-index vocabulary."""
    word_counts = defaultdict(int)
    for word in words:
        word_counts[word] += 1
    
    sorted_words = sorted(word_counts.items(), key=lambda x: -x[1])
    vocab_words = [w for w, c in sorted_words[:max_vocab_size]]
    
    word2idx = {word: idx for idx, word in enumerate(vocab_words)}
    idx2word = {idx: word for word, idx in word2idx.items()}
    
    print(f"📚 Vocabulary size: {len(vocab_words)}")
    return word2idx, idx2word, vocab_words

# ═══════════════════════════════════════════════════════════════════
# WORD IMAGE RENDERING
# ═══════════════════════════════════════════════════════════════════

def render_word_image(word, size=28, font_size=20):
    """Render a word as a grayscale 28x28 image."""
    img = Image.new('L', (size, size), color=255)
    draw = ImageDraw.Draw(img)
    
    try:
        font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", font_size)
    except:
        try:
            font = ImageFont.truetype("arial.ttf", font_size)
        except:
            font = ImageFont.load_default()
    
    bbox = draw.textbbox((0, 0), word, font=font)
    text_width = bbox[2] - bbox[0]
    text_height = bbox[3] - bbox[1]
    
    x = (size - text_width) // 2
    y = (size - text_height) // 2
    
    draw.text((x, y), word, fill=0, font=font)
    
    img_array = np.array(img, dtype=np.float32) / 255.0
    img_array = 1.0 - img_array
    
    return img_array

# ═══════════════════════════════════════════════════════════════════
# DATASET: SEQUENCE OF WORD IMAGES
# ═══════════════════════════════════════════════════════════════════

class SequenceDataset(Dataset):
    """Dataset of word image sequences with next-word prediction."""
    
    def __init__(self, words, word2idx, seq_len=8, max_samples=50000):
        self.samples = []
        self.word2idx = word2idx
        self.seq_len = seq_len
        self.num_classes = len(word2idx)
        
        vocab_set = set(word2idx.keys())
        valid_words = [w for w in words if w in vocab_set]
        
        # Limit samples
        if len(valid_words) > max_samples:
            valid_words = valid_words[:max_samples]
        
        self.words = valid_words
        
        # Build sequence samples
        for i in range(len(valid_words) - seq_len):
            input_words = valid_words[i:i + seq_len]
            target_word = valid_words[i + seq_len]
            
            # Render images
            images = []
            indices = []
            for w in input_words:
                try:
                    img = render_word_image(w)
                except:
                    img = np.zeros((28, 28), dtype=np.float32)
                images.append(torch.FloatTensor(img))
                indices.append(word2idx[w])
            
            target_idx = word2idx[target_word]
            
            self.samples.append({
                'images': images,
                'input_indices': indices,
                'target_idx': target_idx
            })
        
        print(f"🖼️ Sequence dataset: {len(self.samples)} samples (seq_len={seq_len})")
    
    def __len__(self):
        return len(self.samples)
    
    def __getitem__(self, idx):
        sample = self.samples[idx]
        images = torch.stack(sample['images']).unsqueeze(1)  # [seq_len, 1, 28, 28]
        input_indices = torch.LongTensor(sample['input_indices'])
        target_idx = torch.LongTensor([sample['target_idx']])
        
        return images, input_indices, target_idx


def collate_seq_batch(batch):
    """Pad sequences to same length."""
    images_list, input_indices_list, target_idx_list = zip(*batch)
    
    # Stack images (already same size due to fixed seq_len)
    images = torch.stack(images_list)  # [batch, seq_len, 1, 28, 28]
    
    # Input indices (variable length sequences)
    input_indices = torch.stack(input_indices_list)  # [batch, seq_len]
    
    # Target
    target_idx = torch.stack(target_idx_list).squeeze(-1)  # [batch]
    
    return images, input_indices, target_idx

# ═══════════════════════════════════════════════════════════════════
# MODEL: Word-MLP Encoder + LSTM Decoder
# ═══════════════════════════════════════════════════════════════════

class WordMLPEncoder(nn.Module):
    """Conv2D + Dense encoder for word images."""
    
    def __init__(self, embed_dim=128):
        super().__init__()
        
        self.features = nn.Sequential(
            nn.Conv2d(1, 32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2),  # 28 -> 14
            
            nn.Conv2d(32, 64, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2),  # 14 -> 7
            
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2),  # 7 -> 3
        )
        
        self.flatten = nn.Flatten()
        self.embedding = nn.Linear(128 * 3 * 3, embed_dim)
    
    def forward(self, x):
        """x: [batch, seq_len, 1, 28, 28]"""
        batch_size, seq_len = x.size(0), x.size(1)
        
        # Reshape for conv: [batch * seq_len, 1, 28, 28]
        x = x.view(batch_size * seq_len, 1, 28, 28)
        
        # Conv
        x = self.features(x)
        x = self.flatten(x)
        x = self.embedding(x)
        
        # Reshape back: [batch, seq_len, embed_dim]
        x = x.view(batch_size, seq_len, -1)
        
        return x


class SequencePredictor(nn.Module):
    """
    Sequence model with ODE-CCT dynamics:
    - LSTM for trajectory encoding
    - Entropy-based confidence for early stopping
    - Periodic pattern detection
    """
    
    def __init__(self, embed_dim=128, hidden_dim=256, num_classes=2000):
        super().__init__()
        
        self.embed_dim = embed_dim
        self.hidden_dim = hidden_dim
        self.num_classes = num_classes
        
        # Word image encoder
        self.encoder = WordMLPEncoder(embed_dim)
        
        # Sequence encoder (LSTM)
        self.lstm = nn.LSTM(
            input_size=embed_dim,
            hidden_size=hidden_dim,
            num_layers=2,
            batch_first=True,
            dropout=0.2
        )
        
        # Decoder: predict next word
        self.decoder = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(hidden_dim, embed_dim)
        )
        
        # Prediction head
        self.predictor = nn.Linear(embed_dim, num_classes)
        
        # Entropy tracker (ODE-CCT)
        self.register_buffer('entropy_history', torch.zeros(100))
        self.entropy_idx = 0
    
    def forward(self, images, return_confidence=False):
        """Forward pass with optional confidence estimation."""
        # Encode word images
        embeddings = self.encoder(images)  # [batch, seq_len, embed_dim]
        
        # LSTM trajectory
        lstm_out, (hidden, cell) = self.lstm(embeddings)
        
        # Use last hidden state as current trajectory state
        last_hidden = hidden[-1]  # [batch, hidden_dim]
        
        # Decode to prediction space
        decoded = self.decoder(last_hidden)  # [batch, embed_dim]
        
        # Predict next word
        logits = self.predictor(decoded)  # [batch, num_classes]
        
        if return_confidence:
            probs = F.softmax(logits, dim=1)
            entropy = -(probs * torch.log(probs + 1e-8)).sum(dim=1)
            return logits, entropy.mean().item()
        
        return logits, decoded
    
    def detect_periodicity(self, predictions, targets):
        """Detect if model is in a periodic (high-frequency) pattern."""
        # Check if last N predictions were correct (stable cycle)
        if len(predictions) < cfg.periodicity_window:
            return False
        
        recent_correct = predictions[-cfg.periodicity_window:]
        return all(recent_correct)
    
    def update_entropy(self, entropy_value):
        """Update entropy history for ODE-CCT tracking."""
        idx = self.entropy_idx % 100
        self.entropy_history[idx] = entropy_value
        self.entropy_idx += 1
    
    def get_trajectory_stability(self):
        """Get stability measure from entropy history."""
        recent = self.entropy_history[:50].mean().item()
        return recent


class ConditionalCollapsePredictor(nn.Module):
    """
    Enhanced predictor with Conditional Collapse Theory.
    - Adaptively allocates compute based on entropy
    - Uses periodicity detection to skip computation
    """
    
    def __init__(self, embed_dim=128, hidden_dim=256, num_classes=2000):
        super().__init__()
        
        self.base_model = SequencePredictor(embed_dim, hidden_dim, num_classes)
        
        # Collapse threshold
        self.confidence_threshold = 0.7
        
        # Periodic phrase cache
        self.phrase_cache = {}
        
    def forward(self, images, use_collapse=True):
        """
        Forward with optional CCT collapse.
        
        If use_collapse=True:
        - Check if sequence matches known periodic pattern
        - If high confidence, skip LSTM (use cached prediction)
        - If low confidence, run full model
        """
        # Check phrase cache first
        batch_size = images.size(0)
        
        # Create phrase key from word indices
        phrase_key = self._get_phrase_key(images)
        
        if use_collapse and phrase_key in self.phrase_cache:
            # Periodic pattern detected - use cached result
            cached_preds = self.phrase_cache[phrase_key]
            # Run model anyway for training, but could skip in inference
            logits, decoded = self.base_model(images)
            return logits, decoded, True  # Third return = collapsed
        
        # Normal forward pass
        logits, decoded = self.base_model(images)
        
        # Check if we should cache this pattern
        probs = F.softmax(logits, dim=1)
        confidence = probs.max(dim=1)[0].mean().item()
        
        if confidence > self.confidence_threshold:
            # High confidence - cache this pattern
            pred_idx = logits.argmax(dim=1)
            for key, pred in zip(phrase_key, pred_idx.tolist()):
                if key not in self.phrase_cache:
                    self.phrase_cache[key] = pred
        
        return logits, decoded, False
    
    def _get_phrase_key(self, images):
        """Extract discrete key from image sequence (for caching)."""
        # In real implementation, would use embeddings
        # Here we use hash of first/last word positions
        return tuple(images[:, 0, 0, 14, 14].tolist())  # Simplified


# ═══════════════════════════════════════════════════════════════════
# TRAINING
# ═══════════════════════════════════════════════════════════════════

def train_model(model, train_loader, val_loader, epochs=20, lr=0.001):
    """Train sequence prediction model."""
    
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    criterion = nn.CrossEntropyLoss()
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=3, factor=0.5)
    
    model.to(cfg.device)
    
    for epoch in range(epochs):
        model.train()
        train_loss = 0.0
        train_correct = 0
        train_total = 0
        
        entropy_sum = 0.0
        collapse_count = 0
        
        for batch_idx, (images, input_indices, target_idx) in enumerate(train_loader):
            images = images.to(cfg.device)
            target_idx = target_idx.to(cfg.device)
            
            optimizer.zero_grad()
            
            # Forward pass
            logits, entropy = model(images, return_confidence=True)
            
            # Loss
            loss = criterion(logits, target_idx)
            
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            optimizer.step()
            
            # Metrics
            train_loss += loss.item()
            preds = logits.argmax(dim=1)
            train_correct += (preds == target_idx).sum().item()
            train_total += target_idx.size(0)
            
            entropy_sum += entropy
        
        train_acc = 100 * train_correct / train_total
        avg_loss = train_loss / len(train_loader)
        avg_entropy = entropy_sum / len(train_loader)
        
        # Validation
        model.eval()
        val_loss = 0.0
        val_correct = 0
        val_total = 0
        
        with torch.no_grad():
            for images, input_indices, target_idx in val_loader:
                images = images.to(cfg.device)
                target_idx = target_idx.to(cfg.device)
                
                logits, _ = model(images, return_confidence=True)
                loss = criterion(logits, target_idx)
                
                val_loss += loss.item()
                preds = logits.argmax(dim=1)
                val_correct += (preds == target_idx).sum().item()
                val_total += target_idx.size(0)
        
        val_acc = 100 * val_correct / val_total
        avg_val_loss = val_loss / len(val_loader)
        
        scheduler.step(avg_val_loss)
        
        print(f"Epoch {epoch+1}/{epochs} | "
              f"Loss: {avg_loss:.4f} | Train: {train_acc:.2f}% | "
              f"Val Loss: {avg_val_loss:.4f} | Val: {val_acc:.2f}% | "
              f"Entropy: {avg_entropy:.3f}")
    
    return model

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

def test_model(model, test_loader, word2idx, idx2word):
    """Test and analyze ODE-CCT dynamics."""
    
    model.eval()
    criterion = nn.CrossEntropyLoss()
    
    test_loss = 0.0
    test_correct = 0
    test_total = 0
    
    # Perplexity tracking
    entropies = []
    confidence_scores = []
    
    # Sequence analysis
    correct_sequences = []
    all_predictions = []
    
    with torch.no_grad():
        for images, input_indices, target_idx in test_loader:
            images = images.to(cfg.device)
            target_idx = target_idx.to(cfg.device)
            
            logits, entropy = model(images, return_confidence=True)
            loss = criterion(logits, target_idx)
            
            test_loss += loss.item()
            preds = logits.argmax(dim=1)
            test_correct += (preds == target_idx).sum().item()
            test_total += target_idx.size(0)
            
            # Track metrics
            probs = F.softmax(logits, dim=1)
            confidence = probs.max(dim=1)[0]
            
            if isinstance(entropy, float):
                entropies.extend([entropy] * len(target_idx))
            else:
                entropies.extend(entropy.tolist())
            confidence_scores.extend(confidence.tolist())
            
            # Track correct predictions
            for i in range(len(preds)):
                is_correct = preds[i].item() == target_idx[i].item()
                all_predictions.append(is_correct)
    
    test_acc = 100 * test_correct / test_total
    avg_entropy = np.mean(entropies)
    avg_confidence = np.mean(confidence_scores)
    
    # Perplexity
    perplexity = np.exp(avg_entropy)
    
    print("\n" + "="*60)
    print("📊 SEQUENCE PREDICTION TEST RESULTS")
    print("="*60)
    print(f"Test Loss: {test_loss/len(test_loader):.4f}")
    print(f"Test Accuracy: {test_acc:.2f}% ({test_correct}/{test_total})")
    print(f"Average Entropy: {avg_entropy:.4f}")
    print(f"Average Confidence: {avg_confidence:.4f}")
    print(f"Perplexity: {perplexity:.2f}")
    
    # ODE-CCT Analysis
    print("\n🧠 ODE-CCT DYNAMICS ANALYSIS:")
    
    # Entropy distribution
    entropy_bins = [0, 0.5, 1.0, 1.5, 2.0, 3.0, 5.0, float('inf')]
    entropy_counts = [0] * (len(entropy_bins) - 1)
    
    for e in entropies:
        for i in range(len(entropy_bins) - 1):
            if entropy_bins[i] <= e < entropy_bins[i+1]:
                entropy_counts[i] += 1
                break
    
    print("\n  Entropy Distribution (lower = more confident):")
    labels = ['<0.5', '0.5-1.0', '1.0-1.5', '1.5-2.0', '2.0-3.0', '3.0-5.0', '>5.0']
    for label, count in zip(labels, entropy_counts):
        pct = 100 * count / len(entropies)
        bar = '█' * int(pct / 2)
        print(f"    {label:>8}: {bar} {pct:.1f}%")
    
    # Confidence thresholds
    high_conf = sum(1 for c in confidence_scores if c > 0.9)
    med_conf = sum(1 for c in confidence_scores if 0.5 < c <= 0.9)
    low_conf = sum(1 for c in confidence_scores if c <= 0.5)
    
    print(f"\n  Confidence Levels:")
    print(f"    High (>0.9): {100*high_conf/len(confidence_scores):.1f}%")
    print(f"    Medium (0.5-0.9): {100*med_conf/len(confidence_scores):.1f}%")
    print(f"    Low (<0.5): {100*low_conf/len(confidence_scores):.1f}%")
    
    # Sample predictions
    print("\n🔍 Sample Predictions:")
    model.eval()
    sample_count = 0
    
    with torch.no_grad():
        for images, input_indices, target_idx in test_loader:
            for i in range(min(3, len(target_idx))):
                input_words = [idx2word.get(idx.item(), '?') for idx in input_indices[i]]
                true_word = idx2word.get(target_idx[i].item(), '?')
                
                img_batch = images[i:i+1]
                logits, _ = model(img_batch, return_confidence=True)
                pred_idx = logits.argmax(dim=1).item()
                pred_word = idx2word.get(pred_idx, '?')
                
                status = "✅" if pred_idx == target_idx[i].item() else "❌"
                
                print(f"  {status} Input: {' '.join(input_words[-4:])} → | True: '{true_word}' | Pred: '{pred_word}'")
                
                sample_count += 1
                if sample_count >= 8:
                    break
            if sample_count >= 8:
                break
    
    print("="*60)
    
    return test_acc, avg_entropy, perplexity

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

def main():
    print("🧠 Word-MLP Sequence Model (ODE-CCT Language Prediction)")
    print("="*60)
    
    # Step 1: Load data
    print("\n📂 Step 1: Loading markdown files...")
    words = load_markdown_files(cfg.book_folder)
    
    # Step 2: Build vocabulary
    print("\n📚 Step 2: Building vocabulary...")
    word2idx, idx2word, vocab_words = build_vocabulary(words, max_vocab_size=2000)
    cfg.num_classes = len(vocab_words)
    
    # Step 3: Create sequence dataset
    print(f"\n🖼️ Step 3: Creating sequence dataset (seq_len={cfg.seq_len})...")
    dataset = SequenceDataset(words, word2idx, seq_len=cfg.seq_len, max_samples=40000)
    
    # Step 4: Train/Val/Test split
    print("\n📊 Step 4: Splitting data...")
    total = len(dataset)
    train_size = int(0.7 * total)
    val_size = int(0.15 * total)
    test_size = total - train_size - val_size
    
    train_dataset, val_dataset, test_dataset = torch.utils.data.random_split(
        dataset, [train_size, val_size, test_size]
    )
    
    train_loader = DataLoader(train_dataset, batch_size=cfg.batch_size, shuffle=True, collate_fn=collate_seq_batch)
    val_loader = DataLoader(val_dataset, batch_size=cfg.batch_size, shuffle=False, collate_fn=collate_seq_batch)
    test_loader = DataLoader(test_dataset, batch_size=cfg.batch_size, shuffle=False, collate_fn=collate_seq_batch)
    
    print(f"  Train: {len(train_dataset)} | Val: {len(val_dataset)} | Test: {len(test_dataset)}")
    
    # Step 5: Build model
    print("\n🏗️ Step 5: Building model...")
    model = SequencePredictor(
        embed_dim=cfg.embed_dim,
        hidden_dim=cfg.hidden_dim,
        num_classes=cfg.num_classes
    )
    print(model)
    
    # Step 6: Train
    print("\n🚀 Step 6: Training sequence model...")
    model = train_model(model, train_loader, val_loader, epochs=cfg.epochs, lr=cfg.lr)
    
    # Step 7: Test
    print("\n🧪 Step 7: Testing...")
    test_acc, avg_entropy, perplexity = test_model(model, test_loader, word2idx, idx2word)
    
    # Save
    torch.save({
        'model_state_dict': model.state_dict(),
        'word2idx': word2idx,
        'idx2word': idx2word,
        'config': {
            'embed_dim': cfg.embed_dim,
            'hidden_dim': cfg.hidden_dim,
            'num_classes': cfg.num_classes,
            'seq_len': cfg.seq_len
        }
    }, 'word_mlp_sequence.pt')
    
    print("\n💾 Model saved to 'word_mlp_sequence.pt'")
    print(f"Final Test Accuracy: {test_acc:.2f}%")
    print(f"Perplexity: {perplexity:.2f}")

if __name__ == "__main__":
    main()