"""
automl_lib.metrics
~~~~~~~~~~~~~~~~~~
Evaluation metrics and visualization (loss/accuracy curves, confusion matrix).
"""

from __future__ import annotations

import os
from typing import Any

import matplotlib
matplotlib.use("Agg")  # non-interactive backend
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from torch.utils.data import DataLoader

from .utils import ensure_dir, get_logger

logger = get_logger(__name__)


# ═══════════════════════════════════════════════════════════════════════════
# Metric computation
# ═══════════════════════════════════════════════════════════════════════════

def compute_metrics(
    model: nn.Module,
    loader: DataLoader,
    task: str = "classification",
    device: torch.device | None = None,
    num_classes: int = 2,
) -> dict[str, Any]:
    """Evaluate *model* on *loader* and return a metrics dict."""
    if device is None:
        device = next(model.parameters()).device

    model.eval()
    all_preds, all_targets = [], []

    with torch.no_grad():
        for batch in loader:
            if isinstance(batch, dict):
                targets = batch.pop("labels").to(device)
                inputs = {k: v.to(device) for k, v in batch.items()}
                out = model(**inputs)
                out = out.logits if hasattr(out, "logits") else out
            else:
                inputs, targets = batch[0].to(device), batch[1].to(device)
                out = model(inputs)

            if task == "classification":
                preds = out.argmax(dim=1)
            else:
                preds = out.squeeze()

            all_preds.append(preds.cpu().numpy())
            all_targets.append(targets.cpu().numpy())

    preds_np = np.concatenate(all_preds)
    targets_np = np.concatenate(all_targets)

    metrics: dict[str, Any] = {}
    if task == "classification":
        from sklearn.metrics import (
            accuracy_score, f1_score, precision_score,
            recall_score, classification_report, confusion_matrix,
        )
        metrics["accuracy"] = accuracy_score(targets_np, preds_np)
        metrics["f1_macro"] = f1_score(targets_np, preds_np, average="macro", zero_division=0)
        metrics["precision_macro"] = precision_score(targets_np, preds_np, average="macro", zero_division=0)
        metrics["recall_macro"] = recall_score(targets_np, preds_np, average="macro", zero_division=0)
        metrics["classification_report"] = classification_report(targets_np, preds_np, zero_division=0)
        metrics["confusion_matrix"] = confusion_matrix(targets_np, preds_np)
        logger.info(f"Accuracy={metrics['accuracy']:.4f}  F1={metrics['f1_macro']:.4f}")
    else:
        from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
        metrics["mse"] = mean_squared_error(targets_np, preds_np)
        metrics["rmse"] = np.sqrt(metrics["mse"])
        metrics["mae"] = mean_absolute_error(targets_np, preds_np)
        metrics["r2"] = r2_score(targets_np, preds_np)
        logger.info(f"RMSE={metrics['rmse']:.4f}  MAE={metrics['mae']:.4f}  R2={metrics['r2']:.4f}")

    return metrics


# ═══════════════════════════════════════════════════════════════════════════
# Visualization
# ═══════════════════════════════════════════════════════════════════════════

def plot_training_history(
    history: dict[str, list[float]],
    output_dir: str = "automl_output",
) -> str:
    """Plot loss and accuracy curves and save as PNG. Returns file path."""
    ensure_dir(output_dir)

    fig, axes = plt.subplots(1, 2, figsize=(14, 5))
    epochs = range(1, len(history["train_loss"]) + 1)

    # Loss
    axes[0].plot(epochs, history["train_loss"], label="Train Loss", linewidth=2)
    axes[0].plot(epochs, history["val_loss"], label="Val Loss", linewidth=2)
    axes[0].set_xlabel("Epoch")
    axes[0].set_ylabel("Loss")
    axes[0].set_title("Training & Validation Loss")
    axes[0].legend()
    axes[0].grid(True, alpha=0.3)

    # Accuracy
    if any(v > 0 for v in history.get("val_acc", [])):
        axes[1].plot(epochs, history["train_acc"], label="Train Acc", linewidth=2)
        axes[1].plot(epochs, history["val_acc"], label="Val Acc", linewidth=2)
        axes[1].set_xlabel("Epoch")
        axes[1].set_ylabel("Accuracy")
        axes[1].set_title("Training & Validation Accuracy")
        axes[1].legend()
        axes[1].grid(True, alpha=0.3)
    else:
        axes[1].text(0.5, 0.5, "N/A (regression task)",
                     ha="center", va="center", fontsize=14, transform=axes[1].transAxes)
        axes[1].set_title("Accuracy (not applicable)")

    plt.tight_layout()
    path = os.path.join(output_dir, "training_curves.png")
    fig.savefig(path, dpi=150)
    plt.close(fig)
    logger.info(f"Training curves saved -> {path}")
    return path


def plot_confusion_matrix(
    cm: np.ndarray,
    class_names: list[str] | None = None,
    output_dir: str = "automl_output",
) -> str:
    """Plot and save a confusion matrix heatmap."""
    ensure_dir(output_dir)

    fig, ax = plt.subplots(figsize=(8, 6))
    im = ax.imshow(cm, interpolation="nearest", cmap=plt.cm.Blues)
    ax.figure.colorbar(im, ax=ax)

    n = cm.shape[0]
    labels = class_names or [str(i) for i in range(n)]
    ax.set(
        xticks=range(n), yticks=range(n),
        xticklabels=labels, yticklabels=labels,
        xlabel="Predicted", ylabel="True",
        title="Confusion Matrix",
    )
    plt.setp(ax.get_xticklabels(), rotation=45, ha="right")

    # Annotate cells
    thresh = cm.max() / 2.0
    for i in range(n):
        for j in range(n):
            ax.text(j, i, f"{cm[i, j]}",
                    ha="center", va="center",
                    color="white" if cm[i, j] > thresh else "black")

    plt.tight_layout()
    path = os.path.join(output_dir, "confusion_matrix.png")
    fig.savefig(path, dpi=150)
    plt.close(fig)
    logger.info(f"Confusion matrix saved -> {path}")
    return path


def print_dataset_summary(dataset_info: dict[str, Any]):
    """Print a formatted summary of the dataset."""
    logger.info("=" * 60)
    logger.info("Dataset Summary")
    logger.info("=" * 60)
    for key, value in dataset_info.items():
        if key not in ("class_weights", "classes"):
            logger.info(f"  {key}: {value}")
    logger.info("=" * 60)


def compute_feature_importance(model: Any, val_loader: DataLoader, features: list[str]) -> pd.DataFrame:
    """Analyze feature importance using permutation or native scores."""
    df_importance = pd.DataFrame({"feature": features, "importance": 0.0})
    
    # ── Case A: Sklearn native importance ───────────────────────
    if getattr(model, "_is_sklearn", False):
        estimator = getattr(model, "estimator", model)
        if hasattr(estimator, "feature_importances_"):
            df_importance["importance"] = estimator.feature_importances_
            return df_importance.sort_values("importance", ascending=False)
    
    # ── Case B: Permutation Importance (Generic) ────────────────
    # For now, we'll return a stub or simplify it
    # True permutation importance requires multiple passes over val_loader
    return df_importance
