"""Fonctions utilitaires pour exporter les jeux de données associés aux figures."""

from __future__ import annotations

from pathlib import Path
from typing import Any

import numpy as np
import pandas as pd

__all__ = ["export_plot_dataset"]


def export_plot_dataset(data: Any, output_path: str | Path, *, suffix: str = ".csv") -> Path | None:
    """
    Sauvegarde, en regard du fichier image exporté, les données brutes ayant servi à construire la figure.
    """

    if data is None:
        return None

    output_path = Path(output_path)
    dataset_path = output_path.with_suffix(suffix)

    # If the image is exported under a "figures" directory, keep the dataset in
    # an equivalent "data" directory to avoid mixing assets.
    parts = list(dataset_path.parts)
    for idx, part in enumerate(parts):
        if part == "figures":
            parts[idx] = "data"
            dataset_path = Path(*parts)
            break

    dataset_path.parent.mkdir(parents=True, exist_ok=True)

    def _normalize(value: Any, *, default_name: str = "value") -> pd.DataFrame:
        if isinstance(value, pd.DataFrame):
            return value.copy()
        if isinstance(value, pd.Series):
            return value.to_frame(name=value.name or default_name)
        if isinstance(value, np.ndarray):
            return pd.DataFrame(value)
        return pd.DataFrame(value)

    if isinstance(data, dict):
        frames: list[pd.DataFrame] = []
        for key, value in data.items():
            if value is None:
                continue
            frame = _normalize(value, default_name=str(key))
            frame = pd.concat({str(key): frame}, axis=1)
            frames.append(frame)
        if not frames:
            return None
        export_df = pd.concat(frames, axis=1)
    else:
        export_df = _normalize(data)

    export_df.to_csv(dataset_path)
    return dataset_path