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

This file orchestrates the full training loop:
    1. Load and split data
    2. Create model, loss function, and optimizer
    3. Train for multiple epochs with mini-batches
    4. Evaluate on validation/test sets
    5. Generate plots and save results
"""

import os
import sys
import numpy as np

# Add parent directory to path for imports
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_weight_visualization, plot_training_summary)
from common.utils import set_seed, Timer, save_results, print_config

from model import LogisticRegression
from losses import LOSS_FUNCTIONS
from optimizer import GradientDescent


def train_one_epoch(model, loss_fn, optimizer, data_iterator, num_classes):
    """Train for one epoch, returning average loss and accuracy."""
    total_loss = 0.0
    total_correct = 0
    total_samples = 0

    for X_batch, y_batch in data_iterator:
        # Forward pass
        logits = model.forward(X_batch)

        # Compute loss and gradient
        loss, grad_logits = loss_fn(logits, y_batch, num_classes)

        # Update parameters
        optimizer.step(model, grad_logits, X_batch)

        # 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, num_classes, batch_size=256):
    """Evaluate model on a dataset, returning loss, accuracy, and predictions."""
    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(logits, y_batch, num_classes)
        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 logistic regression.

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

    # Create output directory
    output_dir = os.path.join(config.output_dir, config.loss_type)
    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']

    # Quick mode: use subset
    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/validation split
    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 = LogisticRegression(input_dim=784, num_classes=NUM_CLASSES)
    loss_fn = LOSS_FUNCTIONS[config.loss_type]
    optimizer = GradientDescent(learning_rate=config.learning_rate)

    print(f"\nLoss function: {config.loss_type}")
    print(f"Parameters: W {model.W.shape}, b {model.b.shape}")

    # ---- 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):
            # Create data iterator (reshuffles each epoch)
            train_iter = DataIterator(train_imgs, train_lbls,
                                       batch_size=config.batch_size, shuffle=True)

            # Train one epoch
            train_loss, train_acc = train_one_epoch(
                model, loss_fn, optimizer, train_iter, NUM_CLASSES)

            # Evaluate on validation set
            val_loss, val_acc, _ = evaluate(
                model, loss_fn, val_imgs, val_lbls, NUM_CLASSES)

            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, NUM_CLASSES)
    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"Loss Curves ({config.loss_type})",
                     save_path=os.path.join(output_dir, "loss_curves.png"))

    plot_accuracy_curves(train_accs, val_accs,
                         title=f"Accuracy Curves ({config.loss_type})",
                         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"Training Summary ({config.loss_type})",
                          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"Confusion Matrix ({config.loss_type})",
                          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"))

    # Visualize learned weights
    plot_weight_visualization(model.W, CLASS_NAMES,
                              save_path=os.path.join(output_dir, "weight_visualization.png"))

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

    return results


def compare_all_losses(config):
    """
    Train with all four loss functions and generate comparison plots.

    This produces the comparison plot needed for the CW2 report.
    """
    all_results = {}
    loss_types = ["cross_entropy", "hinge", "exponential", "squared"]

    for loss_type in loss_types:
        print(f"\n{'='*60}")
        print(f"  Training with loss: {loss_type}")
        print(f"{'='*60}")
        config.loss_type = loss_type
        results = train(config)
        all_results[loss_type] = results

    # Generate comparison plot
    print(f"\n{'='*60}")
    print("  Generating comparison plots")
    print(f"{'='*60}")

    from common.visualization import plot_loss_comparison
    plot_loss_comparison(all_results,
                         title="Loss Function Comparison on FashionMNIST",
                         save_path=os.path.join(config.output_dir,
                                                "loss_comparison.png"))

    # Print summary table
    print("\n=== Summary: Test Accuracy by Loss Function ===")
    print(f"{'Loss Type':<20s} {'Test Accuracy':>15s}")
    print("-" * 37)
    for name, res in all_results.items():
        print(f"{name:<20s} {res['test_acc']:>15.4f}")

    return all_results
