"""Comprehensive Evaluation, Metrics Computation, Confusion Matrix, and Error Analysis."""

import os
import json
import time
from pathlib import Path
from typing import Dict, Any, List, Optional

import numpy as np
import pandas as pd
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.metrics import (
    accuracy_score,
    precision_recall_fscore_support,
    confusion_matrix,
    classification_report
)

from src.utils.config import (
    EXPERIMENTS_DIR,
    EMOTION_CLASSES,
    ID_TO_EMOTION,
    EMOTION_TO_ID
)


class ExperimentLogger:
    """Standardized Experiment Logger saving configs, metrics, plots, and error analysis."""

    def __init__(self, experiment_id: str, description: str = ""):
        self.experiment_id = experiment_id
        self.exp_dir = EXPERIMENTS_DIR / experiment_id
        self.exp_dir.mkdir(parents=True, exist_ok=True)
        self.model_dir = self.exp_dir / "model"
        self.model_dir.mkdir(parents=True, exist_ok=True)
        self.description = description
        self.start_time = time.time()

    def log_config(self, config: Dict[str, Any]):
        config_path = self.exp_dir / "config.json"
        with open(config_path, "w", encoding="utf-8") as f:
            json.dump(config, f, indent=2)

    def evaluate_and_log(
        self,
        y_true: List[str],
        y_pred: List[str],
        y_proba: Optional[np.ndarray] = None,
        audio_ids: Optional[List[str]] = None,
        split_name: str = "test",
        extra_metrics: Optional[Dict[str, Any]] = None
    ) -> Dict[str, Any]:
        """Compute full metrics suite, plot confusion matrix, and generate error_analysis.csv."""
        training_time = round(time.time() - self.start_time, 2)

        # 1. Macro and weighted metrics
        acc = float(accuracy_score(y_true, y_pred))
        p_macro, r_macro, f1_macro, _ = precision_recall_fscore_support(
            y_true, y_pred, average="macro", zero_division=0
        )
        p_weighted, r_weighted, f1_weighted, _ = precision_recall_fscore_support(
            y_true, y_pred, average="weighted", zero_division=0
        )

        # 2. Per-class metrics
        p_per_class, r_per_class, f1_per_class, support = precision_recall_fscore_support(
            y_true, y_pred, labels=EMOTION_CLASSES, zero_division=0
        )
        per_class_metrics = {}
        for idx, emo in enumerate(EMOTION_CLASSES):
            per_class_metrics[emo] = {
                "precision": float(round(p_per_class[idx], 4)),
                "recall": float(round(r_per_class[idx], 4)),
                "f1_score": float(round(f1_per_class[idx], 4)),
                "support": int(support[idx])
            }

        # 3. Build metrics payload
        metrics = {
            "experiment_id": self.experiment_id,
            "split_evaluated": split_name,
            "training_time_sec": training_time,
            "accuracy": float(round(acc, 4)),
            "accuracy_pct": float(round(acc * 100.0, 2)),
            "macro_f1": float(round(f1_macro, 4)),
            "macro_precision": float(round(p_macro, 4)),
            "macro_recall": float(round(r_macro, 4)),
            "weighted_f1": float(round(f1_weighted, 4)),
            "per_class": per_class_metrics
        }
        if extra_metrics:
            metrics.update(extra_metrics)

        # Save metrics.json
        metrics_path = self.exp_dir / "metrics.json"
        with open(metrics_path, "w", encoding="utf-8") as f:
            json.dump(metrics, f, indent=2)

        # Save classification_report.json
        clf_report_dict = classification_report(y_true, y_pred, labels=EMOTION_CLASSES, output_dict=True, zero_division=0)
        with open(self.exp_dir / "classification_report.json", "w", encoding="utf-8") as f:
            json.dump(clf_report_dict, f, indent=2)

        # 4. Confusion Matrix Plot
        cm = confusion_matrix(y_true, y_pred, labels=EMOTION_CLASSES)
        cm_norm = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]

        fig, ax = plt.subplots(figsize=(8, 6))
        sns.heatmap(
            cm,
            annot=True,
            fmt="d",
            cmap="Blues",
            xticklabels=EMOTION_CLASSES,
            yticklabels=EMOTION_CLASSES,
            ax=ax,
            cbar=True
        )
        ax.set_title(f"{self.experiment_id}: Confusion Matrix (Acc: {acc*100:.1f}%)", fontsize=14, fontweight="bold", pad=12)
        ax.set_ylabel("Actual Emotion", fontsize=12)
        ax.set_xlabel("Predicted Emotion", fontsize=12)
        plt.tight_layout()
        cm_path = self.exp_dir / "confusion_matrix.png"
        fig.savefig(cm_path, dpi=300)
        plt.close(fig)

        # 5. Error Analysis CSV
        if audio_ids is not None:
            error_records = []
            for i, (actual, pred) in enumerate(zip(y_true, y_pred)):
                conf = 0.0
                if y_proba is not None and i < len(y_proba):
                    pred_idx = EMOTION_TO_ID.get(pred, 0)
                    conf = float(round(y_proba[i][pred_idx], 4))

                is_error = (actual != pred)
                error_records.append({
                    "audio_id": audio_ids[i],
                    "actual_emotion": actual,
                    "predicted_emotion": pred,
                    "confidence": conf,
                    "is_misclassified": is_error,
                    "error_pair": f"{actual} -> {pred}" if is_error else "CORRECT"
                })

            err_df = pd.DataFrame(error_records)
            err_csv_path = self.exp_dir / "error_analysis.csv"
            err_df.to_csv(err_csv_path, index=False)

        print(f"\n[EVALUATION] {self.experiment_id} Results:")
        print(f"  Accuracy:  {acc * 100:.2f}%")
        print(f"  Macro F1:  {f1_macro * 100:.2f}%")
        print(f"  Weighted F1: {f1_weighted * 100:.2f}%")
        print(f"  Saved artifacts to: {self.exp_dir}")

        return metrics