import os
import numpy as np
import trees

def load_seg_file(filename,id_file=None,final_zero=True):
    if not id_file:
        lines = [l.strip() for l in open(filename,'r').readlines()]
        lines = [np.array([int(i) for i in l.strip().split()]) for l in lines]
        if final_zero:
            for line in lines:
                line[-1] = 0
        return lines
    
    else:
        treebank, sent_ids = trees.load_trees_with_idx(filename,id_file,strip_top=True)
        bin_lines = []
        for tree in treebank:
            num_toks = len(list(tree.leaves()))
            lbls = [0]*num_toks
            sus = list(tree.children) # SUs are the immediate children of the top level.
            end_idx = 0
            for su in sus:
                num_su_toks = len(list(su.leaves()))
                end_idx += num_su_toks # point at the end of each SU
                if end_idx != num_toks:
                    lbls[end_idx-1] = 1
            bin_lines.append(np.array(lbls))
        return bin_lines
            
def calc_f1(preds,golds):
    true_pos = 0
    all_pos = 0
    all_pred = 0
    for pred,gold in zip(preds,golds):
        all_pos += gold.sum()
        all_pred += pred.sum()

        match_mask = pred==gold
        pos_mask = pred==1
        true_pos += (match_mask&pos_mask).sum()
    #import pdb;pdb.set_trace()

    prec = true_pos/all_pred
    rec = true_pos/all_pos

    print(f'prec: {prec}')
    print(f'rec: {rec}')
    print(f'num pred: {all_pred}')
    f = (2*prec*rec)/(prec+rec)
    return f

def avg_sent_len(preds):
    sent_len = []
    for pred in preds:
        curr_len = 0
        for idx,tok in enumerate(pred):
            if tok == 0 and idx < len(pred)-1:
                curr_len += 1
            else:
                curr_len += 1
                sent_len.append(curr_len)
                curr_len = 0
    return sum(sent_len)/len(sent_len)
        
if __name__=="__main__":
    #data = '/disk/scratch1/s1899827/prosody_nlp/data/input_features/turn_pipeline'
    data = '/disk/scratch1/s1899827/prosody_nlp/data/input_features/turn_pipeline'
    #out = '/disk/scratch1/s1899827/prosody_nlp/data/input_features/turn_opensmile'
    #out =  '/afs/inf.ed.ac.uk/group/project/prosody/prosody_nlp/code/self_attn_speech_parser/output/turn_pipeline'
    out =  '/disk/scratch1/s1899827/prosody_nlp/code/self_attn_speech_parser/output/turn_pipeline_voicing'


    golds = os.path.join(data,'dev.lbl.final_one')
    print(golds)
    golds = load_seg_file(golds)

    pred_list = [#'sp_seg_dev=89.84.pt.dev.pred'
        'nonsp_seg_dev=87.60.pt.dev.pred'
                 #'turn_nonsp_72240_dummy_penalty.5_dev=85.33.pt_dev.pred.lbl'
                 #'sp_predlbl_dev=90.26.pred.reparsed'
                 #'nonsp_seg_dev=87.60.pt.dev.pred',
                 #'sp_seg_dur_only_dev=88.93.pt.dev.pred',
                 #'sp_seg_no_opensmile_dev=89.59.pt.dev.pred',
                 #'sp_seg_pitch_only_dev=89.44.pt.dev.pred',
                 #'sp_seg_dev=89.84.pt.dev.pred',
                 #'sp_seg_fbank_only_dev=89.36.pt.dev.pred',
                 #'sp_seg_pause_only_dev=89.55.pt.dev.pred'
    ]
    

    
    for preds in pred_list:
        print('-'*25)
        print(preds)
        preds = os.path.join(out,preds)
        preds = load_seg_file(preds)
        print(f'num preds: {len(preds)}')
        print(f'num golds: {len(golds)}')
        assert len(preds) == len(golds)
        for pred,gold in zip(preds,golds):
            assert len(pred) == len(gold)
        avg_len = avg_sent_len(preds)
        f = calc_f1(preds,golds)

        print(f)
        print(f'average sent len: {avg_len}')

