"""Tracés orientés calendrier (heatmaps quotidiennes et profils hebdomadaires)."""

from __future__ import annotations

from pathlib import Path
from typing import Sequence

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

from .base import export_plot_dataset
from meteo.variables import Variable

__all__ = ['plot_calendar_heatmap', 'plot_weekday_profiles']


def plot_calendar_heatmap(
    matrix: pd.DataFrame,
    output_path: str | Path,
    *,
    title: str,
    cmap: str = "YlGnBu",
    colorbar_label: str = "",
) -> Path:
    """
    Affiche une heatmap calendrier (lignes = mois, colonnes = jours).
    """
    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    export_plot_dataset(matrix, output_path)

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

    fig, ax = plt.subplots(figsize=(14, 6))
    data = matrix.to_numpy(dtype=float)
    im = ax.imshow(data, aspect="auto", cmap=cmap, interpolation="nearest")

    ax.set_xticks(np.arange(matrix.shape[1]))
    ax.set_xticklabels(matrix.columns, rotation=90)
    ax.set_yticks(np.arange(matrix.shape[0]))
    ax.set_yticklabels(matrix.index)

    ax.set_xlabel("Jour du mois")
    ax.set_ylabel("Mois")
    ax.set_title(title)

    cbar = fig.colorbar(im, ax=ax)
    if colorbar_label:
        cbar.set_label(colorbar_label)

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

def plot_weekday_profiles(
    weekday_df: pd.DataFrame,
    variables: Sequence[Variable],
    output_path: str | Path,
    *,
    title: str,
) -> Path:
    """
    Affiche les moyennes par jour de semaine pour plusieurs variables.
    """
    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    if weekday_df.empty:
        fig, ax = plt.subplots()
        ax.text(0.5, 0.5, "Pas de données hebdomadaires.", 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(weekday_df, output_path)

    weekday_labels = ["Lun", "Mar", "Mer", "Jeu", "Ven", "Sam", "Dim"]
    n_vars = len(variables)
    fig, axes = plt.subplots(n_vars, 1, figsize=(10, 3 * n_vars), sharex=True)
    if n_vars == 1:
        axes = [axes]

    x = np.arange(len(weekday_labels))

    for ax, var in zip(axes, variables):
        if var.column not in weekday_df.columns:
            ax.text(0.5, 0.5, f"Aucune donnée pour {var.label}.", ha="center", va="center")
            ax.set_axis_off()
            continue

        values = weekday_df[var.column].to_numpy(dtype=float)
        ax.plot(x, values, marker="o", label=var.label)
        ax.set_ylabel(f"{var.label} ({var.unit})" if var.unit else var.label)
        ax.grid(True, linestyle=":", alpha=0.5)
        ax.set_xticks(x)
        ax.set_xticklabels(weekday_labels)

    axes[-1].set_xlabel("Jour de semaine")
    axes[0].legend(loc="upper right")
    fig.suptitle(title)
    fig.tight_layout(rect=[0, 0, 1, 0.97])
    fig.savefig(output_path, dpi=150)
    plt.close(fig)
    return output_path.resolve()