"""Visualization helpers for JAO bids."""

from __future__ import annotations

from typing import Any, Tuple

import matplotlib.pyplot as plt

from .parsing import prepare_bid_dataframe, find_offered_capacity, prepare_curtailment_dataframe
from .utils import humanize_label


def plot_bid_stack_step(
    bids_payload: Any,
    offered_capacity: float | None = None,
    title: str = "Bid Stack: Cumulative Step Curve",
) -> Tuple[plt.Figure, plt.Axes]:
    """Plot the cumulative bid stack step curve.

    The chart displays cumulative awarded quantity on the x-axis and price on
    the y-axis, with reference lines and legend annotations.
    """
    df, price_col, y_col, _ = prepare_bid_dataframe(bids_payload)
    if df is None or price_col is None or y_col is None:
        raise ValueError("Could not detect price/quantity columns in bids payload.")

    df = df.sort_values(price_col, ascending=False).reset_index(drop=True)
    df["cum_qty"] = df[y_col].cumsum()

    total_qty = df["cum_qty"].iloc[-1]
    max_price = df[price_col].max()

    awarded_mask = df[y_col] > 0
    if awarded_mask.any():
        min_award_price = df.loc[awarded_mask, price_col].min()
    else:
        min_award_price = df[price_col].min()

    if offered_capacity is None:
        offered_capacity = find_offered_capacity(bids_payload)

    y_label = humanize_label(y_col)
    price_label = humanize_label(price_col)

    plt.style.use("seaborn-v0_8-whitegrid")
    fig, ax = plt.subplots(figsize=(10, 5.5))

    ax.step(df["cum_qty"], df[price_col], where="post", linewidth=2.5, color="#1F77B4")
    ax.fill_between(df["cum_qty"], df[price_col], step="post", alpha=0.15, color="#1F77B4")

    ax.scatter(
        [df["cum_qty"].iloc[0]],
        [max_price],
        color="#1F77B4",
        s=40,
        zorder=3,
        label=f"Max price: {max_price:.2f}",
    )
    ax.scatter(
        [total_qty],
        [min_award_price],
        color="#D62728",
        s=40,
        zorder=3,
        label=f"Min awarded price: {min_award_price:.2f}",
    )

    # Reference lines (no legend entries)
    ax.axvline(total_qty, color="#555555", linestyle="--", linewidth=1, alpha=0.6, label="_nolegend_")
    ax.axhline(min_award_price, color="#D62728", linestyle="--", linewidth=1, alpha=0.7, label="_nolegend_")

    if offered_capacity is not None:
        ax.axvline(
            offered_capacity,
            color="#2CA02C",
            linestyle="-.",
            linewidth=1.5,
            alpha=0.8,
            label=f"Offered capacity: {offered_capacity:.2f}",
        )

    # Legend-only entry for total quantity
    ax.plot([], [], color="none", label=f"Total {y_label}: {total_qty:.2f}")

    ax.set_xlabel(f"Cumulative {y_label}")
    ax.set_ylabel(price_label)
    ax.set_title(title)
    ax.set_xlim(left=0)
    ax.grid(True, alpha=0.25)

    ax.legend(loc="upper right", frameon=True)
    return fig, ax


def plot_curtailment_stacked_bar(
    curtailments: Any,
    title: str = "Curtailments by Reason",
    value_col: str = "curtailedCapacity",
) -> Tuple[plt.Figure, plt.Axes]:
    """Plot a stacked bar chart of curtailments by reason.

    By default, uses curtailed capacity (original - remaining).
    """
    df = prepare_curtailment_dataframe(curtailments)
    if df is None or df.empty:
        raise ValueError("No curtailment data to plot.")

    if value_col not in df.columns:
        if "curtailedCapacity" in df.columns:
            value_col = "curtailedCapacity"
        elif "originalCapacity" in df.columns:
            value_col = "originalCapacity"
        elif "curtailmentFactor" in df.columns:
            value_col = "curtailmentFactor"
        else:
            raise ValueError("No usable numeric curtailment value found.")

    if "periodDate" not in df.columns:
        raise ValueError("Missing periodDate field; cannot group by time.")

    grouped = (
        df.groupby(["periodDate", "reason"], dropna=False)[value_col]
        .sum()
        .unstack(fill_value=0)
        .sort_index()
    )

    plt.style.use("seaborn-v0_8-whitegrid")
    fig, ax = plt.subplots(figsize=(11, 5.5))
    grouped.plot(kind="bar", stacked=True, ax=ax, colormap="tab20")

    ax.set_xlabel("Date")
    ax.set_ylabel(humanize_label(value_col))
    ax.set_title(title)
    ax.grid(True, axis="y", alpha=0.25)
    ax.legend(title="Reason", bbox_to_anchor=(1.02, 1), loc="upper left", frameon=True)
    fig.tight_layout()
    return fig, ax
