import torch
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt
import math
import numpy as np

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# THE ULTIMATE CORPUS: Balanced, Contrastive, and complete with the daughters
corpus = [
    # 1. Direct Gender Contrast (Rips man and woman apart)
    "man is male", "woman is female",
    "king is male", "queen is female",
    "palash is male", "totan is female",
    "boy is male", "girl is female",
    
    # 2. Structural Role Parallels
    "the man is the king", "the woman is the queen",
    "palash is the king", "totan is the queen",
    "palash is a man", "totan is a woman",
    
    # 3. Spousal Contexts
    "palash is husband", "totan is wife",
    "king loves queen", "queen loves king",
    "palash loves totan", "totan loves palash",
    
    # 4. The Daughters (Prapti & Tripti)
    "prapti is girl", "tripti is girl",
    "prapti is princess", "tripti is princess",
    "prapti is daughter", "tripti is daughter",
    "prapti and tripti are sisters",
    
    # 5. Parent-Child Visual Links
    "palash father prapti", "palash father tripti",
    "totan mother prapti", "totan mother tripti",
    "king father princess", "queen mother princess"
]

# Tokenization
words = " ".join(corpus).split()
vocab = list(set(words))
word_to_ix = {w: i for i, w in enumerate(vocab)}
vocab_size = len(vocab)

# Window Context frame generation
data = []
for sentence in corpus:
    tokens = sentence.split()
    for i in range(1, len(tokens) - 1):
        context = [word_to_ix[tokens[i-1]], word_to_ix[tokens[i+1]]]
        target = word_to_ix[tokens[i]]
        data.append((context, target))

x_train = torch.tensor([pair[0] for pair in data], dtype=torch.long).to(device)
y_train = torch.tensor([pair[1] for pair in data], dtype=torch.long).to(device)

# Model
class UltimateEmbeddingModel(nn.Module):
    def __init__(self, vocab_size, d_model=4):
        super(UltimateEmbeddingModel, self).__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.fc_out = nn.Linear(d_model, vocab_size)

    def forward(self, x):
        embedded = self.embedding(x)
        # Unit-sphere norm keeping features perfectly separated
        embedded = embedded / (embedded.norm(dim=-1, keepdim=True) + 1e-8)
        return self.fc_out(embedded.mean(dim=1))

model = UltimateEmbeddingModel(vocab_size=vocab_size, d_model=4).to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.005)

def apply_pca(X):
    X_mean = X - np.mean(X, axis=0)
    cov = np.cov(X_mean.T)
    eigenvalues, eigenvectors = np.linalg.eigh(cov)
    idx = np.argsort(eigenvalues)[::-1]
    return np.dot(X_mean, eigenvectors[:, idx[:2]])

# Plot Config
plt.ion()
fig, ax = plt.subplots(figsize=(12, 10))
target_words = ["man", "woman", "king", "queen", "palash", "totan", "prapti", "tripti", "princess", "girl"]

# Global connections map to look for parallel lines
conn_groups = [
    ("man", "king", 'blue', '--'), ("woman", "queen", 'red', '--'),
    ("palash", "king", 'purple', '-'), ("totan", "queen", 'purple', '-'),
    ("palash", "man", 'blue', ':'), ("totan", "woman", 'red', ':'),
    ("prapti", "princess", 'orange', '-'), ("tripti", "princess", 'orange', '-'),
    ("girl", "princess", 'magenta', '--'),
    ("palash", "prapti", 'green', ':'), ("totan", "tripti", 'green', ':')
]

for epoch in range(400001):
    model.train()
    optimizer.zero_grad()
    
    loss = criterion(model(x_train), y_train)
    loss.backward()
    optimizer.step()
    
    if epoch % 100 == 0:
        ax.clear()
        with torch.no_grad():
            raw_embeds = model.embedding.weight
            norm_embeds = (raw_embeds / (raw_embeds.norm(dim=-1, keepdim=True) + 1e-8)).cpu().numpy()
            embeddings_2d = apply_pca(norm_embeds)
        
        for word in target_words:
            if word in word_to_ix:
                idx = word_to_ix[word]
                vec = embeddings_2d[idx]
                color = 'blue' if word in ['man', 'king', 'palash'] else 'red' if word in ['woman', 'queen', 'totan'] else 'green'
                ax.scatter(vec[0], vec[1], c=color, s=150, edgecolors='k', zorder=5)
                ax.annotate(word, (vec[0], vec[1]), xytext=(7, 7), textcoords='offset points', fontsize=11, fontweight='bold')
        
        for w1, w2, color, style in conn_groups:
            if w1 in word_to_ix and w2 in word_to_ix:
                v1, v2 = embeddings_2d[word_to_ix[w1]], embeddings_2d[word_to_ix[w2]]
                ax.plot([v1[0], v2[0]], [v1[1], v2[1]], color=color, linestyle=style, alpha=0.5, linewidth=2)

        ax.set_title(f"Ultimate Family Semantic Space (PCA 4D->2D)\nEpoch {epoch} | Loss: {loss.item():.4f}")
        ax.grid(True, linestyle='--', alpha=0.4)
        ax.margins(0.2)
        plt.pause(0.01)

plt.ioff()
plt.show()