# AUTOGENERATED! DO NOT EDIT! File to edit: ../../notebooks/utils/00_basic.ipynb.

# %% auto 0
__all__ = ['load_catalog', 'load_params', 'DF_TYPE', 'pmap', 'DataConfig', 'filter_tranges', 'filter_tranges_df', 'pl_norm',
           'partition_data_by_ts', 'partition_data_by_year', 'partition_data_by_year_month', 'partition_data_by_time',
           'concat_df', 'concat_partitions', 'format_timedelta', 'resample', 'df2ts', 'calc_vec_mag', 'check_fgm']

# %% ../../notebooks/utils/00_basic.ipynb 1
from .. import ROOT_DIR
from .kedro import load_context

from functools import partial

from typing import overload

# %% ../../notebooks/utils/00_basic.ipynb 2
load_catalog = partial(load_context, project_path=ROOT_DIR, catalog_only=True)
load_params = partial(load_context, project_path=ROOT_DIR, params_only=True)

# %% ../../notebooks/utils/00_basic.ipynb 4
import polars as pl
import polars.selectors as cs

import pandas as pd
import xarray as xr

import pandas
import numpy as np
from xarray_einstats import linalg

from datetime import timedelta

from loguru import logger

from xarray import DataArray
from typing import Union, Collection, Callable, Optional, Tuple
from typing import Any, Dict


# %% ../../notebooks/utils/00_basic.ipynb 5
from pipe import select
from fastcore.utils import partial

# %% ../../notebooks/utils/00_basic.ipynb 6
def pmap(func, *args, **kwargs):
    """
    map with `partial`
    """
    return select(partial(func, *args, **kwargs))

# %% ../../notebooks/utils/00_basic.ipynb 9
from pydantic import BaseModel
from datetime import datetime, timedelta
from pandas import Timedelta

class DataConfig(BaseModel):
    sat_id: str = None
    start: datetime = None
    end: datetime = None
    ts: timedelta = None
    coord: str = None

# %% ../../notebooks/utils/00_basic.ipynb 11
from fastcore.utils import patch

# %% ../../notebooks/utils/00_basic.ipynb 12
def filter_tranges(time: pl.Series, tranges: Tuple[list, list]):
    """
    - Filter data by time ranges, return the indices of the time that are in the time ranges (left inclusive, right exclusive)
    """

    starts = tranges[0]
    ends = tranges[1]

    start_indices = time.search_sorted(starts)
    end_indices = time.search_sorted(ends)

    return np.concatenate(
        [
            np.arange(start_index, end_index)
            for start_index, end_index in zip(start_indices, end_indices)
        ]
    )

def filter_tranges_df(df: pl.DataFrame, tranges: Tuple[list, list], time_col: str = "time"):
    """
    - Filter data by time ranges
    """

    time = df[time_col]
    filtered_indices = filter_tranges(time, tranges)
    return df[filtered_indices]

# %% ../../notebooks/utils/00_basic.ipynb 13
@patch
def plot(self:pl.DataFrame, *args, **kwargs):
    return self.to_pandas().plot(*args, **kwargs)

# %% ../../notebooks/utils/00_basic.ipynb 14
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

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()

# %% ../../notebooks/utils/00_basic.ipynb 16
def partition_data_by_ts(df: pl.DataFrame, ts: timedelta) -> Dict[str, pl.DataFrame]:
    """Partition the dataset by time

    Args:
        df: Input DataFrame.
        ts: Time interval.

    Returns:
        Partitioned DataFrame.
    """
    return df.with_columns(
        key=pl.col("time").dt.truncate(ts).cast(pl.Utf8)
    ).partition_by("key", include_key=False, as_dict=True)


def partition_data_by_year(df: pl.DataFrame) -> Dict[str, pl.DataFrame]:
    """Partition the dataset by year

    Args:
        df: Input DataFrame.

    Returns:
        Partitioned DataFrame.
    """
    return df.with_columns(year=pl.col("time").dt.year().cast(pl.Utf8)).partition_by(
        "year", include_key=False, as_dict=True
    )


def partition_data_by_year_month(df: pl.DataFrame) -> Dict[str, pl.DataFrame]:
    """Partition the dataset by year

    Args:
        df: Input DataFrame.

    Returns:
        Partitioned DataFrame.
    """
    return df.with_columns(
        year_month=pl.col("time").dt.year().cast(pl.Utf8)
        + "_"
        + pl.col("time").dt.month().cast(pl.Utf8).str.zfill(2),
    ).partition_by("year_month", include_key=False, as_dict=True)
    
def partition_data_by_time(df: pl.LazyFrame | pl.DataFrame, method) -> Dict[str, pl.DataFrame]:
    """Partition the dataset by time

    Args:
        df: Input DataFrame.
        method: The method to partition the data.

    Returns:
        Partitioned DataFrame.
    """
    if isinstance(df, pl.LazyFrame):
        df = df.collect()
    
    if method == "year":
        return partition_data_by_year(df)
    elif method == "year_month":
        return partition_data_by_year_month(df)
    else:
        ts = pd.Timedelta(method)
        return partition_data_by_ts(df, ts)

# %% ../../notebooks/utils/00_basic.ipynb 17
DF_TYPE = Union[pl.DataFrame, pl.LazyFrame, pd.DataFrame]
def concat_df(dfs: list[DF_TYPE]) -> DF_TYPE:
    """Concatenate a list of DataFrames into one DataFrame.
    """
    
    match type(dfs[0]):
        case pl.DataFrame | pl.LazyFrame:
            concat_func = pl.concat
        case pandas.DataFrame:
            concat_func = pandas.concat
        case _:
            raise ValueError(f"Unsupported DataFrame type: {type(dfs[0])}")
    
    return concat_func(dfs)
                     
def concat_partitions(partitioned_input: Dict[str, Callable]):
    """Concatenate input partitions into one DataFrame.

    Args:
        partitioned_input: A dictionary with partition ids as keys and load functions as values.
    """
    partitions_data = [
        partition_load_func() for partition_load_func in partitioned_input.values()
    ]  # load the actual partition data
    
    result = concat_df(partitions_data)
    return result

# %% ../../notebooks/utils/00_basic.ipynb 19
def format_timedelta(time):
    """Format timedelta to `timedelta`"""
    if isinstance(time, timedelta):
        return time
    elif isinstance(time, str):
        return pd.Timedelta(time)
    elif isinstance(time, int):
        return pd.Timedelta(seconds=time)
    else:
        raise TypeError(f"Unsupported type: {type(time)}")

# %% ../../notebooks/utils/00_basic.ipynb 20
@overload
def resample(
    df: pl.DataFrame,
    every: timedelta,
    period: timedelta = None,
    offset: timedelta = None,
    shift: timedelta = None,
    time_column="time",
)-> pl.DataFrame:
    ...

@overload
def resample(
    df: pl.LazyFrame,
    every: timedelta,
    period: timedelta = None,
    offset: timedelta = None,
    shift: timedelta = None,
    time_column="time",
)-> pl.LazyFrame:
    ...

def resample(
    df: pl.LazyFrame | pl.DataFrame,
    every: timedelta,
    period: timedelta = None,
    offset: timedelta = None,
    shift: timedelta = None,
    time_column="time",
):
    """Resample the DataFrame"""
    if period is None:
        period = every
    if shift is None:
        shift = period / 2
    return (
        df.sort(time_column)
        .group_by_dynamic(time_column, every=every, period=period, offset=offset)
        .agg(cs.numeric().mean())
        .with_columns((pl.col(time_column) + shift))
    )

# %% ../../notebooks/utils/00_basic.ipynb 21
def df2ts(
    df: Union[pandas.DataFrame, pl.DataFrame, pl.LazyFrame],
    cols = None,
    time_col = "time",
    attrs=None,
    name=None,
):
    """Convert DataFrame to TimeSeries"""
    
    if cols is None:
        if isinstance(df, (pl.DataFrame, pl.LazyFrame)):
            cols = df.columns
            cols.remove(time_col)
        else:
            cols = df.columns.tolist()
    
    if isinstance(df, pl.LazyFrame):
        df = df.collect()

    # Prepare data
    data = df[cols].to_numpy()

    # Prepare coordinates
    time = df.index if isinstance(df, pandas.DataFrame) else df[time_col]
    if isinstance(cols, str) and len(data.shape) > 1:
        element_len = data.shape[1]
        v_dim = [cols + str(i) for i in range(element_len)]
    else:
        v_dim = cols
        
    # Create the DataArray
    coords = {"time": time, "v_dim": v_dim}

    return xr.DataArray(data, coords=coords, attrs=attrs, name=name)


def calc_vec_mag(vec) -> DataArray:
    return linalg.norm(vec, dims="v_dim")

# %% ../../notebooks/utils/00_basic.ipynb 22
def check_fgm(vec: xr.DataArray):
    # check if time is monotonic increasing
    logger.info("Check if time is monotonic increasing")
    assert vec.time.to_series().is_monotonic_increasing
    # check available time difference
    logger.info(
        f"Available time delta: {vec.time.diff(dim='time').to_series().unique()}"
    )

