# scripts/plot_lagged_correlations.py
from __future__ import annotations

from pathlib import Path
import sys

import argparse

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, VARIABLES_BY_KEY
from meteo.analysis import compute_lagged_correlation
from meteo.plots import plot_lagged_correlation_multi
from meteo.correlation_presets import (
    DEFAULT_ABS_CORRELATION_BANDS,
    DEFAULT_SIGNED_CORRELATION_BANDS,
    CorrelationBand,
)


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


def main() -> None:
    parser = argparse.ArgumentParser(description="Trace les corrélations décalées pour toutes les paires de variables.")
    parser.add_argument(
        "--max-lag",
        type=int,
        default=720,
        help="Décalage maximal en minutes (par défaut : 720 = ±12h).",
    )
    parser.add_argument(
        "--step",
        type=int,
        default=10,
        help="Pas en minutes pour l'évaluation des lags (par défaut : 10).",
    )
    parser.add_argument(
        "--resample",
        default="none",
        help="Fréquence d'agrégation avant calcul (par défaut : 'none') Exemple : '10min'.",
    )
    parser.add_argument(
        "--signed",
        action="store_true",
        default=True,
        help="Affiche les corrélations signées. Utiliser --no-signed pour revenir aux valeurs absolues |r|.",
    )
    parser.add_argument(
        "--no-signed",
        action="store_false",
        dest="signed",
        help="Alias explicite pour repasser en valeurs absolues |r|.",
    )
    parser.add_argument(
        "--thresholds",
        default=None,
        help=(
            "Seuils personnalisés (valeurs séparées par des virgules). "
            "Par défaut, utilise les bandes définies dans meteo.correlation_presets "
            "(abs ou signées selon --signed)."
        ),
    )
    parser.add_argument(
        "--only",
        nargs="*",
        help="Clés de variables à inclure (par défaut toutes les variables numériques).",
    )
    args = parser.parse_args()

    if not CSV_PATH.exists():
        print(f"⚠ Fichier introuvable : {CSV_PATH}")
        return

    df = load_raw_csv(CSV_PATH)
    df = df.select_dtypes(include="number")
    if args.resample and args.resample.lower() != "none":
        df = df.resample(args.resample).mean()
        print(f"Dataset rééchantillonné à {args.resample} pour accélérer le calcul.")
    print(f"Dataset minuté chargé : {CSV_PATH}")
    print(f"  Lignes   : {len(df)}")
    print(f"  Colonnes : {list(df.columns)}")
    print()

    if args.only:
        missing = [k for k in args.only if k not in VARIABLES_BY_KEY]
        if missing:
            raise KeyError(f"Variables inconnues : {', '.join(missing)}")
        variables = [VARIABLES_BY_KEY[k] for k in args.only]
    else:
        variables = list(VARIABLES)
    pairs = [(vx, vy) for i, vx in enumerate(variables) for vy in variables[i + 1 :]]
    print(f"Paires analysées : {len(pairs)} (combinaisons uniques, sans inverses).")

    for var_x, var_y in pairs:
        print(f"→ Corrélation décalée : {var_x.key} → {var_y.key}")

        lag_df_pearson = compute_lagged_correlation(
            df=df,
            var_x=var_x,
            var_y=var_y,
            max_lag_minutes=args.max_lag,
            step_minutes=args.step,
            method="pearson",
        )["correlation"]

        lag_df_spearman = compute_lagged_correlation(
            df=df,
            var_x=var_x,
            var_y=var_y,
            max_lag_minutes=args.max_lag,
            step_minutes=args.step,
            method="spearman",
        )["correlation"]

        use_abs = not args.signed
        if use_abs:
            lag_df_pearson = lag_df_pearson.abs()
            lag_df_spearman = lag_df_spearman.abs()
            ylabel = "Corrélation (|r|)"
            labels = {"Pearson": "Pearson |r|", "Spearman": "Spearman |r|"}
            y_limits = (0.0, 1.0)
            threshold_values = (
                [
                    float(t)
                    for t in (args.thresholds or "").split(",")
                    if t.strip() != ""
                ]
                if args.thresholds
                else []
            )
            bands = list(DEFAULT_ABS_CORRELATION_BANDS)
        else:
            ylabel = "Corrélation"
            labels = {"Pearson": "Pearson", "Spearman": "Spearman"}
            y_limits = (-1.0, 1.0)
            threshold_values: list[float] = []
            if args.thresholds:
                threshold_values = [
                    float(t)
                    for t in args.thresholds.split(",")
                    if t.strip() != ""
                ]
                threshold_values = sorted({v for thr in threshold_values for v in (thr, -thr)})
            bands = list(DEFAULT_SIGNED_CORRELATION_BANDS)

        filename = f"lagcorr_{var_x.key}_to_{var_y.key}.png"
        output_path = OUTPUT_DIR / filename

        plot_lagged_correlation_multi(
            lag_series={
                labels["Pearson"]: lag_df_pearson,
                labels["Spearman"]: lag_df_spearman,
            },
            var_x=var_x,
            var_y=var_y,
            output_path=output_path,
            ylabel=ylabel,
            y_limits=y_limits,
            thresholds=threshold_values,
            bands=bands,
        )

    print("✔ Graphiques de corrélation décalée générés.")


if __name__ == "__main__":
    main()
