# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/07_basic_analysis.ipynb.

# %% auto 0
__all__ = ['custom_describe', 'assign_nearest_research_stage']

# %% ../nbs/07_basic_analysis.ipynb 3
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from typing import List, Any, Dict, Union, Optional
import importlib

# %% ../nbs/07_basic_analysis.ipynb 4
try:
    importlib.import_module("smart_open")
except:
    print("smart_open not installed, please pip install smart-open")

from smart_open import open

# %% ../nbs/07_basic_analysis.ipynb 5
from .config import generate_synthetic_data

# %% ../nbs/07_basic_analysis.ipynb 6
def custom_describe(df: pd.DataFrame) -> pd.DataFrame:
    """
    Generates a custom summary statistics dataframe for mixed data types.
    
    Args:
        df: The input pandas DataFrame
        
    Returns:
        A pandas DataFrame containing the summary statistics
    """
    def describe_column(series: pd.Series) -> Dict[str, Union[int, float]]:
        """
        Generates summary statistics for a given column/series.
        
        Args:
            series: The input pandas Series (column)
            
        Returns:
            A dictionary containing the summary statistics
        """
        if pd.api.types.is_numeric_dtype(series):
            stats = {
                'count': series.count(),
                'unique': series.nunique(),
                'most_frequent': series.mode().iat[0] if series.mode().size > 0 else np.nan,
                'min': series.min(),
                'max': series.max(),
                'mean': series.mean(),
                'median': series.median(),
                'std': series.std(),
            }
        elif pd.api.types.is_datetime64_dtype(series):
            stats = {
                'count': series.count(),
                'unique': series.nunique(),
                'most_frequent': np.nan,
                'min': series.min(),
                'max': series.max(),
                'mean': np.nan,
                'median': np.nan,
                'std': np.nan,
            }
        else:
            stats = {
                'count': series.count(),
                'unique': series.nunique(),
                'most_frequent': series.mode().iat[0] if series.mode().size > 0 else np.nan,
                'min': np.nan,
                'max': np.nan,
                'mean': np.nan,
                'median': np.nan,
                'std': np.nan,
            }
        return stats

    summary = pd.DataFrame({col: describe_column(df[col]) for col in df.columns}).transpose()
    summary = summary[['count', 'unique', 'most_frequent', 'min', 'max', 'mean', 'median', 'std']]
    return summary.T

# %% ../nbs/07_basic_analysis.ipynb 8
def assign_nearest_research_stage(dataset: pd.DataFrame, 
                                  population: pd.DataFrame, 
                                  max_days: int = 60, 
                                  stages: List[str] = ['visit'], 
                                  agg: Union[str, None] = 'first') -> pd.DataFrame:
    """
    Assign the nearest research stage to each record in a dataset.
    
    Args:
        dataset (pd.DataFrame): The dataset containing records to be assigned research stages.
        population (pd.DataFrame): The population data with participant_id, cohort, research_stage, and research_stage_date.
        max_days (int, optional): The maximum number of days allowed between the collection date and research stage date. Defaults to 60.
        stages (List[str], optional): The list of types of research stages to consider. Defaults to ['visit'].
        agg (Union[str, None], optional): The aggregation function to be used when (optionally) aggregating multiple rows
                                          from the same research stage. The rows are already sorted by distance from the
                                          date of the research stage. Can be 'first' (closest), 'last' (farthest), 'mean',
                                          'min', 'max', or None. Defaults to 'first'.

    Returns:
        pd.DataFrame: The dataset with the nearest research stage assigned to each record.
    """
    data_wo_stage = dataset.sort_values('collection_date').reset_index()\
        .drop(columns=['research_stage'], errors='ignore')
    population = population.loc[
        population.index.get_level_values('research_stage').str.contains('|'.join(stages))]\
        .sort_values('research_stage_date').reset_index()\
        [['participant_id', 'cohort', 'research_stage', 'research_stage_date']]

    data_w_stage = pd.merge_asof(
        data_wo_stage,
        population,
        by=['participant_id', 'cohort'],
        left_on='collection_date',
        right_on='research_stage_date',
        direction='nearest')\
        .assign(delta=lambda x: (x['collection_date'] - x['research_stage_date']).dt.days.abs())
    data_w_stage = data_w_stage.loc[data_w_stage['delta'] < max_days].sort_values('delta')\
        .drop(columns=['delta', 'research_stage_date'])

    if agg is None:
        # return data_wo_stage, data_w_stage
        return data_wo_stage.merge(
            data_w_stage[['participant_id', 'cohort', 'collection_date', 'research_stage']].drop_duplicates(),
            how='left')\
            .set_index(dataset.index.names)

    return data_w_stage.groupby(['participant_id', 'cohort', 'research_stage']).agg(agg)\
        .drop(columns=['array_index'], errors='ignore')

