# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/11_meta_loader.ipynb.

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

# %% ../nbs/11_meta_loader.ipynb 3
import os
import re
from typing import List, Any, Dict, Union
import warnings

import numpy as np
import pandas as pd
import dask.dataframe as dd

# %% ../nbs/11_meta_loader.ipynb 4
from pheno_utils.config import (
    DATASETS_PATH, 
    COHORT, 
    ERROR_ACTION
    )
from .pheno_loader import PhenoLoader

# %% ../nbs/11_meta_loader.ipynb 5
class MetaLoader:
    """
    Class to load multiple dictionaries and allows to easily access the relevant fields.

    Args:
    
        base_path (str, optional): The base path where the data is stored. Defaults to DATASETS_PATH.
        cohort (str, optional): The name of the cohort within the dataset. Defaults to COHORT.
        flexible_field_search (bool, optional): Whether to allow regex field search. Defaults to False.
        errors (str, optional): Whether to raise an error or issue a warning if missing data is encountered.
            Possible values are 'raise', 'warn' and 'ignore'. Defaults to 'raise'.
        **kwargs: Additional keyword arguments to pass to a DataLoader class.

    Attributes:
    
        dicts (pd.DataFrame): A dictionary of data dictionaries (dataframes) of all availbale datasets in the base_path.
        fields (list): A list of all fields.
        cohort (str): The name of the cohort being used.
        base_path (str): The base path where the data is stored.
        flexible_field_search (bool): Whether to allow regex field search.
        errors (str): Whether to raise an error or issue a warning if missing data is encountered.
        kwargs (dict): Additional keyword arguments to pass to a DataLoader class.
    """

    def __init__(
        self,
        base_path: str = DATASETS_PATH,
        cohort: str = COHORT,
        flexible_field_search: bool = False,
        errors: str = ERROR_ACTION,
        **kwargs,
    ) -> None:
        self.cohort = cohort
        self.base_path = base_path
        self.dataset_path = self.__get_dataset_path__()
        
        self.flexible_field_search = flexible_field_search
        self.errors = errors
        self.kwargs = kwargs

        self.__load_dictionaries__()

    def load(self, fields: Union[str,List[str]], flexible: bool=None, prop: str='tabular_field_name') -> pd.DataFrame:
        """
        Return a dataframe containing the fields from the respective datasets.

        Args:
            fields (Union[str,List[str]]): Fields to return
            flexible (bool, optional): Whether to use fuzzy matching to find fields. Defaults to None, which uses the DataLoader's flexible_field_search attribute.
            prop (str, optional): The property to use for searching. Defaults to 'tabular_field_name'.

        Returns:
            pd.DataFrame: Dataframe containing the fields from the respective datasets.
        """
        found_fields = self.get(fields, flexible, prop)
        found_fields.columns = found_fields.columns.str.split('/').str[1]
        dup_fields = found_fields.columns.value_counts()\
            .to_frame('count').query('count > 1').index

        loaded_fields = []
        for ds, f in found_fields.T.groupby('dataset'):
            df = PhenoLoader(ds, base_path=self.base_path, cohort=self.cohort,
                             age_sex_dataset=None, **self.kwargs)\
                [f.index.tolist()]
            # rename duplicate fields
            df = df.rename(columns=pd.Series(dup_fields + f'_{ds}', index=dup_fields))

            if not len(loaded_fields):
                loaded_fields = df
                continue

            loaded_fields = loaded_fields.join(df, how='outer')

        return loaded_fields

    def get(self, fields: Union[str,List[str]], flexible: bool=None, prop='tabular_field_name') -> pd.DataFrame:
        """
        Return metadata for the specified fields from all tables.

        Args:
            fields (List[str]): Fields to return
            flexible (bool, optional): Whether to use fuzzy matching to find fields. Defaults to None, which uses the DataLoader's flexible_field_search attribute.
            prop (str, optional): The property to use for searching. Defaults to 'tabular_field_name'.

        Returns:
            pd.DataFrame: Data for the specified fields from all tables
        """
        if flexible is None:
            flexible = self.flexible_field_search
        if isinstance(fields, str):
            fields = [fields]
        fields = [f.lower() for f in fields]

        data = pd.DataFrame()
        for dataset, df in self.dicts.items():
            if prop == 'tabular_field_name':
                search_in = pd.Series(df.columns, index=df.columns).str.lower()
            else:
                search_in = df.loc[prop].dropna().str.lower()
            if flexible:
                # use fuzzy matching including regex to find fields
                fields_in_col = np.unique([col for f in fields for col, text in search_in.items()
                                           if type(text) is str and re.search(f, text)])
            else:
                fields_in_col = search_in[search_in.isin(fields)].index
            if len(fields_in_col):
                this_data = df[fields_in_col]
                this_data.columns = dataset + '/' + this_data.columns
                data = self.__concat__(data, this_data)

        return data

    def __repr__(self):
        """
        Return string representation of object

        Returns:
            str: String representation of object
        """
        return self.__str__()

    def __str__(self):
        """
        Return string representation of object

        Returns:
            str: String representation of object
        """
        ds_list = str(list(self.dicts.keys())).replace(',', '\n')
        return f'MetaLoader for: {self.dataset_path}\n' + \
            f'with {len(self.fields)} fields\n{len(self.dicts)} datasets:\n{ds_list}'

    def __getitem__(self, fields: Union[str,List[str]]):
        """
        Return data for the specified fields from all tables

        Args:
            fields (Union[str, List[str]]): Fields to return

        Returns:
            pd.DataFrame: Data for the specified fields from all tables
        """
        return self.get(fields)

    def __concat__(self, df1, df2):
        if df1.empty:
            return df2
        if df2.empty:
            return df1
        return df1.join(df2, how='outer')

    def __load_dictionaries__(self) -> None:
        """
        Load all dictionaries in the base_path.
        """
        dicts = dd.read_csv(os.path.join(self.dataset_path, '*_dict*.csv'),
                            include_path_column=True, dtype={'parent_dataframe': 'object', 'sampling_rate': 'object'}).compute()
        if self.cohort is None:
            dataset_ind = -2
        else:
            dataset_ind = -3
        dicts['dataset'] = dicts['path'].str.split('/').str[dataset_ind]
        dicts = dicts.drop(columns=['path'])
        self.fields = dicts['tabular_field_name'].unique()

        self.dicts = {}
        col_order = ['dataset'] + dicts.columns.drop('dataset').tolist()
        for dataset in dicts['dataset'].unique():
            self.dicts[dataset] = dicts.loc[dicts['dataset'] == dataset, col_order].set_index('tabular_field_name').T

    def __get_dataset_path__(self):
        """
        Get the dataset path.

        Args:
            dataset (str): the name of the dataset

        Returns:
            str: the path to the dataset
        """
        if self.cohort is not None:
            return os.path.join(self.base_path, '*', self.cohort)
        return os.path.join(self.base_path, '*')

