# AUTOGENERATED! DO NOT EDIT! File to edit: ../../notebooks/utils/02_analysis_utils.ipynb.

# %% auto 0
__all__ = ['AVG_SATS', 'link_coord2dim', 'filter_tranges_ds', 'filter_before_jupiter', 'n2_normalize']

# %% ../../notebooks/utils/02_analysis_utils.ipynb 2
import polars as pl
import pandas as pd

# %% ../../notebooks/utils/02_analysis_utils.ipynb 3
def link_coord2dim(df: pl.DataFrame, dim="time", coord: str = "radial_distance"):
    """Link the coord to a dimension across different subgroups

    Note: this idea is borrowed from the `xarray.DataArray.coords`.
    """
    base_df = df.filter(sat="JNO").select(dim, coord).rename({coord: f"ref_{coord}"})
    return df.join(base_df, on=dim, how="left")

# %% ../../notebooks/utils/02_analysis_utils.ipynb 4
from ..datasets import IDsDataset
from .basic import filter_tranges_df

# %% ../../notebooks/utils/02_analysis_utils.ipynb 5
def filter_tranges_ds(ds: IDsDataset, tranges: tuple[list, list]):
    """Filter a dataset by a list of time ranges"""
    new_ds = ds.copy()
    new_ds.candidates = filter_tranges_df(ds.candidates, tranges)
    new_ds.data = filter_tranges_df(ds.data.collect(), tranges).lazy()
    
    return new_ds

# %% ../../notebooks/utils/02_analysis_utils.ipynb 6
def filter_before_jupiter(df: pl.DataFrame):
    return df.filter(pl.col("time") < pd.Timestamp("2016-05-01"))

# %% ../../notebooks/utils/02_analysis_utils.ipynb 7
AVG_SATS = ["STA", "THB", "Wind"]


def n2_normalize(df: pl.DataFrame, cols, avg_sats: list = AVG_SATS):
    exprs = [pl.col(f"{col}").mean().alias(f"{col}_n2_factor") for col in cols]

    avg_df = df.filter(pl.col("sat").is_in(avg_sats)).group_by("time").agg(exprs)

    exprs = [
        (pl.col(f"{col}") / pl.col(f"{col}_n2_factor")).alias(f"{col}_n2")
        for col in cols
    ]

    return df.join(avg_df, on="time").with_columns(exprs)
