"""HF Datasets bridge utilities.

This module provides bridge functions for converting between PLAID datasets/samples
and Hugging Face Datasets format. It includes utilities for feature type conversion,
dataset generation from PLAID objects, and sample reconstruction.
"""

# -*- coding: utf-8 -*-
#
# This file is subject to the terms and conditions defined in
# file 'LICENSE.txt', which is part of this source code package.
#
#

from functools import partial
from typing import Any, Callable, Generator, Optional

import datasets
import numpy as np
import pyarrow as pa
from datasets import Features, Sequence, Value

from plaid import Sample
from plaid.storage.common.preprocessor import build_sample_dict
from plaid.types import IndexType


def convert_dtype_to_hf_feature(feature_type: dict[str, Any]):
    """Convert a PLAID feature type dict to Hugging Face Feature.

    Args:
        feature_type (dict): Dictionary with 'dtype' and 'ndim' keys.

    Returns:
        Features or Sequence: The corresponding HF feature type.
    """
    base_dtype = feature_type["dtype"]
    ndim = feature_type["ndim"]

    feature = Value(base_dtype)
    for _ in range(ndim):
        feature = Sequence(feature)
    return feature


def convert_to_hf_feature(variable_schema: dict[str, dict]):
    """Convert a PLAID variable schema to Hugging Face Features.

    Args:
        variable_schema (dict[str, dict]): Mapping of variable names to type dicts.

    Returns:
        Features: The HF Features object.
    """
    return Features(
        {k: convert_dtype_to_hf_feature(v) for k, v in variable_schema.items()}
    )


def generator_to_datasetdict(
    generators: dict[str, Callable[..., Generator[Sample, None, None]]],
    variable_schema: dict,
    cache_dir: str,
    gen_kwargs: Optional[dict[str, dict[str, list[IndexType]]]] = None,
    processes_number: int = 1,
    writer_batch_size: int = 1,
) -> datasets.DatasetDict:
    """Convert PLAID dataset generators into a Hugging Face `datasets.DatasetDict`.

    This function takes generator functions that yield PLAID samples and converts them
    into a Hugging Face DatasetDict. Each generator corresponds to a split (e.g., "train",
    "test") and the function processes samples by flattening their structure and converting
    them to the Hugging Face format based on the provided variable schema.

    Args:
        generators (dict[str, Callable[..., Generator[Sample, None, None]]]):
            Mapping from split names (e.g., "train", "test") to generator functions.
            Each generator function must yield PLAID Sample objects that will be
            converted to the Hugging Face format.
        variable_schema (dict[str, dict]):
            Dictionary defining the schema of variables/features in the dataset.
            Maps feature names to their type information (dtype and ndim).
        cache_dir (str):
            Directory path used as cache directory for the Hugging Face
            dataset generation process.
        gen_kwargs (dict[str, dict[str, list[IndexType]]], optional):
            Optional mapping from split names to dictionaries of keyword arguments
            to be passed to each generator function. Useful for passing split-specific
            parameters like sample indices. Default is None, which creates empty
            kwargs for each split.
        processes_number (int, optional):
            Number of parallel processes to use when materializing the dataset from
            the generators. Default is 1 (no parallelization).
        writer_batch_size (int, optional):
            Batch size used when writing samples to disk in Hugging Face format.
            Default is 1.

    Returns:
        datasets.DatasetDict:
            A Hugging Face DatasetDict containing one Dataset per split, where each
            dataset contains the samples generated by the corresponding generator.

    Example:
        >>> def train_generator():
        ...     for sample in train_samples:
        ...         yield sample
        >>> def test_generator():
        ...     for sample in test_samples:
        ...         yield sample
        >>> variable_schema = {
        ...     "velocity_x": {"dtype": "float32", "ndim": 2},
        ...     "velocity_y": {"dtype": "float32", "ndim": 2}
        ... }
        >>> ds_dict = generator_to_datasetdict(
        ...     generators={"train": train_generator, "test": test_generator},
        ...     variable_schema=variable_schema,
        ...     cache_dir="/tmp/hf_cache",
        ...     processes_number=4,
        ...     writer_batch_size=10
        ... )
        >>> print(ds_dict)
        DatasetDict({
            train: Dataset({
                features: ['velocity_x', 'velocity_y'],
                num_rows: ...
            }),
            test: Dataset({
                features: ['velocity_x', 'velocity_y'],
                num_rows: ...
            })
        })
    """
    hf_features = convert_to_hf_feature(variable_schema)

    all_features_keys = list(variable_schema.keys())

    def generator_fn(gen_func, all_features_keys, **kwargs):
        for sample in gen_func(**kwargs):
            hf_sample, _, _ = build_sample_dict(sample)
            yield {path: hf_sample.get(path, None) for path in all_features_keys}

    _dict = {}
    for split_name, gen_func in generators.items():
        gen = partial(generator_fn, all_features_keys=all_features_keys)
        gen_kwargs_ = gen_kwargs or {split_name: {} for split_name in generators.keys()}
        _dict[split_name] = datasets.Dataset.from_generator(
            generator=gen,
            gen_kwargs={"gen_func": gen_func, **gen_kwargs_[split_name]},
            features=hf_features,
            cache_dir=cache_dir,
            num_proc=processes_number,
            writer_batch_size=writer_batch_size,
            split=datasets.splits.NamedSplit(split_name),
        )
    return datasets.DatasetDict(_dict)


def to_var_sample_dict(
    ds: datasets.Dataset,
    i: int,
    features: Optional[list[str]] = None,
    enforce_shapes: bool = True,
) -> dict[str, Optional[np.ndarray]]:
    """Convert a Hugging Face dataset row to a variable sample dict containing the features that vary in the dataset.

    Args:
        ds (datasets.Dataset): The Hugging Face dataset.
        i (int): The row index.
        features: Iterable of feature names (keys) to extract from the dataset.
        enforce_shapes (bool): Whether to enforce consistent shapes.

    Returns:
        dict[str, Optional[np.ndarray]]: The variable sample dictionary.
    """
    table = ds.data

    if features is None:
        features = table.column_names
    else:
        missing = set(features) - set(table.column_names)
        if missing:  # pragma: no cover
            raise KeyError(f"Missing features in hf_dataset: {sorted(missing)}")

    var_sample_dict = {}
    if not enforce_shapes:
        for name in features:
            value = table[name][i].values
            if value is None:
                var_sample_dict[name] = None  # pragma: no cover
            else:
                var_sample_dict[name] = value.to_numpy(zero_copy_only=False)
    else:
        for name in features:
            if isinstance(table[name][i], pa.NullScalar):
                var_sample_dict[name] = None  # pragma: no cover
            else:
                value = table[name][i].values
                if value is None:
                    var_sample_dict[name] = None  # pragma: no cover
                else:
                    if isinstance(value, pa.ListArray):
                        var_sample_dict[name] = np.stack(
                            value.to_numpy(zero_copy_only=False)
                        )
                    elif isinstance(value, pa.StringArray):  # pragma: no cover
                        var_sample_dict[name] = value.to_numpy(zero_copy_only=False)
                    else:
                        var_sample_dict[name] = value.to_numpy(zero_copy_only=True)

    return var_sample_dict


def sample_to_var_sample_dict(
    hf_sample: dict[str, Any],
) -> dict[str, Any]:
    """Convert a Hugging Face sample dict to variable sample dict.

    Args:
        hf_sample (dict): The HF sample dictionary.

    Returns:
        dict: The processed variable sample dictionary.
    """
    var_sample_dict = {}
    for name, value in hf_sample.items():
        if value is None:
            var_sample_dict[name] = None
        else:
            var_sample_dict[name] = np.array(value)
    return var_sample_dict
