src/models/ModelsP3_experiment.py

"""
GRU Model Training - GTSRB Traffic Sign Classification
"""

import torch
from torch.utils.data import DataLoader
from pathlib import Path

from src.config import TrainingConfig
from src.detection.sequence_dataset import create_sequence_dataloaders
from src.models.base_sequential_models import GRUModel, RNNModel, LSTMModel, TransformerModel
from src.models.unified_trainer import UnifiedTrainer


FEATURES_DIR   = "./cache/vgg16_sequence_features"
RESULTS_DIR    = "./results"
CHECKPOINT_DIR = "./checkpoints"
BATCH_SIZE     = 32
NUM_WORKERS    = 0
NUM_EPOCHS     = 30


# Wrapper that casts sequences to float32 without breaking len()
class Float32Loader:
    def __init__(self, loader):
        self.loader = loader

    def __len__(self):
        return len(self.loader)

    def __iter__(self):
        for sequences, labels in self.loader:
            yield sequences.float(), labels


def load_data(
    features_dir = FEATURES_DIR,
    batch_size   = BATCH_SIZE,
    num_workers  = NUM_WORKERS,
):
    train_loader, val_loader, test_loader = create_sequence_dataloaders(
        features_dir = features_dir,
        batch_size   = batch_size,
        num_workers  = num_workers,
        augment      = True,
        seed         = 42,
    )

    train_loader = Float32Loader(train_loader)
    val_loader   = Float32Loader(val_loader)
    test_loader  = Float32Loader(test_loader)

    print(f"Train batches : {len(train_loader)}")
    print(f"Val   batches : {len(val_loader)}")
    print(f"Test  batches : {len(test_loader)}")

    return train_loader, val_loader, test_loader


def build_model(
    hidden_size   = 256,
    num_layers    = 2,
    bidirectional = True,
    dropout       = 0.3,
    model_type     = "gru",
):
    device = "cuda" if torch.cuda.is_available() else "cpu"

    if model_type.lower() == "gru":
        model = GRUModel(
            input_size    = 512,
            hidden_size   = hidden_size,
            num_layers    = num_layers,
            output_size   = 43,
            dropout       = dropout,
            bidirectional = bidirectional,
            device        = device,
        )
    elif model_type.lower() == "rnn":
        model = RNNModel(
            input_size    = 512,
            hidden_size   = hidden_size,
            num_layers    = num_layers,
            output_size   = 43,
            dropout       = dropout,
            bidirectional = bidirectional,
            device        = device,
        )
    elif model_type.lower() == "lstm":
        model = LSTMModel(
            input_size    = 512,
            hidden_size   = hidden_size,
            num_layers    = num_layers,
            output_size   = 43,
            dropout       = dropout,
            bidirectional = bidirectional,
            device        = device,
        )
    elif model_type.lower() == "transformer":
        model = TransformerModel(
            input_size    = 512,
            output_size   = 43,
            device        = device,
        )
    else:
        raise ValueError(f"Unsupported model type: {model_type}")

    print(f"Model      : {model.get_model_name()}")
    print(f"Parameters : {model.get_num_parameters():,}")
    print(f"Device     : {device}")

    return model


def train_model(
    model,
    train_loader,
    val_loader,
    num_epochs     = NUM_EPOCHS,
    learning_rate  = 1e-3,
    weight_decay   = 1e-4,
    checkpoint_dir = CHECKPOINT_DIR,
):
    device = "cuda" if torch.cuda.is_available() else "cpu"

    Path(checkpoint_dir).mkdir(parents=True, exist_ok=True)

    config = TrainingConfig()
    config.LEARNING_RATE            = learning_rate
    config.WEIGHT_DECAY             = weight_decay
    config.NUM_EPOCHS               = num_epochs
    config.EARLY_STOPPING_PATIENCE  = 10
    config.SCHEDULER_TYPE           = "cosine"
    config.SAVE_BEST_ONLY           = True
    config.DEVICE                   = device

    trainer = UnifiedTrainer(
        model        = model,
        train_loader = train_loader,
        val_loader   = val_loader,
        config       = config,
        device       = device,
        save_dir     = checkpoint_dir,
    )

    trainer.train(num_epochs=num_epochs)

    return trainer


def evaluate_model(
    trainer,
    test_loader,
    results_dir = RESULTS_DIR,
):
    Path(results_dir).mkdir(parents=True, exist_ok=True)

    metrics = trainer.evaluate(test_loader)
    trainer.save_training_curves(results_dir)
    trainer.save_metrics_json(metrics, results_dir)

    print(f"\nResults saved to: {results_dir}/")

    return metrics


def main():
    """Train all 4 models (RNN, GRU, LSTM, Transformer) sequentially."""
    import json
    from sklearn.metrics import confusion_matrix
    import matplotlib.pyplot as plt
    
    device = "cuda" if torch.cuda.is_available() else "cpu"
    
    print("=" * 80)
    print("ModelsP3 FINAL TRAINING EXPERIMENT")
    print("=" * 80)
    print(f"Device: {device}")
    
    # Load data once (reuse across all models)
    print("\n[Loading data...]")
    train_loader, val_loader, test_loader = load_data(
        features_dir=FEATURES_DIR,
        batch_size=BATCH_SIZE,
        num_workers=NUM_WORKERS,
    )
    
    model_types = ["rnn", "gru", "lstm", "transformer"]
    results_summary = {}
    
    for model_type in model_types:
        print("\n" + "=" * 80)
        print(f"TRAINING {model_type.upper()}")
        print("=" * 80)
        
        # Build model
        model = build_model(model_type=model_type)
        
        # Train model
        trainer = train_model(
            model=model,
            train_loader=train_loader,
            val_loader=val_loader,
            num_epochs=NUM_EPOCHS,
            checkpoint_dir=f"{CHECKPOINT_DIR}/{model_type}",
        )
        
        # Evaluate model
        metrics = evaluate_model(
            trainer=trainer,
            test_loader=test_loader,
            results_dir=RESULTS_DIR,
        )
        
        # Save confusion matrix PNG
        if trainer.last_eval_targets and trainer.last_eval_preds:
            cm = confusion_matrix(trainer.last_eval_targets, trainer.last_eval_preds)
            fig, ax = plt.subplots(figsize=(10, 8))
            im = ax.imshow(cm, cmap="Blues")
            fig.colorbar(im, ax=ax)
            ax.set_title(f"{model_type.upper()} Confusion Matrix")
            ax.set_xlabel("Predicted")
            ax.set_ylabel("True")
            output_path = Path(RESULTS_DIR) / f"{model_type.upper()}_confusion_matrix.png"
            fig.tight_layout()
            fig.savefig(output_path, dpi=180, bbox_inches="tight")
            plt.close(fig)
            print(f"Saved confusion matrix: {output_path}")
        
        # Store results for summary
        results_summary[model_type.upper()] = {
            "accuracy": metrics.get("accuracy"),
            "precision": metrics.get("precision"),
            "recall": metrics.get("recall"),
            "f1": metrics.get("f1"),
            "loss": metrics.get("loss"),
        }
    
    # Print final summary
    print("\n" + "=" * 80)
    print("FINAL RESULTS SUMMARY")
    print("=" * 80)
    for model_name, model_metrics in results_summary.items():
        print(f"\n{model_name}:")
        for metric_name, metric_value in model_metrics.items():
            print(f"  {metric_name}: {metric_value:.4f}")
    
    # Save summary to JSON
    summary_path = Path(RESULTS_DIR) / "training_summary.json"
    with open(summary_path, "w") as f:
        json.dump(results_summary, f, indent=2)
    print(f"\nSummary saved to: {summary_path}")


if __name__ == "__main__":
    main()