# AUTOGENERATED! DO NOT EDIT! File to edit: ../../notebooks/01_ids_detection.ipynb.

# %% auto 0
__all__ = ['INDEX_STD_THRESHOLD', 'INDEX_FLUC_THRESHOLD', 'INDEX_DIFF_THRESHOLD', 'SPARSE_THRESHOLD', 'pl_format_time',
           'compute_std', 'compute_combinded_std', 'add_neighbor_std', 'compute_index_std', 'compute_index_fluctuation',
           'pl_dvec', 'compute_index_diff', 'compute_indices', 'filter_indices', 'detect_events']

# %% ../../notebooks/01_ids_detection.ipynb 3
from datetime import timedelta
import polars as pl

from fastcore.utils import *
from fastcore.test import *
from ..utils.polars import pl_norm

# %% ../../notebooks/01_ids_detection.ipynb 4
from ..utils.basic import _expand_selectors

# %% ../../notebooks/01_ids_detection.ipynb 5
# some helper functions
def pl_format_time(df: pl.LazyFrame | pl.DataFrame, tau: timedelta):
    return df.with_columns(
        tstart=pl.col("time"),
        tstop=(pl.col("time") + tau),
        time=(pl.col("time") + tau / 2),
    )

# %% ../../notebooks/01_ids_detection.ipynb 6
def compute_std(
    df: pl.DataFrame | pl.LazyFrame,
    period: timedelta,  # period to group by
    b_cols=["BX", "BY", "BZ"],
    every: timedelta = None,  # every to group by (default: period / 2)
):
    if every is None:
        every = period / 2
    b_std_cols = [col_name + "_std" for col_name in b_cols]

    std_df = (
        df.group_by_dynamic("time", every=every, period=period)
        .agg(
            pl.count(),
            pl.col(b_cols).std(ddof=0).map_alias(lambda col_name: col_name + "_std"),
        )
        .with_columns(
            pl_norm(b_std_cols).alias("B_std"),
        )
        .drop(b_std_cols)
    )
    return std_df


def compute_combinded_std(df: pl.DataFrame | pl.LazyFrame, tau, cols) -> pl.DataFrame:
    combined_std_cols = [col_name + "_combined_std" for col_name in cols]
    offsets = [0 * tau, tau / 2]
    combined_std_dfs = []
    for offset in offsets:
        truncated_df = df.select(
            (pl.col("time") - offset).dt.truncate(tau, offset=offset).alias("time"),
            pl.col(cols),
        )

        prev_df = truncated_df.select(
            (pl.col("time") + tau),
            pl.col(cols),
        )

        next_df = truncated_df.select(
            (pl.col("time") - tau),
            pl.col(cols),
        )

        temp_combined_std_df = (
            pl.concat([prev_df, next_df])
            .group_by("time")
            .agg(
                pl.col(cols)
                .std(ddof=0)
                .map_alias(lambda col_name: col_name + "_combined_std"),
            )
            .with_columns(B_std_combined=pl_norm(combined_std_cols))
            .drop(combined_std_cols)
            .sort("time")
        )

        combined_std_dfs.append(temp_combined_std_df)

    combined_std_df = pl.concat(combined_std_dfs)
    return combined_std_df

# %% ../../notebooks/01_ids_detection.ipynb 8
def add_neighbor_std(
    df: pl.LazyFrame, tau: timedelta, join_strategy="inner"
):  # noqa: F811
    """
    Get the neighbor standard deviations

    Parameters
    ----------
    - df (pl.LazyFrame): The input DataFrame.
    - tau : The time interval value.

    Notes
    -----
    Simply shift would not work correctly if data is missing, like `std_next = pl.col("B_std").shift(-2)`.

    """

    # Calculate the standard deviation index
    prev_std_df = df.select(
        pl.col("time") + tau,
        B_std_prev=pl.col("B_std"),
        count_prev=pl.col("count"),
    )

    next_std_df = df.select(
        pl.col("time") - tau,
        B_std_next=pl.col("B_std"),
        count_next=pl.col("count"),
    )

    return df.join(prev_std_df, on="time", how=join_strategy).join(
        next_std_df, on="time", how=join_strategy
    )

# %% ../../notebooks/01_ids_detection.ipynb 9
def compute_index_std(df: pl.LazyFrame):  # noqa: F811
    """
    Compute the standard deviation index based on the given DataFrame

    Parameters
    ----------
    - df (pl.LazyFrame): The input DataFrame.

    Returns
    -------
    - pl.LazyFrame: DataFrame with calculated 'index_std' column.

    Examples
    --------
    >>> index_std_df = compute_index_std_pl(df)
    >>> index_std_df
    """

    return df.with_columns(
        index_std=pl.col("B_std") / pl.max_horizontal("B_std_prev", "B_std_next"),
    )

# %% ../../notebooks/01_ids_detection.ipynb 11
def compute_index_fluctuation(df: pl.LazyFrame, base_col="B_std"):
    std_combined = pl.col("B_std_combined")
    std_added = pl.sum_horizontal("B_std_prev", "B_std_next")
    return df.with_columns(index_fluctuation=std_combined / std_added)

# %% ../../notebooks/01_ids_detection.ipynb 13
def pl_dvec(columns, *more_columns):
    all_columns = _expand_selectors(columns, *more_columns)
    return [
        (pl.col(column).first() - pl.col(column).last()).alias(f"d{column}_vec")
        for column in all_columns
    ]

# %% ../../notebooks/01_ids_detection.ipynb 14
def compute_index_diff(df: pl.LazyFrame, period: timedelta, cols):
    db_cols = ["d" + col + "_vec" for col in cols]

    index_diff = (
        df.with_columns(B=pl_norm(cols))
        .group_by_dynamic("time", every=period / 2, period=period)
        .agg(
            pl.col("B").mean().alias("B_mean"),
            *pl_dvec(cols),
        )
        .with_columns(dB_vec=pl_norm(db_cols))
        .with_columns(index_diff=pl.col("dB_vec") / pl.col("B_mean"))
        .drop(db_cols)
    )

    return index_diff

# %% ../../notebooks/01_ids_detection.ipynb 15
def _compute_indices(
    df: pl.LazyFrame, tau: timedelta, cols: list[str] = ["BX", "BY", "BZ"]
) -> pl.LazyFrame:
    join_strategy = "inner"

    std_df = compute_std(df, tau, cols)
    stds_df = add_neighbor_std(std_df, tau)

    combined_std_df = compute_combinded_std(df, tau, cols)
    index_diff = compute_index_diff(df, tau, cols)

    indices = (
        stds_df.join(index_diff, on="time")
        .join(combined_std_df, on="time", how=join_strategy)
        .pipe(compute_index_std)
        .pipe(compute_index_fluctuation)
        .drop(["B_std_prev", "B_std_next", "B_added_std", "B_std_combined"])
    )

    return indices


def compute_indices(
    df: pl.DataFrame, tau: timedelta, bcols: list[str] = ["BX", "BY", "BZ"]
) -> pl.DataFrame:
    """
    Compute all index based on the given DataFrame and tau value.

    Parameters
    ----------
    df : pl.DataFrame
        Input DataFrame.
    tau : datetime.timedelta
        Time interval value.

    Returns
    -------
    tuple :
        Tuple containing DataFrame results for fluctuation index,
        standard deviation index, and 'index_num'.

    Examples
    --------
    >>> indices = compute_indices(df, tau)

    Notes
    -----
    - This is a wrapper for `_compute_indices` with `pl.LazyFrame` input.
    - Simply shift to calculate index_std would not work correctly if data is missing,
        like `std_next = pl.col("B_std").shift(-2)`.
    - Drop null though may lose some IDs (using the default `join_strategy`).
        Because we could not tell if it is a real ID or just a partial wave
        from incomplete data without previous or/and next std.
        Hopefully we can pick up the lost ones with smaller tau.
    - TODO: Can be optimized further, but this is already fast enough.
        - TEST: if `join` can be improved by shift after filling the missing values.
        - TEST: if `list` in `polars` really fast?
    """
    return _compute_indices(df.lazy(), tau, bcols).collect()

# %% ../../notebooks/01_ids_detection.ipynb 17
# from discontinuitypy import PARAMS

# INDEX_STD_THRESHOLD = PARAMS['detection']['index_std_threshold']
# INDEX_FLUC_THRESHOLD = PARAMS['detection']['index_fluc_threshold']
# INDEX_DIFF_THRESHOLD = PARAMS['detection']['index_diff_threshold']
INDEX_STD_THRESHOLD = 2
INDEX_FLUC_THRESHOLD = 1
INDEX_DIFF_THRESHOLD = 0.1
SPARSE_THRESHOLD = 15

# %% ../../notebooks/01_ids_detection.ipynb 18
def filter_indices(
    df: pl.DataFrame | pl.LazyFrame,
    index_std_threshold : float = INDEX_STD_THRESHOLD,
    index_fluc_threshold : float =INDEX_FLUC_THRESHOLD,
    index_diff_threshold : float =INDEX_DIFF_THRESHOLD,
    sparse_num : int = SPARSE_THRESHOLD,
) -> pl.DataFrame | pl.LazyFrame:
    # filter indices to get possible IDs

    return df.filter(
        pl.col("index_std") > index_std_threshold,
        pl.col("index_fluctuation") > index_fluc_threshold,
        pl.col("index_diff") > index_diff_threshold,
        pl.col("index_std").is_finite(), # for cases where neighboring groups have std=0
        pl.col("count") > sparse_num, 
        pl.col("count_prev") > sparse_num, # filter out sparse intervals, which may give unreasonable results.
        pl.col("count_next") > sparse_num # filter out sparse intervals, which may give unreasonable results.
    ).drop(["count_prev", "count_next"])

# %% ../../notebooks/01_ids_detection.ipynb 19
def detect_events(data: pl.DataFrame, tau: timedelta, ts: timedelta, bcols):
    indices = compute_indices(data, tau, bcols)
    sparse_num = tau / ts // 3
    events = indices.pipe(filter_indices, sparse_num=sparse_num).pipe(
        pl_format_time, tau
    )
    return events
