data_loader.py

"""
Chargement des données brutes depuis Reflets-bspp et encodage en numpy arrays.
Aucune dépendance sur Streamlit. Aucune modification du simulateur séquentiel.

Conventions d'encodage:
  Type véhicule  : POMPE=0, VSAV=1, PSE=2
  Type intervention : POMPE=0, VSAV=1
"""

from __future__ import annotations

import datetime
import math
from pathlib import Path
from typing import Tuple

import numpy as np
import pandas as pd

# ---------------------------------------------------------------------------
# Chemins absolus vers les données
# ---------------------------------------------------------------------------
DATA_DIR = Path(__file__).parent.parent / "datas"

INTER_PARQUET   = DATA_DIR / "interventions" / "inter_23_modifie.parquet"
CMA_CSV         = DATA_DIR / "interventions" / "cma_inter.csv"
SECTEURS_GEOJSON= DATA_DIR / "geo" / "secteurs_cs.geojson"
LSO_CSV         = DATA_DIR / "utils" / "lso.csv"
ENGINS_CSV      = DATA_DIR / "utils" / "liste_engins_fix.csv"
MODIF_CSV       = DATA_DIR / "utils" / "modifications_engins.csv"
INDISPO_PARQUET = DATA_DIR / "utils" / "hour_cs_vsav.parquet"

# ---------------------------------------------------------------------------
# Encodage des types
# ---------------------------------------------------------------------------
POMPE_TYPE, VSAV_TYPE, PSE_TYPE = 0, 1, 2

VHL_TYPE_MAP: dict[str, int] = {
    "POMPE": POMPE_TYPE,
    "VSAV":  VSAV_TYPE,
    "PSE":   PSE_TYPE,
}

INTER_TYPE_MAP: dict[str, int] = {
    "POMPE": 0,
    "VSAV":  1,
}

# Pour chaque type d'intervention, quels types de véhicules peuvent y répondre
# index [inter_type, vehicle_type] → bool
# POMPE inter: POMPE(0) et PSE(2) ; VSAV inter: VSAV(1) et PSE(2)
TYPE_ELIGIBLE = np.array([
    [True,  False, True],   # POMPE intervention
    [False, True,  True],   # VSAV  intervention
], dtype=bool)


# ---------------------------------------------------------------------------
# Loaders
# ---------------------------------------------------------------------------

def _load_interventions_raw() -> pd.DataFrame:
    # La colonne date s'appelle "selection" dans le parquet
    # cstc, depart, trajet sont déjà dans le parquet (pas besoin de cma_inter pour ça)
    df = pd.merge(
        pd.read_parquet(INTER_PARQUET),
        pd.read_csv(CMA_CSV),
        on="IdMMASelection",
    ).dropna()
    df = df.loc[df["cstc"].ne("STEC")]

    for col in ["x", "y", "x_final", "y_final"]:
        df[col] = pd.to_numeric(df[col])

    # "selection" est la colonne datetime d'intervention
    df["date"] = pd.to_datetime(df["selection"], format="%Y-%m-%d %H:%M")
    df = df.sort_values("date", ascending=True).reset_index(drop=True)
    df["id"] = range(len(df))

    df["traitement"] = pd.to_timedelta(df["traitement"], unit="s")
    df = df[df["traitement"] > datetime.timedelta()]
    df["depart"]  = pd.to_timedelta(df["depart"],  unit="s")
    df["trajet"]  = pd.to_timedelta(df["trajet"],  unit="s")

    # Sanitiser les départs aberrants
    dep_s = df["depart"].dt.total_seconds()
    mean_dep, std_dep = dep_s.mean(), dep_s.std()
    aberrant = (dep_s <= 0) | (dep_s > mean_dep + 2 * std_dep)
    df.loc[aberrant, "depart"] = pd.to_timedelta(mean_dep, unit="s")

    return df


def _load_secteurs_raw() -> pd.DataFrame:
    import geopandas as gpd
    gdf = gpd.read_file(SECTEURS_GEOJSON).dropna()
    lso = pd.read_csv(LSO_CSV)
    gdf = gdf.merge(lso, how="left", left_on="nom", right_on="cs")
    gdf = gdf.dropna(subset=["x", "y"])
    gdf = gdf.rename(columns={"nom": "id"})
    return gdf


def _load_engins_raw() -> pd.DataFrame:
    df = pd.merge(
        pd.read_csv(ENGINS_CSV),
        pd.read_csv(LSO_CSV),
        on="cs",
    ).dropna()

    # Appliquer les modifications de CS
    modifs = pd.read_csv(MODIF_CSV).set_index("id")["cs"]
    df["cs"] = df["id"].map(modifs).fillna(df["cs"])
    return df


def _load_indispos_raw() -> pd.DataFrame:
    df = pd.read_parquet(INDISPO_PARQUET)
    return df[df["delta"] < 0].copy()


# ---------------------------------------------------------------------------
# Fonction principale
# ---------------------------------------------------------------------------

def load_data() -> Tuple[
    dict,      # interventions_np
    dict,      # vehicles_np
    list,      # cs_list  (index → cs_name)
    np.ndarray,# cs_coords [N_cs, 2] float64  (x, y)
    np.ndarray,# vehicle_unavail [V, N_inter] bool
]:
    """
    Charge toutes les données et les encode en numpy arrays GPU-prêts.

    Returns
    -------
    interventions_np : dict avec clés
        timestamps  [N]  float64  unix timestamp de l'intervention
        x, y        [N]  float64  coordonnées du site
        x_final     [N]  float64  coordonnées fin (retour)
        y_final     [N]  float64
        type        [N]  int32    0=POMPE 1=VSAV
        cstc_idx    [N]  int32    index du CS compétent (-1 si hors liste)
        duration    [N]  float32  durée sur place (s)
        depart      [N]  float32  temps de préparation au départ (s)
        cstc_name   [N]  object   nom du CS compétent (pour debug)
        proc        [N]  object   "R" ou autre (procédure)

    vehicles_np : dict avec clés
        ids            [V]  int/str    identifiants originaux
        type           [V]  int32      0=POMPE 1=VSAV 2=PSE
        initial_cs_idx [V]  int32      CS initial (allocation de référence)
        x_base         [V]  float64    coordonnée x de la caserne initiale
        y_base         [V]  float64    coordonnée y de la caserne initiale

    cs_list  : list[str]   – cs_list[i] = nom du CS d'index i
    cs_coords: ndarray [N_cs, 2]  – (x, y) par CS dans l'ordre de cs_list

    vehicle_unavail : ndarray [V, N_inter] bool
        True si véhicule v est indisponible à l'heure de l'intervention i
    """
    if _cache_is_valid():
        print("[data_loader] Cache trouve — chargement rapide...")
        try:
            result = _load_cache()
            print(f"[data_loader] Charge depuis cache: {len(result[0]['timestamps'])} inter, {len(result[1]['ids'])} vhl, {len(result[2])} CS")
            return result
        except Exception as e:
            print(f"[data_loader] Cache invalide ({e}), recalcul...")

    print("[data_loader] Chargement des données brutes...")
    df_inter = _load_interventions_raw()
    df_sect  = _load_secteurs_raw()
    df_eng   = _load_engins_raw()
    df_ind   = _load_indispos_raw()

    # ---- Index CS ---------------------------------------------------------
    cs_list = sorted(df_sect["id"].unique().tolist())
    cs_to_idx: dict[str, int] = {cs: i for i, cs in enumerate(cs_list)}
    N_cs = len(cs_list)

    lso_map: dict[str, tuple] = {
        row.cs: (row.x, row.y) for row in pd.read_csv(LSO_CSV).itertuples()
    }
    cs_coords = np.array([
        lso_map.get(cs, (0.0, 0.0)) for cs in cs_list
    ], dtype=np.float64)

    # ---- Interventions ----------------------------------------------------
    N = len(df_inter)
    timestamps = np.zeros(N, dtype=np.float64)
    ix = np.zeros(N, dtype=np.float64)
    iy = np.zeros(N, dtype=np.float64)
    ix_f = np.zeros(N, dtype=np.float64)
    iy_f = np.zeros(N, dtype=np.float64)
    itype = np.zeros(N, dtype=np.int32)
    ictc = np.full(N, -1, dtype=np.int32)
    idur = np.zeros(N, dtype=np.float32)
    idep = np.zeros(N, dtype=np.float32)
    cstc_names = np.empty(N, dtype=object)
    proc_arr   = np.empty(N, dtype=object)

    for i, row in enumerate(df_inter.itertuples()):
        timestamps[i] = int(
            row.date.to_pydatetime()
            .replace(tzinfo=datetime.timezone.utc)
            .timestamp()
        )
        ix[i]   = row.x
        iy[i]   = row.y
        ix_f[i] = row.x_final
        iy_f[i] = row.y_final
        itype[i] = INTER_TYPE_MAP.get(row.fem_mma, 0)
        ictc[i]  = cs_to_idx.get(row.cstc, -1)
        idur[i]  = row.traitement.total_seconds()
        idep[i]  = row.depart.total_seconds()
        cstc_names[i] = row.cstc
        proc_arr[i]   = getattr(row, "proc", "B")

    interventions_np = dict(
        timestamps=timestamps,
        x=ix, y=iy, x_final=ix_f, y_final=iy_f,
        type=itype, cstc_idx=ictc, duration=idur, depart=idep,
        cstc_name=cstc_names, proc=proc_arr,
    )

    # ---- Véhicules --------------------------------------------------------
    # Déterminer le type de chaque engin
    def _eng_type(row) -> int:
        if str(row.Type_VHL) == "PSE":
            return PSE_TYPE
        if str(row.Interventions) == "VSAV":
            return VSAV_TYPE
        return POMPE_TYPE

    df_eng = df_eng.reset_index(drop=True)
    V = len(df_eng)
    vids  = df_eng["id"].values
    vtype = np.array([_eng_type(r) for r in df_eng.itertuples()], dtype=np.int32)
    vcs   = np.array([cs_to_idx.get(str(r.cs), 0) for r in df_eng.itertuples()], dtype=np.int32)
    vx    = df_eng["x"].values.astype(np.float64)
    vy    = df_eng["y"].values.astype(np.float64)

    # ---- Modularité PSE↔VSAV -----------------------------------------------
    # vsav_pse_pair_idx[v] = index du PSE pair du VSAV v  (-1 si pas couplé)
    # pse_vsav_pair_idx[v] = index du VSAV pair du PSE v  (-1 si pas couplé)
    vid_to_v_int: dict[int, int] = {int(vid): v for v, vid in enumerate(vids)}
    vsav_pse_pair_idx = np.full(V, -1, dtype=np.int32)
    pse_vsav_pair_idx = np.full(V, -1, dtype=np.int32)

    for v, row in enumerate(df_eng.itertuples()):
        mod_id = int(row.modularite) if row.modularite != 0 else 0
        if mod_id == 0:
            continue
        paired_v = vid_to_v_int.get(mod_id, -1)
        if paired_v == -1:
            continue
        if vtype[v] == VSAV_TYPE:
            vsav_pse_pair_idx[v] = paired_v   # VSAV → index de son PSE
        elif vtype[v] == PSE_TYPE:
            pse_vsav_pair_idx[v] = paired_v   # PSE  → index de son VSAV

    # ---- Plage interdite nuit → [N_inter] bool ----------------------------
    # Vectorisation de plage_interdite_nuit() depuis les datetime pandas
    _hours   = df_inter["date"].dt.hour + df_inter["date"].dt.minute / 60.0
    _weekday = df_inter["date"].dt.dayofweek.values   # 0=lun, 6=dim
    _h = _hours.values

    night_restricted = np.zeros(N, dtype=bool)
    # Lun–Ven (0-4) : 0h–8h interdit
    night_restricted |= ((_weekday >= 0) & (_weekday <= 4)) & (_h < 8.0)
    # Dim(6) ou Lun–Jeu(0-3) : 23h–minuit interdit
    night_restricted |= ((_weekday == 6) | ((_weekday >= 0) & (_weekday <= 3))) & (_h >= 23.0)
    # Sam–Dim (5-6) : 0h–7h interdit
    night_restricted |= ((_weekday >= 5) & (_weekday <= 6)) & (_h < 7.0)

    # Procédure rouge [N_inter] bool
    inter_rouge = (proc_arr == "R")

    vehicles_np = dict(
        ids=vids,
        type=vtype,
        initial_cs_idx=vcs,
        x_base=vx,
        y_base=vy,
        vsav_pse_pair_idx=vsav_pse_pair_idx,
        pse_vsav_pair_idx=pse_vsav_pair_idx,
    )
    interventions_np["night_restricted"] = night_restricted
    interventions_np["rouge"]            = inter_rouge

    # ---- Indisponibilités → [V, N_inter] bool ----------------------------
    print("[data_loader] Construction de la matrice d'indisponibilité...")
    # Indexer les véhicules par id pour lookup rapide
    vid_to_v = {vid: v for v, vid in enumerate(vids)}

    # Agréger les intervalles d'indispo par véhicule
    # Format: vehicle_intervals[v] = list of (start_ts, end_ts)
    vehicle_intervals: list[list] = [[] for _ in range(V)]

    # df_ind contient des colonnes : date, cs, delta (négatif = manque VSAV)
    # delta = nombre d'heures manquantes → on met |delta| véhicules VSAV indispos
    cs_to_vsav_ids: dict[str, list] = {}
    for v, row in enumerate(df_eng.itertuples()):
        if _eng_type(row) == VSAV_TYPE:
            cs = str(row.cs)
            cs_to_vsav_ids.setdefault(cs, []).append(v)

    for row in df_ind.itertuples():
        cs = str(row.cs)
        delta_h = abs(float(row.delta))
        date_dt = pd.to_datetime(row.date).to_pydatetime().replace(
            tzinfo=datetime.timezone.utc
        )
        start_ts = int(date_dt.timestamp())
        end_ts   = start_ts + int(delta_h * 3600)

        vsav_vs = cs_to_vsav_ids.get(cs, [])
        if vsav_vs:
            # On rend 1 VSAV indispo (le premier disponible – approximation)
            vehicle_intervals[vsav_vs[0]].append((start_ts, end_ts))

    # Construire la matrice booléenne par vectorisation
    vehicle_unavail = np.zeros((V, N), dtype=bool)
    for v, intervals in enumerate(vehicle_intervals):
        if not intervals:
            continue
        # Fusionne les intervalles
        intervals_sorted = sorted(intervals)
        merged: list = []
        for s, e in intervals_sorted:
            if not merged or s > merged[-1][1]:
                merged.append([s, e])
            else:
                merged[-1][1] = max(merged[-1][1], e)
        starts = np.array([s for s, _ in merged], dtype=np.int64)
        ends   = np.array([e for _, e in merged], dtype=np.int64)
        # Pour chaque intervention, vérifier si l'horodatage est dans un intervalle
        ts = timestamps.astype(np.int64)
        for s, e in zip(starts, ends):
            vehicle_unavail[v] |= (ts >= s) & (ts < e)

    print(f"[data_loader] Chargé: {N} interventions, {V} véhicules, {N_cs} CS")
    _save_cache(interventions_np, vehicles_np, cs_list, cs_coords, vehicle_unavail)
    return interventions_np, vehicles_np, cs_list, cs_coords, vehicle_unavail


_CACHE_PATH = Path(__file__).parent / "cache" / "data_loader_cache.pkl"

def _cache_is_valid():
    if not _CACHE_PATH.exists():
        return False
    cache_mtime = _CACHE_PATH.stat().st_mtime
    for src in [INTER_PARQUET, CMA_CSV, SECTEURS_GEOJSON, LSO_CSV, ENGINS_CSV, MODIF_CSV, INDISPO_PARQUET]:
        if src.exists() and src.stat().st_mtime > cache_mtime:
            return False
    return True

def _save_cache(interventions_np, vehicles_np, cs_list, cs_coords, vehicle_unavail):
    import pickle
    _CACHE_PATH.parent.mkdir(parents=True, exist_ok=True)
    with open(_CACHE_PATH, "wb") as f:
        pickle.dump({"interventions_np": interventions_np, "vehicles_np": vehicles_np,
                     "cs_list": cs_list, "cs_coords": cs_coords, "vehicle_unavail": vehicle_unavail},
                    f, protocol=pickle.HIGHEST_PROTOCOL)

def _load_cache():
    import pickle
    with open(_CACHE_PATH, "rb") as f:
        d = pickle.load(f)
    return d["interventions_np"], d["vehicles_np"], d["cs_list"], d["cs_coords"], d["vehicle_unavail"]