import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim.lr_scheduler import ReduceLROnPlateau
import os
from tqdm import tqdm
from sklearn.metrics import accuracy_score, classification_report
import argparse

from model import SimpleRNN
from data_loader import create_dataloaders


class EarlyStopping:
    """Early stopping to stop training when validation loss doesn't improve"""
    
    def __init__(self, patience=5, min_delta=0, restore_best_weights=True):
        self.patience = patience
        self.min_delta = min_delta
        self.restore_best_weights = restore_best_weights
        self.best_loss = None
        self.counter = 0
        self.best_weights = None
        
    def __call__(self, val_loss, model):
        if self.best_loss is None:
            self.best_loss = val_loss
            self.save_checkpoint(model)
        elif val_loss < self.best_loss - self.min_delta:
            self.best_loss = val_loss
            self.counter = 0
            self.save_checkpoint(model)
        else:
            self.counter += 1
            
        if self.counter >= self.patience:
            if self.restore_best_weights:
                model.load_state_dict(self.best_weights)
            return True
        return False
    
    def save_checkpoint(self, model):
        """Save model checkpoint"""
        self.best_weights = model.state_dict().copy()


def train_epoch(model, train_loader, criterion, optimizer, device):
    """Train for one epoch"""
    model.train()
    total_loss = 0
    all_preds = []
    all_labels = []
    
    pbar = tqdm(train_loader, desc="Training")
    for texts, labels in pbar:
        texts = texts.to(device)
        labels = labels.to(device)
        
        # Forward pass
        optimizer.zero_grad()
        outputs = model(texts)
        loss = criterion(outputs, labels)
        
        # Backward pass
        loss.backward()
        # Gradient clipping to prevent exploding/vanishing gradients
        # Use larger max_norm for simple RNN to avoid over-clipping
        grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=5.0)
        optimizer.step()
        
        # Log gradient norm occasionally for debugging
        if pbar.n % 100 == 0:
            pbar.set_postfix({'loss': loss.item(), 'grad_norm': f'{grad_norm:.2f}'})
        else:
            pbar.set_postfix({'loss': loss.item()})
        
        # Metrics
        total_loss += loss.item()
        preds = torch.argmax(outputs, dim=1)
        all_preds.extend(preds.cpu().detach().numpy())
        all_labels.extend(labels.cpu().detach().numpy())
    
    avg_loss = total_loss / len(train_loader)
    accuracy = accuracy_score(all_labels, all_preds)
    
    return avg_loss, accuracy


def evaluate(model, val_loader, criterion, device):
    """Evaluate on validation set"""
    model.eval()
    total_loss = 0
    all_preds = []
    all_labels = []
    
    with torch.no_grad():
        for texts, labels in tqdm(val_loader, desc="Evaluating"):
            texts = texts.to(device)
            labels = labels.to(device)
            
            outputs = model(texts)
            loss = criterion(outputs, labels)
            
            total_loss += loss.item()
            preds = torch.argmax(outputs, dim=1)
            all_preds.extend(preds.cpu().detach().numpy())
            all_labels.extend(labels.cpu().detach().numpy())
    
    avg_loss = total_loss / len(val_loader)
    accuracy = accuracy_score(all_labels, all_preds)
    
    return avg_loss, accuracy, all_preds, all_labels


def train_model(
    dataset_name,
    embedding_dim=128,
    hidden_dim=256,
    num_layers=2,
    num_hidden_nodes=128,
    dropout=0.3,
    batch_size=32,
    max_length=128,
    learning_rate=0.001,
    num_epochs=20,
    patience=5,
    vocab_min_freq=2,
    device=None
):
    """Main training function"""
    
    # Set device
    if device is None:
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Using device: {device}")
    
    # Create dataloaders
    print(f"\nLoading {dataset_name} dataset...")
    train_loader, val_loader, vocab, num_classes, train_labels = create_dataloaders(
        dataset_name, batch_size=batch_size, max_length=max_length, min_freq=vocab_min_freq
    )
    
    # Calculate class weights for imbalanced datasets (especially emotion)
    if dataset_name == 'emotion':
        from collections import Counter
        import numpy as np
        label_counts = Counter(train_labels)
        total = len(train_labels)
        # Use sqrt scaling for smoother weights (less extreme than inverse)
        # This prevents over-weighting rare classes
        class_weights = torch.tensor([
            np.sqrt(total / label_counts.get(i, 1)) if label_counts.get(i, 0) > 0 else 1.0
            for i in range(num_classes)
        ], dtype=torch.float32)
        # Normalize so weights don't dominate the loss
        class_weights = class_weights / class_weights.mean()
        print(f"Class distribution: {dict(label_counts)}")
        print(f"Class weights (sqrt-scaled): {class_weights.numpy()}")
    else:
        class_weights = None
    
    # Create model
    model = SimpleRNN(
        vocab_size=len(vocab),
        embedding_dim=embedding_dim,
        hidden_dim=hidden_dim,
        num_layers=num_layers,
        num_classes=num_classes,
        dropout=dropout,
        num_hidden_nodes=num_hidden_nodes
    ).to(device)
    
    # Initialize weights - simpler and more stable for Simple RNN
    def init_weights(m):
        if isinstance(m, nn.Linear):
            torch.nn.init.xavier_uniform_(m.weight, gain=1.0)
            if m.bias is not None:
                # For output layer with many classes, use small negative bias
                # This helps prevent initial collapse to one class
                if hasattr(m, 'out_features') and m.out_features == num_classes and num_classes > 4:
                    m.bias.data.fill_(-0.05)  # Small negative bias for multi-class
                else:
                    m.bias.data.fill_(0.0)
        elif isinstance(m, nn.RNN):
            # Simple RNN initialization - smaller scale to prevent instability
            for name, param in m.named_parameters():
                if 'weight_ih' in name:
                    # Input-to-hidden: use smaller initialization
                    torch.nn.init.xavier_uniform_(param.data, gain=0.5)
                elif 'weight_hh' in name:
                    # Hidden-to-hidden: use orthogonal with smaller scale
                    torch.nn.init.orthogonal_(param.data, gain=0.5)
                elif 'bias' in name:
                    # Initialize biases to zero (no forget gate in simple RNN)
                    param.data.fill_(0.0)
        elif isinstance(m, nn.Embedding):
            # Embedding initialization - uniform distribution
            torch.nn.init.normal_(m.weight, mean=0.0, std=0.01)
            # Set padding to zero
            if m.padding_idx is not None:
                m.weight.data[m.padding_idx].fill_(0)
    
    model.apply(init_weights)
    
    print(f"\nModel architecture:")
    print(model)
    print(f"\nTotal parameters: {sum(p.numel() for p in model.parameters()):,}")
    
    # Loss function and optimizer
    # Use class weights for emotion dataset to handle imbalance
    if class_weights is not None:
        class_weights = class_weights.to(device)
        criterion = nn.CrossEntropyLoss(weight=class_weights)
        print("Using weighted CrossEntropyLoss for class imbalance")
    else:
        criterion = nn.CrossEntropyLoss()
    
    # Reduced weight decay for simple RNN (they need more flexibility)
    optimizer = optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=0.001, betas=(0.9, 0.999))
    scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.7, patience=3, verbose=True, min_lr=1e-5)
    
    # Early stopping
    early_stopping = EarlyStopping(patience=patience)
    
    # Training loop
    print(f"\nStarting training for {num_epochs} epochs...")
    best_val_loss = float('inf')
    
    for epoch in range(num_epochs):
        print(f"\n{'='*60}")
        print(f"Epoch {epoch+1}/{num_epochs}")
        print(f"{'='*60}")
        
        # Train
        train_loss, train_acc = train_epoch(model, train_loader, criterion, optimizer, device)
        
        # Validate
        val_loss, val_acc, val_preds, val_labels = evaluate(model, val_loader, criterion, device)
        
        # Learning rate scheduling
        scheduler.step(val_loss)
        
        # Print metrics
        print(f"\nTrain Loss: {train_loss:.4f}, Train Acc: {train_acc:.4f}")
        print(f"Val Loss: {val_loss:.4f}, Val Acc: {val_acc:.4f}")
        print(f"Learning Rate: {optimizer.param_groups[0]['lr']:.6f}")
        
        # Early stopping
        if early_stopping(val_loss, model):
            print(f"\nEarly stopping triggered after {epoch+1} epochs")
            break
        
        # Save best model
        if val_loss < best_val_loss:
            best_val_loss = val_loss
            os.makedirs('checkpoints', exist_ok=True)
            torch.save({
                'epoch': epoch,
                'model_state_dict': model.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'val_loss': val_loss,
                'val_acc': val_acc,
                'vocab': vocab,
                'num_classes': num_classes,
                'model_config': {
                    'embedding_dim': embedding_dim,
                    'hidden_dim': hidden_dim,
                    'num_layers': num_layers,
                    'num_hidden_nodes': num_hidden_nodes,
                    'dropout': dropout,
                }
            }, f'checkpoints/best_model_{dataset_name}.pt')
            print(f"Saved best model (val_loss: {val_loss:.4f})")
    
    # Final evaluation
    print(f"\n{'='*60}")
    print("Final Evaluation")
    print(f"{'='*60}")
    final_loss, final_acc, final_preds, final_labels = evaluate(model, val_loader, criterion, device)
    print(f"\nFinal Validation Accuracy: {final_acc:.4f}")
    print(f"\nClassification Report:")
    print(classification_report(final_labels, final_preds))
    
    return model, vocab


def main():
    parser = argparse.ArgumentParser(description='Train RNN on text classification datasets')
    parser.add_argument('--dataset', type=str, choices=['emotion', 'ag_news'], required=True,
                       help='Dataset to use: emotion or ag_news')
    parser.add_argument('--embedding_dim', type=int, default=128, help='Embedding dimension')
    parser.add_argument('--hidden_dim', type=int, default=256, help='RNN hidden dimension')
    parser.add_argument('--num_layers', type=int, default=2, help='Number of RNN layers')
    parser.add_argument('--num_hidden_nodes', type=int, default=128, help='Number of nodes in hidden layer')
    parser.add_argument('--dropout', type=float, default=0.3, help='Dropout rate')
    parser.add_argument('--batch_size', type=int, default=32, help='Batch size')
    parser.add_argument('--max_length', type=int, default=128, help='Maximum sequence length')
    parser.add_argument('--learning_rate', type=float, default=0.001, help='Learning rate')
    parser.add_argument('--num_epochs', type=int, default=20, help='Number of epochs')
    parser.add_argument('--patience', type=int, default=5, help='Early stopping patience')
    parser.add_argument('--vocab_min_freq', type=int, default=2, help='Minimum word frequency for vocabulary')
    
    args = parser.parse_args()
    
    train_model(
        dataset_name=args.dataset,
        embedding_dim=args.embedding_dim,
        hidden_dim=args.hidden_dim,
        num_layers=args.num_layers,
        num_hidden_nodes=args.num_hidden_nodes,
        dropout=args.dropout,
        batch_size=args.batch_size,
        max_length=args.max_length,
        learning_rate=args.learning_rate,
        num_epochs=args.num_epochs,
        patience=args.patience,
        vocab_min_freq=args.vocab_min_freq
    )


if __name__ == '__main__':
    main()