# meteo/plots.py
from __future__ import annotations

from pathlib import Path

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

from .variables import Variable


def plot_scatter_pair(
    df: pd.DataFrame,
    var_x: Variable,
    var_y: Variable,
    output_path: str | Path,
    *,
    sample_step: int = 10,
) -> Path:
    """
    Trace un nuage de points (scatter) pour une paire de variables.

    - On sous-échantillonne les données avec `sample_step` (par exemple,
      1 point sur 10) pour éviter un graphique illisible.
    """
    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    # On ne garde que les colonnes pertinentes et les lignes complètes
    df_pair = df[[var_x.column, var_y.column]].dropna()

    if sample_step > 1:
        df_pair = df_pair.iloc[::sample_step, :]

    plt.figure()
    plt.scatter(df_pair[var_x.column], df_pair[var_y.column], s=5, alpha=0.5)
    plt.xlabel(f"{var_x.label} ({var_x.unit})")
    plt.ylabel(f"{var_y.label} ({var_y.unit})")
    plt.title(f"{var_y.label} en fonction de {var_x.label}")
    plt.tight_layout()
    plt.savefig(output_path, dpi=150)
    plt.close()

    return output_path.resolve()


def plot_lagged_correlation(
    lag_df: pd.DataFrame,
    var_x: Variable,
    var_y: Variable,
    output_path: str | Path,
) -> Path:
    """
    Trace la corrélation en fonction du lag (en minutes) entre deux variables.
    """
    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    plt.figure()
    plt.plot(lag_df.index, lag_df["correlation"])
    plt.axvline(0, linestyle="--")  # lag = 0
    plt.xlabel("Décalage (minutes)\n(lag > 0 : X précède Y)")
    plt.ylabel("Corrélation")
    plt.title(f"Corrélation décalée : {var_x.label} → {var_y.label}")
    plt.grid(True)
    plt.tight_layout()
    plt.savefig(output_path, dpi=150)
    plt.close()

    return output_path.resolve()



def plot_correlation_heatmap(
    corr: pd.DataFrame,
    variables: Sequence[Variable],
    output_path: str | Path,
    *,
    annotate: bool = True,
) -> Path:
    """
    Trace une heatmap de la matrice de corrélation.

    Paramètres
    ----------
    corr :
        Matrice de corrélation (index et colonnes doivent correspondre
        aux noms de colonnes des variables).
    variables :
        Liste de Variable, dans l'ordre où elles doivent apparaître.
    output_path :
        Chemin du fichier image à écrire.
    annotate :
        Si True, affiche la valeur numérique dans chaque case.
    """
    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    columns = [v.column for v in variables]
    labels = [v.label for v in variables]

    # On aligne la matrice sur l'ordre désiré
    corr = corr.loc[columns, columns]

    data = corr.to_numpy()

    fig, ax = plt.subplots()
    im = ax.imshow(data, vmin=-1.0, vmax=1.0)

    # Ticks et labels
    ax.set_xticks(np.arange(len(labels)))
    ax.set_yticks(np.arange(len(labels)))
    ax.set_xticklabels(labels, rotation=45, ha="right")
    ax.set_yticklabels(labels)

    # Axe en haut/bas selon préférence (ici on laisse en bas)
    ax.set_title("Matrice de corrélation (coef. de Pearson)")

    # Barre de couleur
    cbar = plt.colorbar(im, ax=ax)
    cbar.set_label("Corrélation")

    # Annotation des cases
    if annotate:
        n = data.shape[0]
        for i in range(n):
            for j in range(n):
                if i == j:
                    text = "—"
                else:
                    val = data[i, j]
                    if np.isnan(val):
                        text = ""
                    else:
                        text = f"{val:.2f}"
                ax.text(
                    j,
                    i,
                    text,
                    ha="center",
                    va="center",
                )

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

    return output_path.resolve()