"""
Viz — Matplotlib and Plotly chart helpers for business data visualization.
"""

from __future__ import annotations

from typing import Any, Literal

import numpy as np
import pandas as pd


def bar(
    df: pd.DataFrame,
    x: str,
    y: str,
    title: str = "",
    color: str = "#8b5cf6",
    output: str | None = None,
    backend: Literal["matplotlib", "plotly"] = "plotly",
) -> Any:
    """Bar chart."""
    if backend == "plotly":
        import plotly.express as px
        fig = px.bar(df, x=x, y=y, title=title, color_discrete_sequence=[color])
        _plotly_style(fig)
        if output:
            fig.write_html(output)
        return fig
    else:
        import matplotlib.pyplot as plt
        fig, ax = plt.subplots(figsize=(10, 5))
        ax.bar(df[x], df[y], color=color)
        ax.set_title(title)
        _mpl_style(ax)
        if output:
            fig.savefig(output, dpi=150, bbox_inches="tight")
        return fig


def line(
    df: pd.DataFrame,
    x: str,
    y: str | list[str],
    title: str = "",
    output: str | None = None,
    backend: Literal["matplotlib", "plotly"] = "plotly",
) -> Any:
    """Line chart. y can be a single column or a list of columns."""
    ys = [y] if isinstance(y, str) else y
    if backend == "plotly":
        import plotly.graph_objects as go
        colors = ["#8b5cf6", "#34d399", "#f59e0b", "#ef4444", "#60a5fa"]
        fig = go.Figure()
        for i, col in enumerate(ys):
            fig.add_trace(go.Scatter(x=df[x], y=df[col], name=col, line=dict(color=colors[i % len(colors)], width=2)))
        fig.update_layout(title=title)
        _plotly_style(fig)
        if output:
            fig.write_html(output)
        return fig
    else:
        import matplotlib.pyplot as plt
        fig, ax = plt.subplots(figsize=(10, 5))
        for col in ys:
            ax.plot(df[x], df[col], label=col)
        ax.set_title(title)
        ax.legend()
        _mpl_style(ax)
        if output:
            fig.savefig(output, dpi=150, bbox_inches="tight")
        return fig


def scatter(
    df: pd.DataFrame,
    x: str,
    y: str,
    color_col: str | None = None,
    size_col: str | None = None,
    title: str = "",
    output: str | None = None,
    backend: Literal["matplotlib", "plotly"] = "plotly",
) -> Any:
    """Scatter plot with optional color and size encoding."""
    if backend == "plotly":
        import plotly.express as px
        fig = px.scatter(df, x=x, y=y, color=color_col, size=size_col, title=title,
                         color_discrete_sequence=["#8b5cf6", "#34d399", "#f59e0b", "#ef4444"])
        _plotly_style(fig)
        if output:
            fig.write_html(output)
        return fig
    else:
        import matplotlib.pyplot as plt
        fig, ax = plt.subplots(figsize=(8, 6))
        ax.scatter(df[x], df[y], alpha=0.6, color="#8b5cf6")
        ax.set_xlabel(x)
        ax.set_ylabel(y)
        ax.set_title(title)
        _mpl_style(ax)
        if output:
            fig.savefig(output, dpi=150, bbox_inches="tight")
        return fig


def histogram(
    series: pd.Series,
    bins: int = 30,
    title: str = "",
    output: str | None = None,
    backend: Literal["matplotlib", "plotly"] = "plotly",
) -> Any:
    """Histogram of a numeric series."""
    if backend == "plotly":
        import plotly.express as px
        fig = px.histogram(x=series, nbins=bins, title=title, color_discrete_sequence=["#8b5cf6"])
        _plotly_style(fig)
        if output:
            fig.write_html(output)
        return fig
    else:
        import matplotlib.pyplot as plt
        fig, ax = plt.subplots(figsize=(8, 5))
        ax.hist(series.dropna(), bins=bins, color="#8b5cf6", edgecolor="none", alpha=0.8)
        ax.set_title(title)
        _mpl_style(ax)
        if output:
            fig.savefig(output, dpi=150, bbox_inches="tight")
        return fig


def heatmap(
    df: pd.DataFrame,
    title: str = "Correlation Heatmap",
    output: str | None = None,
    backend: Literal["matplotlib", "plotly"] = "plotly",
) -> Any:
    """Correlation heatmap."""
    corr = df.select_dtypes(include="number").corr()
    if backend == "plotly":
        import plotly.figure_factory as ff
        fig = ff.create_annotated_heatmap(
            z=corr.values,
            x=list(corr.columns),
            y=list(corr.index),
            annotation_text=corr.round(2).values.astype(str),
            colorscale="Purples",
        )
        fig.update_layout(title=title)
        _plotly_style(fig)
        if output:
            fig.write_html(output)
        return fig
    else:
        import matplotlib.pyplot as plt
        fig, ax = plt.subplots(figsize=(max(6, len(corr.columns)), max(5, len(corr))))
        im = ax.imshow(corr.values, cmap="Purples", vmin=-1, vmax=1)
        ax.set_xticks(range(len(corr.columns)))
        ax.set_yticks(range(len(corr)))
        ax.set_xticklabels(corr.columns, rotation=45, ha="right")
        ax.set_yticklabels(corr.index)
        plt.colorbar(im, ax=ax)
        ax.set_title(title)
        if output:
            fig.savefig(output, dpi=150, bbox_inches="tight")
        return fig


def cohort_heatmap(
    cohort_df: pd.DataFrame,
    title: str = "Cohort Retention",
    output: str | None = None,
) -> Any:
    """Cohort retention heatmap using Plotly."""
    import plotly.graph_objects as go
    fig = go.Figure(go.Heatmap(
        z=cohort_df.values,
        x=list(cohort_df.columns),
        y=list(cohort_df.index),
        colorscale="Purples",
        text=cohort_df.values.round(1),
        texttemplate="%{text}%",
        showscale=True,
    ))
    fig.update_layout(title=title, xaxis_title="Period Offset", yaxis_title="Cohort")
    _plotly_style(fig)
    if output:
        fig.write_html(output)
    return fig


def montecarlo_chart(
    bins_df: pd.DataFrame,
    p10: float, p50: float, p90: float,
    title: str = "Monte Carlo Distribution",
    output: str | None = None,
) -> Any:
    """Monte Carlo histogram with P10/P50/P90 vertical lines."""
    import plotly.graph_objects as go

    colors = ["#ef4444" if row["is_negative"] else "#8b5cf6" for _, row in bins_df.iterrows()]

    fig = go.Figure()
    fig.add_trace(go.Bar(
        x=bins_df["center"],
        y=bins_df["count"],
        marker_color=colors,
        name="Simulations",
        width=(bins_df["bin_max"] - bins_df["bin_min"]).mean() * 0.9,
    ))
    for val, name, color in [(p10, "P10", "#f87171"), (p50, "P50 (Median)", "#a78bfa"), (p90, "P90", "#34d399")]:
        fig.add_vline(x=val, line_dash="dash", line_color=color, annotation_text=name)

    fig.update_layout(title=title, xaxis_title="Cumulative Cash Flow", yaxis_title="Count")
    _plotly_style(fig)
    if output:
        fig.write_html(output)
    return fig


# ── Style helpers ─────────────────────────────────────────────────────────────
def _plotly_style(fig):
    fig.update_layout(
        paper_bgcolor="#0f0e17",
        plot_bgcolor="#0f0e17",
        font=dict(color="#e0dff0", family="system-ui"),
        title_font=dict(color="#a78bfa", size=16),
        legend=dict(bgcolor="rgba(0,0,0,0)"),
    )
    fig.update_xaxes(gridcolor="rgba(255,255,255,0.05)", zeroline=False)
    fig.update_yaxes(gridcolor="rgba(255,255,255,0.05)", zeroline=False)


def _mpl_style(ax):
    ax.set_facecolor("#0f0e17")
    ax.figure.patch.set_facecolor("#0f0e17")
    ax.tick_params(colors="#9ca3af")
    ax.xaxis.label.set_color("#9ca3af")
    ax.yaxis.label.set_color("#9ca3af")
    ax.title.set_color("#a78bfa")
    for spine in ax.spines.values():
        spine.set_color("#374151")


__all__ = ["bar", "line", "scatter", "histogram", "heatmap", "cohort_heatmap", "montecarlo_chart"]
