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]

#NODE = 'FRAG'
NODE = 'DUMMY'



data = f'/disk/scratch1/s1899827/prosody_nlp/data/input_features/turn_pipeline_{NODE}'
data_os = f'/disk/scratch1/s1899827/prosody_nlp/data/input_features/turn_opensmile'
out = f'/disk/scratch1/s1899827/prosody_nlp/code/self_attn_speech_parser/output/turn_pipeline_no_voicing_redux'

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

parser_training = 'silver'


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
print(f'pred_tree_file: {pred_tree_file}')
pred_score_file = args.predpath+'.scores'


split=args.predpath.split('DUMMY')[-1].strip('.')
if not split=='test':
    split='dev'
print(split)
    


model_prefix = os.path.basename(pred_tree_file).replace(f'.{parser_training}.parsed','')
model_prefix = model_prefix.replace(f'.{NODE}','')
model_prefix = model_prefix.replace('.wo_opensmile','')
model_prefix = model_prefix.replace('.no_voice_parser','')
model_prefix = model_prefix.replace('.test','')

print(f'model prefix: {model_prefix}')

id_file = os.path.join(data,f'{model_prefix}.silver.{split}.ids')
print()
print(id_file)
print(pred_tree_file)
print()
gold_id_file = os.path.join(data_os,f'turn_{split}.ids')
gold_tree_file = os.path.join(data_os,f'turn_{split}.trees')

pred_ids = load_lines(id_file)
pred_trees = load_lines(pred_tree_file)
pred_scores = [float(line.split()[0]) for line in load_lines(pred_score_file)]
modsent_id2score = dict(zip(pred_ids,pred_scores))

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)

scores_file = pred_tree_file+'.scores'

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:
        #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)

gold_id2score = {}
pred2turntree = {}
for turn in gold_ids:
    modsents = turn2modsent[turn]
    if len(modsents)==1:
        pred = f'(TURN {modsent2pred[modsents[0]]})'
        pred2turntree[turn] = pred
        try:
            gold_id2score[turn] = [modsent_id2score[turn2modsent[turn][0]]]
        except:
            import pdb;pdb.set_trace()
    if len(modsents) > 1:
        pred = f'(TURN {" ".join([modsent2pred[sent] for sent in modsents])})'
        pred2turntree[turn] = pred
        gold_id2score[turn] = [modsent_id2score[sent] for sent in modsents]

# TODO at this point, each gold turn should be mapped to a list of len >= 1, which indicates the scores.


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+f'.{parser_training}.turns'
scores_path = pred_tree_file+f'.{parser_training}.turns.scores'
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')

with open(scores_path,'w') as f:
    for turn in gold_ids:
        f.write(' '.join([str(num) for num in gold_id2score[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()
