"""Tracés dédiés aux analyses du vent (roses et vecteurs agrégés)."""

from __future__ import annotations

from pathlib import Path
from typing import Sequence

import matplotlib.dates as mdates
import matplotlib.pyplot as plt
from matplotlib.ticker import FuncFormatter
import numpy as np
import pandas as pd

from .base import export_plot_dataset

__all__ = ['plot_wind_rose', 'plot_wind_vector_series']


def plot_wind_rose(
    frequencies: pd.DataFrame,
    speed_bin_labels: Sequence[str],
    output_path: str | Path,
    *,
    sector_size_deg: float,
    cmap: str = "viridis",
) -> Path:
    """
    Trace une rose des vents empilée par classes de vitesses (en % du temps).
    """
    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    if frequencies.empty:
        fig, ax = plt.subplots(subplot_kw={"projection": "polar"})
        ax.text(0.5, 0.5, "Données de vent insuffisantes.", ha="center", va="center")
        ax.set_axis_off()
        fig.savefig(output_path, dpi=150, bbox_inches="tight")
        plt.close(fig)
        return output_path.resolve()

    dataset = frequencies.copy()
    dataset.insert(0, "sector_start_deg", frequencies.index)
    dataset.insert(1, "sector_center_deg", frequencies.index + sector_size_deg / 2.0)
    export_plot_dataset(dataset, output_path)

    fig, ax = plt.subplots(subplot_kw={"projection": "polar"}, figsize=(6, 6))
    cmap_obj = plt.get_cmap(cmap, len(speed_bin_labels))
    colors = cmap_obj(np.linspace(0.2, 0.95, len(speed_bin_labels)))

    angles = np.deg2rad(frequencies.index.to_numpy(dtype=float) + sector_size_deg / 2.0)
    width = np.deg2rad(sector_size_deg)
    bottom = np.zeros_like(angles, dtype=float)

    for label, color in zip(speed_bin_labels, colors):
        values = frequencies[label].to_numpy(dtype=float)
        bars = ax.bar(
            angles,
            values,
            width=width,
            bottom=bottom,
            color=color,
            edgecolor="white",
            linewidth=0.5,
            align="center",
        )
        bottom += values

    ax.set_theta_zero_location("N")
    ax.set_theta_direction(-1)
    ax.set_xticks(np.deg2rad(np.arange(0, 360, 45)))
    ax.set_xticklabels(["N", "NE", "E", "SE", "S", "SO", "O", "NO"])
    max_radius = np.max(bottom)
    ax.set_ylim(0, max(max_radius * 1.1, 1))
    ax.yaxis.set_major_formatter(FuncFormatter(lambda val, _pos: f"{val:.0f}%"))
    ax.set_title("Rose des vents (fréquence en %)")
    legend_handles = [
        plt.Line2D([0], [0], color=color, linewidth=6, label=label) for label, color in zip(speed_bin_labels, colors)
    ]
    ax.legend(
        handles=legend_handles,
        loc="lower center",
        bbox_to_anchor=(0.5, -0.15),
        ncol=2,
        title="Vitesses (km/h)",
    )

    fig.tight_layout()
    fig.savefig(output_path, dpi=150, bbox_inches="tight")
    plt.close(fig)
    return output_path.resolve()

def plot_wind_vector_series(
    vector_df: pd.DataFrame,
    output_path: str | Path,
    *,
    title: str = "Vecteurs moyens du vent",
) -> Path:
    """
    Représente les composantes moyennes du vent sous forme de flèches (u/v).
    """
    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    if vector_df.empty:
        fig, ax = plt.subplots()
        ax.text(0.5, 0.5, "Pas de données de vent.", ha="center", va="center")
        ax.set_axis_off()
        fig.savefig(output_path, dpi=150, bbox_inches="tight")
        plt.close(fig)
        return output_path.resolve()

    export_plot_dataset(vector_df, output_path)

    times = vector_df.index
    x = mdates.date2num(times)
    u = vector_df["u"].to_numpy(dtype=float)
    v = vector_df["v"].to_numpy(dtype=float)
    speed = vector_df["speed"].to_numpy(dtype=float)

    fig, ax = plt.subplots(figsize=(12, 4))
    q = ax.quiver(
        x,
        np.zeros_like(x),
        u,
        v,
        speed,
        angles="xy",
        scale_units="xy",
        scale=1,
        cmap="viridis",
    )
    ax.axhline(0, color="black", linewidth=0.5)
    ax.set_ylim(-max(abs(v)) * 1.2 if np.any(v) else -1, max(abs(v)) * 1.2 if np.any(v) else 1)
    ax.xaxis.set_major_locator(mdates.AutoDateLocator())
    ax.xaxis.set_major_formatter(mdates.ConciseDateFormatter(ax.xaxis.get_major_locator()))
    ax.set_ylabel("Composante nord (v)")
    ax.set_xlabel("Date")
    ax.set_title(title)
    cbar = fig.colorbar(q, ax=ax)
    cbar.set_label("Vitesse moyenne (km/h)")

    fig.tight_layout()
    fig.savefig(output_path, dpi=150)
    plt.close(fig)
    return output_path.resolve()