# AUTOGENERATED! DO NOT EDIT! File to edit: ../../../notebooks/detection/02_TS_criterion.ipynb.

# %% auto 0
__all__ = ['detect_gradient']

# %% ../../../notebooks/detection/02_TS_criterion.ipynb 1
import polars as pl
from datetime import timedelta
from beforerr.polars import pl_norm, format_time

# %% ../../../notebooks/detection/02_TS_criterion.ipynb 2
def detect_gradient(
    df: pl.LazyFrame,
    cols: list[str],
    time: str = "time",
    avg_interval=timedelta(minutes=1),
    window=timedelta(minutes=3),
):
    """
    Identifies discontinuities in the averaged vector field and specified criteria:
    |ΔB| > max(|B_i|, |B_{i-window}|) / 2

    Parameters
    ----------
    df : pl.LazyFrame
        Dataframe containing the raw vector components.
    time : str
        Column name for time.
    avg_interval : timedelta
        Time interval for averaging (e.g., '1m' for 1 minute).
    window : timedelta
        Time interval to look back for computing ΔB.
    """
    # Step 1: 1-Minute Averaging
    averaged_df = df.group_by_dynamic(time, every=avg_interval).mean()

    # Step 2: Shifted Columns
    suffix = "_shifted"
    prev_df = averaged_df.with_columns(pl.col(time) + window).pipe(format_time)
    combined_df = averaged_df.join(prev_df, on=time, suffix=suffix)
    shifted_cols = [col + suffix for col in cols]

    # Step 3: Compute ΔB Components
    delta_cols = [pl.col(col) - pl.col(col + suffix) for col in cols]

    # Step 4: Compute Magnitudes |B_i|, |B_{i-window}|, and |ΔB|
    magnitude_Bi = pl_norm(cols)
    magnitude_Biw = pl_norm(shifted_cols)
    magnitude_deltaB = pl_norm(delta_cols)

    # Step 5: Compute B_L = max(|B_i|, |B_{i-window}|)
    B_L = pl.max_horizontal(magnitude_Bi, magnitude_Biw)

    # Step 6: Apply Transformations and Filter
    return (
        combined_df.with_columns(
            magnitude_deltaB.alias("|ΔB|"),
            B_L.alias("B_L"),
        )
        .filter(pl.col("|ΔB|") > pl.col("B_L") / 2)
        .drop(cols + shifted_cols)
        .with_columns(
            tstart=pl.col("time") - window,
            tstop=pl.col("time") + avg_interval,
            time=pl.col("time") + (avg_interval - window) / 2,
        )
        .pipe(format_time)
        .collect()
    )
