# AUTOGENERATED! DO NOT EDIT! File to edit: ../notebooks/20_datasets.ipynb.

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

# %% ../notebooks/20_datasets.ipynb 1
import polars as pl
import pandas as pd
from loguru import logger
from random import sample

from pydantic import BaseModel, Field

# %% ../notebooks/20_datasets.ipynb 4
from .utils.basic import df2ts
from .utils.plot import plot_candidate as _plot_candidate
from .integration import combine_features, calc_combined_features

# %% ../notebooks/20_datasets.ipynb 5
class IDsDataset(BaseModel):
    class Config:
        arbitrary_types_allowed = True
        extra = "allow"

    mag_data: pl.LazyFrame = Field(alias="data")
    candidates: pl.DataFrame = None
    plasma_data: pl.LazyFrame = None

    bcols: list[str] = None
    vec_cols: list[str] = None

    def update_candidates_with_plasma_data(self):
        df_combined = combine_features(
            self.candidates, self.plasma_data.collect()
        ).lazy()

        self.candidates = calc_combined_features(
            df_combined, vec_cols=self.vec_cols
        ).collect()

    def get_candidate(self, index=None, predicates=None, **kwargs):
        if index is not None:
            candidate = self.candidates.row(index, named=True)
        elif predicates is not None:
            candidate = self.candidates.filter(predicates).row(0, named=True)
        return candidate

    def get_candidate_data(self, candidate=None, index=None, predicates=None, **kwargs):
        if candidate is None:
            candidate = self.get_candidate(index, predicates, **kwargs)

        _data = self.mag_data.filter(
            pl.col("time").is_between(candidate["tstart"], candidate["tstop"])
        )
        return df2ts(_data, self.bcols)

    def plot_candidate(self, candidate=None, index=None, predicates=None, **kwargs):
        if candidate is None:
            candidate = self.get_candidate(index, predicates, **kwargs)
        sat_fgm = self.get_candidate_data(candidate)

        return _plot_candidate(candidate, sat_fgm, **kwargs)

    def plot_candidates(self, indices=None, num=4, random=True, **kwargs):
        if indices is None:
            indices = sample(range(len(self.candidates)), num) if random else range(num)
            logger.info(f"Candidates indices: {indices}")

        return [self.plot_candidate(index=i, **kwargs) for i in indices]
