# AUTOGENERATED! DO NOT EDIT! File to edit: ../notebooks/11_ids_config.ipynb.

# %% auto 0
__all__ = ['split_timerange', 'IDsConfig', 'get_vars', 'SpeasyIDsConfig']

# %% ../notebooks/11_ids_config.ipynb 0
from datetime import datetime
from beforerr.project import produce_or_load_file
from .datasets import IDsDataset
from space_analysis.meta import Dataset
from space_analysis.utils.speasy import Variables
import polars as pl
from functools import cached_property

from tqdm.auto import tqdm

# %% ../notebooks/11_ids_config.ipynb 1
def split_timerange(timerange: list[datetime], split: int = 1):
    """
    Split a timerange into multiple timeranges.
    """
    from sunpy.time import TimeRange

    trs: list[TimeRange] = TimeRange(timerange).split(split)
    return [[tr.start.value, tr.end.value] for tr in trs]


class IDsConfig(IDsDataset):
    """
    Extend the IDsDataset class to provide additional functionalities:

    - Split data to handle large datasets
    """

    timerange: list[datetime] = None
    split: int = 1

    @property
    def timeranges(self):
        return split_timerange(self.timerange, self.split)

    @property
    def file_prefix(self):
        if self.timerange is None:
            return super().file_prefix
        else:
            tr_str = "-".join(t.strftime("%Y%m%d") for t in self.timerange)
            return super().file_prefix + f"_tr={tr_str}"

# %% ../notebooks/11_ids_config.ipynb 2
def get_vars(self, vars: str, timerange: list[datetime] = None):
    meta: Dataset = getattr(self, f"{vars}_meta")
    timerange = timerange or self.timerange or meta.timerange
    return Variables(
        timerange=timerange,
        provider=self.provider,
        **meta.model_dump(exclude_unset=True),
    )


class SpeasyIDsConfig(IDsConfig):
    """Based on `speasy` Variables to get the data"""

    provider: str = "cda"

    def get_vars(self, *args, **kwargs):
        return get_vars(self, *args, **kwargs)

    def get_vars_df(self, vars: str, **kwargs):
        return get_vars(self, vars, **kwargs).to_polars()

    # Variables
    @cached_property
    def mag_vars(self):
        return self.get_vars("mag")

    @cached_property
    def plasma_vars(self):
        return self.get_vars("plasma")

    def get_data(self):
        # TODO: directly get columns from the data without loading them
        self.plasma_meta.density_col = (
            self.plasma_meta.density_col or self.plasma_vars.data[0].columns[0]
        )
        self.plasma_meta.velocity_cols = (
            self.plasma_meta.velocity_cols or self.plasma_vars.data[1].columns
        )
        if self.data is None:
            self.data = self.get_vars_df("mag")
        if self.plasma_data is None:
            self.plasma_data = self.get_vars_df("plasma")
        return self

    def produce_or_load(self, **kwargs):
        if self.split == 1:
            self.file.exists() or self.get_data()
            return super().produce_or_load(**kwargs)
        else:
            updates = [{"timerange": tr, "split": 1} for tr in self.timeranges]
            configs = [self.model_copy(update=update, deep=True) for update in updates]
            datas, files = zip(
                config.produce_or_load(**kwargs) for config in tqdm(configs)
            )

            return produce_or_load_file(
                f=pl.concat,
                c=dict(items=datas),
                file=self.file,
            )
