import os
import csv

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.metrics import (
    accuracy_score,
    f1_score,
    classification_report,
    confusion_matrix,
    cohen_kappa_score,
)

from utils import CLASS_NAMES


def evaluate_model(model: nn.Module, test_loader: DataLoader, device: torch.device) -> dict:
    """Run inference on test_loader and return a dict of all metrics."""
    model.eval()
    all_preds, all_labels = [], []

    with torch.no_grad():
        for images, labels in test_loader:
            images = images.to(device, non_blocking=True)
            outputs = model(images)
            preds = outputs.argmax(dim=1)
            all_preds.extend(preds.cpu().numpy())
            all_labels.extend(labels.numpy())

    all_preds = np.array(all_preds)
    all_labels = np.array(all_labels)

    return {
        'accuracy': accuracy_score(all_labels, all_preds),
        'weighted_f1': f1_score(all_labels, all_preds, average='weighted', zero_division=0),
        'macro_f1': f1_score(all_labels, all_preds, average='macro', zero_division=0),
        'classification_report': classification_report(
            all_labels, all_preds, target_names=CLASS_NAMES, zero_division=0
        ),
        'confusion_matrix': confusion_matrix(all_labels, all_preds),
        'cohen_kappa': cohen_kappa_score(all_labels, all_preds, weights='quadratic'),
    }


def plot_confusion_matrix(cm: np.ndarray, save_path: str) -> None:
    """Save a seaborn confusion matrix heatmap to save_path."""
    fig, ax = plt.subplots(figsize=(8, 6))
    sns.heatmap(
        cm, annot=True, fmt='d', cmap='Blues',
        xticklabels=CLASS_NAMES, yticklabels=CLASS_NAMES,
        ax=ax
    )
    ax.set_xlabel('Predicted')
    ax.set_ylabel('True')
    ax.set_title('Confusion Matrix')
    plt.tight_layout()
    fig.savefig(save_path, dpi=150)
    plt.close(fig)


def plot_training_curves(
    train_losses: list, val_losses: list, val_f1s: list, save_path: str
) -> None:
    """Save loss and F1 training curves to save_path."""
    epochs = range(1, len(train_losses) + 1)
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))

    ax1.plot(epochs, train_losses, label='Train loss')
    ax1.plot(epochs, val_losses, label='Val loss')
    ax1.set_xlabel('Epoch')
    ax1.set_ylabel('Loss')
    ax1.set_title('Training and Validation Loss')
    ax1.legend()

    ax2.plot(epochs, val_f1s, label='Val macro F1', color='green')
    ax2.set_xlabel('Epoch')
    ax2.set_ylabel('Macro F1')
    ax2.set_title('Validation Macro F1')
    ax2.legend()

    plt.tight_layout()
    fig.savefig(save_path, dpi=150)
    plt.close(fig)


def generate_summary(all_results: dict, save_dir: str) -> None:
    """Print comparison table, save as CSV, and generate bar chart."""
    os.makedirs(save_dir, exist_ok=True)

    rows = []
    for key, r in all_results.items():
        model_name, clahe_str = key.rsplit('_clahe=', 1)
        rows.append({
            'model': model_name,
            'clahe': clahe_str,
            'weighted_f1': r['weighted_f1'],
            'macro_f1': r['macro_f1'],
            'accuracy': r['accuracy'],
            'cohen_kappa': r['cohen_kappa'],
            'train_time_s': round(r['train_time'], 1),
            'epochs_trained': r['epochs_trained'],
        })

    # Print table
    header = f"{'Model':<18} {'CLAHE':<6} {'W-F1':>6} {'M-F1':>6} {'Acc':>6} {'Kappa':>6} {'Time(s)':>8} {'Epochs':>7}"
    print('\n' + '=' * len(header))
    print(header)
    print('-' * len(header))
    for row in rows:
        print(
            f"{row['model']:<18} {row['clahe']:<6} "
            f"{row['weighted_f1']:>6.4f} {row['macro_f1']:>6.4f} "
            f"{row['accuracy']:>6.4f} {row['cohen_kappa']:>6.4f} "
            f"{row['train_time_s']:>8} {row['epochs_trained']:>7}"
        )
    print('=' * len(header))

    # Save CSV
    csv_path = os.path.join(save_dir, 'summary.csv')
    with open(csv_path, 'w', newline='') as f:
        writer = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
        writer.writeheader()
        writer.writerows(rows)
    print(f"\nSummary saved to {csv_path}")

    # Comparison bar chart — show both macro and weighted F1
    x_labels = [f"{r['model']}\nCLAHE={r['clahe']}" for r in rows]
    w_f1_vals = [r['weighted_f1'] for r in rows]
    m_f1_vals = [r['macro_f1'] for r in rows]

    x = np.arange(len(x_labels))
    bar_width = 0.35
    fig, ax = plt.subplots(figsize=(12, 5))
    bars_w = ax.bar(x - bar_width / 2, w_f1_vals, bar_width, label='Weighted F1', color='#4C72B0')
    bars_m = ax.bar(x + bar_width / 2, m_f1_vals, bar_width, label='Macro F1', color='#DD8452')
    ax.set_ylabel('F1-Score')
    ax.set_title('F1 Comparison Across Experiments')
    ax.set_xticks(x)
    ax.set_xticklabels(x_labels)
    ax.set_ylim(0, 1.0)
    ax.legend()
    for bar, val in zip(bars_w, w_f1_vals):
        ax.text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 0.01,
                f'{val:.3f}', ha='center', va='bottom', fontsize=8)
    for bar, val in zip(bars_m, m_f1_vals):
        ax.text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 0.01,
                f'{val:.3f}', ha='center', va='bottom', fontsize=8)
    plt.tight_layout()
    chart_path = os.path.join(save_dir, 'comparison_chart.png')
    fig.savefig(chart_path, dpi=150)
    plt.close(fig)
    print(f"Comparison chart saved to {chart_path}")