# scripts/plot_calendar_overview.py
from __future__ import annotations

from pathlib import Path

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

from meteo.dataset import load_raw_csv
from meteo.analysis import compute_daily_rainfall_totals
from meteo.plots import plot_calendar_heatmap, plot_weekday_profiles
from meteo.variables import VARIABLES_BY_KEY


CSV_PATH = Path("data/weather_minutely.csv")
OUTPUT_DIR = Path("figures/calendar")

WEEKDAY_VARIABLE_KEYS = ["temperature", "humidity", "wind_speed", "illuminance"]


def _format_calendar_matrix(series: pd.Series, year: int, agg_label: str) -> pd.DataFrame:
    """
    Transforme une série quotidienne en matrice mois x jours (1-31).
    """
    start = pd.Timestamp(year=year, month=1, day=1, tz=series.index.tz)
    end = pd.Timestamp(year=year, month=12, day=31, tz=series.index.tz)
    filtered = series.loc[(series.index >= start) & (series.index <= end)]

    matrix = pd.DataFrame(
        np.nan,
        index=[calendar.month_name[m][:3] for m in range(1, 13)],
        columns=list(range(1, 32)),
    )

    for timestamp, value in filtered.items():
        month = timestamp.month
        day = timestamp.day
        matrix.at[calendar.month_name[month][:3], day] = value

    matrix.index.name = f"{agg_label} ({year})"
    return matrix


def compute_daily_mean(df: pd.DataFrame, column: str) -> pd.Series:
    return df[column].resample("1D").mean()


def plot_combined_calendar(
    matrices: dict[str, pd.DataFrame],
    output_path: Path,
    *,
    title: str,
) -> None:
    if not matrices:
        return

    n = len(matrices)
    fig, axes = plt.subplots(n, 1, figsize=(14, 4 * n), sharex=True)
    if n == 1:
        axes = [axes]

    for ax, (label, matrix) in zip(axes, matrices.items()):
        data = matrix.to_numpy(dtype=float)
        im = ax.imshow(data, aspect="auto", interpolation="nearest", cmap=matrix.attrs.get("cmap", "viridis"))
        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_ylabel(label)
        cbar = fig.colorbar(im, ax=ax)
        if matrix.attrs.get("colorbar_label"):
            cbar.set_label(matrix.attrs["colorbar_label"])

    axes[-1].set_xlabel("Jour du mois")
    fig.suptitle(title)
    fig.tight_layout(rect=[0, 0, 1, 0.97])
    fig.savefig(output_path, dpi=150)
    plt.close(fig)


def main() -> None:
    if not CSV_PATH.exists():
        print(f"⚠ Fichier introuvable : {CSV_PATH}")
        return

    df = load_raw_csv(CSV_PATH)
    if df.empty:
        print("⚠ Dataset vide.")
        return

    if not isinstance(df.index, pd.DatetimeIndex):
        print("⚠ Le dataset doit avoir un index temporel.")
        return

    print(f"Dataset minuté chargé : {CSV_PATH}")
    print(f"  Lignes   : {len(df)}")
    print(f"  Colonnes : {list(df.columns)}")
    print()

    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

    latest_year = df.index.year.max()
    print(f"Année retenue pour le calendrier : {latest_year}")

    daily_totals = compute_daily_rainfall_totals(df=df)
    daily_rain = daily_totals["daily_total"]
    rain_matrix = _format_calendar_matrix(daily_rain, latest_year, "Pluie (mm)")
    rain_matrix.attrs["cmap"] = "Blues"
    rain_matrix.attrs["colorbar_label"] = "mm"
    rain_path = OUTPUT_DIR / f"calendar_rain_{latest_year}.png"
    plot_calendar_heatmap(
        matrix=rain_matrix,
        output_path=rain_path,
        title=f"Pluie quotidienne - {latest_year}",
        cmap="Blues",
        colorbar_label="mm",
    )
    print(f"✔ Heatmap pluie {latest_year} : {rain_path}")

    daily_temp = compute_daily_mean(df, "temperature")
    temp_matrix = _format_calendar_matrix(daily_temp, latest_year, "Température (°C)")
    temp_matrix.attrs["cmap"] = "coolwarm"
    temp_matrix.attrs["colorbar_label"] = "°C"
    temp_path = OUTPUT_DIR / f"calendar_temperature_{latest_year}.png"
    plot_calendar_heatmap(
        matrix=temp_matrix,
        output_path=temp_path,
        title=f"Température moyenne quotidienne - {latest_year}",
        cmap="coolwarm",
        colorbar_label="°C",
    )
    print(f"✔ Heatmap température {latest_year} : {temp_path}")

    matrices_for_combined = {
        "Pluie (mm)": rain_matrix,
        "Température (°C)": temp_matrix,
    }

    if "pressure" in df.columns:
        daily_pressure = compute_daily_mean(df, "pressure")
        pressure_matrix = _format_calendar_matrix(daily_pressure, latest_year, "Pression (hPa)")
        pressure_matrix.attrs["cmap"] = "Greens"
        pressure_matrix.attrs["colorbar_label"] = "hPa"
        pressure_path = OUTPUT_DIR / f"calendar_pressure_{latest_year}.png"
        plot_calendar_heatmap(
            matrix=pressure_matrix,
            output_path=pressure_path,
            title=f"Pression moyenne quotidienne - {latest_year}",
            cmap="Greens",
            colorbar_label="hPa",
        )
        print(f"✔ Heatmap pression {latest_year} : {pressure_path}")
        matrices_for_combined["Pression (hPa)"] = pressure_matrix

    if "illuminance" in df.columns:
        daily_lux = compute_daily_mean(df, "illuminance")
        lux_matrix = _format_calendar_matrix(daily_lux, latest_year, "Illuminance (lux)")
        lux_matrix.attrs["cmap"] = "YlOrBr"
        lux_matrix.attrs["colorbar_label"] = "lux"
        lux_path = OUTPUT_DIR / f"calendar_illuminance_{latest_year}.png"
        plot_calendar_heatmap(
            matrix=lux_matrix,
            output_path=lux_path,
            title=f"Illuminance moyenne quotidienne - {latest_year}",
            cmap="YlOrBr",
            colorbar_label="lux",
        )
        print(f"✔ Heatmap illuminance {latest_year} : {lux_path}")
        matrices_for_combined["Illuminance (lux)"] = lux_matrix

    if "wind_speed" in df.columns:
        daily_wind = compute_daily_mean(df, "wind_speed")
        wind_matrix = _format_calendar_matrix(daily_wind, latest_year, "Vent (km/h)")
        wind_matrix.attrs["cmap"] = "Purples"
        wind_matrix.attrs["colorbar_label"] = "km/h"
        wind_path = OUTPUT_DIR / f"calendar_wind_{latest_year}.png"
        plot_calendar_heatmap(
            matrix=wind_matrix,
            output_path=wind_path,
            title=f"Vitesse moyenne du vent - {latest_year}",
            cmap="Purples",
            colorbar_label="km/h",
        )
        print(f"✔ Heatmap vent {latest_year} : {wind_path}")
        matrices_for_combined["Vent (km/h)"] = wind_matrix

    combined_path = OUTPUT_DIR / f"calendar_combined_{latest_year}.png"
    plot_combined_calendar(
        matrices=matrices_for_combined,
        output_path=combined_path,
        title=f"Calendrier combiné {latest_year}",
    )
    print(f"✔ Calendrier combiné : {combined_path}")

    hourly = df[WEEKDAY_VARIABLE_KEYS].resample("1h").mean()
    weekday_stats = hourly.groupby(hourly.index.dayofweek).mean()
    weekday_path = OUTPUT_DIR / "weekday_profiles.png"
    variables = [VARIABLES_BY_KEY[key] for key in WEEKDAY_VARIABLE_KEYS]
    plot_weekday_profiles(
        weekday_df=weekday_stats,
        variables=variables,
        output_path=weekday_path,
        title="Profils moyens par jour de semaine",
    )
    print(f"✔ Profils hebdomadaires : {weekday_path}")

    print("✔ Graphiques calendrier générés.")


if __name__ == "__main__":
    main()