"""
CW3: Training and Evaluation Pipeline (PROVIDED — do not modify).

This file orchestrates the MLP training loop with explicit backpropagation:
    1. model.forward(X)       — forward pass through all layers
    2. loss.forward(logits, y) — compute loss
    3. loss.backward()         — gradient of loss w.r.t. logits
    4. model.backward(grad)    — backpropagate through all layers
    5. optimizer.step()        — update all parameters

Compare with CW2's trainer: the flow is the same, but now the backward
pass goes through MULTIPLE layers instead of just one.
"""

import os
import sys
import numpy as np

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

from common.data_loader import (load_numpy_data, train_val_split, subset_data,
                                 DataIterator, CLASS_NAMES, NUM_CLASSES)
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_hyperparameter_search, plot_training_summary)
from common.utils import set_seed, Timer, save_results, print_config

from model import MLP
from losses import CrossEntropyLoss
from optimizer import SGD


def train_one_epoch(model, loss_fn, optimizer, data_iterator):
    """Train for one epoch with explicit backpropagation."""
    total_loss = 0.0
    total_correct = 0
    total_samples = 0

    for X_batch, y_batch in data_iterator:
        # Step 1: Forward pass through the network
        logits = model.forward(X_batch)

        # Step 2: Compute loss
        loss = loss_fn.forward(logits, y_batch)

        # Step 3: Compute gradient of loss w.r.t. logits
        grad_logits = loss_fn.backward()

        # Step 4: Backpropagate through all layers (chain rule!)
        model.backward(grad_logits)

        # Step 5: Update parameters
        optimizer.step()

        # Track metrics
        predictions = np.argmax(logits, axis=1)
        total_loss += loss * len(X_batch)
        total_correct += np.sum(predictions == y_batch)
        total_samples += len(X_batch)

    avg_loss = total_loss / total_samples
    avg_acc = total_correct / total_samples
    return avg_loss, avg_acc


def evaluate(model, loss_fn, images, labels, batch_size=256):
    """Evaluate model on a dataset."""
    data_iter = DataIterator(images, labels, batch_size=batch_size, shuffle=False)

    total_loss = 0.0
    total_correct = 0
    total_samples = 0
    all_predictions = []

    for X_batch, y_batch in data_iter:
        logits = model.forward(X_batch)
        loss = loss_fn.forward(logits, y_batch)
        predictions = np.argmax(logits, axis=1)

        total_loss += loss * len(X_batch)
        total_correct += np.sum(predictions == y_batch)
        total_samples += len(X_batch)
        all_predictions.append(predictions)

    avg_loss = total_loss / total_samples
    avg_acc = total_correct / total_samples
    all_predictions = np.concatenate(all_predictions)

    return avg_loss, avg_acc, all_predictions


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

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

    # Create output directory
    hidden_str = "-".join(map(str, config.hidden_dims))
    output_dir = os.path.join(config.output_dir,
                              f"mlp_{hidden_str}_{config.activation}")
    os.makedirs(output_dir, exist_ok=True)

    # ---- Data Loading ----
    data = load_numpy_data(config.data_dir)
    train_images, train_labels = data['train_images'], data['train_labels']
    test_images, test_labels = data['test_images'], data['test_labels']

    if config.quick:
        print(f"\n[Quick mode] Using {config.quick_num_samples} samples, "
              f"{config.quick_epochs} epochs")
        train_images, train_labels = subset_data(
            train_images, train_labels, config.quick_num_samples, config.seed)
        num_epochs = config.quick_epochs
    else:
        num_epochs = config.num_epochs

    train_imgs, train_lbls, val_imgs, val_lbls = train_val_split(
        train_images, train_labels, config.val_ratio, config.seed)

    print(f"Train: {len(train_imgs)}, Val: {len(val_imgs)}, Test: {len(test_labels)}")

    # ---- Model, Loss, Optimizer ----
    model = MLP(
        input_dim=784,
        hidden_dims=config.hidden_dims,
        num_classes=NUM_CLASSES,
        activation=config.activation
    )
    loss_fn = CrossEntropyLoss()
    optimizer = SGD(
        parameters=model.parameters(),
        learning_rate=config.learning_rate,
        momentum=config.momentum,
        weight_decay=config.weight_decay
    )

    print(f"\nArchitecture: 784 -> {' -> '.join(map(str, config.hidden_dims))} -> 10")
    print(f"Activation: {config.activation}")
    print(f"Total parameters: {model.num_parameters():,}")

    # ---- Training Loop ----
    train_losses, val_losses = [], []
    train_accs, val_accs = [], []

    print(f"\nTraining for {num_epochs} epochs...")
    with Timer("Training"):
        for epoch in range(1, num_epochs + 1):
            train_iter = DataIterator(train_imgs, train_lbls,
                                       batch_size=config.batch_size, shuffle=True)

            train_loss, train_acc = train_one_epoch(
                model, loss_fn, optimizer, train_iter)

            val_loss, val_acc, _ = evaluate(model, loss_fn, val_imgs, val_lbls)

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

            if epoch % 5 == 0 or epoch == 1:
                print(f"  Epoch {epoch:3d}/{num_epochs}: "
                      f"Train Loss={train_loss:.4f}, Train Acc={train_acc:.4f}, "
                      f"Val Loss={val_loss:.4f}, Val Acc={val_acc:.4f}")

    # ---- Test Evaluation ----
    print("\n=== Test Set Evaluation ===")
    test_loss, test_acc, test_preds = evaluate(
        model, loss_fn, test_images, test_labels)
    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"MLP [{hidden_str}] Loss Curves",
                     save_path=os.path.join(output_dir, "loss_curves.png"))

    plot_accuracy_curves(train_accs, val_accs,
                         title=f"MLP [{hidden_str}] 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"MLP [{hidden_str}] 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"MLP [{hidden_str}] Confusion Matrix",
                          save_path=os.path.join(output_dir, "confusion_matrix.png"))

    plot_sample_predictions(test_images, test_labels, test_preds, CLASS_NAMES,
                            save_path=os.path.join(output_dir, "sample_predictions.png"))

    # ---- Save Results ----
    results = {
        'hidden_dims': config.hidden_dims,
        'activation': config.activation,
        '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': model.num_parameters(),
        'config': {
            'learning_rate': config.learning_rate,
            'momentum': config.momentum,
            'batch_size': config.batch_size,
            'num_epochs': num_epochs,
        }
    }
    save_results(results, os.path.join(output_dir, "results.json"))

    return results


def sweep_hyperparameter(config, param_name, param_values):
    """
    Run a hyperparameter sweep and generate comparison plot.

    Args:
        config: Base Config
        param_name: Name of hyperparameter to sweep
        param_values: List of values to try
    """
    all_train_accs = []
    all_val_accs = []

    for val in param_values:
        print(f"\n{'='*60}")
        print(f"  Sweep: {param_name} = {val}")
        print(f"{'='*60}")

        # Update config
        if param_name == "hidden_dim":
            # Sweep single hidden layer width
            config.hidden_dims = [val]
        elif param_name == "num_layers":
            # Sweep depth with fixed width
            config.hidden_dims = [128] * val
        elif param_name == "learning_rate":
            config.learning_rate = val
        else:
            setattr(config, param_name, val)

        results = train(config)
        all_train_accs.append(results['train_accs'][-1])
        all_val_accs.append(results['val_accs'][-1])

    # Generate sweep plot
    plot_hyperparameter_search(
        param_name, param_values, all_train_accs, all_val_accs,
        save_path=os.path.join(config.output_dir, f"sweep_{param_name}.png"))

    # Print summary
    print(f"\n=== Sweep Results: {param_name} ===")
    print(f"{'Value':<15s} {'Train Acc':>12s} {'Val Acc':>12s}")
    print("-" * 41)
    for val, ta, va in zip(param_values, all_train_accs, all_val_accs):
        print(f"{str(val):<15s} {ta:>12.4f} {va:>12.4f}")
