# scripts/plot_correlation_heatmap.py
from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
import sys
import argparse
import pandas as pd
import numpy as np


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")


@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 (coef. de Pearson)",
        colorbar_label="Coefficient de corrélation",
        cmap="viridis",
        vmin=0.0,
        vmax=1.0,
    ),
    "spearman": HeatmapConfig(
        filename="correlation_heatmap_spearman.png",
        title="Corrélations (coef. de Spearman)",
        colorbar_label="Coefficient de corrélation",
        cmap="viridis",
        vmin=0.0,
        vmax=1.0,
    ),
}


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:
    parser = argparse.ArgumentParser(description="Trace des matrices de corrélation instantanées (signées, absolues ou r²).")
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=DOC_DIR / "figures",
        help="Dossier de sortie pour les heatmaps.",
    )
    parser.add_argument(
        "--transform",
        choices=["identity", "absolute", "square"],
        default="absolute",
        help="Transformation de la matrice (signée, |r| ou r²). Par défaut : |r|.",
    )
    parser.add_argument(
        "--upper-only",
        action="store_true",
        help="Masque la partie inférieure de la matrice pour alléger la lecture.",
    )
    args = parser.parse_args()

    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)
    df = df[[v.column for v in VARIABLES]]
    print(f"Dataset minuté chargé : {CSV_PATH}")
    print(f"  Lignes   : {len(df)}")
    print(f"  Colonnes : {list(df.columns)}")
    print()

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

    args.output_dir.mkdir(parents=True, exist_ok=True)

    for method, corr in matrices.items():
        if args.upper_only:
            mask = np.tril(np.ones_like(corr, dtype=bool), k=-1)
            corr = corr.mask(mask)

        print(f"Matrice de corrélation (méthode={method}, transform={transform}) :")
        print(corr)
        print()

        config = _get_heatmap_config(method)
        filename = config.filename
        title = config.title
        if transform == "absolute":
            title = f"{title} (|r|)"
            stem, suffix = filename.rsplit(".", 1)
            filename = f"{stem}_abs.{suffix}"
        elif transform == "square":
            title = f"{title} (r²)"
            stem, suffix = filename.rsplit(".", 1)
            filename = f"{stem}_r2.{suffix}"
            config = HeatmapConfig(
                filename=filename,
                title=title,
                colorbar_label="Coefficient de corrélation r²",
                cmap="viridis",
                vmin=0.0,
                vmax=1.0,
            )
        elif transform == "identity":
            config = HeatmapConfig(
                filename=filename,
                title=title,
                colorbar_label="Coefficient de corrélation r",
                cmap="coolwarm",
                vmin=-1.0,
                vmax=1.0,
            )

        output_path = plot_correlation_heatmap(
            corr=corr,
            variables=VARIABLES,
            output_path=args.output_dir / filename,
            annotate=True,
            title=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()
