# runs the detection tests

from constat import ConStat, perform_test, load_result
import os
import seaborn as sns
import matplotlib.pyplot as plt
import pandas as pd
from tqdm import tqdm
import numpy as np
from loguru import logger

base_path = 'lm-evaluation-harness/output'
base_path_data = 'data/contamination'

ref_models = [
    "microsoft/phi-2",
    "meta-llama/Llama-2-7b-chat-hf",
    "meta-llama/Llama-2-7b-hf",
    "meta-llama/Llama-2-13b-chat-hf",
    "meta-llama/Llama-2-13b-hf",
    "meta-llama/Llama-3-70b-chat-hf",
    "meta-llama/Meta-Llama-3-70B",
    "meta-llama/Llama-2-70b-chat-hf",
    # "mistralai/Mistral-7B-v0.1",
    "mistralai/Mistral-7B-Instruct-v0.2",
    "mistralai/Mistral-7B-Instruct-v0.1",
    "tiiuae/falcon-7b-instruct",
    "tiiuae/falcon-7b",
    "google/gemma-1.1-2b-it",
    "google/gemma-1.1-7b-it",
    "meta-llama/Meta-Llama-3-8B",
    "meta-llama/Meta-Llama-3-8B-Instruct",
    "microsoft/Phi-3-mini-4k-instruct",
    "mistralai/Mixtral-8x22B-Instruct-v0.1",
    "allenai/OLMo-7B-Instruct",
    "mistralai/Mixtral-8x7B-Instruct-v0.1"

]

extra_ref_models = [
    "mistralai/Mistral-7B-v0.1"
]

test_models = [
  "yam-peleg/Experiment26-7B",
  "BarraHome/Mistroll-7B-v2.2",
  "MTSAIR/multi_verse_model",
  "Qwen/Qwen1.5-1.8B-Chat",
  "internlm/internlm2-7b",
  "internlm/internlm2-1_8b",
  "internlm/internlm2-math-7b",
  "internlm/internlm2-math-base-7b",
  "Qwen/Qwen1.5-4B-Chat",
  "Qwen/Qwen1.5-7B-Chat",
  "Qwen/Qwen1.5-14B-Chat",
  "Qwen/Qwen1.5-72B-Chat",
  "Qwen/Qwen1.5-110B-Chat",
  "stabilityai/stablelm-base-alpha-7b-v2",
  "stabilityai/stablelm-2-12b-chat",
  "stabilityai/stablelm-2-12b",
  "stabilityai/stablelm-2-1_6b",
  "stabilityai/stablelm-zephyr-3b",
  "stabilityai/stablelm-2-1_6b-chat",
  "mistralai/Mistral-7B-v0.1",
  "zero-one-ai/Yi-34B",
  "zero-one-ai/Yi-6B",
  "gpt-3.5-turbo",
  "gpt-4-turbo",
  "gpt-4o"
]

default_metric_map = {
    'gsm8k': 'flexible_extract',
    'mathqa': 'acc_norm',
    'arc': 'acc_norm',
    'mmlu': 'acc',
    'sciq': 'acc_norm',
    'hellaswag': 'acc_norm',
    'mmlu_combination': 'acc',
    'lambada_openai': 'acc',
    'mmlu_synthetic': 'acc'
}

default_reference_benchmark = {
    'gsm8k': 'mathqa',
    'hellaswag': 'lambada_openai',
    'arc': 'sciq',
    'mmlu': 'mmlu_combination'
}

random_performances = {
    'gsm8k': 0,
    'mathqa': 0.25,
    'arc': 0.25,
    'mmlu': 0.25,
    'sciq': 0.25,
    'hellaswag': 0.25,
    'mmlu_combination': 0.25,
    'lambada_openai': 0,
    'mmlu_synthetic': 0.25
}

def get_contaminated_models(benchmark, username):
    models = ["default", "rephrase", "lr_1e-4", "lr_1e-5", "repeat_1", "default_other", "default_no", 
              "default_with_ref", "rephrase_with_ref"]
    full_model_names = [
        f"{username}/contamination-models-{benchmark}-meta-llama-Llama-2-7b-chat-hf-{model}" for model in models
    ]
    full_model_names += [
        f"{username}/contamination-models-{benchmark}-{model}" for model in models
    ]
    return full_model_names



def perform_full_test(benchmark, ref_benchmark, ref_models, to_test_models, 
                 add_extended_models=True, n_bootstrap=1000, add_no_cont_result=True, 
                 extra_ref_models=[]):
    logger.info(benchmark + ref_benchmark)
    random_performance_no_dash = random_performances.get(benchmark[:benchmark.index('_')], None)
    random_performance_benchmark = random_performances.get(benchmark, random_performance_no_dash)
    random_performance_ref_benchmark = random_performances.get(ref_benchmark, random_performance_benchmark)
    metric_no_dash = default_metric_map.get(benchmark[:benchmark.index('_')], None)
    metric = default_metric_map.get(benchmark, metric_no_dash)
    metric_reference_benchmark = default_metric_map.get(ref_benchmark, metric)
    test = ConStat(
        add_extended_models=add_extended_models,
        n_bootstrap=n_bootstrap,
        random_performance=(random_performance_benchmark, random_performance_ref_benchmark),
        p_value_delta=0
    )
    no_cont_benchmark = benchmark[:benchmark.index('_')] + '_no_cont'
    results = []
    for model in tqdm(to_test_models):
        ref_models_here = ref_models[:]
        if model in ref_models:
            ref_models_here.remove(model)
            ref_models_here += extra_ref_models
        try:
            add_no_cont_results_test = 'contamination-models-' not in model
            test_result = perform_test(
                model, benchmark, ref_benchmark, test, ref_models, 
                base_path, metric, metric_reference_benchmark, add_no_cont_results=add_no_cont_results_test, 
                base_path_data=base_path_data
            )
        except Exception as e:
            logger.warning(f"{e} for {model} at {benchmark} with {ref_benchmark}")
            continue
        test_result['model'] = model
        if add_no_cont_result:
            result_no_cont = load_result(
                base_path, model, no_cont_benchmark, metric, base_path_data=base_path_data
            )
            test_result['no_cont'] = np.mean(result_no_cont)
            test_result['no_cont_std'] = np.sqrt((1 - test_result['no_cont']) * test_result['no_cont'] / len(result_no_cont))
        results.append(test_result)

    # convert to dataframe
    df = pd.DataFrame(results)
    return df

benchmarks = [
    'gsm8k', 'arc', 'mmlu', 'hellaswag'
]

import argparse

parser = argparse.ArgumentParser(description='Run detection tests')
parser.add_argument('--benchmark', type=str)
parser.add_argument('--type', type=str)
parser.add_argument('--normal-type', type=str, default='normal')
parser.add_argument('--username', type=str, default='anonymous')

args = parser.parse_args()

os.makedirs('tables', exist_ok=True)

n_bootstrap = 10000
benchmark = args.benchmark

if args.type == 'reference':
    other_benchmark = default_reference_benchmark[benchmark]
else:
    other_benchmark = benchmark + f'_{args.type}'

res_rephrase = perform_full_test(benchmark + f'_{args.normal_type}', other_benchmark, ref_models, 
                                    ref_models + test_models + get_contaminated_models(benchmark, args.username), add_no_cont_result=True, 
                                    n_bootstrap=n_bootstrap)
res_rephrase.to_csv(f'tables/{benchmark}_{args.type}.csv')