import torch
import torchaudio
import torchaudio.transforms as T
import matplotlib.pyplot as plt
# from Ipython.display import Audio
import numpy as np
import seaborn as sns
from sklearn.metrics import (
    roc_curve,
    auc,
    precision_recall_curve,
    average_precision_score,
    confusion_matrix as sk_confusion_matrix,
)
from typing import List
import wandb

__all__ = ["Plotter", "plot_waveform", "plot_spectrogram", "plot_fbank"]

class Plotter:

    @staticmethod
    def roc_auc_plt(
        y_true: np.ndarray,
        y_probs: np.ndarray,
        prefix: str,
        epoch: int,
        wandb_handler,
    ) -> tuple:
        fpr, tpr, thresholds = roc_curve(y_true, y_probs)
        roc_auc = auc(fpr, tpr)

        j_scores       = tpr - fpr
        best_idx       = int(np.argmax(j_scores))
        best_threshold = float(thresholds[best_idx])
        best_fpr       = fpr[best_idx]
        best_tpr       = tpr[best_idx]

        fig, ax = plt.subplots(figsize=(6, 6))
        ax.plot(fpr, tpr, lw=2, label=f"AUC = {roc_auc:.4f}")
        ax.plot([0, 1], [0, 1], linestyle="--", color="grey", lw=1)
        ax.scatter(best_fpr, best_tpr, color="red", s=60, zorder=5)
        ax.annotate(
            f"Best thr={best_threshold:.2f}",
            xy=(best_fpr, best_tpr),
            xytext=(best_fpr + 0.04, best_tpr - 0.08),
            fontsize=9,
            arrowprops=dict(arrowstyle="->", color="red"),
        )

        # Annotate a spread of threshold values along the curve
        for idx in np.linspace(0, len(thresholds) - 1, 5, dtype=int):
            ax.annotate(
                f"{thresholds[idx]:.2f}",
                xy=(fpr[idx], tpr[idx]),
                fontsize=7,
                color="dimgray",
            )

        ax.set_title(f"{prefix.capitalize()} ROC Curve (Epoch {epoch})", fontsize=12)
        ax.set_xlabel("False Positive Rate")
        ax.set_ylabel("True Positive Rate")
        ax.legend(loc="lower right")
        plt.tight_layout()

        wandb_handler.log_metrics({
            f"{prefix}/roc_auc":            roc_auc,
            f"{prefix}/roc_best_threshold": best_threshold,
            f"{prefix}/roc_curve":          wandb.Image(fig),
        })

        plt.close(fig)
        return roc_auc, best_threshold

    @staticmethod
    def pr_curve(
        y_true: np.ndarray,
        y_probs: np.ndarray,
        prefix: str,
        epoch: int,
        wandb_handler,
    ) -> tuple:
        precision, recall, thresholds = precision_recall_curve(y_true, y_probs)
        pr_auc = average_precision_score(y_true, y_probs)

        f1_scores      = 2 * precision * recall / np.maximum(precision + recall, 1e-10)
        best_idx       = int(np.argmax(f1_scores))
        best_threshold = float(thresholds[best_idx]) if best_idx < len(thresholds) else 0.5

        fig, ax = plt.subplots(figsize=(6, 6))
        ax.plot(recall, precision, lw=2, label=f"AP = {pr_auc:.4f}")
        ax.scatter(recall[best_idx], precision[best_idx], color="red", s=60, zorder=5)
        ax.annotate(
            f"Best thr={best_threshold:.2f}",
            xy=(recall[best_idx], precision[best_idx]),
            xytext=(recall[best_idx] + 0.04, precision[best_idx] - 0.08),
            fontsize=9,
            arrowprops=dict(arrowstyle="->", color="red"),
        )

        ax.set_title(f"{prefix.capitalize()} PR Curve (Epoch {epoch})", fontsize=12)
        ax.set_xlabel("Recall")
        ax.set_ylabel("Precision")
        ax.legend(loc="lower left")
        plt.tight_layout()

        wandb_handler.log_metrics({
            f"{prefix}/pr_auc":            pr_auc,
            f"{prefix}/pr_best_threshold": best_threshold,
            f"{prefix}/pr_curve":          wandb.Image(fig),
        })

        plt.close(fig)
        return pr_auc, best_threshold
    @staticmethod
    def confusion_matrix(
        y_true: np.ndarray,
        y_pred: np.ndarray,
        class_names: List[str],
        prefix: str,
        epoch: int,
        wandb_handler,
    ) -> np.ndarray:
        cm = sk_confusion_matrix(y_true, y_pred)

        fig, ax = plt.subplots(figsize=(5, 4))
        sns.heatmap(
            cm,
            annot=True,
            fmt="d",
            cmap="Blues",
            xticklabels=class_names,
            yticklabels=class_names,
            ax=ax,
        )
        ax.set_title(f"{prefix.capitalize()} Confusion Matrix (Epoch {epoch})", fontsize=12)
        ax.set_xlabel("Predicted Label")
        ax.set_ylabel("True Label")
        plt.tight_layout()

        wandb_handler.log_metrics({
            f"{prefix}/confusion_matrix": wandb.Image(fig)
        })

        plt.close(fig)
        return cm

# def audio_from_waveform(waveform, sample_rate):
#     return Audio(waveform.numpy(), sample_rate)

def plot_waveform(waveform, sr, title="Waveform", ax=None):
    waveform = waveform.numpy()

    num_channels, num_frames = waveform.shape
    time_axis = torch.arange(0, num_frames) / sr

    if ax is None:
        _, ax = plt.subplots(num_channels, 1)
    ax.plot(time_axis, waveform[0], linewidth=1)
    ax.grid(True)
    ax.set_xlim([0, time_axis[-1]])
    ax.set_title(title)

def plot_spectrogram(specgram, title=None, ylabel="freq_bin", ax=None):
    if ax is None:
        _, ax = plt.subplots(1, 1)
    if title is not None:
        ax.set_title(title)
    ax.set_ylabel(ylabel)
    power_to_db = T.AmplitudeToDB("power", 80.0)
    ax.imshow(power_to_db(specgram), origin="lower", aspect="auto", interpolation="nearest")


def plot_fbank(fbank, title=None):
    fig, axs = plt.subplots(1, 1)
    axs.set_title(title or "Filter bank")
    axs.imshow(fbank, aspect="auto")
    axs.set_ylabel("frequency bin")
    axs.set_xlabel("mel bin")
