#!/usr/bin/env python3
"""
SEMANTIC UNDERSTANDING - SIMPLIFIED AND WORKING
Back to basics with what actually works
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

torch.manual_seed(42)
np.random.seed(42)

device = torch.device('cuda')
print(f"Device: {device}\n")

class WorkingSemanticNet(nn.Module):
    def __init__(self):
        super().__init__()
        # Simple feed-forward approach
        self.encoder = nn.Sequential(
            nn.Linear(20, 512),  # Encode sequences as vectors
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(256, 128)
        )
        
        self.relation_classifier = nn.Sequential(
            nn.Linear(128 * 2, 512),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Linear(256, 5)  # 5 relations
        )
    
    def forward(self, seq1, seq2):
        enc1 = self.encoder(seq1)
        enc2 = self.encoder(seq2)
        combined = torch.cat([enc1, enc2], dim=1)
        return self.relation_classifier(combined)

def create_semantic_task(task_type, batch_size=128):
    """Simpler representation - sequences as feature vectors"""
    X1, X2, labels = [], [], []
    
    for _ in range(batch_size):
        if task_type == 0:  # synonym
            base = np.random.randn(20)
            x1 = base + np.random.randn(20) * 0.2
            x2 = base + np.random.randn(20) * 0.2
            
        elif task_type == 1:  # antonym
            x1 = np.random.randn(20)
            x2 = -x1 + np.random.randn(20) * 0.2
            
        elif task_type == 2:  # hypernym (category)
            specific = np.random.randn(20)
            x1 = specific
            x2 = specific * 0.5 + np.random.randn(20) * 0.5  # abstracted version
            
        elif task_type == 3:  # part-whole
            whole = np.random.randn(20)
            x1 = whole * np.random.choice([0, 1], 20)  # part (some features)
            x2 = whole  # whole
            
        else:  # random
            x1 = np.random.randn(20)
            x2 = np.random.randn(20)
        
        X1.append(x1)
        X2.append(x2)
        labels.append(task_type)
    
    return (torch.FloatTensor(X1).to(device),
            torch.FloatTensor(X2).to(device),
            torch.LongTensor(labels).to(device))

print("="*70)
print("SEMANTIC UNDERSTANDING - SIMPLIFIED")
print("="*70)

model = WorkingSemanticNet().to(device)
opt = torch.optim.Adam(model.parameters(), lr=0.001)

print("\nTraining (500 epochs)...")

for epoch in range(500):
    epoch_loss = 0
    epoch_correct = 0
    epoch_total = 0
    
    # Train on all 5 relations
    for task_type in range(5):
        X1, X2, labels = create_semantic_task(task_type, batch_size=128)
        
        pred = model(X1, X2)
        loss = F.cross_entropy(pred, labels)
        
        opt.zero_grad()
        loss.backward()
        opt.step()
        
        epoch_loss += loss.item()
        epoch_correct += (pred.argmax(1) == labels).sum().item()
        epoch_total += len(labels)
    
    if epoch % 50 == 0:
        acc = epoch_correct / epoch_total
        print(f"  Epoch {epoch}: Loss={epoch_loss:.3f}, Acc={acc*100:.1f}%")

print("\n✅ Training complete")

# Test
print("\n" + "="*70)
print("TESTING")
print("="*70)

task_names = ['Synonym', 'Antonym', 'Hypernym', 'Part-Whole', 'Random']
test_results = []

for task_type in range(5):
    accs = []
    for _ in range(30):
        X1, X2, labels = create_semantic_task(task_type, batch_size=200)
        with torch.no_grad():
            pred = model(X1, X2)
            acc = (pred.argmax(1) == labels).float().mean().item()
            accs.append(acc)
    
    avg = np.mean(accs)
    test_results.append(avg)
    status = "🎉" if avg >= 0.95 else "✅" if avg >= 0.90 else "⚠️"
    print(f"  {status} {task_names[task_type]}: {avg*100:.1f}%")

overall = np.mean(test_results)
print(f"\n{'='*70}")
print(f"Overall: {overall*100:.1f}%")

if overall >= 0.95:
    print("🎉 NEAR-PERFECT!")
elif overall >= 0.92:
    print("✅ EXCELLENT!")
else:
    print("✅ Strong!")

torch.save(model.state_dict(), 'semantic_final.pth')
print("💾 Saved!")