"""Exploratory data analysis and visualization utilities.

All plot functions accept flexible arguments (column names, style, save_path)
and return (fig, ax) for display or further composition. Optional **kwargs
are passed through to the underlying matplotlib/seaborn/plotly calls.
"""

from __future__ import annotations

from collections import Counter
from typing import TYPE_CHECKING, Any, Sequence

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
from sklearn.metrics import confusion_matrix as sklearn_confusion_matrix
from sklearn.metrics import roc_curve, auc

if TYPE_CHECKING:
    from matplotlib.axes import Axes
    from matplotlib.figure import Figure


# ---------------------------------------------------------------------------
# Inspection
# ---------------------------------------------------------------------------


def initial_inspection(
    df: pd.DataFrame,
    sample_n: int = 5,
    columns: Sequence[str] | None = None,
    include_dtypes: bool = True,
    include_nulls: bool = True,
    return_dict: bool = False,
) -> dict[str, Any] | None:
    """Shape, dtypes, null counts, and sample rows. Prints by default."""
    subset = df[columns] if columns is not None else df
    out: dict[str, Any] = {
        "shape": subset.shape,
        "columns": list(subset.columns),
    }
    if include_dtypes:
        out["dtypes"] = subset.dtypes.to_dict()
    if include_nulls:
        out["null_counts"] = subset.isnull().sum().to_dict()
    out["sample"] = subset.head(sample_n).to_dict("records")

    if return_dict:
        return out
    print("Shape:", out["shape"])
    print("Columns:", out["columns"])
    if include_dtypes:
        print("Dtypes:\n", pd.Series(out["dtypes"]))
    if include_nulls:
        print("Null counts:\n", pd.Series(out["null_counts"]))
    print("Sample rows:\n", subset.head(sample_n))
    return None


# ---------------------------------------------------------------------------
# Helpers: figure/axis and save
# ---------------------------------------------------------------------------


def _get_fig_ax(
    ax: Axes | None,
    figsize: tuple[float, float] = (10, 5),
) -> tuple[Figure, Axes]:
    from matplotlib.figure import Figure

    if ax is not None:
        return ax.get_figure(), ax
    fig, ax = plt.subplots(figsize=figsize)
    return fig, ax


def _save_and_show(
    fig: Figure,
    save_path: str | None = None,
    show: bool = True,
    dpi: int = 100,
) -> None:
    if save_path:
        fig.savefig(save_path, bbox_inches="tight", dpi=dpi)
    if show:
        plt.show()


# ---------------------------------------------------------------------------
# Generic plotters
# ---------------------------------------------------------------------------


def plot_bar(
    data: pd.DataFrame | pd.Series | dict[str, Any] | Sequence[Any],
    x_col: str | None = None,
    y_col: str | None = None,
    order: Sequence[Any] | None = None,
    figsize: tuple[float, float] = (10, 5),
    title: str | None = None,
    xlabel: str | None = None,
    ylabel: str | None = None,
    palette: str | Sequence[str] | None = None,
    ax: Axes | None = None,
    save_path: str | None = None,
    show: bool = True,
    **kwargs: Any,
) -> tuple[Figure, Axes]:
    """Bar chart from counts or from a DataFrame column."""
    fig, ax = _get_fig_ax(ax, figsize)

    if isinstance(data, pd.DataFrame):
        if x_col is None:
            raise ValueError("x_col is required when data is a DataFrame")
        counts = data[x_col].value_counts()
        if order is not None:
            counts = counts.reindex(order).dropna()
        x = counts.index.tolist()
        y = counts.values
    elif isinstance(data, dict):
        if order is not None:
            x = list(order)
            y = [data.get(k, 0) for k in order]
        else:
            x = list(data.keys())
            y = list(data.values())
    elif isinstance(data, pd.Series):
        counts = data.value_counts()
        if order is not None:
            counts = counts.reindex(order).dropna()
        x = counts.index.tolist()
        y = counts.values.tolist()
    else:
        raise TypeError("data must be DataFrame, Series, or dict")

    bars = ax.bar(x, y, color=None, **kwargs)
    if palette:
        colors = sns.color_palette(palette, n_colors=len(x)) if isinstance(palette, str) else palette
        for bar, c in zip(bars, colors[: len(x)]):
            bar.set_facecolor(c)

    if title:
        ax.set_title(title)
    if xlabel:
        ax.set_xlabel(xlabel)
    if ylabel:
        ax.set_ylabel(ylabel)
    ax.tick_params(axis="x", rotation=45)
    fig.tight_layout()
    _save_and_show(fig, save_path, show)
    return fig, ax


def plot_histogram(
    data: pd.DataFrame | np.ndarray | Sequence[float],
    column: str | None = None,
    n_bins: int = 50,
    kde: bool = False,
    figsize: tuple[float, float] = (10, 5),
    title: str | None = None,
    xlabel: str | None = None,
    ylabel: str | None = "Count",
    ax: Axes | None = None,
    save_path: str | None = None,
    show: bool = True,
    **kwargs: Any,
) -> tuple[Figure, Axes]:
    """Histogram (and optional KDE) for a single numeric variable."""
    fig, ax = _get_fig_ax(ax, figsize)
    values = data[column] if isinstance(data, pd.DataFrame) and column else np.asarray(data)
    values = np.asarray(values).flatten()
    values = values[~np.isnan(values)]

    sns.histplot(values, bins=n_bins, kde=kde, ax=ax, **kwargs)
    if title:
        ax.set_title(title)
    if xlabel:
        ax.set_xlabel(xlabel)
    if ylabel:
        ax.set_ylabel(ylabel)
    fig.tight_layout()
    _save_and_show(fig, save_path, show)
    return fig, ax


def plot_heatmap(
    data: pd.DataFrame | np.ndarray,
    xlabel: str | None = None,
    ylabel: str | None = None,
    title: str | None = None,
    annot: bool = True,
    fmt: str = ".2f",
    cmap: str = "Blues",
    figsize: tuple[float, float] = (8, 6),
    ax: Axes | None = None,
    save_path: str | None = None,
    show: bool = True,
    **kwargs: Any,
) -> tuple[Figure, Axes]:
    """2D heatmap from a matrix or DataFrame."""
    fig, ax = _get_fig_ax(ax, figsize)
    if isinstance(data, pd.DataFrame):
        plot_data = data
    else:
        plot_data = np.asarray(data)
        if plot_data.ndim != 2:
            raise ValueError("data must be 2D")
    sns.heatmap(plot_data, annot=annot, fmt=fmt, cmap=cmap, ax=ax, **kwargs)
    if title:
        ax.set_title(title)
    if xlabel:
        ax.set_xlabel(xlabel)
    if ylabel:
        ax.set_ylabel(ylabel)
    fig.tight_layout()
    _save_and_show(fig, save_path, show)
    return fig, ax


def plot_line(
    data: pd.DataFrame | np.ndarray,
    x_col: str | None = None,
    y_col: str | list[str] | None = None,
    figsize: tuple[float, float] = (10, 5),
    title: str | None = None,
    xlabel: str | None = None,
    ylabel: str | None = None,
    ax: Axes | None = None,
    save_path: str | None = None,
    show: bool = True,
    **kwargs: Any,
) -> tuple[Figure, Axes]:
    """Line plot (x vs y or multiple series)."""
    fig, ax = _get_fig_ax(ax, figsize)
    if isinstance(data, pd.DataFrame):
        if x_col is None or y_col is None:
            raise ValueError("x_col and y_col required when data is DataFrame")
        x = data[x_col].values
        y_cols = [y_col] if isinstance(y_col, str) else y_col
        for col in y_cols:
            ax.plot(x, data[col].values, label=col, **kwargs)
        if len(y_cols) > 1:
            ax.legend()
    else:
        arr = np.asarray(data)
        if arr.ndim == 1:
            ax.plot(arr, **kwargs)
        else:
            for i in range(arr.shape[1]):
                ax.plot(arr[:, i], label=str(i), **kwargs)
            ax.legend()
    if title:
        ax.set_title(title)
    if xlabel:
        ax.set_xlabel(xlabel)
    if ylabel:
        ax.set_ylabel(ylabel)
    fig.tight_layout()
    _save_and_show(fig, save_path, show)
    return fig, ax


def plot_scatter(
    data: pd.DataFrame,
    x_col: str,
    y_col: str,
    hue: str | None = None,
    size: str | None = None,
    figsize: tuple[float, float] = (8, 6),
    title: str | None = None,
    xlabel: str | None = None,
    ylabel: str | None = None,
    ax: Axes | None = None,
    save_path: str | None = None,
    show: bool = True,
    **kwargs: Any,
) -> tuple[Figure, Axes]:
    """Scatter plot (x vs y, optional color/size by column)."""
    fig, ax = _get_fig_ax(ax, figsize)
    sns.scatterplot(
        data=data,
        x=x_col,
        y=y_col,
        hue=hue,
        size=size,
        ax=ax,
        **kwargs,
    )
    if title:
        ax.set_title(title)
    if xlabel:
        ax.set_xlabel(xlabel)
    if ylabel:
        ax.set_ylabel(ylabel)
    fig.tight_layout()
    _save_and_show(fig, save_path, show)
    return fig, ax


# ---------------------------------------------------------------------------
# EDA helpers
# ---------------------------------------------------------------------------


def plot_class_distribution(
    df: pd.DataFrame,
    target_col: str,
    order: Sequence[Any] | None = None,
    normalize: bool = False,
    figsize: tuple[float, float] = (10, 5),
    title: str | None = None,
    xlabel: str | None = None,
    ylabel: str | None = None,
    palette: str | Sequence[str] | None = None,
    ax: Axes | None = None,
    save_path: str | None = None,
    show: bool = True,
    **kwargs: Any,
) -> tuple[Figure, Axes]:
    """Counts of a categorical column (bar chart)."""
    counts = df[target_col].value_counts(normalize=normalize)
    if order is not None:
        counts = counts.reindex(order).dropna()
    fig, ax = plot_bar(
        counts,
        figsize=figsize,
        title=title or f"Distribution of {target_col}",
        xlabel=xlabel or target_col,
        ylabel=ylabel or ("Proportion" if normalize else "Count"),
        palette=palette,
        ax=ax,
        save_path=save_path,
        show=show,
        **kwargs,
    )
    return fig, ax


def plot_text_length_distribution(
    df: pd.DataFrame,
    text_col: str,
    kind: str = "char",
    n_bins: int = 50,
    kde: bool = True,
    figsize: tuple[float, float] = (10, 5),
    title: str | None = None,
    xlabel: str | None = None,
    ax: Axes | None = None,
    save_path: str | None = None,
    show: bool = True,
    **kwargs: Any,
) -> tuple[Figure, Axes]:
    """Distribution of text lengths (char or word count)."""
    texts = df[text_col].dropna().astype(str)
    if kind == "char":
        lengths = texts.str.len().values
        xlabel = xlabel or "Character count"
    elif kind == "word":
        lengths = texts.str.split().str.len().values
        xlabel = xlabel or "Word count"
    else:
        raise ValueError("kind must be 'char' or 'word'")
    lengths = lengths[~np.isnan(lengths)]
    fig, ax = plot_histogram(
        lengths,
        n_bins=n_bins,
        kde=kde,
        figsize=figsize,
        title=title or f"Text length distribution ({kind})",
        xlabel=xlabel,
        ax=ax,
        save_path=save_path,
        show=show,
        **kwargs,
    )
    return fig, ax


def _tokenize_ngrams(texts: Sequence[str], n: int) -> Counter:
    """Return Counter of n-grams (space tokenized)."""
    counter: Counter = Counter()
    for text in texts:
        tokens = str(text).lower().split()
        for i in range(len(tokens) - n + 1):
            ngram = tuple(tokens[i : i + n])
            counter[ngram] += 1
    return counter


def plot_ngram_frequency(
    texts: Sequence[str],
    n: int = 1,
    top_k: int = 20,
    title: str | None = None,
    figsize: tuple[float, float] = (10, 5),
    backend: str = "matplotlib",
    save_path: str | None = None,
    show: bool = True,
    **kwargs: Any,
) -> Figure | tuple[Figure, Axes]:
    """Top-k n-grams as bar chart. Returns Plotly fig if backend='plotly', else (fig, ax)."""
    counter = _tokenize_ngrams(texts, n)
    top = counter.most_common(top_k)
    if not top:
        # empty figure
        fig, ax = plt.subplots(figsize=figsize)
        if title:
            ax.set_title(title)
        if save_path:
            fig.savefig(save_path, bbox_inches="tight")
        return fig, ax
    labels = [" ".join(t) for t, _ in top]
    counts = [c for _, c in top]

    if backend == "plotly":
        try:
            import plotly.express as px
            import plotly.graph_objects as go
        except ImportError:
            backend = "matplotlib"
        else:
            fig = go.Figure(go.Bar(x=counts, y=labels, orientation="h"))
            fig.update_layout(
                title=title or f"Top {top_k} {n}-grams",
                xaxis_title="Count",
                yaxis_title="N-gram",
                height=figsize[1] * 80,
                width=figsize[0] * 80,
                **kwargs,
            )
            if save_path:
                fig.write_image(save_path) if hasattr(fig, "write_image") else fig.write_html(save_path.replace(".png", ".html"))
            return fig

    fig, ax = _get_fig_ax(None, figsize)
    ax.barh(labels[::-1], counts[::-1], **kwargs)
    ax.set_title(title or f"Top {top_k} {n}-grams")
    ax.set_xlabel("Count")
    fig.tight_layout()
    _save_and_show(fig, save_path, show)
    return fig, ax


def plot_wordcloud(
    texts: Sequence[str],
    label: str | None = None,
    max_words: int = 100,
    width: int = 800,
    height: int = 400,
    figsize: tuple[float, float] = (12, 6),
    save_path: str | None = None,
    show: bool = True,
    **kwargs: Any,
) -> tuple[Figure, Axes]:
    """Word cloud from a list of text strings. Requires optional dependency: wordcloud."""
    try:
        from wordcloud import WordCloud
    except ImportError as e:
        raise ImportError("plot_wordcloud requires the 'wordcloud' package. Install with: pip install wordcloud") from e
    combined = " ".join(str(t) for t in texts)
    wc = WordCloud(max_words=max_words, width=width, height=height, **kwargs).generate(combined)
    fig, ax = plt.subplots(figsize=(figsize[0] / 80, figsize[1] / 80))
    ax.imshow(wc, interpolation="bilinear")
    ax.axis("off")
    if label:
        ax.set_title(label)
    if save_path:
        fig.savefig(save_path, bbox_inches="tight", pad_inches=0)
    if show:
        plt.show()
    return fig, ax


# ---------------------------------------------------------------------------
# Evaluation helpers
# ---------------------------------------------------------------------------


def plot_confusion_matrix(
    y_true: Sequence[Any],
    y_pred: Sequence[Any],
    labels: Sequence[Any] | None = None,
    title: str | None = "Confusion matrix",
    figsize: tuple[float, float] = (8, 6),
    annot: bool = True,
    cmap: str = "Blues",
    ax: Axes | None = None,
    save_path: str | None = None,
    show: bool = True,
    **kwargs: Any,
) -> tuple[Figure, Axes]:
    """Confusion matrix as heatmap."""
    y_true = np.asarray(y_true)
    y_pred = np.asarray(y_pred)
    cm = sklearn_confusion_matrix(y_true, y_pred)
    if labels is not None:
        # assume labels order matches 0,1,...,n-1
        ticklabels = list(labels)
    else:
        ticklabels = list(range(cm.shape[0]))
    fig, ax = _get_fig_ax(ax, figsize)
    sns.heatmap(cm, annot=annot, fmt="d", cmap=cmap, ax=ax, xticklabels=ticklabels, yticklabels=ticklabels, **kwargs)
    ax.set_xlabel("Predicted")
    ax.set_ylabel("True")
    if title:
        ax.set_title(title)
    fig.tight_layout()
    _save_and_show(fig, save_path, show)
    return fig, ax


def plot_roc_curves(
    y_true: Sequence[int | float],
    y_scores_dict: dict[str, np.ndarray | Sequence[float]],
    figsize: tuple[float, float] = (8, 6),
    title: str | None = "ROC curves",
    ax: Axes | None = None,
    save_path: str | None = None,
    show: bool = True,
    **kwargs: Any,
) -> tuple[Figure, Axes]:
    """Multiple ROC curves (one per model)."""
    y_true = np.asarray(y_true)
    fig, ax = _get_fig_ax(ax, figsize)
    ax.plot([0, 1], [0, 1], "k--", label="Chance")
    for name, scores in y_scores_dict.items():
        scores = np.asarray(scores).flatten()
        fpr, tpr, _ = roc_curve(y_true, scores)
        roc_auc = auc(fpr, tpr)
        ax.plot(fpr, tpr, label=f"{name} (AUC = {roc_auc:.2f})", **kwargs)
    ax.set_xlabel("False Positive Rate")
    ax.set_ylabel("True Positive Rate")
    if title:
        ax.set_title(title)
    ax.legend()
    ax.set_xlim(0, 1)
    ax.set_ylim(0, 1)
    fig.tight_layout()
    _save_and_show(fig, save_path, show)
    return fig, ax


def plot_model_comparison(
    results_dict: dict[str, dict[str, float]],
    metrics: Sequence[str] | None = None,
    figsize: tuple[float, float] = (10, 5),
    title: str | None = "Model comparison",
    backend: str = "matplotlib",
    save_path: str | None = None,
    show: bool = True,
    **kwargs: Any,
) -> Figure | tuple[Figure, Axes]:
    """Grouped bar chart comparing metrics across models."""
    df = pd.DataFrame(results_dict).T
    if metrics is not None:
        df = df[[m for m in metrics if m in df.columns]]
    df = df.reset_index().rename(columns={"index": "model"})
    df_melt = df.melt(id_vars="model", var_name="metric", value_name="value")

    if backend == "plotly":
        try:
            import plotly.express as px
        except ImportError:
            backend = "matplotlib"
        else:
            fig = px.bar(df_melt, x="model", y="value", color="metric", barmode="group", title=title or "Model comparison", **kwargs)
            if save_path:
                fig.write_image(save_path) if hasattr(fig, "write_image") else fig.write_html(save_path.replace(".png", ".html"))
            return fig

    fig, ax = _get_fig_ax(None, figsize)
    sns.barplot(data=df_melt, x="model", y="value", hue="metric", ax=ax, **kwargs)
    if title:
        ax.set_title(title)
    ax.tick_params(axis="x", rotation=45)
    fig.tight_layout()
    _save_and_show(fig, save_path, show)
    return fig, ax
