"""
Word-MLP: Conv2D + Dense for Visual Word Embeddings
Skip-gram context head for stronger semantic relationships
"""

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"
    
    # Image
    img_size = 28
    font_size = 20
    
    # Model
    embed_dim = 128
    num_classes = None
    num_negatives = 5  # Negative samples for skip-gram
    
    # Training
    batch_size = 64
    epochs = 20
    lr = 0.001
    
    # Loss weights
    cls_weight = 0.5
    skipgram_weight = 0.5
    
    # 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()
        
        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
    """
    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
    
    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 WITH SKIP-GRAM PAIRS
# ═══════════════════════════════════════════════════════════════════

class WordDataset(Dataset):
    """Dataset with word images + skip-gram context pairs."""
    
    def __init__(self, words, word2idx, context_window=2, max_samples=50000):
        self.samples = []
        self.word2idx = word2idx
        self.num_classes = len(word2idx)
        
        vocab_set = set(word2idx.keys())
        valid_words = [w for w in words if w in vocab_set]
        
        if len(valid_words) > max_samples:
            valid_words = valid_words[:max_samples]
        
        # Pre-compute all embeddings cache for skip-gram loss
        self.words = valid_words
        self.word2idx_list = [word2idx[w] for w in valid_words]
        
        # Build samples
        for i, word in enumerate(valid_words):
            idx = word2idx[word]
            
            # Context words within window
            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:
                    context_indices.append(word2idx[valid_words[j]])
            
            # Random negative samples
            neg_indices = []
            for _ in range(cfg.num_negatives):
                neg_word = random.choice(list(word2idx.keys()))
                neg_indices.append(word2idx[neg_word])
            
            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),
            torch.LongTensor([sample['word_idx']]),
            torch.LongTensor(sample['context_indices']),
            torch.LongTensor(sample['neg_indices'])
        )


def collate_word_batch(batch):
    """Pad variable-length context and negative lists."""
    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 with Skip-Gram Head
# ═══════════════════════════════════════════════════════════════════

class WordMLP(nn.Module):
    """
    Conv2D + Dense model with two heads:
    1. Classification head (word identity)
    2. Skip-gram context head (predict context words)
    """
    
    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
        )
        
        self.flatten = nn.Flatten()
        
        # Shared embedding layer
        self.embedding = nn.Linear(128 * 3 * 3, embed_dim)
        
        # === HEAD 1: Classification ===
        # Predicts which word the image represents
        self.classifier = nn.Linear(embed_dim, num_classes)
        
        # === HEAD 2: Skip-Gram Context ===
        # Predicts context words from the embedding
        # Uses same output space as embeddings (word2vec style)
        self.context_predictor = nn.Sequential(
            nn.Linear(embed_dim, embed_dim),
            nn.ReLU(),
            nn.Linear(embed_dim, embed_dim)
        )
        
        # Word embeddings lookup (for skip-gram training)
        self.word_embeddings = nn.Embedding(num_classes, embed_dim)
        
        # Initialize embeddings
        nn.init.uniform_(self.word_embeddings.weight, -0.5/embed_dim, 0.5/embed_dim)
    
    def forward(self, x, return_context_embeds=False):
        # Conv features
        x = self.features(x)
        x = self.flatten(x)
        
        # Shared embedding
        embed = self.embedding(x)
        
        # Classification logits
        cls_logits = self.classifier(embed)
        
        # Context prediction (skip-gram)
        context_pred = self.context_predictor(embed)
        
        if return_context_embeds:
            # Return word embeddings for skip-gram loss
            word_embeds = self.word_embeddings.weight
            return cls_logits, embed, context_pred, word_embeds
        
        return cls_logits, embed, context_pred

# ═══════════════════════════════════════════════════════════════════
# SKIP-GRAM LOSS (NEGATIVE SAMPLING)
# ═══════════════════════════════════════════════════════════════════

def skipgram_loss(embeddings, context_pred, word_embeddings, word_idx, context_indices, neg_indices):
    """
    Skip-gram loss with negative sampling.
    
    Maximize: sim(target, context_word) for positive pairs
    Minimize: sim(target, random_word) for negative pairs
    
    Uses in-batch negatives for efficiency.
    """
    batch_size = embeddings.size(0)
    embed_dim = embeddings.size(1)
    device = embeddings.device
    
    # Normalize embeddings for cosine similarity
    target = F.normalize(embeddings, dim=1)
    context_pred_norm = F.normalize(context_pred, dim=1)
    
    # Word embeddings for context words
    word_embeds = F.normalize(word_embeddings.weight, dim=1)
    
    # Collect positive and negative pairs
    pos_pairs = []
    neg_pairs = []
    
    for i in range(batch_size):
        target_vec = target[i]  # [embed_dim]
        
        # Positive context words
        for ctx_idx in context_indices[i]:
            ctx_idx = int(ctx_idx.item())
            if ctx_idx >= 0:
                ctx_vec = word_embeds[ctx_idx]  # [embed_dim]
                pos_pairs.append((target_vec, ctx_vec))
        
        # Negative samples
        for neg_idx in neg_indices[i]:
            neg_idx = int(neg_idx.item())
            if neg_idx >= 0:
                neg_vec = word_embeds[neg_idx]
                neg_pairs.append((target_vec, neg_vec))
    
    if len(pos_pairs) == 0:
        return torch.tensor(0.0, device=device)
    
    # Compute positive scores
    pos_loss = 0.0
    for tgt, ctx in pos_pairs:
        score = torch.dot(tgt, ctx)
        pos_loss += -torch.log(torch.sigmoid(score) + 1e-8)
    
    # Compute negative scores
    neg_loss = 0.0
    for tgt, neg in neg_pairs:
        score = torch.dot(tgt, neg)
        neg_loss += -torch.log(torch.sigmoid(-score) + 1e-8)
    
    total_loss = (pos_loss + neg_loss) / max(len(pos_pairs) + len(neg_pairs), 1)
    
    return total_loss


def inbatch_skipgram_loss(embeddings, context_pred, word_embeddings, word_idx_batch, num_classes):
    """
    In-batch skip-gram loss with efficient negative sampling.
    
    For each target, the other words in the batch (with different labels)
    serve as negative samples.
    """
    batch_size = embeddings.size(0)
    embed_dim = embeddings.size(1)
    device = embeddings.device
    
    # Normalize
    target = F.normalize(embeddings, dim=1)  # [batch, embed]
    context_pred = F.normalize(context_pred, dim=1)  # [batch, embed]
    
    # Word embeddings matrix
    word_embeds = F.normalize(word_embeddings.weight, dim=1)  # [vocab, embed]
    
    # Get embeddings for context words (from in-batch)
    context_embeds = word_embeds[word_idx_batch]  # [batch, embed]
    context_embeds = F.normalize(context_embeds, dim=1)
    
    # Positive: target dot context prediction
    pos_scores = torch.sum(target * context_pred, dim=1)  # [batch]
    pos_loss = -torch.log(torch.sigmoid(pos_scores) + 1e-8).mean()
    
    # Negative: target dot random words in batch (different labels)
    neg_loss = 0.0
    num_negatives = 5
    
    for i in range(batch_size):
        tgt = target[i]
        label = word_idx_batch[i].item()
        
        # Find negatives: other words in batch with different label
        neg_indices = []
        for j in range(batch_size):
            if word_idx_batch[j].item() != label:
                neg_indices.append(j)
        
        if len(neg_indices) > 0:
            # Sample random negatives
            neg_indices = random.sample(neg_indices, min(num_negatives, len(neg_indices)))
            
            for j in neg_indices:
                neg_vec = target[j]  # Use target embedding as negative sample
                score = torch.dot(tgt, neg_vec)
                neg_loss += torch.log(torch.sigmoid(score) + 1e-8)
    
    if batch_size > 1:
        neg_loss = -neg_loss / (batch_size * num_negatives)
    else:
        neg_loss = torch.tensor(0.0, device=device)
    
    return pos_loss + neg_loss

# ═══════════════════════════════════════════════════════════════════
# TRAINING WITH COMBINED LOSS
# ═══════════════════════════════════════════════════════════════════

def train_model(model, train_loader, val_loader, epochs=20, lr=0.001):
    """Train with classification + skip-gram losses."""
    
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    cls_criterion = nn.CrossEntropyLoss()
    
    model.to(cfg.device)
    
    for epoch in range(epochs):
        # Training
        model.train()
        train_loss = 0.0
        train_cls_loss = 0.0
        train_sg_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)
            context_idx = context_idx.to(cfg.device)
            neg_idx = neg_idx.to(cfg.device)
            
            optimizer.zero_grad()
            
            # Forward pass
            cls_logits, embeddings, context_pred = model(images)
            
            # === Loss 1: Classification ===
            cls_loss = cls_criterion(cls_logits, word_idx)
            
            # === Loss 2: Skip-gram ===
            sg_loss = skipgram_loss(
                embeddings,
                context_pred,
                model.word_embeddings,
                word_idx,
                context_idx,
                neg_idx,
            )
            
            # Combined loss
            loss = cfg.cls_weight * cls_loss + cfg.skipgram_weight * sg_loss
            
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            optimizer.step()
            
            train_loss += loss.item()
            train_cls_loss += cls_loss.item()
            train_sg_loss += sg_loss.item()
            
            preds = cls_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)
        avg_cls_loss = train_cls_loss / len(train_loader)
        avg_sg_loss = train_sg_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)
                
                cls_logits, _, _ = model(images)
                
                loss = cls_criterion(cls_logits, word_idx)
                val_loss += loss.item()
                
                preds = cls_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"Loss: {avg_train_loss:.4f} (Cls: {avg_cls_loss:.4f} + SG: {avg_sg_loss:.4f}) | "
              f"Train: {train_acc:.2f}% | Val: {val_acc:.2f}%")
    
    return model

# ═══════════════════════════════════════════════════════════════════
# TESTING WITH SEMANTIC ANALYSIS
# ═══════════════════════════════════════════════════════════════════

def test_model(model, test_loader, word2idx, idx2word):
    """Test and analyze learned semantic relationships."""
    
    model.eval()
    criterion = nn.CrossEntropyLoss()
    
    test_loss = 0.0
    test_correct = 0
    test_total = 0
    
    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)
            
            cls_logits, _, _ = model(images)
            
            loss = criterion(cls_logits, word_idx)
            test_loss += loss.item()
            
            preds = cls_logits.argmax(dim=1)
            test_correct += (preds == word_idx).sum().item()
            test_total += word_idx.size(0)
            
            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})")
    
    # 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("\n📈 Top 5:")
    for idx, acc in sorted_acc[:5]:
        word = idx2word.get(idx, f"<idx:{idx}>")
        print(f"  {word}: {acc*100:.1f}%")
    
    print("\n📉 Bottom 5:")
    for idx, acc in sorted_acc[-5:]:
        word = idx2word.get(idx, f"<idx:{idx}>")
        print(f"  {word}: {acc*100:.1f}%")
    
    # Semantic similarity analysis
    print("\n🔗 SKIP-GRAM SEMANTIC CLUSTERS:")
    print("Testing if words appearing in similar contexts cluster together...\n")
    
    model.eval()
    with torch.no_grad():
        # Test word categories
        test_groups = {
            "articles": ['the', 'a', 'an'],
            "conjunctions": ['and', 'or', 'but'],
            "prepositions": ['in', 'on', 'at', 'to', 'from'],
            "verbs_be": ['is', 'are', 'was', 'were', 'be', 'been'],
            "common": ['and', 'the', 'is', 'to', 'of']
        }
        
        for group_name, test_words in test_groups.items():
            valid_words = [w for w in test_words if w in word2idx]
            if len(valid_words) < 2:
                continue
            
            print(f"  [{group_name}]: {valid_words}")
            
            # Get embeddings
            embeds = []
            for word in valid_words:
                img = torch.FloatTensor(render_word_image(word)).unsqueeze(0).unsqueeze(0).to(cfg.device)
                _, embed, _ = model(img)
                embeds.append(F.normalize(embed, dim=1).squeeze().cpu().numpy())
            
            # Compute pairwise similarities
            for i, w1 in enumerate(valid_words):
                for j, w2 in enumerate(valid_words):
                    if j > i:
                        sim = np.dot(embeds[i], embeds[j])
                        status = "✅" if sim > 0.5 else "⚠️"
                        print(f"    {status} '{w1}' <-> '{w2}': {sim:.3f}")
    
    # Find semantically similar words (skip-gram space)
    print("\n🔍 Top Similar Words (Skip-Gram Embedding Space):")
    
    test_words = ['and', 'the', 'is', 'of', 'to']
    for word in test_words:
        if word not in word2idx:
            continue
        
        img = torch.FloatTensor(render_word_image(word)).unsqueeze(0).unsqueeze(0).to(cfg.device)
        _, embed, _ = model(img)
        embed = F.normalize(embed, dim=1)
        
        # Get context prediction
        _, _, context_pred = model(img)
        context_pred = F.normalize(context_pred, dim=1)
        
        # Compare with word embeddings
        word_embeds = F.normalize(model.word_embeddings.weight, dim=1)
        similarities = torch.mm(context_pred, word_embeds.T).squeeze()
        
        # Get top-k similar
        top_k = 5
        _, top_indices = similarities.topk(top_k + 1)
        
        similar_words = []
        for idx in top_indices:
            idx = idx.item()
            if idx != word2idx[word]:
                similar_words.append(idx2word.get(idx, f"idx:{idx}"))
        
        print(f"  '{word}' → {similar_words[:top_k]}")
    
    print("="*60)
    
    return test_acc

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

def main():
    print("🧠 Word-MLP with Skip-Gram Context Head")
    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 (Classification + Skip-Gram)...")
    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_skipgram.pt')
    
    print("\n💾 Model saved to 'word_mlp_skipgram.pt'")
    print(f"Final Test Accuracy: {test_acc:.2f}%")

if __name__ == "__main__":
    main()