"""
CW4: PyTorch Training Pipeline (PROVIDED — do not modify).

Compare this with CW3's trainer:
    CW3 (manual):                    CW4 (PyTorch):
    ─────────────                    ──────────────
    logits = model.forward(X)        output = model(x)
    loss = loss_fn.forward(logits,y) loss = criterion(output, target)
    grad = loss_fn.backward()        loss.backward()  ← autograd!
    model.backward(grad)             (included above)
    optimizer.step()                 optimizer.step()

PyTorch's loss.backward() automatically does what you implemented
by hand in CW3: propagating gradients through ALL layers via chain rule.
"""

import os
import sys
import numpy as np
import torch
import torch.nn as nn
from tqdm import tqdm

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

from common.metrics import accuracy, confusion_matrix, classification_report
from common.visualization import (plot_loss_curves, plot_accuracy_curves,
                                   plot_confusion_matrix, plot_sample_predictions,
                                   plot_training_summary)
from common.utils import set_seed, Timer, save_results, print_config

from dataset import get_dataloaders, CLASS_NAMES, NUM_CLASSES
from model import get_model


def train_one_epoch(model, dataloader, criterion, optimizer, device):
    """Train for one epoch."""
    model.train()
    total_loss = 0.0
    total_correct = 0
    total_samples = 0

    for images, labels in tqdm(dataloader, desc="  Training", leave=False):
        images, labels = images.to(device), labels.to(device)

        # Forward pass
        output = model(images)
        loss = criterion(output, labels)

        # Backward pass (PyTorch autograd does the chain rule!)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # Track metrics
        predictions = output.argmax(dim=1)
        total_loss += loss.item() * len(labels)
        total_correct += (predictions == labels).sum().item()
        total_samples += len(labels)

    return total_loss / total_samples, total_correct / total_samples


@torch.no_grad()
def evaluate(model, dataloader, criterion, device):
    """Evaluate model on a dataset."""
    model.eval()
    total_loss = 0.0
    total_correct = 0
    total_samples = 0
    all_preds = []
    all_labels = []

    for images, labels in tqdm(dataloader, desc="  Evaluating", leave=False):
        images, labels = images.to(device), labels.to(device)

        output = model(images)
        loss = criterion(output, labels)

        predictions = output.argmax(dim=1)
        total_loss += loss.item() * len(labels)
        total_correct += (predictions == labels).sum().item()
        total_samples += len(labels)
        all_preds.append(predictions.cpu().numpy())
        all_labels.append(labels.cpu().numpy())

    all_preds = np.concatenate(all_preds)
    all_labels = np.concatenate(all_labels)

    return (total_loss / total_samples,
            total_correct / total_samples,
            all_preds, all_labels)


def train(config):
    """
    Full training pipeline for CNN.

    Args:
        config: Config dataclass
    """
    print_config(config)
    set_seed(config.seed)

    # Device
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"\nDevice: {device}")

    # Output directory
    output_dir = os.path.join(config.output_dir, config.model_type)
    os.makedirs(output_dir, exist_ok=True)

    # ---- Data ----
    train_loader, val_loader, test_loader = get_dataloaders(config)

    # ---- Model ----
    model = get_model(config.model_type, NUM_CLASSES).to(device)
    num_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"\nModel: {config.model_type}")
    print(f"Total parameters: {num_params:,}")
    print(model)

    # ---- Loss and Optimizer ----
    criterion = nn.CrossEntropyLoss()

    if config.optimizer_type == "adam":
        optimizer = torch.optim.Adam(model.parameters(), lr=config.learning_rate)
    elif config.optimizer_type == "sgd":
        optimizer = torch.optim.SGD(model.parameters(), lr=config.learning_rate,
                                     momentum=0.9)
    else:
        raise ValueError(f"Unknown optimizer: {config.optimizer_type}")

    # ---- Training ----
    num_epochs = config.quick_epochs if config.quick else config.num_epochs
    if config.quick:
        print(f"\n[Quick mode] Training for {num_epochs} epochs")

    train_losses, val_losses = [], []
    train_accs, val_accs = [], []
    best_val_acc = 0.0

    print(f"\nTraining for {num_epochs} epochs...")
    with Timer("Training"):
        for epoch in range(1, num_epochs + 1):
            train_loss, train_acc = train_one_epoch(
                model, train_loader, criterion, optimizer, device)

            val_loss, val_acc, _, _ = evaluate(
                model, val_loader, criterion, device)

            train_losses.append(train_loss)
            val_losses.append(val_loss)
            train_accs.append(train_acc)
            val_accs.append(val_acc)

            # Save best model
            if val_acc > best_val_acc:
                best_val_acc = val_acc
                torch.save(model.state_dict(),
                           os.path.join(output_dir, "best_model.pth"))

            print(f"  Epoch {epoch:2d}/{num_epochs}: "
                  f"Train Loss={train_loss:.4f}, Train Acc={train_acc:.4f}, "
                  f"Val Loss={val_loss:.4f}, Val Acc={val_acc:.4f}")

    # ---- Load Best Model and Test ----
    model.load_state_dict(torch.load(os.path.join(output_dir, "best_model.pth"),
                                      weights_only=True))

    print("\n=== Test Set Evaluation ===")
    test_loss, test_acc, test_preds, test_labels = evaluate(
        model, test_loader, criterion, device)
    print(f"Test Loss: {test_loss:.4f}, Test Accuracy: {test_acc:.4f}")
    print(f"\n{classification_report(test_preds, test_labels, CLASS_NAMES)}")

    # ---- Generate Plots ----
    print("\nGenerating plots...")
    plot_loss_curves(train_losses, val_losses,
                     title=f"CNN [{config.model_type}] Loss Curves",
                     save_path=os.path.join(output_dir, "loss_curves.png"))

    plot_accuracy_curves(train_accs, val_accs,
                         title=f"CNN [{config.model_type}] Accuracy Curves",
                         save_path=os.path.join(output_dir, "accuracy_curves.png"))

    # Combined training summary (recommended for report)
    plot_training_summary(train_losses, val_losses, train_accs, val_accs,
                          title=f"CNN [{config.model_type}] Training Summary",
                          save_path=os.path.join(output_dir, "training_summary.png"))

    cm = confusion_matrix(test_preds, test_labels)
    plot_confusion_matrix(cm, CLASS_NAMES,
                          title=f"CNN [{config.model_type}] Confusion Matrix",
                          save_path=os.path.join(output_dir, "confusion_matrix.png"))

    # Get some test images for visualization (denormalize for display)
    test_images_for_plot = []
    for images, _ in test_loader:
        test_images_for_plot = images[:16].numpy().squeeze(1)  # (16, 28, 28)
        break
    plot_sample_predictions(
        test_images_for_plot, test_labels[:16], test_preds[:16], CLASS_NAMES,
        save_path=os.path.join(output_dir, "sample_predictions.png"))

    # ---- Save Results ----
    results = {
        'model_type': config.model_type,
        'train_losses': train_losses,
        'val_losses': val_losses,
        'train_accs': train_accs,
        'val_accs': val_accs,
        'test_loss': test_loss,
        'test_acc': test_acc,
        'num_parameters': num_params,
        'best_val_acc': best_val_acc,
        'config': {
            'learning_rate': config.learning_rate,
            'optimizer': config.optimizer_type,
            'batch_size': config.batch_size,
            'num_epochs': num_epochs,
            'augmentation': config.use_augmentation,
        }
    }
    save_results(results, os.path.join(output_dir, "results.json"))

    return results
