# scripts/plot_correlation_heatmap.py
from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
import sys


PROJECT_ROOT = Path(__file__).resolve().parents[3]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from meteo.dataset import load_raw_csv
from meteo.variables import VARIABLES
from meteo.analysis import compute_correlation_matrices_for_methods
from meteo.plots import plot_correlation_heatmap


CSV_PATH = Path("data/weather_minutely.csv")
DOC_DIR = Path(__file__).resolve().parent.parent

CORRELATION_METHODS: tuple[str, ...] = ("pearson", "spearman")
CORRELATION_TRANSFORM = "square"


@dataclass(frozen=True)
class HeatmapConfig:
    filename: str
    title: str
    colorbar_label: str
    cmap: str = "viridis"
    vmin: float = 0.0
    vmax: float = 1.0


HEATMAP_CONFIGS: dict[str, HeatmapConfig] = {
    "pearson": HeatmapConfig(
        filename="correlation_heatmap.png",
        title="Corrélations R² (coef. de Pearson)",
        colorbar_label="Coefficient de corrélation R²",
    ),
    "spearman": HeatmapConfig(
        filename="correlation_heatmap_spearman.png",
        title="Corrélations R² (coef. de Spearman)",
        colorbar_label="Coefficient de corrélation R²",
    ),
}


def _get_heatmap_config(method: str) -> HeatmapConfig:
    if method in HEATMAP_CONFIGS:
        return HEATMAP_CONFIGS[method]

    # Valeurs par défaut pour un scénario non prévu.
    return HeatmapConfig(
        filename=f"correlation_heatmap_{method}.png",
        title=f"Matrice de corrélation ({method})",
        colorbar_label="Coefficient de corrélation",
        cmap="viridis" if CORRELATION_TRANSFORM == "square" else "coolwarm",
        vmin=0.0 if CORRELATION_TRANSFORM == "square" else -1.0,
        vmax=1.0,
    )


def main() -> None:
    if not CSV_PATH.exists():
        print(f"⚠ Fichier introuvable : {CSV_PATH}")
        print("   Assurez-vous d'avoir généré le dataset minuté.")
        return

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

    matrices = compute_correlation_matrices_for_methods(
        df=df,
        variables=VARIABLES,
        methods=CORRELATION_METHODS,
        transform=CORRELATION_TRANSFORM,
    )

    for method, corr in matrices.items():
        print(f"Matrice de corrélation (méthode={method}, transform={CORRELATION_TRANSFORM}) :")
        print(corr)
        print()

        config = _get_heatmap_config(method)
        output_path = plot_correlation_heatmap(
            corr=corr,
            variables=VARIABLES,
            output_path=DOC_DIR / "figures" / config.filename,
            annotate=True,
            title=config.title,
            cmap=config.cmap,
            vmin=config.vmin,
            vmax=config.vmax,
            colorbar_label=config.colorbar_label,
        )

        print(f"✔ Heatmap de corrélation sauvegardée dans : {output_path}")


if __name__ == "__main__":
    main()
