from PYEVALB import scorer,parser
import itertools
import random
import os
import pickle
import time
import string
import tempfile
from trees import load_trees
from evaluate import seg_fscore_lists,tree2num,seg_counts_for_numstr_example

num_resamples = 10**5

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 get_p_value(baseline_f,gold_f,experiment_f,ids,mode='parse'):



    
    if mode=='parse':
        print(experiment_f.split('/')[-1])

        gold_lines = [l.strip() for l in open(gold_f).readlines()]
        experiment_lines = [l.strip() for l in open(experiment_f).readlines()]
        baseline_lines = [l.strip() for l in open(baseline_f).readlines()]
        num_lines = len(gold_lines)
        assert len(gold_lines)==len(ids)
        assert len(gold_lines)==len(experiment_lines)
        assert len(gold_lines)==len(baseline_lines)

        scr = scorer.Scorer()
        print('Calculate baseline...')
        fullset_experiment_f1 = scr.get_f1_from_list(gold_lines,experiment_lines)
        fullset_baseline_f1 = scr.get_f1_from_list(gold_lines,baseline_lines)
    elif mode=='seg':
        gold_lines = [tree2num(gtree,output_format='np') for gtree in load_trees(gold_f)]
        experiment_lines = [tree2num(gtree,output_format='np') for gtree in load_trees(experiment_f)]
        baseline_lines = [tree2num(gtree,output_format='np') for gtree in load_trees(baseline_f)]
        print('Calculate baseline...')
        fullset_experiment_f1 = seg_fscore_lists(gold_lines,experiment_lines,input_format='np').fscore
        fullset_baseline_f1 = seg_fscore_lists(gold_lines,baseline_lines,input_format='np').fscore
        print(f'baseline f1: {fullset_baseline_f1}')
        print(f'experiment f1: {fullset_experiment_f1}')
    
    fullset_delta = abs(fullset_experiment_f1 - fullset_baseline_f1)

        
    big_diffs = 0

    id2matched = {'experiment':{},
                  'baseline':{}}
    id2gold = {'experiment':{},
               'baseline':{}}
    id2test = {'experiment':{},
               'baseline':{}}

    model2output = {'experiment':experiment_lines,
                    'baseline':baseline_lines}


    print('Store bracket scores ...')
    for model in ('experiment','baseline'):
        test_lines = model2output[model]
        for i,turn in enumerate(ids):
            gold_line = gold_lines[i]
            test_line = test_lines[i]
            if mode=='parse':
                assert tree2str(gold_line)==tree2str(test_line)
                mat,gol,tes = scr.get_bracket_counts_from_tree(gold_line,test_line)
            elif mode=='seg':
                mat,gol,tes = seg_counts_for_numstr_example(gold_line,test_line)
            id2matched[model][turn] = mat
            id2gold[model][turn] = gol
            id2test[model][turn] = tes
        
        
    idxs = [i for i in range(len(gold_lines))]
    print('Resample ...')
    for i in range(num_resamples):
        if i%100000==1:
            print(i)
        gold_resamples = []
        experiment_resamples = []
        baseline_resamples = []
        resample_idx = random.choices(idxs,k=len(idxs))
        resampled_turns = []
        for idx in resample_idx:
            resampled_turns.append(ids[idx])
            
        gold = sum([id2gold['experiment'][turn] for turn in resampled_turns])
        experiment_matched = sum([id2matched['experiment'][turn] for turn in resampled_turns])
        experiment_test = sum([id2test['experiment'][turn] for turn in resampled_turns])
        experiment_rec = experiment_matched/gold
        experiment_prec = experiment_matched/experiment_test
        experiment_f1 = ((2*experiment_rec*experiment_prec)/(experiment_rec+experiment_prec))*100

        baseline_matched = sum([id2matched['baseline'][turn] for turn in resampled_turns])
        baseline_test = sum([id2test['baseline'][turn] for turn in resampled_turns])
        baseline_rec = baseline_matched/gold
        baseline_prec = baseline_matched/baseline_test
        baseline_f1 = ((2*baseline_rec*baseline_prec)/(baseline_rec+baseline_prec))*100

        #if experiment_f1 - baseline_f1 > (2*fullset_delta):
        curr_delta = abs(experiment_f1 - baseline_f1)
        if curr_delta > (2*fullset_delta):
            big_diffs += 1
    print(f'p-value estimate: {big_diffs/num_resamples}')




def main():
    node = 'DUMMY'
    #output_dir = '/disk/scratch1/s1899827/prosody_nlp/code/self_attn_speech_parser/output/turn_pipeline_no_voicing_redux'
    output_dir = '/disk/scratch1/s1899827/prosody_nlp/code/self_attn_speech_parser/output/turn_no_voicing'
    #data_dir = f'/disk/scratch1/s1899827/prosody_nlp/data/input_features/turn_pipeline_{NODE}'
    data_dir = f'/disk/scratch1/s1899827/prosody_nlp/data/input_features/turn_opensmile'
    gold_f = os.path.join(data_dir,'turn_dev.trees')
    id_f = os.path.join(data_dir,'turn_dev.ids')
    
    # Basic pros/nonpros comparison
    #  nonsp_seg_5154_dev=87.24.pt.silver.parsed.DUMMY.silver.turns   nonsp_seg_73784_dev=87.61.pt.silver.parsed.DUMMY.silver.turns
    #nonsp_seg_38976_dev=87.36.pt.silver.parsed.DUMMY.silver.turns  nonsp_seg_72240_dev=87.60.pt.silver.parsed.DUMMY.silver.turns

    #baseline_f = os.path.join(output_dir,'nonsp_parse_dp.3_26229_dev=85.94.pred')
    #baseline_f = os.path.join(output_dir,'nonsp_parse_dp.3_38976_dev=86.11.pred')
    #baseline_f = os.path.join(output_dir,'turn_pause_only_correct_eval_72240_dev=85.41.pred')
    #baseline_f = os.path.join(output_dir,'turn_dur_only_correct_eval_72240_dev=85.53.pred')
    baseline_f = os.path.join(output_dir,'turn_fbank_only_correct_eval_72240_dev=86.00.pred')
    
    #baseline_f = os.path.join(output_dir,'sp_seg_duration_only_dev=88.93.pt.silver.parsed.DUMMY.silver.turns.defrag')
    #baseline_f = os.path.join(output_dir,'sp_seg_fbank_only_dev=89.36.pt.silver.parsed.DUMMY.silver.turns.defrag')
    #baseline_f = os.path.join(output_dir,'sp_seg_pause_only_dev=89.55.pt.silver.parsed.DUMMY.silver.turns.defrag')
    #baseline_f = os.path.join(output_dir,'sp_seg_pitch_only_dev=89.44.pt.silver.parsed.DUMMY.silver.turns.defrag')
                              
    #baseline_f = os.path.join(output_dir,'turn_dur_only_correct_eval_72240_dev=85.53.pred')
    #baseline_f = os.path.join(output_dir,'nonsp_parse_dp.3_72240_dev=85.33.pred')
    experiments = ['turn_fbank_only_correct_eval_72240_dev=86.00.pred',
                   'turn_pitch_only_correct_eval_72240_dev=85.81.pred',
                   'turn_pause_only_correct_eval_72240_dev=85.41.pred',
                   'turn_dur_only_correct_eval_72240_dev=85.53.pred']
    
    ids = [l.strip() for l in open(id_f).readlines()]

    for ex in [experiments[1]]:
        print('baseline: ',baseline_f)
        print(ex)
        ex = os.path.join(output_dir,ex)
        get_p_value(baseline_f,gold_f,ex,ids,mode='parse')

if __name__=="__main__":
    main()
    """
    COMB = FALSE
    
    if COMB:
        temp_dir = tempfile.TemporaryDirectory(prefix="bootstrap-")
        print("Temporary dir", temp_dir)
        baseline_f = os.path.join(temp_dir.name, 'baseline.txt')
        experiment_f = os.path.join(temp_dir.name, 'experiment.txt')
        gold_out_f = os.path.join(temp_dir.name, 'gold.txt')
        baseline_files = ['nonsp_parse_dp.3_26229_dev=85.94.pred',
                      'nonsp_parse_dp.3_38976_dev=86.11.pred',
                      'nonsp_parse_dp.3_5154_dev=85.76.pred',
                      'nonsp_parse_dp.3_72240_dev=85.33.pred',
                      'nonsp_parse_dp.3_73784_dev=85.36.pred']


        experiments = ['sp_glove_dp.3_26229_dev=86.34.pred',
                       'sp_glove_dp.3_38976_dev=86.24.pred',
                       'sp_glove_dp.3_5154_dev=85.99.pred',
                       'sp_glove_dp.3_72240_dev=86.48.pred',
                       'sp_glove_dp.3_73784_dev=85.98.pred']

        bases = []
        exps = []
        ids = []
        golds = []
        for base,exp in zip(baseline_files,experiments):
            base = os.path.join(output_dir,base)
            exp = os.path.join(output_dir,exp)
            bases.extend([line.strip() for line in open(base,'r').readlines()])
            exps.extend([line.strip() for line in open(exp,'r').readlines()])
            ids.extend([line.strip() for line in open(id_f,'r').readlines()])
            golds.extend([line.strip() for line in open(gold_f,'r').readlines()])
        with open(baseline_f,'w') as f:
            f.write('\n'.join(bases))
        with open(experiment_f,'w') as f:
            f.write('\n'.join(exps))
        with open(gold_out_f,'w') as f:
            f.write('\n'.join(golds))
        get_p_value(baseline_f,gold_out_f,experiment_f,ids)

    else: # if not combining files:
    """

