# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/10_polars.ipynb.

# %% auto 0
__all__ = ['convert_to_pd_dataframe', 'sort', 'pl_norm', 'decompose_vector', 'filter_series_by_ranges_i', 'filter_df_by_ranges']

# %% ../nbs/10_polars.ipynb 2
import polars as pl
from typing import Any, Collection

# %% ../nbs/10_polars.ipynb 4
def convert_to_pd_dataframe(
    df: pl.DataFrame | pl.LazyFrame,  # original DataFrame or LazyFrame
):
    """
    Convert a Polars DataFrame or LazyFrame into a pandas-like DataFrame.
    """
    if isinstance(df, pl.LazyFrame):
        df = df.collect()
    elif not isinstance(df, pl.DataFrame):
        raise TypeError("Input must be a Polars DataFrame or LazyFrame")

    data = df.to_pandas(use_pyarrow_extension_array=True)

    return data

# %% ../nbs/10_polars.ipynb 6
def sort(df: pl.DataFrame, col="time"):
    if df.get_column(col).is_sorted():
        return df.set_sorted(col)
    else:
        return df.sort(col)

# %% ../nbs/10_polars.ipynb 7
def _expand_selectors(items: Any, *more_items: Any) -> list[Any]:
    """
    See `_expand_selectors` in `polars`.
    """
    expanded: list[Any] = []
    for item in (
        *(
            items
            if isinstance(items, Collection) and not isinstance(items, str)
            else [items]
        ),
        *more_items,
    ):
        expanded.append(item)
    return expanded

# %% ../nbs/10_polars.ipynb 8
def pl_norm(columns, *more_columns) -> pl.Expr:
    """
    Computes the square root of the sum of squares for the given columns.

    Args:
    *columns (str): Names of the columns.

    Returns:
    pl.Expr: Expression representing the square root of the sum of squares.
    """
    all_columns = _expand_selectors(columns, *more_columns)
    squares = [pl.col(column).pow(2) for column in all_columns]

    return sum(squares).sqrt()

# %% ../nbs/10_polars.ipynb 9
def decompose_vector(
    df: pl.DataFrame, vector_col, name=None, suffixes: list = ["_x", "_y", "_z"]
):
    """
    Decompose a vector column in a DataFrame into separate columns for each component with custom suffixes.

    Parameters:
    - df (pl.DataFrame): The input DataFrame.
    - vector_col (str): The name of the vector column to decompose.
    - name (str, optional): Base name for the decomposed columns. If None, uses `vector_col` as the base name.
    - suffixes (list, optional): A list of suffixes to use for the decomposed columns.
      If None or not enough suffixes are provided, defaults to '_0', '_1', etc.

    Returns:
    - pl.DataFrame: A DataFrame with the original vector column decomposed into separate columns.
    """

    if name is None:
        name = vector_col

    # Determine the maximum length of vectors in the column to handle dynamic vector lengths
    max_length = df.select(pl.col(vector_col).list.len()).max()[0, 0]

    if suffixes is None or len(suffixes) < max_length:
        if suffixes is None:
            suffixes = []
        # Extend or create the list of suffixes with default values
        suffixes.extend([f"_{i}" for i in range(len(suffixes), max_length)])

    # Create column expressions for each element in the vector
    column_expressions = [
        pl.col(vector_col).list.get(i).alias(name).name.suffix(suffixes[i])
        for i in range(max_length)
    ]

    return df.with_columns(column_expressions)

# %% ../nbs/10_polars.ipynb 11
def filter_series_by_ranges_i(data: pl.Series, starts: list, stops: list):
    starts_index = data.search_sorted(starts)
    ends_index = data.search_sorted(stops, side="right")

    return pl.concat(
        pl.arange(*range, eager=True) for range in zip(starts_index, ends_index)
    ).unique()


def filter_df_by_ranges(data: pl.DataFrame, starts: list, stops: list, col="time"):
    """
    Filter a DataFrame from ranges
    """

    indices_unique = filter_series_by_ranges_i(data[col], starts, stops)
    return data[indices_unique]
