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(baselines,
                experiments,
                output_dir,
                gold_f,
                ids):
    base = {}
    exp = {}
    for i,baseline in enumerate(baselines):
        baseline = os.path.join(output_dir,baseline)
        base[i] = [line.strip() for line in open(baseline,'r').readlines()]
    for i,experiment in enumerate(experiments):
        experiment = os.path.join(output_dir,experiment)
        exp[i] = [line.strip() for line in open(experiment,'r').readlines()]
    golds = [line.strip() for line in open(gold_f,'r').readlines()]
    

    
    all_gold_lines = golds*len(base)
    all_ids = ids*len(base)
    all_experiment_lines = []
    for ex in exp:
        all_experiment_lines.extend(exp[ex])
    all_baseline_lines = []
    for bas in base:
        all_baseline_lines.extend(base[bas])

    num_lines = len(all_gold_lines)
    assert len(all_gold_lines)==len(all_ids)
    assert len(all_gold_lines)==len(all_experiment_lines)
    assert len(all_gold_lines)==len(all_baseline_lines)

    scr = scorer.Scorer()
    print('Calculate baseline...')
    fullset_experiment_f1 = scr.get_f1_from_list(all_gold_lines,all_experiment_lines)
    fullset_baseline_f1 = scr.get_f1_from_list(all_gold_lines,all_baseline_lines)
    
    fullset_delta = abs(fullset_experiment_f1 - fullset_baseline_f1)
        
    big_diffs = 0

    id2matched = {'experiment':{0:{},
                                1:{},
                                2:{},
                                3:{},
                                4:{}},
                  'baseline':{0:{},
                              1:{},
                              2:{},
                              3:{},
                              4:{}}}
    id2gold = {'experiment':{0:{},
                             1:{},
                             2:{},
                             3:{},
                             4:{}},
               'baseline':{0:{},
                           1:{},
                           2:{},
                           3:{},
                           4:{}}}
    id2test = {'experiment':{0:{},
                             1:{},
                             2:{},
                             3:{},
                             4:{}},
               'baseline':{0:{},
                           1:{},
                           2:{},
                           3:{},
                           4:{}}}



    print('Store bracket scores ...')
    for model in base:
        baseline_lines = base[model]
        experiment_lines = exp[model]
        
        for i,turn in enumerate(ids):
            gold_line = golds[i]
            baseline_line = baseline_lines[i]
            experiment_line = experiment_lines[i]
            assert tree2str(gold_line)==tree2str(baseline_line)
            assert tree2str(gold_line)==tree2str(experiment_line)
            bmat,bgol,btes = scr.get_bracket_counts_from_tree(gold_line,baseline_line)
            id2matched['baseline'][model][turn] = bmat
            id2gold['baseline'][model][turn] = bgol
            id2test['baseline'][model][turn] = btes
            emat,egol,etes = scr.get_bracket_counts_from_tree(gold_line,experiment_line)
            id2matched['experiment'][model][turn] = emat
            id2gold['experiment'][model][turn] = egol
            id2test['experiment'][model][turn] = etes
        
        
    idxs = [i for i in range(len(golds))]
    print('Resample ...')
    for i in range(num_resamples):
        if i%10000==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])

        base_resample_model = random.choices(list(range(len(base))),k=len(idxs))
        exp_resample_model = random.choices(list(range(len(exp))),k=len(idxs))
            
        gold = sum([id2gold['experiment'][model][turn] for turn,model in zip(resampled_turns,exp_resample_model)])
        experiment_matched = sum([id2matched['experiment'][model][turn] for turn,model in zip(resampled_turns,exp_resample_model)])
        experiment_test = sum([id2test['experiment'][model][turn] for turn,model in zip(resampled_turns,exp_resample_model)])
        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'][model][turn] for turn,model in zip(resampled_turns,base_resample_model)])
        baseline_test = sum([id2test['baseline'][model][turn] for turn,model in zip(resampled_turns,base_resample_model)])
        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'
    split = 'test'
    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_opensmile'
    gold_f = os.path.join(data_dir,f'turn_{split}.trees')
    id_f = os.path.join(data_dir,f'turn_{split}.ids')
    

    """
    baselines = ['nonsp_parse_dp.3_26229_dev=85.94.pt_test.pred',
                 'nonsp_parse_dp.3_5154_dev=85.76.pt_test.pred',
                 'nonsp_parse_dp.3_73784_dev=85.36.pt_test.pred',
                 'nonsp_parse_dp.3_38976_dev=86.11.pt_test.pred',
                 'nonsp_parse_dp.3_72240_dev=85.33.pt_test.pred']
    
    experiments = ['sp_glove_dp.3_26229_dev=86.34.pt_test.pred',
                   'sp_glove_dp.3_5154_dev=85.99.pt_test.pred',
                   'sp_glove_dp.3_73784_dev=85.98.pt_test.pred',
                   'sp_glove_dp.3_38976_dev=86.24.pt_test.pred',
                   'sp_glove_dp.3_72240_dev=86.48.pt_test.pred']
    """





    
    baselines = ['nonsp_seg_26229_dev=87.61.pt.silver.parsed.DUMMY.test.silver.turns.defrag',
                 'nonsp_seg_38976_dev=87.36.pt.silver.parsed.DUMMY.test.silver.turns.defrag',
                 'nonsp_seg_5154_dev=87.24.pt.silver.parsed.DUMMY.test.silver.turns.defrag',
                 'nonsp_seg_72240_dev=87.60.pt.silver.parsed.DUMMY.test.silver.turns.defrag',
                 'nonsp_seg_73784_dev=87.61.pt.silver.parsed.DUMMY.test.silver.turns.defrag']
    
    experiments = ['sp_seg_26229_dev=89.54.pt.silver.parsed.DUMMY.test.silver.turns.defrag',
                   'sp_seg_38976_dev=89.46.pt.silver.parsed.DUMMY.test.silver.turns.defrag',
                   'sp_seg_5154_dev=89.62.pt.silver.parsed.DUMMY.test.silver.turns.defrag',
                   'sp_seg_72240_dev=89.59.pt.silver.parsed.DUMMY.test.silver.turns.defrag',
                   'sp_seg_73784_dev=89.99.pt.silver.parsed.DUMMY.test.silver.turns.defrag']
    
    #    """
    ids = [l.strip() for l in open(id_f).readlines()]

    get_p_value(baselines,
                experiments,
                output_dir,
                gold_f,
                ids)

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:
    """

