# meteo/season.py
from __future__ import annotations

from typing import Iterable, Sequence

import numpy as np
import pandas as pd


SEASON_LABELS = np.array(["winter", "spring", "summer", "autumn"])
MONTH_TO_SEASON_INDEX = {
    12: 0,
    1: 0,
    2: 0,
    3: 1,
    4: 1,
    5: 1,
    6: 2,
    7: 2,
    8: 2,
    9: 3,
    10: 3,
    11: 3,
}


def _ensure_datetime_index(index: pd.Index) -> pd.DatetimeIndex:
    if not isinstance(index, pd.DatetimeIndex):
        raise TypeError("Cette fonction nécessite un DatetimeIndex.")
    return index


def _season_indices_for_month(months: np.ndarray, hemisphere: str) -> np.ndarray:
    base_indices = np.vectorize(MONTH_TO_SEASON_INDEX.get)(months)
    if hemisphere == "south":
        return (base_indices + 2) % len(SEASON_LABELS)
    return base_indices


def compute_season_series(
    index: pd.Index,
    *,
    hemisphere: str = "north",
    column_name: str = "season",
) -> pd.Series:
    """
    Retourne une série catégorielle indiquant la saison météorologique pour chaque timestamp.
    """
    hemisphere = hemisphere.lower()
    if hemisphere not in {"north", "south"}:
        raise ValueError("hemisphere doit valoir 'north' ou 'south'.")

    dt_index = _ensure_datetime_index(index)
    month_array = dt_index.month.to_numpy()
    season_indices = _season_indices_for_month(month_array, hemisphere)
    labels = SEASON_LABELS[season_indices]
    return pd.Series(labels, index=dt_index, name=column_name)


def add_season_column(
    df: pd.DataFrame,
    *,
    hemisphere: str = "north",
    column_name: str = "season",
) -> pd.DataFrame:
    """
    Ajoute une colonne 'season' (winter/spring/summer/autumn) au DataFrame.
    """
    series = compute_season_series(df.index, hemisphere=hemisphere, column_name=column_name)
    df[column_name] = series
    return df


def sort_season_labels(
    labels: Iterable[str],
    *,
    order: Sequence[str] | None = None,
) -> list[str]:
    """
    Trie la liste fournie en respectant l'ordre saisonnier par défaut.
    """
    reference = [str(season) for season in (order if order is not None else SEASON_LABELS)]
    label_set = {str(label) for label in labels if label}
    return [season for season in reference if season in label_set]