import numpy as np
import pandas as pd
from linearmodels import RandomEffects
import statsmodels.tools.tools as sm
import os
from PYEVALB import scorer



# EXAMPLE:
"""
from linearmodels.datasets import jobtraining
data = jobtraining.load()
year = pd.Categorical(data.year)
data = data.set_index(['fcode', 'year'])
data['year'] = year

exog_vars = ['grant', 'employ']
exog = sm.add_constant(data[exog_vars])
mod = RandomEffects(data.clscrap, exog)
re_res = mod.fit()
print(re_res)
"""
###################

def load_lines(filename):
    return [l.strip() for l in open(filename,'r').readlines()]




scr = scorer.Scorer(unlbl=False)

output_dir = '/disk/scratch1/s1899827/prosody_nlp/code/self_attn_speech_parser/output/turn_no_voicing'
data_dir = '/disk/scratch1/s1899827/prosody_nlp/data/input_features/turn_opensmile'

pros_preds = ['sp_glove_dp.3_26229_dev=86.34.pred',
              'sp_glove_dp.3_38976_dev=86.24.pred',
              'sp_glove_dp.3_5154_dev=85.99.pred',
              'sp_glove_dp.3_72240_dev=86.48.pred',
              'sp_glove_dp.3_73784_dev=85.98.pred'
]
nonpros_preds = ['nonsp_parse_dp.3_26229_dev=85.94.pred',
                 'nonsp_parse_dp.3_38976_dev=86.11.pred',
                 'nonsp_parse_dp.3_5154_dev=85.76.pred',
                 'nonsp_parse_dp.3_72240_dev=85.33.pred',
                 'nonsp_parse_dp.3_73784_dev=85.36.pred'
]
gold_file = 'turn_dev.trees'
id_file = 'turn_dev.ids'

gold_f = os.path.join(data_dir, gold_file)
id_f = os.path.join(data_dir, id_file)

gold_lines = load_lines(gold_f)
ids = load_lines(id_f)

def get_results(pred_lines,gold_lines,ids,pros=1):

    result_ids = []
    result_mat = []
    result_gol = []
    result_tes = []
    result_parse_f1 = []
    result_seg_f1 = []
    result_pros = []

    for gline,pline,idnum in zip(gold_lines,pred_lines,ids):
        mat,gol,tes = scr.get_bracket_counts_from_tree(gline,pline)
        pp = mat/tes
        rr = mat/gol
        if mat==0:
            ff = 0
        else:
            ff = (2*pp*rr)/(pp+rr)
        result_pros.append(pros)
        result_mat.append(mat)
        result_gol.append(gol)
        result_tes.append(tes)
        result_parse_f1.append(ff)
        result_ids.append(idnum)
    return result_pros,result_parse_f1,result_ids,result_mat,result_gol,result_tes


pros = []
parse_f1 = []
idnums = []
mat = []
gol = []
tes = []
years = []

base_year = 1900

for pred_file in pros_preds:
    pred_f = os.path.join(output_dir, pred_file)
    print(pred_f)
    pred_lines = load_lines(pred_f)
    result_pros,result_parse_f1,result_ids,result_mat,result_gol,result_tes = get_results(pred_lines,gold_lines,ids,pros=1)
    pros.extend(result_pros)
    parse_f1.extend(result_parse_f1)
    idnums.extend(result_ids)
    mat.extend(result_mat)
    gol.extend(result_gol)
    tes.extend(result_tes)
    years.extend([base_year+i for i in range(len(result_mat))])
for pred_file in nonpros_preds:
    pred_f = os.path.join(output_dir, pred_file)
    print(pred_f)
    pred_lines = load_lines(pred_f)
    result_pros,result_parse_f1,result_ids,result_mat,result_gol,result_tes = get_results(pred_lines,gold_lines,ids,pros=0)
    pros.extend(result_pros)
    parse_f1.extend(result_parse_f1)
    idnums.extend(result_ids)
    mat.extend(result_mat)
    gol.extend(result_gol)
    tes.extend(result_tes)
    years.extend([base_year+i for i in range(len(result_mat))])
    
df = pd.DataFrame(list(zip(pros,parse_f1,idnums,mat,gol,tes,years)),
                  columns =['pros','parse_f1','id','mat','gol','tes','year'])


print(len(df))
#import pdb;pdb.set_trace()
year = pd.Categorical(df.year)
df = df.set_index(['id', 'year'])
df['year'] = year

exog_vars = ['pros']
exog = sm.add_constant(df[exog_vars])
print(df)
#import pdb;pdb.set_trace()
mod = RandomEffects(df.parse_f1, exog)
#mod = RandomEffects(df.mat, exog)


re_res = mod.fit()
print(re_res) 


import pdb;pdb.set_trace()


"""
from linearmodels.datasets import jobtraining
data = jobtraining.load()
year = pd.Categorical(data.year)
data = data.set_index(['fcode', 'year'])
data['year'] = year

exog_vars = ['grant', 'employ']
exog = sm.add_constant(data[exog_vars])
mod = RandomEffects(data.clscrap, exog)
re_res = mod.fit()
print(re_res)
"""


