from dacy.score import score
import spacy
import os 
import pandas as pd
import math


def logratio(x:float ,y:float):
    return round(math.log2(x/ y), 2)

def get_means_sds(df):
    '''
    Compute mean F1 and SD for each augmenter of a single model.

    input:
     - df with model performance of a single model. 

    output: 
    -  
    '''
    df["f1"] = df["ents_excl_MISC_ents_f"]*100
    df = df[df["augmenter"].isin(["Female names", "Male names", "Muslim male names", "Muslim female names"])]
    augs, means, sds = [],[],[]
    for aug in set(df["augmenter"]):
        augs.append(aug)
        sub = df[df["augmenter"]==aug]
        means.append(round(sub["f1"].mean(),4))
        sds.append(round(sub["f1"].std(),2))
    # create df
    out_df = pd.DataFrame({"Group": augs, "F1": means, "SD": sds})
    # create df with mean and SD in one column
    combinations = [str(str(a)+" ("+str(b)+")") for (a,b) in zip(means,sds)]
    out_df_condensed = pd.DataFrame({"Group": augs, "F1 (SD)": combinations})

    # Compute log-ratio
    f = out_df[out_df["Group"]=="Female names"].iloc[0,1]
    m = out_df[out_df["Group"]=="Male names"].iloc[0,1]
    ratio = logratio(m,f)  

    return out_df_condensed, ratio


def eval_ner_simple(df, l_ratio, model_name):
    data = {
        f'Simple Output {model_name}': ['', 'Augmented DaNe', ''],
        'Gender Effect Size': ['', l_ratio,''],
        'Name Augmentation': ['Female', '', 'Male'],
        'Macro F1 (SD)': [df[df["Group"]=="Female names"].iloc[0,1], '', 
                     df[df["Group"]=="Male names"].iloc[0,1]],
    }
    df_out = pd.DataFrame(data).T
    df_out.columns = df_out.iloc[0]
    df_out = df_out.iloc[1:]
    return df_out

def eval_ner_detailed(df, l_ratio, model_name):
    data = {
        f'Detailed Output for {model_name}': ['', 'Augmented DaNe', '', ''],
        'Gender Effect Size': ['', l_ratio,'', ''],
        'Name Augmentation': ['Female', '', 'Male', ''],
        'Macro F1 (SD)': [df[df["Group"]=="Female names"].iloc[0,1],'', df[df["Group"]=="Male names"].iloc[0,1],''],
        'Name Augmentation ': ['Majority Female', 'Minority Female', 'Majority Male', 'Minority Male'],
        'Macro F1 (SD) ': [df[df["Group"]=="Female names"].iloc[0,1], 
                      df[df["Group"]=="Muslim female names"].iloc[0,1], 
                      df[df["Group"]=="Male names"].iloc[0,1], 
                      df[df["Group"]=="Muslim male names"].iloc[0,1]]
    }
    out_df = pd.DataFrame(data).T
    out_df.columns = out_df.iloc[0]
    out_df = out_df.iloc[1:]
    return out_df


def eval_model_augmentation(mdl, model_name, n, augmenters, dataset):
    '''
    Return CSV file of model performance on NER task with different name augmentations
    using DACY score function
    input:
    - model_dict : dictionary of models to be run
    - dataset : test dataset for model eval
    - augmenters 
    output: 
    - CSV file in folder "robustness" (creates directory if it does not exist)
    '''

    # loop over all models in model_dict 
    scores = []
    i = 0
    for aug, nam, k in augmenters:
        print(f"\t Running augmenter: {nam} with N iterations: {k}.")
        scores_ = score(corpus=dataset, apply_fn=mdl, augmenters=aug, k=k)
        scores_["model"] = model_name
        scores_["augmenter"] = nam
        scores_["i"] = i
        scores.append(scores_)
        i += 1
    scores_df = pd.concat(scores)

    # summarise scores over augmentations, compute gender effect
    scores_summed, l_ratio = get_means_sds(scores_df)
    print("ratio")
    print(l_ratio)
    # compute simple and detailed output
    out_detailed = eval_ner_detailed(scores_summed, l_ratio, model_name)

    out_simple = eval_ner_simple(scores_summed, l_ratio, model_name)

    return [out_simple, out_detailed]

