"""Graphiques consacrés aux cumuls de pluie et à leur répartition temporelle."""

from __future__ import annotations

from pathlib import Path

import matplotlib.dates as mdates
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

from .base import export_plot_dataset

__all__ = ['plot_daily_rainfall_hyetograph', 'plot_rainfall_by_season']


def plot_daily_rainfall_hyetograph(
    daily_rain: pd.DataFrame,
    output_path: str | Path,
) -> Path:
    """
    Affiche les cumuls quotidiens de pluie (barres) et le cumul annuel (ligne).
    """
    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    if daily_rain.empty:
        fig, ax = plt.subplots()
        ax.text(0.5, 0.5, "Pas de données de précipitations disponibles.", ha="center", va="center")
        ax.set_axis_off()
        fig.savefig(output_path, dpi=150, bbox_inches="tight")
        plt.close(fig)
        return output_path.resolve()

    export_plot_dataset(daily_rain, output_path)

    fig, ax1 = plt.subplots(figsize=(12, 5))
    ax1.bar(
        daily_rain.index,
        daily_rain["daily_total"],
        width=0.8,
        color="tab:blue",
        alpha=0.7,
        label="Pluie quotidienne",
    )
    ax1.set_ylabel("Pluie quotidienne (mm)")
    ax1.set_xlabel("Date")
    ax1.grid(True, axis="y", linestyle=":", alpha=0.5)

    ax2 = ax1.twinx()
    ax2.plot(
        daily_rain.index,
        daily_rain["cumulative_total"],
        color="tab:red",
        linewidth=2,
        label="Cumul annuel",
    )
    ax2.set_ylabel("Pluie cumulée (mm)")

    locator = mdates.AutoDateLocator()
    formatter = mdates.ConciseDateFormatter(locator)
    ax1.xaxis.set_major_locator(locator)
    ax1.xaxis.set_major_formatter(formatter)

    lines_labels = [
        (ax1.get_legend_handles_labels()),
        (ax2.get_legend_handles_labels()),
    ]
    lines, labels = [sum(lol, []) for lol in zip(*lines_labels)]
    ax1.legend(lines, labels, loc="upper left")

    fig.tight_layout()
    fig.savefig(output_path, dpi=150)
    plt.close(fig)
    return output_path.resolve()

def plot_rainfall_by_season(
    rainfall_df: pd.DataFrame,
    output_path: str | Path,
    *,
    title: str = "Pluie cumulée par saison",
) -> Path:
    """
    Affiche la pluie cumulée par saison ainsi que le nombre d'heures pluvieuses.
    """
    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    if rainfall_df.empty:
        fig, ax = plt.subplots()
        ax.text(0.5, 0.5, "Pas de données de pluie saisonnière.", ha="center", va="center")
        ax.set_axis_off()
        fig.savefig(output_path, dpi=150, bbox_inches="tight")
        plt.close(fig)
        return output_path.resolve()

    export_plot_dataset(rainfall_df, output_path)

    seasons = rainfall_df.index.tolist()
    x = np.arange(len(seasons))
    totals = rainfall_df["total_rain_mm"].to_numpy(dtype=float)

    fig, ax1 = plt.subplots(figsize=(9, 4))
    bars = ax1.bar(x, totals, color="tab:blue", alpha=0.7, label="Pluie cumulée")
    ax1.set_ylabel("Pluie cumulée (mm)")
    ax1.set_xlabel("Saison")
    ax1.set_xticks(x)
    ax1.set_xticklabels([season.capitalize() for season in seasons])
    ax1.grid(True, axis="y", linestyle=":", alpha=0.5)

    for rect, value in zip(bars, totals):
        height = rect.get_height()
        ax1.text(rect.get_x() + rect.get_width() / 2, height, f"{value:.0f}", ha="center", va="bottom", fontsize=8)

    lines = []
    labels = []

    if "rainy_hours" in rainfall_df.columns:
        ax2 = ax1.twinx()
        rainy_hours = rainfall_df["rainy_hours"].to_numpy(dtype=float)
        line = ax2.plot(
            x,
            rainy_hours,
            color="tab:red",
            marker="o",
            label="Heures pluvieuses",
        )[0]
        ax2.set_ylabel("Heures pluvieuses")
        lines.append(line)
        labels.append("Heures pluvieuses")

    handles, lbls = ax1.get_legend_handles_labels()
    handles.extend(lines)
    lbls.extend(labels)
    if handles:
        ax1.legend(handles, lbls, loc="upper left")

    ax1.set_title(title)
    fig.tight_layout()
    fig.savefig(output_path, dpi=150)
    plt.close(fig)
    return output_path.resolve()