import os
import pandas as pd
import trees
import numpy as np
from math import isnan
from PYEVALB import scorer,parser
from nltk.tree import Tree
from evaluate import seg_fscore_lists,evalb

SCR = scorer.Scorer()
ISCR = scorer.InteractiveScorer()

def load_lines(filename,split=True):
    lines = [line.strip() for line in open(filename,'r').readlines()]
    if split:
        return [[int(x) for x in line.split()] for line in lines]
    else:
        return lines

def seg_scores(gold,pred):
    gold[-1] = 0
    pred[-1] = 0
    gold = np.array(gold)
    pred = np.array(pred)
    pred_pos = np.sum(pred)
    match = gold==pred
    pred1 = pred==1
    true_pos = np.sum(match&pred1)
    gold_pos = np.sum(gold)
    if pred_pos > 0:
        prec = 100*(true_pos/pred_pos)
    else:
        prec = float('nan')
    if gold_pos > 0:
        rec = 100*(true_pos/gold_pos)
    else:
        rec = float('nan')
    if gold_pos == pred_pos == 0:
        f1 = 100
    elif isnan(prec) or isnan(rec) or rec+prec==0:
        f1 = float('nan')
    else:
        f1 = (2*prec*rec)/(prec+rec)
    #return prec,rec,f1

    return f1

def parse_scores(gold,pred):
    gold = parser.create_from_bracket_string(gold)
    pred = parser.create_from_bracket_string(pred)
    result = SCR.score_trees(gold,pred)
    prec = result.prec
    rec = result.recall
    if rec == prec == 0:
        f1 = 0
    else:
        f1 = 100*((2*prec*rec)/(prec+rec))
    return f1


def count_nodes(tree):
    tree = parser.create_from_bracket_string(tree)
    nts = [nt for nt in tree.non_terminals if nt.value!='TURN']
    return len(nts)


def count_missing(gold,tree):
    gold = parser.create_from_bracket_string(gold)
    tree = parser.create_from_bracket_string(tree)
    res,missing,extra = ISCR.score_trees(gold,tree)
    return len(missing)

def count_extra(gold,tree):
    gold = parser.create_from_bracket_string(gold)
    tree = parser.create_from_bracket_string(tree)
    res,missing,extra = ISCR.score_trees(gold,tree)
    return len(extra)



def examine_parses(df):
    df = df.sort_values(by=['pipe_parse_f1'])
    min_len = min([len(parser.create_from_bracket_string(gg).sentence) for gg in df.gold])
    print(min_len)
    for i,row in df.iterrows():
        #for ii,gg,pp,ee in zip(ids,golds,pipes,e2es):
        gg = parser.create_from_bracket_string(row.gold)
        pp = parser.create_from_bracket_string(row.pipeline)
        ee = parser.create_from_bracket_string(row.e2e)
        res_e2e,missing_e2e,extra_e2e = ISCR.score_trees(gg,ee)
        res_pipe,missing_pipe,extra_pipe = ISCR.score_trees(gg,pp)
        print('='*50)
        print(row.ids)
        print('='*50)
        print('GOLD')
        print(gg)
        print(row.gold_lbls)
        print('PIPELINE')
        print(pp)
        print(row.pipe_lbls)
        print('pipeline seg f1:',row.pipe_seg_f1)
        print('pipeline parse f1:',row.pipe_parse_f1)
        print('E2E')
        print(ee)
        print(row.e2e_lbls)
        print('e2e seg f1:',row.e2e_seg_f1)
        print('e2e parse f1:',row.e2e_parse_f1)
        print('GOLD NODES')
        print(gg.non_terminal_labels)
        print('='*50)
        print('MISSING')
        print(f'pipeline: {missing_pipe}')
        print(f'e2e: {missing_e2e}')
        print('='*50)
        print('EXTRA')
        print(f'pipeline: {extra_pipe}')
        print(f'e2e: {extra_e2e}')
        import pdb;pdb.set_trace()
        
def main():

    data = '/disk/scratch1/s1899827/prosody_nlp/data/input_features/turn_opensmile'
    pipeout = '/disk/scratch1/s1899827/prosody_nlp/code/self_attn_speech_parser/output/turn_pipeline_no_voicing_redux'
    e2eout = '/disk/scratch1/s1899827/prosody_nlp/code/self_attn_speech_parser/output/turn_no_voicing'

    gold_tree_f = os.path.join(data,'turn_dev.trees')
    gold_lbl_f = os.path.join(data,'turn_dev.lbl')
    gold_id_f = os.path.join(data,'turn_dev.ids')

    pipe_tree_f = os.path.join(pipeout,'sp_seg_72240_dev=89.59.pt.silver.parsed.DUMMY.silver.turns.defrag')
    pipe_lbl_f = os.path.join(pipeout,'sp_seg_72240_dev=89.59.pt.dev.pred')
    pipe_id_f = os.path.join(data,'turn_dev.ids')

    e2e_tree_f = os.path.join(e2eout,'sp_glove_dp.3_72240_dev=86.48.pred')
    e2e_lbl_f = os.path.join(e2eout,'sp_glove_dp.3_72240_dev=86.48.pred.lbl')
    e2e_id_f = os.path.join(data,'turn_dev.ids')

    gold_trees,ids = trees.load_trees_with_idx(gold_tree_f,gold_id_f,strip_top=True)
    pipe_trees,_ = trees.load_trees_with_idx(pipe_tree_f,pipe_id_f,strip_top=True)
    e2e_trees,_ = trees.load_trees_with_idx(e2e_tree_f,e2e_id_f,strip_top=True)

    gold_trees_txt = [tree.linearize() for tree in gold_trees]
    pipe_trees_txt = [tree.linearize() for tree in pipe_trees]
    e2e_trees_txt = [tree.linearize() for tree in e2e_trees]
    
    gold_lbls = load_lines(gold_lbl_f)
    pipe_lbls = load_lines(pipe_lbl_f)
    e2e_lbls = load_lines(e2e_lbl_f)

    
    df = pd.DataFrame(list(zip(ids,gold_trees_txt,pipe_trees_txt,e2e_trees_txt,gold_lbls,pipe_lbls,e2e_lbls,gold_trees,pipe_trees,e2e_trees)),
                      columns=['ids','gold','pipeline','e2e','gold_lbls','pipe_lbls','e2e_lbls','gold_trees','pipe_trees','e2e_trees'])

    df['pipe_seg_f1'] = df.apply(lambda x: seg_scores(x.gold_lbls,x.pipe_lbls),axis=1)
    df['e2e_seg_f1'] = df.apply(lambda x: seg_scores(x.gold_lbls,x.e2e_lbls),axis=1)

    df['pipe_parse_f1'] = df.apply(lambda x: parse_scores(x.gold,x.pipeline),axis=1)
    df['e2e_parse_f1'] = df.apply(lambda x: parse_scores(x.gold,x.e2e),axis=1)    

    df['pipe_node_count'] = df.apply(lambda x: count_nodes(x.pipeline),axis=1)
    df['e2e_node_count'] = df.apply(lambda x: count_nodes(x.e2e),axis=1)

    df['pipe_missing'] = df.apply(lambda x: count_missing(x.gold,x.pipeline),axis=1)
    df['e2e_missing'] = df.apply(lambda x: count_missing(x.gold,x.e2e),axis=1)
     
    df['pipe_extra'] = df.apply(lambda x: count_extra(x.gold,x.pipeline),axis=1)
    df['e2e_extra'] = df.apply(lambda x: count_extra(x.gold,x.e2e),axis=1)

    # I want to see places where e2e outdoes pipelne on parse, but does worse on seg
    # Trying just low parse scores from pipeline
    #mask = (df['e2e_parse_f1']>df['pipe_parse_f1'])&(df['e2e_seg_f1']<df['pipe_seg_f1'])
    mask = (df['e2e_parse_f1']>df['pipe_parse_f1'])&(df.apply(lambda x: sum(x.gold_lbls)>0,axis=1))&(df['e2e_seg_f1']<df['pipe_seg_f1'])&(df.apply(lambda x: sum(x.e2e_lbls)>sum(x.pipe_lbls),axis=1))
    masked = df[mask]
    others = df[~mask]
    
    """
    overall_pipe_seg_f1= seg_fscore_lists(df['gold_lbls'].tolist(),df['pipe_lbls'].tolist()).fscore
    overall_e2e_seg_f1= seg_fscore_lists(df['gold_lbls'].tolist(),df['e2e_lbls'].tolist()).fscore
    masked_pipe_seg_f1 = seg_fscore_lists(masked['gold_lbls'].tolist(),masked['pipe_lbls'].tolist()).fscore
    masked_e2e_seg_f1 = seg_fscore_lists(masked['gold_lbls'].tolist(),masked['e2e_lbls'].tolist()).fscore
    others_pipe_seg_f1 = seg_fscore_lists(others['gold_lbls'].tolist(),others['pipe_lbls'].tolist()).fscore
    others_e2e_seg_f1 = seg_fscore_lists(others['gold_lbls'].tolist(),others['e2e_lbls'].tolist()).fscore
    
    overall_pipe_parse_f1= evalb('/tmp',df['gold_trees'].tolist(),df['pipe_trees'].tolist()).fscore
    overall_e2e_parse_f1= evalb('/tmp',df['gold_trees'].tolist(),df['e2e_trees'].tolist()).fscore
    masked_pipe_parse_f1 = evalb('/tmp',masked['gold_trees'].tolist(),masked['pipe_trees'].tolist()).fscore
    masked_e2e_parse_f1 = evalb('/tmp',masked['gold_trees'].tolist(),masked['e2e_trees'].tolist()).fscore
    others_pipe_parse_f1 = evalb('/tmp',others['gold_trees'].tolist(),others['pipe_trees'].tolist()).fscore
    others_e2e_parse_f1 = evalb('/tmp',others['gold_trees'].tolist(),others['e2e_trees'].tolist()).fscore
    import pdb;pdb.set_trace()
    """
    # Examine parses:
    examine_parses(masked)

    print('done.')
if __name__=="__main__":
    main()
    
