# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/01_data.ipynb.

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

# %% ../nbs/01_data.ipynb 3
from .utils import rawdata
import pandas as pd
from typing import Optional, List, Callable

class AData:
    def __init__(self, 
                 all: bool = True, # If True, all data is loaded. If False, only a subset is loaded
                 tables: Optional[List] = None # If not None, returns only the dataframes in the list
                 ) -> None:
        "Does stuff"
        self.all = all
        self.tables = tables
        
        self.raw_data = rawdata(all=all, tables=tables)
        self.cleaned_data = {}

        # Map of tables to functions to clean them
        # - order of list is order the cleaning functions is applied
        self.table_func_map = {
            'employee_details': [
                {'func':self.remove_missing_by_col,
                 'kwargs':{'colname': 'emp_id'}
                },
                {'func':self.remove_duplicate_by_col,
                 'kwargs':{'colname': 'emp_id'}
                }
            ],
            'manager_details': [
                {'func':self.remove_missing_by_col,
                 'kwargs':{'colname': 'mgr_id'}
                },
                {'func':lambda x: x.groupby('mgr_id')['exp_as_mgr'].mean().reset_index(),
                 'kwargs':{}
                }
            ],
            'manager_mapping': [
                {'func':self.remove_missing_by_col,
                 'kwargs':{'colname': 'mgr_id'}
                },
                {'func':self.remove_duplicate_by_col,
                 'kwargs':{'colname': 'emp_id'}
                },
            ]
        }

        self.clean_data()
    
    def __repr__(self) -> str:
        return f"AData(all={self.all}, tables={self.tables})"
    
    def __str__(self) -> str:
        return f"AData(all={self.all}, tables={self.tables})"

    def _clean_data(self, 
                    func: Callable, # Function to clean data, should 
                    tables: List[str], # List of tables to clean
                    **kwargs # Additional arguments to pass to the function
                    ) -> None:
        "Cleans data using the function passed"
        for table in tables:
            if table not in self.raw_data.keys():
                raise ValueError(f"Table {table} not found in raw data")
            
            if table not in self.cleaned_data.keys():
                self.cleaned_data[table] = func(self.raw_data[table], **kwargs)
            else:
                self.cleaned_data[table] = func(self.cleaned_data[table], **kwargs)

    def clean_data(self) -> None:
        "Cleans data"
        
        for table_name in self.raw_data.keys():
            # if no cleaning functions are specified for the table, add the raw data to cleaned data
            if table_name not in self.table_func_map.keys():
                self.cleaned_data[table_name] = self.raw_data[table_name]
                continue

            for func in self.table_func_map[table_name]:
                self._clean_data(func=func['func'], tables=[table_name], **func['kwargs'])
    
    @staticmethod
    def remove_missing_by_col(
        data: pd.DataFrame, # Dataframe to remove records from 
        colname: str, # Column name to remove records by
        ):
        "Removes records with missing values in the specified column"
        return data[~data.loc[:, colname].isna()]

    @staticmethod
    def remove_duplicate_by_col(
        data: pd.DataFrame, # Dataframe to remove records from 
        colname: str, # Column name to remove records by
        ):
        "Removes duplicate records in the specified column"
        return data.drop_duplicates(subset=colname)


