# AUTOGENERATED! DO NOT EDIT! File to edit: ../notebooks/03_mag_plasma.ipynb.

# %% auto 0
__all__ = ['combine_features', 'vector_project', 'vector_project_pl', 'calc_rotation_angle_pl', 'compute_inertial_length',
           'compute_Alfven_speed', 'compute_Alfven_current', 'calc_combined_features']

# %% ../notebooks/03_mag_plasma.ipynb 1
import polars as pl
import polars.selectors as cs

# %% ../notebooks/03_mag_plasma.ipynb 2
def combine_features(candidates: pl.DataFrame, states_data: pl.DataFrame):
    # change time format: see issue: https://github.com/pola-rs/polars/issues/12023
    states_data = states_data.with_columns(
        cs.datetime().dt.cast_time_unit("ns"),
    )
    candidates = candidates.with_columns(
        cs.datetime().dt.cast_time_unit("ns"),
    )

    return candidates.sort("time").join_asof(states_data.sort("time"), on="time")

# %% ../notebooks/03_mag_plasma.ipynb 4
import astropy.units as u
from astropy.constants import mu0, e
from plasmapy.formulary.lengths import inertial_length
from plasmapy.formulary.speeds import Alfven_speed
from xarray_einstats import linalg
from .utils.basic import df2ts
from .core.propeties import calc_rotation_angle
import xarray as xr

# %% ../notebooks/03_mag_plasma.ipynb 5
def vector_project(v1,v2, dim="v_dim"):
    return xr.dot(v1 , v2, dims=dim) / linalg.norm(v2, dims=dim)

def vector_project_pl(df: pl.LazyFrame, v1_cols, v2_cols, name=None):
    
    v1 = df2ts(df, v1_cols).assign_coords(v_dim=["x","y","z"])
    v2 = df2ts(df, v2_cols).assign_coords(v_dim=["x","y","z"])  
    result = vector_project(v1, v2, dim="v_dim")
    
    return df.with_columns(
        pl.Series(result.data).alias(name or "v_proj")
    )

# %% ../notebooks/03_mag_plasma.ipynb 6
def calc_rotation_angle_pl(ldf: pl.LazyFrame, v1_cols, v2_cols, name):
    df = ldf.collect()
    v1 = df.select(v1_cols).to_numpy()
    v2 = df.select(v2_cols).to_numpy()
    
    result = calc_rotation_angle(v1, v2)
    
    return ldf.with_columns(
        pl.Series(result).alias(name)
    )

# %% ../notebooks/03_mag_plasma.ipynb 8
def compute_inertial_length(ldf: pl.LazyFrame, density_col = "plasma_density"):
    df = ldf.collect()

    density = df[density_col].to_numpy() * u.cm ** (-3)
    result = inertial_length(density, "H+").to(u.km)

    return df.with_columns(ion_inertial_length=pl.Series(result.value)).lazy()

# %% ../notebooks/03_mag_plasma.ipynb 10
def compute_Alfven_speed(
    ldf: pl.LazyFrame,
    density_col = "plasma_density",
):
    df = ldf.collect()

    B = df["B"] if "B" in df.columns else df["b_mag"]  # backwards compatiblity
    density = df[density_col].to_numpy() * u.cm ** (-3)
    result = Alfven_speed(B.to_numpy() * u.nT, density=density, ion="p+").to(u.km / u.s)

    return df.with_columns(Alfven_speed=pl.Series(result.value)).lazy()


def compute_Alfven_current(
    ldf: pl.LazyFrame,
    density_col = "plasma_density",
):
    df = ldf.collect()

    Alfven_speed = df["Alfven_speed"].to_numpy() * u.km / u.s
    density = df[density_col].to_numpy() * u.cm ** (-3)

    result = (e.si * Alfven_speed * density)
    result = result.to(u.nA / u.m**2)

    return df.with_columns(j_Alfven=pl.Series(result.value)).lazy()

# %% ../notebooks/03_mag_plasma.ipynb 12
def calc_combined_features(
    df: pl.LazyFrame,
    vec_cols = ["v_x", "v_y", "v_z"],  # plasma velocity vector in any fixed coordinate system,
    normal_cols = ["k_x", "k_y", "k_z"], # normal vector of the discontinuity plane
    b_cols = None, # ["B_background_x", "B_background_y", "B_background_z"]
):
    
    Vl_cols = ["Vl_x", "Vl_y", "Vl_z"]
    Vn_cols = ["Vn_x", "Vn_y", "Vn_z"]
    
    j_factor = ((u.nT / u.s) * (1 / mu0 / (u.km / u.s))).to(u.nA / u.m**2)

    result = df.with_columns( duration = pl.col("d_tstop") - pl.col("d_tstart"))

    if b_cols is not None:
        result = result.pipe(calc_rotation_angle_pl, b_cols, normal_cols, name="theta_n_b")
        
    result = (
        result
        .pipe(vector_project_pl, vec_cols, Vl_cols, name="v_l")
        .pipe(vector_project_pl, vec_cols, Vn_cols, name="v_n")
        .pipe(vector_project_pl, vec_cols, normal_cols, name="v_k")
        .with_columns(
            pl.col("v_n").abs(),
            pl.col("v_k").abs(),
            v_mn=(pl.col("plasma_speed") ** 2 - pl.col("v_l") ** 2).sqrt(),
        )
        .with_columns(
            L_n=pl.col("v_n") * pl.col("duration").dt.nanoseconds() / 1e9,
            L_mn=pl.col("v_mn") * pl.col("duration").dt.nanoseconds() / 1e9,
            L_k=pl.col("v_k") * pl.col("duration").dt.nanoseconds() / 1e9,
            j0=pl.col("d_star") / pl.col("v_mn"),
            j0_k=pl.col("d_star") / pl.col("v_k"),
        )
        .pipe(compute_inertial_length)
        .pipe(compute_Alfven_speed)
        .pipe(compute_Alfven_current)
        .with_columns(
            cs.by_name("j0", "j0_k") * j_factor.value,
        )
    )
    return result
