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

# %% auto 0
__all__ = ['get_candidate_data', 'get_candidates', 'calc_candidate_duration', 'get_data_at_times', 'calc_rotation_angle',
           'calc_events_rotation_angle', 'calc_normal_direction', 'calc_events_normal_direction',
           'calc_events_vec_change', 'IDsPdPipeline', 'process_events']

# %% ../../notebooks/02_ids_properties.ipynb 1
#| code-summary: "Import all the packages needed for the project"
import polars as pl
import xarray as xr
from fastcore.all import patch

try:
    import modin.pandas as pd
    import modin.pandas as mpd
    from modin.config import ProgressBar
    ProgressBar.enable()
except ImportError:
    import pandas as pd
import pandas
    
import numpy as np

from datetime import timedelta

from loguru import logger

import pdpipe as pdp
from pdpipe.util import out_of_place_col_insert

# %% ../../notebooks/02_ids_properties.ipynb 2
def get_candidate_data(
    candidate: dict, data: xr.DataArray, neighbor: int = 0
) -> xr.DataArray:
    duration = candidate["tstop"] - candidate["tstart"]
    offset = neighbor * duration
    temp_tstart = candidate["tstart"] - offset
    temp_tstop = candidate["tstop"] + offset

    return data.sel(time=slice(temp_tstart, temp_tstop))


def get_candidates(candidates: pd.DataFrame, candidate_type=None, num: int = 4):
    if candidate_type is not None:
        _candidates = candidates[candidates["type"] == candidate_type]
    else:
        _candidates = candidates

    # Sample a specific number of candidates if num is provided and it's less than the total number
    if num < len(_candidates):
        logger.info(
            f"Sampling {num} {candidate_type} candidates out of {len(_candidates)}"
        )
        return _candidates.sample(num)
    else:
        return _candidates

# %% ../../notebooks/02_ids_properties.ipynb 4
from ..propeties.duration import calc_duration

# %% ../../notebooks/02_ids_properties.ipynb 5
def calc_candidate_duration(candidate, data, **kwargs):
    try:
        candidate_data = get_candidate_data(candidate, data)
        result = calc_duration(candidate_data, **kwargs)
        return pandas.Series(result)
    except Exception as e:
        logger.debug(
            f"Error for candidate {candidate} at {candidate['time']}: {str(e)}"
        )
        raise e

# %% ../../notebooks/02_ids_properties.ipynb 7
from ..propeties.mva import calc_candidate_mva_features

# %% ../../notebooks/02_ids_properties.ipynb 9
def get_data_at_times(data: xr.DataArray, times) -> np.ndarray:
    """
    Select data at specified times.
    """
    # Use xarray's selection capability if data supports it
    return data.sel(time=times, method="nearest").to_numpy()

# %% ../../notebooks/02_ids_properties.ipynb 10
def calc_rotation_angle(v1, v2):
    """
    Computes the rotation angle between two vectors.
    
    Parameters:
    - v1: The first vector(s).
    - v2: The second vector(s).
    """
    
    if v1.shape != v2.shape:
        raise ValueError("Vectors must have the same shape.")
    
    # Normalize the vectors
    v1_u = v1 / np.linalg.norm(v1, axis=-1, keepdims=True)
    v2_u = v2 / np.linalg.norm(v2, axis=-1, keepdims=True)
    
    # Calculate the cosine of the angle for each time step
    cosine_angle = np.sum(v1_u * v2_u, axis=-1)
    
    # Clip the values to handle potential floating point errors
    cosine_angle = np.clip(cosine_angle, -1, 1)
    
    angle = np.arccos(cosine_angle)
    
    # Convert the angles from radians to degrees
    return np.degrees(angle)


# %% ../../notebooks/02_ids_properties.ipynb 11
def calc_events_rotation_angle(events, data: xr.DataArray):
    """
    Computes the rotation angle(s) at two different time steps.
    """
    tstart = events['t.d_start'].to_numpy()
    tstop = events['t.d_end'].to_numpy()

    vecs_before = get_data_at_times(data, tstart)
    vecs_after = get_data_at_times(data, tstop)

    rotation_angles = calc_rotation_angle(vecs_before, vecs_after)
    return rotation_angles

# %% ../../notebooks/02_ids_properties.ipynb 13
def calc_normal_direction(v1, v2, normalize=True) -> np.ndarray:
    """
    Computes the normal direction of two vectors.

    Parameters
    ----------
    v1 : array_like 
        The first vector(s).
    v2 : array_like 
        The second vector(s).
    """
    c = np.cross(v1, v2)
    return c / np.linalg.norm(c, axis=-1, keepdims=True)


# %% ../../notebooks/02_ids_properties.ipynb 14
def calc_events_normal_direction(events, data: xr.DataArray):
    """
    Computes the normal directions(s) at two different time steps.
    """
    tstart = events['t.d_start'].to_numpy()
    tstop = events['t.d_end'].to_numpy()

    vecs_before = get_data_at_times(data, tstart)
    vecs_after = get_data_at_times(data, tstop)

    normal_directions = calc_normal_direction(vecs_before, vecs_after)
    # need to convert to list first, as only 1D array is supported
    return normal_directions.tolist()


# %% ../../notebooks/02_ids_properties.ipynb 15
def calc_events_vec_change(events, data: xr.DataArray):
    """
    Utils function to calculate features related to the change of the magnetic field
    """
    tstart = events['t.d_start'].to_numpy()
    tstop = events['t.d_end'].to_numpy()
    
    vecs_before = get_data_at_times(data, tstart)
    vecs_after = get_data_at_times(data, tstop)
    return (vecs_after - vecs_before).tolist()
    

# %% ../../notebooks/02_ids_properties.ipynb 18
@patch
def _transform(self: pdp.ApplyToRows, X, verbose):
    new_cols = X.apply(self._func, axis=1)
    if isinstance(new_cols, (pd.Series, pandas.Series)):
        loc = len(X.columns)
        if self._follow_column:
            loc = X.columns.get_loc(self._follow_column) + 1
        return out_of_place_col_insert(
            X=X, series=new_cols, loc=loc, column_name=self._colname
        )
    if isinstance(new_cols, (mpd.DataFrame, pandas.DataFrame)):
        sorted_cols = sorted(list(new_cols.columns))
        new_cols = new_cols[sorted_cols]
        if self._follow_column:
            inter_X = X
            loc = X.columns.get_loc(self._follow_column) + 1
            for colname in new_cols.columns:
                inter_X = out_of_place_col_insert(
                    X=inter_X,
                    series=new_cols[colname],
                    loc=loc,
                    column_name=colname,
                )
                loc += 1
            return inter_X
        assign_map = {
            colname: new_cols[colname] for colname in new_cols.columns
        }
        return X.assign(**assign_map)
    raise TypeError(  # pragma: no cover
        "Unexpected type generated by applying a function to a DataFrame."
        " Only Series and DataFrame are allowed."
    )

# %% ../../notebooks/02_ids_properties.ipynb 20
class IDsPdPipeline:
    @staticmethod
    def calc_duration(data: xr.DataArray, **kwargs):
        return pdp.ApplyToRows(
            lambda df: calc_candidate_duration(df, data, **kwargs),
            func_desc="calculating pre-duration parameters",
        )

    @staticmethod
    def calc_mva_features(data, **kwargs):
        return pdp.ApplyToRows(
            lambda df: calc_candidate_mva_features(df, data, **kwargs),
            func_desc="calculating MVA features",
        )

    @staticmethod
    def calc_vec_change(data, **kwargs):
        return pdp.ColByFrameFunc(
            "dB",
            lambda df: calc_events_vec_change(df, data, **kwargs),
            func_desc="calculating compound change",
        )

    @staticmethod
    def calc_rotation_angle(data, **kwargs):
        return pdp.ColByFrameFunc(
            "rotation_angle",
            lambda df: calc_events_rotation_angle(df, data, **kwargs),
            func_desc="calculating rotation angle",
        )

    @staticmethod
    def calc_normal_direction(data, name="normal_direction", **kwargs):
        return pdp.ColByFrameFunc(
            name,
            lambda df: calc_events_normal_direction(df, data, **kwargs),
            func_desc="calculating normal direction",
        )

# %% ../../notebooks/02_ids_properties.ipynb 21
from beforerr.polars import convert_to_pd_dataframe, decompose_vector

# %% ../../notebooks/02_ids_properties.ipynb 22
from typing import Literal


def process_events(
    candidates_pl: pl.DataFrame,  # potential candidates DataFrame
    sat_fgm: xr.DataArray,  # satellite FGM data
    data_resolution: timedelta,  # time resolution of the data
    modin=True,
    method: Literal["fit", "derivative"]= "fit",
    **kwargs,
) -> pl.DataFrame:
    "Process candidates DataFrame"

    candidates = pd.DataFrame(convert_to_pd_dataframe(candidates_pl, modin=modin))
    
    if method == "fit":
        duration_method = "distance"  
    else:
        duration_method = "derivative"

    candidates = (
        IDsPdPipeline.calc_duration(sat_fgm, method=duration_method, **kwargs)
        .apply(candidates)
        .dropna()
    )  # Remove candidates with NaN values)

    ids = (
        IDsPdPipeline.calc_mva_features(sat_fgm, method=method, **kwargs)
        + IDsPdPipeline.calc_vec_change(sat_fgm)
        + IDsPdPipeline.calc_rotation_angle(sat_fgm)
        + IDsPdPipeline.calc_normal_direction(sat_fgm, name="k")
    ).apply(candidates)

    if isinstance(ids, mpd.DataFrame):
        ids = ids._to_pandas()

    vectors2decompose = ["dB", "dB_lmn", "k", "Vl", "Vn"]

    df = pl.DataFrame(
        ids.dropna(), schema_overrides={vec: pl.List for vec in vectors2decompose}
    )  # ArrowInvalid: Could not convert [0.9799027968348948, -0.17761542644940076, -0.07309766783111293] with type list: tried to convert to double

    if method == "fit":
        duration_expr = pl.col("fit.vars.sigma") * 2
    else:
        duration_expr = (
            pl.col("t.d_end") - pl.col("t.d_start")
        ).dt.total_nanoseconds() / 1e9  # convert to seconds

    for vec in vectors2decompose:
        df = decompose_vector(df, vec)

    return df.with_columns(duration=duration_expr).drop(vectors2decompose)
    # ValueError: Data type fixed_size_list[pyarrow] not supported by interchange protocol
