import os
import pickle
import numpy as np

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

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

def load_pickle(filename):
    with open(filename,'rb') as f:
        return pickle.load(f)

golds = load_lines(os.path.join(data,'turn_dev.lbl'))
preds = load_lines(os.path.join(out,'turn_sp_seg_72240_dev=99.71.pt_turn_dev_predicted.txt'))
ids = load_lines(os.path.join(data,'turn_dev.ids'))
turn2sent = load_pickle(os.path.join(data,'..','..','turn2sent.pickle'))

### STEP 1 #################################################
# correctly segmented turns -> add sent IDS to list
# incorrectly segmented turns 
#    -> combine sent ids and add to list
#    -> split sent ids and add to list
############################################################

wrong_turn_ids = []
usable_sent_ids = []
for idnum,gold,pred in zip(ids,golds,preds):
    if gold==pred:
        usable_sent_ids.extend(turn2sent[idnum])
    else:
        g_list = np.array([int(i) for i in gold.split()])
        p_list = np.array([int(i) for i in pred.split()])
        if np.sum(g_list) > np.sum(p_list):
            print(f'{idnum}: UNDERSEG')
            sents = turn2sent[idnum]
            idx = 0
            curr_sent = ''
            for g,p in zip(g_list,p_list):
                if g==1 and p==0:
                    curr_sent = curr_sent+sents[idx]+':'
                    idx += 1
                elif g==1 and p==1:
                    usable_sent_ids.append(curr_sent+sents[idx])
                    curr_sent = ''
                    idx += 1
            if curr_sent:
                usable_sent_ids.append(curr_sent+sents[-1])
        elif np.sum(g_list) < np.sum(p_list):
            print(f'{idnum}: OVERSEG')
            idx = 0
            sents = turn2sent[idnum]
            subsent = 0
            unfinished_split = False
            for g,p in zip(g_list,p_list):

                if g==0 and p==1:
                    usable_sent_ids.append(f'{sents[idx]}.{subsent}')
                    subsent += 1
                    unfinished_split = True
                elif g==1 and p==1:
                    if unfinished_split:
                        usable_sent_ids.append(f'{sents[idx]}.{subsent}')
                        subsent += 1
                    usable_sent_ids.append(sents[idx])
                    idx += 1
                    unfinished_split = False
                elif g==0 and p==0:
                    pass
                else:
                    import pdb;pdb.set_trace()
            if unfinished_split:
                usable_sent_ids.append(f'{sents[idx]}.{subsent}')
                import pdb;pdb.set_trace()
        else:
            print(f'{idnum}: wrongseg')


import pdb;pdb.set_trace()

##### STEP 2 ##############################
# Put sent ids and words into new files
###########################################


##### STEP 3 ##############################
# Remake acoustic features
###########################################

