import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
from sklearn.datasets import make_classification
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, f1_score
import numpy as np

# ==========================================
# 1. Tabular ResNet Implementation
# ==========================================

class ResNetBlock(nn.Module):
    def __init__(self, dim, dropout=0.1):
        super().__init__()
        self.block = nn.Sequential(
            nn.Linear(dim, dim),
            nn.ReLU(),
            nn.Dropout(dropout)
        )

    def forward(self, x):
        return x + self.block(x)

class TabularResNet(nn.Module):
    def __init__(self, input_dim, hidden_dim, num_blocks, output_dim, dropout=0.1):
        super().__init__()
        self.input_proj = nn.Linear(input_dim, hidden_dim)
        self.res_blocks = nn.Sequential(
            *[ResNetBlock(hidden_dim, dropout) for _ in range(num_blocks)]
        )
        self.output_head = nn.Linear(hidden_dim, output_dim)

    def forward(self, x):
        x = torch.relu(self.input_proj(x))
        x = self.res_blocks(x)
        return self.output_head(x)

# ==========================================
# 2. FT-Transformer Implementation
# ==========================================

class FeatureTokenizer(nn.Module):
    def __init__(self, num_features, embed_dim):
        super().__init__()
        self.tokenizers = nn.ModuleList([nn.Linear(1, embed_dim) for _ in range(num_features)])

    def forward(self, x):
        # Project each feature independently: (batch, num_features) -> (batch, num_features, embed_dim)
        tokens = [tokenizer(x[:, i].unsqueeze(1)) for i, tokenizer in enumerate(self.tokenizers)]
        return torch.stack(tokens, dim=1)

class FTTransformer(nn.Module):
    def __init__(self, num_features, embed_dim, num_heads, num_layers, output_dim, dropout=0.1):
        super().__init__()
        self.tokenizer = FeatureTokenizer(num_features, embed_dim)
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=embed_dim, nhead=num_heads, dim_feedforward=embed_dim * 4, 
            dropout=dropout, batch_first=True
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        self.mlp = nn.Sequential(
            nn.Linear(embed_dim, embed_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(embed_dim, output_dim)
        )

    def forward(self, x):
        x = self.tokenizer(x)
        x = self.transformer(x)
        x = torch.mean(x, dim=1) # Global Average Pooling
        return self.mlp(x)

# ==========================================
# 3. Training & Evaluation Pipeline
# ==========================================

def train_and_eval():
    # Data Setup
    X, y = make_classification(n_samples=2000, n_features=20, random_state=42)
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
    
    scaler = StandardScaler()
    X_train = torch.FloatTensor(scaler.fit_transform(X_train))
    X_test = torch.FloatTensor(scaler.transform(X_test))
    y_train, y_test = torch.LongTensor(y_train), torch.LongTensor(y_test)
    
    train_loader = DataLoader(TensorDataset(X_train, y_train), batch_size=64, shuffle=True)
    test_loader = DataLoader(TensorDataset(X_test, y_test), batch_size=64)

    # Model Initialization
    models = {
        "ResNet": TabularResNet(20, 64, 3, 2),
        "FT-Transformer": FTTransformer(20, 32, 4, 3, 2)
    }

    for name, model in models.items():
        optimizer = optim.Adam(model.parameters(), lr=1e-3)
        criterion = nn.CrossEntropyLoss()
        
        # Simple Training Loop
        model.train()
        for epoch in range(20):
            for bx, by in train_loader:
                optimizer.zero_grad()
                criterion(model(bx), by).backward()
                optimizer.step()
        
        # Evaluation
        model.eval()
        all_preds, all_targets = [], []
        with torch.no_grad():
            for bx, by in test_loader:
                all_preds.extend(torch.argmax(model(bx), dim=1).numpy())
                all_targets.extend(by.numpy())
        
        print(f"{name} -> Accuracy: {accuracy_score(all_targets, all_preds):.4f}")

if __name__ == '__main__':
    train_and_eval()