import numpy as np
from functools import partial

import pandas as pd


def predict_memory_states(my_collection, group):
    states = my_collection.states(*group.name)
    group['stability'] = float(states[0])
    group['difficulty'] = float(states[1])
    group['count'] = len(group)
    return pd.DataFrame({
        'r_history': [group.name[1]],
        't_history': [group.name[0]],
        'stability': [round(float(states[0]), 2)],
        'difficulty': [round(float(states[1]), 2)],
        'count': [len(group)]
        })


def get_my_memory_states(proj_dir, dataset, my_collection):
    prediction = dataset.groupby(by=['t_history', 'r_history']).progress_apply(
        partial(predict_memory_states, my_collection))
    prediction.reset_index(drop=True, inplace=True)
    prediction.sort_values(by=['r_history'], inplace=True)
    prediction.to_csv(proj_dir / "prediction.tsv", sep='\t', index=None)
    print("prediction.tsv saved.")
    prediction['difficulty'] = prediction['difficulty'].map(lambda x: int(round(x)))
    difficulty_distribution = prediction.groupby(by=['difficulty'])['count'].sum() / prediction['count'].sum()
    print(difficulty_distribution)
    difficulty_distribution_padding = np.zeros(10)
    for i in range(10):
        if i + 1 in difficulty_distribution.index:
            difficulty_distribution_padding[i] = difficulty_distribution.loc[i + 1]
    return difficulty_distribution_padding, difficulty_distribution