import pandas as pd
import numpy as np

## Production Function

def get_production_func(df,task):
    """Obtain production function on second level

    Computes for every second the expected points each subject obtains.

    Parameters
    ------------

    df : pd.DataFrame
        Information in standard (subject-round)-format.
    task : str
        `math` : math task.
        `mem` : memory.

    Returns
    ------------
    df_res : pd.DataFrame
        Expected value in with seconds as index.

    """
    df_res = pd.Series(dtype=float)
    if task == 'math':
        max_num = 8
    elif task == 'mem':
        max_num = 10
    for sec in range(0,61):
        tab = (df<=sec).mean()
        for num in reversed(range(1,max_num)):
            weight = np.array(tab[f'{task}_t_{num}']-tab[f'{task}_t_{num+1}'])
            if (num+1) == max_num:
                # Treat the last round extra
                exp_val = ((num+1) * np.array(tab[f'{task}_t_{num+1}']))
            exp_val += (num * weight)
        df_res.loc[sec] =  exp_val
    #df_res.index = tab.index
    return df_res

def get_optimal_allocation(df,p_vars):
    """
    Find the optimal allocation of time between 'mem' and 'math' activities.

    Parameters:
    df (pd.DataFrame): DataFrame containing scores for 'mem' and 'math' activities

    Returns:
    tuple: A tuple of two elements:
        rslt (pd.DataFrame): DataFrame with columns 'switch_point', 'math_score', 'mem_score', and 'total_score'
        detail_dict (dict): A dictionary with individual level results for each subject

    The function iterates over all seconds, calculates scores for 'mem' and 'math' activities,
    and finds the optimal allocation that results in the highest total score.
    The results are returned as a DataFrame with columns 'switch_point', 'math_score', 'mem_score', and 'total_score' and as a dictionary with individual level results.
    """
    allotted_time=60
    xs = np.ones((3,(allotted_time+1)))*np.nan
    for sec in range(0,allotted_time+1):
        mem = df.loc[allotted_time-sec]['mem']
        math = df.loc[sec]['math']
        xs[0,sec] = mem
        xs[1,sec] = math
        xs[2,sec] = mem+math
    # Optimal allocation
    alloc = np.argmax(xs[2,:])
    p_vars['switch_point'] = int(alloc)
    p_vars['opt_math_score'] = round(float(xs[1,alloc]),3)
    p_vars['opt_mem_score'] = round(float(xs[0,alloc]),3)
    p_vars['opt_total_score'] = round(float(xs[2,alloc]),3)

    def comp_opt_alloc(self):
        p_vars = self.participant.vars
        # Simplifies code later
        math_cols = ['math_t_1', 'math_t_2', 'math_t_3', 'math_t_4', 'math_t_5', 'math_t_6', 'math_t_7', 'math_t_8']
        mem_cols = ['mem_t_1' ,'mem_t_2', 'mem_t_3', 'mem_t_4', 'mem_t_5', 'mem_t_6', 'mem_t_7', 'mem_t_8', 'mem_t_9', 'mem_t_10']
        # Create df of only mem data
        memdict={col:p_vars[col] for col in mem_cols}
        df_mem = pd.DataFrame(memdict)
        # Create df of only math data
        mathdict={col:p_vars[col] for col in math_cols}
        df_math = pd.DataFrame(mathdict)
        # Create respective production function
        rslt_mem = get_production_func(df_mem,'mem')
        rslt_math = get_production_func(df_math,'math')
        # Find optimal allocation
        df_both = pd.DataFrame({'mem': rslt_mem, 'math': rslt_math})
        get_optimal_allocation(df_both,p_vars)

p_vars = {'math_t_1': [2.847,2.249,2.143,1.985,2.306,2.512,2.396,2.346,1.861,2.031],
 'mem_t_1': [6.164,21.833,6.722,4.212,20.728,5.739,21.099,21.069,8.834,7.728],
 'math_t_2': [5.919,4.097,4.351,3.593,4.258,4.664,3.988,4.586,3.884,4.191],
 'mem_t_2': [13.955,27.291,15.569,8.987,24.423,8.166,28.79,22.869,10.834,17.64],
 'math_t_3': [8.855,6.979,8.526,7.137,6.713,8.119,8.052,7.362,8.172,6.719],
 'mem_t_3': [16.963,34.89,20.169,20.582,30.439,14.493,32.446,29.319,12.658,19.312],
 'math_t_4': [11.383,9.843,12.625,10.92,9.233,11.135,11.971,10.402,10.7,11.039],
 'mem_t_4': [32.645,48.714,22.601,22.438,35.095,18.973,38.797,33.199,14.73,23.392],
 'math_t_5': [17.846,15.643,17.553,16.32,16.473,14.695,17.115,15.778,17.244,16.086],
 'mem_t_5': [37.445,56.674,26.857,25.301,42.998,20.861,43.749,35.111,25.081,31.239],
 'math_t_6': [28.873,21.658,22.857,23.424,19.833,19.935,22.715,20.329,21.092,23.758],
 'mem_t_6': [41.741,67.185,33.064,33.661,46.494,27.164,47.637,41.607,38.436,36.921],
 'math_t_7': [52.111,35.466,32.032,39.298,30.523,30.782,38.938,33.232,32.451,35.133],
 'mem_t_7': [49.877,69.073,37.28,43.837,48.942,30.748,51.549,47.759,41.804,42.233],
 'math_t_8': [75.686,47.249,39.359,55.129,38.786,43.637,51.105,65.673,46.25,44.237],
 'mem_t_8': [54.244,70.632,41.992,54.38,51.422,32.852,53.773,49.494,44.627,47.193],
 'mem_t_9': [56.068,72.08,46.008,58.332,53.494,34.892,55.989,50.862,47.004,49.169],
 'mem_t_10': [57.188,73.129,47.592,59.556,54.966,36.1,57.437,52.198,48.195,50.513]}