# AUTOGENERATED! DO NOT EDIT! File to edit: ../notebooks/21_datasets_kedro.ipynb.

# %% auto 0
__all__ = ['E1IDsDataset', 'cIDsDataset']

# %% ../notebooks/21_datasets_kedro.ipynb 1
from .datasets import IDsDataset

from datetime import timedelta

import polars as pl
import polars.selectors as cs

from kedro.io import DataCatalog
from .utils.basic import concat_partitions

# %% ../notebooks/21_datasets_kedro.ipynb 2
class E1IDsDataset(IDsDataset):
    sat_id: str
    tau: timedelta
    ts: timedelta = timedelta(seconds=1)

# %% ../notebooks/21_datasets_kedro.ipynb 3
class cIDsDataset(E1IDsDataset):
    catalog: DataCatalog = None

    _load_data_format = "{sat}.MAG.primary_data_{ts}"
    _load_events_format = "events.{sat}_{ts}_{tau}"
    or_df: pl.DataFrame | None = None  # occurence rate
    or_df_normalized: pl.DataFrame | None = None  # normalized occurence rate

    def __init__(self, **data):
        super().__init__(**data)

        tau_str = f"tau_{self.tau.seconds}s"
        ts_mag_str = f"ts_{self.ts.seconds}s"

        self._tau_str = tau_str
        self._ts_mag_str = ts_mag_str

        self.events_format = self._load_events_format.format(
            sat=self.sat_id, ts=ts_mag_str, tau=tau_str
        )

        if data.get("data_format") is None:
            self.data_format = self._load_data_format.format(
                sat=self.sat_id, ts=ts_mag_str
            )

        if self.candidates is None:
            self.load_events()
        if self.mag_data is None:
            self.load_data()

    def load_events(self):
        data_format = self.events_format
        self.candidates = (
            self.catalog.load(data_format)
            .fill_nan(None)
            .with_columns(
                cs.float().cast(pl.Float64),
                sat=pl.lit(self.sat_id),
            )
            .collect()
        )

    def load_data(self):
        data_format = self.data_format
        self.mag_data = concat_partitions(self.catalog.load(data_format))
