data.py

import datetime
import math
from typing import List, Tuple

import geopandas as gpd
import pandas as pd
import streamlit as st
from classes import (
    POMPE,
    PSE,
    VSAV,
    Engin,
    Indisponibilite,
    Intervention,
    POMPE_Engin,
    PSE_Engin,
    Secteur,
    VSAV_Engin,
)


@st.cache_data
def get_raw_interventions() -> gpd.GeoDataFrame:
    df: pd.DataFrame = pd.merge(
        pd.read_parquet("datas/interventions/inter_23_modifie.parquet"),
        pd.read_csv("datas/interventions/cma_inter.csv"),
        on="IdMMASelection",
    ).dropna()
    df = df.loc[df["cstc"].ne("STEC")]
    df = df.rename(
        columns={
            "inter": "inter_id",
            "selection": "date",
            "traitement": "traitement",  # Durée sur place + temps de retour
            "x": "x",
            "y": "y",
            "fem_mma": "fem_mma",
            "cstc": "cstc",  # Centre de secours territoriellement competent
            "depart": "depart",
            "trajet": "trajet",
            "cs": "cs",
            "proc": "proc",
            "x_final": "x_final",
            "y_final": "y_final",
            "retour": "retour",
            "fam_engin": "fam_engin",
            # Inutile
            "engagement": "_engagement",
            "IdMMASelection": "_IdMMASelection",
            "grpt": "_grpt",
            "IdInterventionSolution": "_IdInterventionSolution",
            "CodeClasseFamilleMateriel": "_CodeClasseFamilleMateriel",
            "n_red": "_n_red",
        }
    )

    cols = ["x", "y", "x_final", "y_final"]
    df[cols] = df[cols].apply(pd.to_numeric)

    gdf = gpd.GeoDataFrame(
        df, geometry=gpd.points_from_xy(df["x"], df["y"]), crs="EPSG:4326"
    )

    # Transforme les temps pour pandas
    gdf["date"] = pd.to_datetime(gdf["date"], format="%Y-%m-%d %H:%M")
    gdf: gpd.GeoDataFrame = gdf.sort_values("date", ascending=True).reset_index(
        drop=True
    )
    gdf["id"] = range(0, len(gdf))
    gdf["traitement"] = pd.to_timedelta(gdf["traitement"], unit="s")
    gdf = gdf[gdf["traitement"] > datetime.timedelta()]
    gdf["depart"] = pd.to_timedelta(gdf["depart"], unit="s")
    gdf["trajet"] = pd.to_timedelta(gdf["trajet"], unit="s")

    # Sanitisation du temps de départ : si aberrant, on remplace par la moyenne
    depart_seconds = gdf["depart"].dt.total_seconds()
    mean_dep = depart_seconds.mean()
    std_dep = depart_seconds.std()
    aberrant = (depart_seconds <= 0) | (depart_seconds > mean_dep + 2 * std_dep)
    gdf.loc[aberrant, "depart"] = pd.to_timedelta(int(round(mean_dep)), unit="s")

    # Ajoute des catégories
    if "cs_secteur" in gdf.columns:
        gdf["cs_secteur"] = gdf["cs_secteur"].astype("category")
    if "proc" in gdf.columns:
        gdf["proc"] = gdf["proc"].astype("category")
        if "type" in gdf.columns:
            gdf["type"] = gdf["type"].astype("category")

    return gdf


def parse_custom_interventions_csv(df: pd.DataFrame) -> list[Intervention]:
    """
    Parse un DataFrame CSV custom en liste d'objets Intervention.
    Colonnes attendues : selection (datetime), x, y, x_final, y_final,
    traitement (s), depart (s), trajet (s), retour (s), cstc, cs, proc, fem_mma.
    Les colonnes manquantes recoivent des valeurs par defaut.
    """
    df = df.copy()

    # Renommage souple
    col_map = {"selection": "date", "inter": "inter_id"}
    df = df.rename(columns={k: v for k, v in col_map.items() if k in df.columns})

    # Colonnes numeriques
    for c in ["x", "y", "x_final", "y_final"]:
        if c in df.columns:
            df[c] = pd.to_numeric(df[c], errors="coerce")
        else:
            df[c] = 0.0
    if "x_final" not in df.columns or df["x_final"].isna().all():
        df["x_final"] = df["x"]
        df["y_final"] = df["y"]

    # Date
    if "date" in df.columns:
        df["date"] = pd.to_datetime(df["date"], errors="coerce")
    else:
        df["date"] = pd.Timestamp("2023-01-01")

    df = df.sort_values("date").reset_index(drop=True)
    df["id"] = range(len(df))

    # Durees en timedelta
    for col, default_s in [("traitement", 300), ("depart", 120), ("trajet", 300), ("retour", 600)]:
        if col in df.columns:
            vals = pd.to_numeric(df[col], errors="coerce").fillna(default_s)
            df[col] = pd.to_timedelta(vals, unit="s")
        else:
            df[col] = pd.to_timedelta(default_s, unit="s")

    # Filtrer traitement > 0
    df = df[df["traitement"] > datetime.timedelta()]

    # Sanitiser depart
    dep_s = df["depart"].dt.total_seconds()
    mean_dep = dep_s.mean() if len(dep_s) > 0 else 120
    std_dep = dep_s.std() if len(dep_s) > 1 else 60
    aberrant = (dep_s <= 0) | (dep_s > mean_dep + 2 * std_dep)
    df.loc[aberrant, "depart"] = pd.to_timedelta(mean_dep, unit="s")

    # Colonnes texte avec defaults
    for col, default in [("cstc", "ANTO"), ("cs", "ANTO"), ("proc", "B"),
                         ("fem_mma", "VSAV"), ("inter_id", "custom")]:
        if col not in df.columns:
            df[col] = default

    df = df.dropna(subset=["x", "y"])

    # GeoDataFrame pour Intervention.df
    gdf = gpd.GeoDataFrame(df, geometry=gpd.points_from_xy(df["x"], df["y"]), crs="EPSG:4326")
    Intervention.df = gdf

    interventions = []
    for inter in gdf.itertuples():
        interventions.append(
            Intervention(
                proc=getattr(inter, "proc", "B"),
                id=getattr(inter, "id", 0),
                x=inter.x,
                y=inter.y,
                x_final=inter.x_final,
                y_final=inter.y_final,
                retour=inter.retour if isinstance(inter.retour, datetime.timedelta)
                       else datetime.timedelta(seconds=float(inter.retour)),
                date=inter.date.to_pydatetime() if hasattr(inter.date, "to_pydatetime")
                     else inter.date,
                traitement=inter.traitement if isinstance(inter.traitement, datetime.timedelta)
                           else datetime.timedelta(seconds=float(inter.traitement)),
                trajet=inter.trajet if isinstance(inter.trajet, datetime.timedelta)
                       else datetime.timedelta(seconds=float(inter.trajet)),
                fem_mma=getattr(inter, "fem_mma", "VSAV"),
                cstc=getattr(inter, "cstc", "ANTO"),
                inter_id_original=str(getattr(inter, "inter_id", f"custom_{inter.id}")),
                cs_vhl=getattr(inter, "cs", None),
                depart=inter.depart if isinstance(inter.depart, datetime.timedelta)
                       else datetime.timedelta(seconds=float(inter.depart)),
            )
        )

    return interventions


def get_interventions() -> list[Intervention]:
    # Pas de cache ici car on doit assigner Intervention.df (attribut de classe)
    # Mais get_raw_interventions() est caché, donc c'est rapide quand même
    df = get_raw_interventions()
    Intervention.df = df

    interventions = []

    for inter in df.itertuples():
        interventions.append(
            Intervention(
                proc=getattr(inter, "proc"),
                id=getattr(inter, "id"),
                x=getattr(inter, "x"),
                y=getattr(inter, "y"),
                x_final=getattr(inter, "x_final"),
                y_final=getattr(inter, "y_final"),
                retour=datetime.timedelta(seconds=inter.retour),
                date=getattr(inter, "date").to_pydatetime(),
                traitement=getattr(inter, "traitement").to_pytimedelta(),
                trajet=getattr(inter, "trajet").to_pytimedelta(),
                fem_mma=getattr(inter, "fem_mma"),
                cstc=getattr(inter, "cstc"),
                inter_id_original=getattr(inter, "inter_id"),
                cs_vhl=getattr(inter, "cs", None),
                depart=getattr(inter, "depart").to_pytimedelta(),
            )
        )

    return interventions


@st.cache_data
def get_raw_secteurs() -> gpd.GeoDataFrame:
    df = gpd.read_file("datas/geo/secteurs_cs.geojson").dropna()
    df_utils = pd.read_csv("datas/utils/lso.csv")
    df = df.merge(df_utils, how="left", left_on="nom", right_on="cs")
    # Supprimer les secteurs sans coordonnées (ex: PCDG, ORLY absents de lso.csv)
    df = df.dropna(subset=["x", "y"])
    df.drop(columns=["cs"], inplace=True)
    df = df.rename(
        columns={
            "nom": "id",
            "geometry": "geometry",
            "x": "x",
            "y": "y",
            # Inutile
            "rg_cle": "_rg_cle",
            "node": "_node",
            "compagnie": "_compagnie",
        }
    )

    # Donne l'emplacement du centre de secours
    if "x" in df.columns and "y" in df.columns:
        df["cs_geometry"] = gpd.points_from_xy(df["x"], df["y"])

    return df


def get_secteurs() -> dict[str, Secteur]:
    # Pas de cache ici car on doit assigner Secteur.df (attribut de classe)
    # Mais get_raw_secteurs() est caché, donc c'est rapide quand même
    df = get_raw_secteurs()
    Secteur.df = df

    secteurs = {}

    for secteur in df.itertuples():
        secteurs[secteur.id] = Secteur(
            id=getattr(secteur, "id"), x=getattr(secteur, "x"), y=getattr(secteur, "y")
        )

    return secteurs


@st.cache_data
def get_raw_engins() -> gpd.GeoDataFrame:
    # Récupère les données et les regroupe
    df = pd.merge(
        pd.read_csv("datas/utils/liste_engins_fix.csv"),
        pd.read_csv("datas/utils/lso.csv"),
        on="cs",
    ).dropna()

    df["cs"] = (
        df["id"]
        .map(pd.read_csv("datas/utils/modifications_engins.csv").set_index("id")["cs"])
        .fillna(df["cs"])
    )

    df = df.rename(
        columns={
            "id": "id",
            "Interventions": "intervention_fem_mma",
            "cs": "cs",
            "x": "x",
            "y": "y",
            "Type_VHL": "Type_VHL",
            "modularite": "modularite",
            # Inutile
            "Regime": "_Regime",
            "ordre": "_ordre",
            "id_cs": "_id_cs",
            "disponible": "_disponible",
            "node": "_node",
        }
    )

    # GeoDataFrame with initial availability
    df = gpd.GeoDataFrame(
        df, geometry=gpd.points_from_xy(df["x"], df["y"]), crs="EPSG:4326"
    )

    # Met "id" en tant qu'index
    if "id" in df.columns and df["id"].is_unique:
        df = df.set_index("id", drop=False)

    # Ajoute des catégories
    if "cs" in df.columns:
        df["cs"] = df["cs"].astype("category")
    if "type_interventions" in df.columns:
        df["type_interventions"] = df["type_interventions"].astype("category")

    return df


def get_engins() -> dict[int, Engin]:
    # Pas de cache ici car on doit assigner Engin.df (attribut de classe)
    # Mais get_raw_engins() est caché, donc c'est rapide quand même
    df = get_raw_engins()
    Engin.df = df

    dict_engins = {}
    modularites: List[Tuple[Engin, int]] = []  # Il associe l engin a son modulaire

    for engin in df.itertuples():
        type_engin: PSE_Engin | POMPE_Engin | VSAV_Engin | None = None
        if engin.Type_VHL == PSE:
            type_engin = PSE_Engin
        else:
            if engin.intervention_fem_mma == VSAV:
                type_engin = VSAV_Engin
            elif engin.intervention_fem_mma == POMPE:
                type_engin = POMPE_Engin

        dict_engins[engin.id] = type_engin(engin.id, engin.cs, engin.x, engin.y)
        modularites.append((dict_engins[engin.id], engin.modularite))

    for engin, engin_id in modularites:
        if engin_id != 0:
            engin.set_modularite(dict_engins[engin_id])

    return dict_engins


def get_data():
    # Note: Pas de cache ici car retourne des objets custom (non-sérialisables)
    # Mais les fonctions sous-jacentes sont cachées, donc c'est rapide quand même
    interventions = get_interventions()
    secteurs = get_secteurs()
    engins = get_engins()
    indisponibilites = get_indisponibilites()
    return interventions, secteurs, engins, indisponibilites


@st.cache_data
def get_indisponibilites():
    df = pd.read_parquet("datas/utils/hour_cs_vsav.parquet")
    df = df[df["delta"] < 0]

    indisponibilites = []
    for row in df.itertuples():
        delta = -row.delta
        indisponibilite = Indisponibilite(
            row.date,
            row.cs,
            datetime.timedelta(hours=delta)
        )
        indisponibilites.append(indisponibilite)
    return indisponibilites