"""Fonctions spécifiques aux analyses de vent (roses et composantes)."""

from __future__ import annotations

from typing import Sequence

import numpy as np
import pandas as pd

from .core import _ensure_datetime_index

__all__ = ['compute_wind_rose_distribution', 'compute_mean_wind_components']


def _format_speed_bin_labels(speed_bins: Sequence[float]) -> list[str]:
    labels: list[str] = []
    for i in range(len(speed_bins) - 1):
        low = speed_bins[i]
        high = speed_bins[i + 1]
        if np.isinf(high):
            labels.append(f"≥{low:g}")
        else:
            labels.append(f"{low:g}–{high:g}")
    return labels

def compute_wind_rose_distribution(
    df: pd.DataFrame,
    *,
    direction_sector_size: int = 30,
    speed_bins: Sequence[float] = (0, 10, 20, 30, 50, float("inf")),
) -> tuple[pd.DataFrame, list[str], float]:
    """
    Regroupe la distribution vent/direction en secteurs angulaires et classes de vitesse.
    Retourne un DataFrame indexé par le début du secteur (en degrés) et colonnes = classes de vitesse (%).
    """
    if direction_sector_size <= 0 or direction_sector_size > 180:
        raise ValueError("direction_sector_size doit être compris entre 1 et 180 degrés.")

    if "wind_speed" not in df.columns or "wind_direction" not in df.columns:
        raise KeyError("Le DataFrame doit contenir 'wind_speed' et 'wind_direction'.")

    data = df[["wind_speed", "wind_direction"]].dropna()
    if data.empty:
        return pd.DataFrame(), [], float(direction_sector_size)

    n_sectors = int(360 / direction_sector_size)
    direction = data["wind_direction"].to_numpy(dtype=float) % 360.0
    sector_indices = np.floor(direction / direction_sector_size).astype(int) % n_sectors

    bins = list(speed_bins)
    if not np.isinf(bins[-1]):
        bins.append(float("inf"))
    labels = _format_speed_bin_labels(bins)

    speed_categories = pd.cut(
        data["wind_speed"],
        bins=bins,
        right=False,
        include_lowest=True,
        labels=labels,
    )

    counts = (
        pd.crosstab(sector_indices, speed_categories)
        .reindex(range(n_sectors), fill_value=0)
        .reindex(columns=labels, fill_value=0)
    )

    total = counts.values.sum()
    frequencies = counts / total * 100.0 if total > 0 else counts.astype(float)
    frequencies.index = frequencies.index * direction_sector_size
    return frequencies, labels, float(direction_sector_size)

def compute_mean_wind_components(
    df: pd.DataFrame,
    *,
    freq: str = "1M",
) -> pd.DataFrame:
    """
    Calcule les composantes zonale (u) et méridienne (v) du vent pour une fréquence donnée.
    Retourne également la vitesse moyenne.
    """
    if "wind_speed" not in df.columns or "wind_direction" not in df.columns:
        raise KeyError("Les colonnes 'wind_speed' et 'wind_direction' sont requises.")

    _ensure_datetime_index(df)
    subset = df[["wind_speed", "wind_direction"]].dropna()
    if subset.empty:
        return pd.DataFrame(columns=["u", "v", "speed"])

    radians = np.deg2rad(subset["wind_direction"].to_numpy(dtype=float))
    speed = subset["wind_speed"].to_numpy(dtype=float)

    u = speed * np.sin(radians) * -1  # composante est-ouest (positive vers l'est)
    v = speed * np.cos(radians) * -1  # composante nord-sud (positive vers le nord)

    vector_df = pd.DataFrame(
        {
            "u": u,
            "v": v,
            "speed": speed,
        },
        index=subset.index,
    )

    actual_freq = "1ME" if freq == "1M" else freq
    grouped = vector_df.resample(actual_freq).mean()
    return grouped.dropna(how="all")