import os
import pickle
from PYEVALB import scorer
import tempfile
import string
import argparse

def tree2str(tree):
    return ' '.join(''.join(ch for ch in tree if not ch.isupper()).replace('(','').replace(')','').translate(str.maketrans('', '', string.punctuation)).strip().split())

def load_lines(filename):
    with open(filename,'r') as f:
        lines = f.readlines()
    return [l.strip() for l in lines]



data = '/afs/inf.ed.ac.uk/group/project/prosody/prosody_nlp/data/input_features/turn_pipeline'
out = '/afs/inf.ed.ac.uk/group/project/prosody/prosody_nlp/code/self_attn_speech_parser/output/turn_pipeline'

sent2turn = pickle.load(open(os.path.join(data,'..','sent2turn.pickle'),'rb'))
turn2sent = pickle.load(open(os.path.join(data,'..','turn2sent.pickle'),'rb'))



parser = argparse.ArgumentParser()
parser.add_argument("--predpath",type=str,default="sp_seg_dev=87.61.pt.dev.pred")
args = parser.parse_args()
pred_tree_file = args.predpath

model_prefix = os.path.basename(pred_tree_file).replace('.silver.parsed','')

#id_file = os.path.join(data,'dev.silver.ids')
id_file = os.path.join(data,f'{model_prefix}.silver.dev.ids')

#pred_tree_file = os.path.join(out,'turn_nonsp_seg_72240_focal_loss_no_alpha_gamma3_final_one_dev=87.61.pt.silver.parsed')
#gold_id_file = os.path.join(data,'..','turn_opensmile','dev.ids')
#gold_tree_file = os.path.join(data,'..','turn_opensmile','dev.trees')
gold_id_file = os.path.join(data,'dev.ids')
gold_tree_file = os.path.join(data,'dev.trees')

pred_ids = load_lines(id_file)
pred_trees = load_lines(pred_tree_file)
gold_ids = load_lines(gold_id_file)
gold_trees = load_lines(gold_tree_file)

print(f'ids {len(pred_ids)}')
print(f'trees {len(pred_trees)}')
assert len(pred_ids)==len(pred_trees)


modsent2pred = dict(zip(pred_ids,pred_trees))
turn2gold = dict(zip(gold_ids,gold_trees))

modsent2turn = {}
turn2modsent = {}
for modsent in pred_ids:
    if modsent in sent2turn:
        turn = sent2turn[modsent]
        modsent2turn[modsent] = turn
        if turn in turn2modsent:
            turn2modsent[turn].append(modsent)
        else:
            turn2modsent[turn] = [modsent]
            """
    elif ':' in modsent:
        print(modsent)

        sent_id = modsent.split(':')[0]
        turn = sent2turn[sent_id]
        modsent2turn[modsent] = turn
        if turn in turn2modsent:
            turn2modsent[turn].append(modsent)
        else:
            turn2modsent[turn] = [modsent]
            """
    elif '+' in modsent:
        #sent_id = modsent.split('+')[0]
        turn = modsent.split('+')[0] #sent2turn[sent_id]
        modsent2turn[modsent] = turn
        if turn in turn2modsent:
            turn2modsent[turn].append(modsent)
        else:
            turn2modsent[turn] = [modsent] 

for turn in turn2modsent:
    sents = turn2modsent[turn]
    if '+' in sents[0]:
        
        turn2modsent[turn] = sorted(sents, key=lambda sent: int(sent.split('+')[1]))

    else:
        turn2modsent[turn] = sorted(sents)

    
pred2turntree = {}
for turn in gold_ids:
    modsents = turn2modsent[turn]
    if len(modsents)==1:
        pred = f'(TURN {modsent2pred[modsents[0]]})'
        pred2turntree[turn] = pred
    if len(modsents) > 1:
        pred = f'(TURN {" ".join([modsent2pred[sent] for sent in modsents])})'
        pred2turntree[turn] = pred


import string
for turn in pred2turntree:

    pred = tree2str(pred2turntree[turn])
    gold = tree2str(turn2gold[turn])
    if not pred==gold:
        print(turn)
        print(f'pred: {pred}')
        print('---')
        print(f'gold: {gold}')
        print('ERROR: non-matching sentences')
        import pdb;pdb.set_trace()
                
 

temp_dir = tempfile.TemporaryDirectory(prefix="evalb-")
predicted_path = pred_tree_file+'.turns'
output_path = os.path.join(temp_dir.name,'out.txt')

with open(predicted_path,'w') as f:
    for turn in gold_ids:
        f.write(pred2turntree[turn])
        f.write('\n')

scr = scorer.Scorer()
scr.evalb(gold_tree_file,predicted_path,output_path)
print(f'scoring done; find output at {output_path}')
import pdb;pdb.set_trace()
temp_dir.cleanup()
