"""
Word-MLP: Conv2D + Dense for Visual Word Embeddings
Loads markdown files → Renders words as 28x28 images → Trains → Tests
"""

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
from PIL import Image, ImageDraw, ImageFont
from collections import defaultdict
from pathlib import Path

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

class Config:
    # Paths
    book_folder = "./books"  # Folder containing .md files
    
    # Image
    img_size = 28
    font_size = 20
    
    # Model
    embed_dim = 128
    num_classes = None  # Set after building vocab
    
    # Training
    batch_size = 64
    epochs = 20
    lr = 0.001
    
    # Device
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

cfg = Config()

# ═══════════════════════════════════════════════════════════════════
# WORD EXTRACTION FROM MARKDOWN
# ═══════════════════════════════════════════════════════════════════

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()
        
        # Extract words: lowercase, alphabetic only
        extracted = re.findall(r'[a-zA-Z]{2,}', content)
        words.extend([w.lower() for w in extracted])
    
    if len(words) == 0:
        print(f"⚠️ No words found in {folder_path}. Using sample text.")
        return sample_words()
    
    print(f"✅ Extracted {len(words)} words from markdown files.")
    return words

def sample_words():
    """Fallback sample words if no markdown files found."""
    sample_text = """
    intelligence artificial machine learning neural network deep 
    consciousness awareness thinking reasoning planning algorithm 
    data information knowledge wisdom understanding insight logic 
    mathematics physics chemistry biology science research study
    theory hypothesis experiment observation analysis synthesis
    computer digital electronic computational programming code
    network connection relationship system structure function
    energy power strength force motion dynamics change growth
    """
    return re.findall(r'[a-zA-Z]{2,}', sample_text.lower())

# ═══════════════════════════════════════════════════════════════════
# VOCABULARY BUILDING
# ═══════════════════════════════════════════════════════════════════

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
    
    # Sort by frequency
    sorted_words = sorted(word_counts.items(), key=lambda x: -x[1])
    
    # Limit vocabulary
    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."""
    # Create white background
    img = Image.new('L', (size, size), color=255)
    draw = ImageDraw.Draw(img)
    
    # Try to use a nice font, fallback to default
    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()
    
    # Calculate text position (centered)
    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 (black on white)
    draw.text((x, y), word, fill=0, font=font)
    
    # Convert to numpy array
    img_array = np.array(img, dtype=np.float32) / 255.0
    
    # Invert: white background becomes 0, black text becomes 1
    img_array = 1.0 - img_array
    
    return img_array

# ═══════════════════════════════════════════════════════════════════
# DATASET
# ═══════════════════════════════════════════════════════════════════

class WordDataset(Dataset):
    """Dataset of word images with context co-occurrence."""
    
    def __init__(self, words, word2idx, context_window=2, max_samples=50000):
        self.samples = []
        self.word2idx = word2idx
        
        # Build context pairs
        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]
        
        # Build co-occurrence pairs
        for i, word in enumerate(valid_words):
            idx = word2idx[word]
            
            # Get context words
            start = max(0, i - context_window)
            end = min(len(valid_words), i + context_window + 1)
            
            context_indices = []
            for j in range(start, end):
                if j != i and valid_words[j] in word2idx:
                    context_indices.append(word2idx[valid_words[j]])
            
            # Also add some negative samples (random words)
            neg_indices = []
            for _ in range(min(5, len(context_indices))):
                neg_word = random.choice(list(word2idx.keys()))
                neg_indices.append(word2idx[neg_word])
            
            # Render word image
            try:
                img = render_word_image(word)
            except:
                img = np.zeros((28, 28), dtype=np.float32)
            
            self.samples.append({
                'image': img,
                'word_idx': idx,
                'context_indices': context_indices,
                'neg_indices': neg_indices
            })
        
        print(f"🖼️ Dataset size: {len(self.samples)} samples")
    
    def __len__(self):
        return len(self.samples)
    
    def __getitem__(self, idx):
        sample = self.samples[idx]
        return (
            torch.FloatTensor(sample['image']).unsqueeze(0),  # [1, 28, 28]
            torch.LongTensor([sample['word_idx']]),            # [1]
            torch.LongTensor(sample['context_indices']),       # [N]
            torch.LongTensor(sample['neg_indices'])            # [N]
        )


def collate_word_batch(batch):
    """Pad variable-length context and negative lists within a batch."""
    images, word_idx, context_idx, neg_idx = zip(*batch)

    images = torch.stack(images, dim=0)
    word_idx = torch.stack(word_idx, dim=0)

    if any(t.numel() > 0 for t in context_idx):
        context_idx = pad_sequence(context_idx, batch_first=True, padding_value=-1)
    else:
        context_idx = torch.empty((len(batch), 0), dtype=torch.long)

    if any(t.numel() > 0 for t in neg_idx):
        neg_idx = pad_sequence(neg_idx, batch_first=True, padding_value=-1)
    else:
        neg_idx = torch.empty((len(batch), 0), dtype=torch.long)

    return images, word_idx, context_idx, neg_idx

# ═══════════════════════════════════════════════════════════════════
# MODEL: Conv2D + Dense
# ═══════════════════════════════════════════════════════════════════

class WordMLP(nn.Module):
    """Conv2D + Dense model for word embeddings."""
    
    def __init__(self, num_classes, embed_dim=128):
        super().__init__()
        
        # Convolutional feature extraction
        self.features = nn.Sequential(
            # Block 1
            nn.Conv2d(1, 32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2),  # 28 -> 14
            
            # Block 2
            nn.Conv2d(32, 64, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2),  # 14 -> 7
            
            # Block 3
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2),  # 7 -> 3
        )
        
        # Calculate flattened size: 128 * 3 * 3 = 1152
        self.flatten = nn.Flatten()
        
        # Embedding layer
        self.embedding = nn.Linear(128 * 3 * 3, embed_dim)
        
        # Classification head
        self.classifier = nn.Linear(embed_dim, num_classes)
    
    def forward(self, x):
        # Conv features
        x = self.features(x)
        x = self.flatten(x)
        
        # Embedding
        embed = self.embedding(x)
        
        # Classification
        logits = self.classifier(embed)
        
        return logits, embed

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

def train_model(model, train_loader, val_loader, epochs=20, lr=0.001):
    """Train the model with classification + context loss."""
    
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    criterion = nn.CrossEntropyLoss()
    
    model.to(cfg.device)
    
    for epoch in range(epochs):
        # Training
        model.train()
        train_loss = 0.0
        train_correct = 0
        train_total = 0
        
        for batch_idx, (images, word_idx, context_idx, neg_idx) in enumerate(train_loader):
            images = images.to(cfg.device)
            word_idx = word_idx.squeeze(-1).to(cfg.device)
            
            optimizer.zero_grad()
            
            logits, embeddings = model(images)
            
            # Classification loss
            cls_loss = criterion(logits, word_idx)
            
            # Context loss (contrastive)
            if context_idx.shape[1] > 0:
                context_loss = compute_context_loss(embeddings, word_idx, context_idx.to(cfg.device), neg_idx.to(cfg.device))
            else:
                context_loss = torch.tensor(0.0).to(cfg.device)
            
            # Combined loss
            loss = cls_loss + 0.2 * context_loss
            
            loss.backward()
            optimizer.step()
            
            train_loss += loss.item()
            preds = logits.argmax(dim=1)
            train_correct += (preds == word_idx).sum().item()
            train_total += word_idx.size(0)
        
        train_acc = 100 * train_correct / train_total
        avg_train_loss = train_loss / len(train_loader)
        
        # Validation
        model.eval()
        val_loss = 0.0
        val_correct = 0
        val_total = 0
        
        with torch.no_grad():
            for images, word_idx, context_idx, neg_idx in val_loader:
                images = images.to(cfg.device)
                word_idx = word_idx.squeeze(-1).to(cfg.device)
                
                logits, embeddings = model(images)
                
                loss = criterion(logits, word_idx)
                val_loss += loss.item()
                
                preds = logits.argmax(dim=1)
                val_correct += (preds == word_idx).sum().item()
                val_total += word_idx.size(0)
        
        val_acc = 100 * val_correct / val_total
        avg_val_loss = val_loss / len(val_loader)
        
        print(f"Epoch {epoch+1}/{epochs} | "
              f"Train Loss: {avg_train_loss:.4f} Acc: {train_acc:.2f}% | "
              f"Val Loss: {avg_val_loss:.4f} Acc: {val_acc:.2f}%")
    
    return model

def compute_context_loss(embeddings, word_idx, context_idx, neg_idx, margin=1.0):
    """
    Contrastive loss over in-batch matches for context and negative word ids.
    """
    batch_size = embeddings.size(0)
    
    # Normalize embeddings
    embeddings = F.normalize(embeddings, dim=1)
    device = embeddings.device

    total_loss = torch.tensor(0.0, device=device)
    label_to_positions = defaultdict(list)
    for i, label in enumerate(word_idx.tolist()):
        label_to_positions[int(label)].append(i)
    
    for i in range(batch_size):
        emb = embeddings[i]  # [embed_dim]
        
        # Positive context loss
        pos_loss = torch.tensor(0.0, device=device)
        for ctx_idx in context_idx[i]:
            ctx_label = int(ctx_idx.item())
            if ctx_label < 0:
                continue
            for j in label_to_positions.get(ctx_label, []):
                if j == i:
                    continue
                pos_loss = pos_loss + F.pairwise_distance(
                    emb.unsqueeze(0),
                    embeddings[j].unsqueeze(0),
                ).squeeze()
        
        # Negative margin loss
        neg_loss = torch.tensor(0.0, device=device)
        for neg_label in neg_idx[i]:
            neg_label = int(neg_label.item())
            if neg_label < 0:
                continue
            for j in label_to_positions.get(neg_label, []):
                if j == i:
                    continue
                neg_loss = neg_loss + F.relu(
                    margin - F.pairwise_distance(
                        emb.unsqueeze(0),
                        embeddings[j].unsqueeze(0),
                    ).squeeze()
                )
        
        total_loss += pos_loss + neg_loss
    
    return total_loss / max(batch_size, 1)

# ═══════════════════════════════════════════════════════════════════
# TESTING
# ═══════════════════════════════════════════════════════════════════

def test_model(model, test_loader, word2idx, idx2word):
    """Test the model and show qualitative results."""
    
    model.eval()
    criterion = nn.CrossEntropyLoss()
    
    test_loss = 0.0
    test_correct = 0
    test_total = 0
    
    # Per-class accuracy
    class_correct = defaultdict(int)
    class_total = defaultdict(int)
    
    with torch.no_grad():
        for images, word_idx, context_idx, neg_idx in test_loader:
            images = images.to(cfg.device)
            word_idx = word_idx.squeeze(-1).to(cfg.device)
            
            logits, embeddings = model(images)
            
            loss = criterion(logits, word_idx)
            test_loss += loss.item()
            
            preds = logits.argmax(dim=1)
            test_correct += (preds == word_idx).sum().item()
            test_total += word_idx.size(0)
            
            # Per-class tracking
            for pred, true in zip(preds.cpu(), word_idx.cpu()):
                class_total[true.item()] += 1
                if pred == true:
                    class_correct[true.item()] += 1
    
    test_acc = 100 * test_correct / test_total
    avg_test_loss = test_loss / len(test_loader)
    
    print("\n" + "="*60)
    print("📊 TEST RESULTS")
    print("="*60)
    print(f"Test Loss: {avg_test_loss:.4f}")
    print(f"Test Accuracy: {test_acc:.2f}% ({test_correct}/{test_total})")
    
    # Top/Bottom classes
    print("\n📈 Per-Class Accuracy:")
    class_acc = {k: class_correct[k]/max(class_total[k],1) for k in class_total}
    sorted_acc = sorted(class_acc.items(), key=lambda x: -x[1])
    
    print("  Top 5:")
    for idx, acc in sorted_acc[:5]:
        word = idx2word.get(idx, f"<idx:{idx}>")
        print(f"    {word}: {acc*100:.1f}% ({class_correct[idx]}/{class_total[idx]})")
    
    print("  Bottom 5:")
    for idx, acc in sorted_acc[-5:]:
        word = idx2word.get(idx, f"<idx:{idx}>")
        print(f"    {word}: {acc*100:.1f}% ({class_correct[idx]}/{class_total[idx]})")
    
    # Qualitative examples
    print("\n🔍 Sample Predictions:")
    model.eval()
    sample_count = 0
    
    with torch.no_grad():
        for images, word_idx, context_idx, neg_idx in test_loader:
            for i in range(min(5, images.size(0))):
                img = images[i:i+1].to(cfg.device)
                true_idx = word_idx[i].item()
                
                logits, _ = model(img)
                pred_idx = logits.argmax(dim=1).item()
                
                true_word = idx2word.get(true_idx, f"<idx:{true_idx}>")
                pred_word = idx2word.get(pred_idx, f"<idx:{pred_idx}>")
                status = "✅" if pred_idx == true_idx else "❌"
                
                print(f"  {status} True: '{true_word}' | Pred: '{pred_word}'")
                sample_count += 1
                if sample_count >= 10:
                    break
            if sample_count >= 10:
                break
    
    # Embedding similarity test
    print("\n🔗 Embedding Similarity (semantically similar words):")
    test_words = ['the', 'and', 'is', 'are', 'was', 'were']
    available_words = [w for w in test_words if w in word2idx]
    
    if len(available_words) >= 2:
        model.eval()
        with torch.no_grad():
            for word in available_words[:4]:
                img = torch.FloatTensor(render_word_image(word)).unsqueeze(0).unsqueeze(0).to(cfg.device)
                _, emb = model(img)
                
                # Find closest words
                distances = []
                for w, idx in word2idx.items():
                    if w != word:
                        img2 = torch.FloatTensor(render_word_image(w)).unsqueeze(0).unsqueeze(0).to(cfg.device)
                        _, emb2 = model(img2)
                        dist = F.cosine_similarity(emb, emb2).item()
                        distances.append((w, dist))
                
                distances.sort(key=lambda x: -x[1])
                similar = [w for w, d in distances[:3]]
                print(f"  '{word}' → {similar}")
    
    print("="*60)
    
    return test_acc

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

def main():
    print("🧠 Word-MLP: Conv2D + Dense Word Embeddings")
    print("="*60)
    
    # Step 1: Load markdown files
    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 dataset
    print("\n🖼️ Step 3: Creating dataset...")
    dataset = WordDataset(words, word2idx, context_window=2, max_samples=30000)
    
    # 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_word_batch)
    val_loader = DataLoader(val_dataset, batch_size=cfg.batch_size, shuffle=False, collate_fn=collate_word_batch)
    test_loader = DataLoader(test_dataset, batch_size=cfg.batch_size, shuffle=False, collate_fn=collate_word_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 = WordMLP(num_classes=cfg.num_classes, embed_dim=cfg.embed_dim)
    print(model)
    
    # Step 6: Train
    print("\n🚀 Step 6: Training...")
    model = train_model(model, train_loader, val_loader, epochs=cfg.epochs, lr=cfg.lr)
    
    # Step 7: Test
    print("\n🧪 Step 7: Testing...")
    test_acc = test_model(model, test_loader, word2idx, idx2word)
    
    # Save model
    torch.save({
        'model_state_dict': model.state_dict(),
        'word2idx': word2idx,
        'idx2word': idx2word,
        'config': {
            'num_classes': cfg.num_classes,
            'embed_dim': cfg.embed_dim
        }
    }, 'word_mlp_model.pt')
    
    print("\n💾 Model saved to 'word_mlp_model.pt'")
    print(f"Final Test Accuracy: {test_acc:.2f}%")

if __name__ == "__main__":
    main()