# 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, timedelta
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, get_data, get_time_resolution
from space_analysis.ds.spz.io import spzvars2pldf
import polars as pl
from functools import cached_property
from abc import abstractmethod

from tqdm.auto import tqdm
from loguru import logger

# %% ../../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 (thus often requiring getting data lazily)
    """

    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}"

    @abstractmethod
    def get_data(self):
        pass

    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, _ = zip(
                *(config.produce_or_load(**kwargs) for config in tqdm(configs))
            )

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

# %% ../../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()

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

    def get_data(self):
        return self._get_mag_data()._get_plasma_data()

    def _get_plasma_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.plasma_data is None:
            self.plasma_data = self.get_vars_df("plasma")
        return self

    def _get_mag_data(self):
        if self.data is None:
            data = get_data(self.mag_meta, self.provider, self.timerange)
            self.mag_meta.data = data
            if self.mag_meta.ts is None:
                ts = timedelta(seconds=get_time_resolution(data[0])["median"])
                logger.info(f"Setting time resolution to {ts}")
                self.mag_meta.ts = ts
            self.data = spzvars2pldf(data)
        return self
