import math
import os
import re
import subprocess
import sys
from math import sqrt
from statistics import mean, variance
from typing import List, Dict

import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
from scipy.stats import wilcoxon
from tqdm import tqdm

ROOT_DIR = os.path.dirname(__file__)
IN_SOURCE_DIR = os.path.join(ROOT_DIR, "benchmark")
IN_TEST_DIR = os.path.join(ROOT_DIR, "tests")
IN_FUZZER_TEST_DIR = os.path.join(ROOT_DIR, "fuzzer_tests")
OUT_DIR = os.path.join(ROOT_DIR, "out")
MUT_PY_PATH = os.path.join(ROOT_DIR, 'env37', 'bin', 'mut.py')
REPS: int = 10


def cohen_d(d1: List[float], d2: List[float]) -> float:
    pooled_sd = sqrt(((len(d1) - 1) * variance(d1) + (len(d2) - 1) * variance(d2)) /
                     (len(d1) + len(d2) - 2))

    if pooled_sd == 0:
        return math.inf

    return (mean(d1) - mean(d2)) / pooled_sd


def effect_size(eff: float) -> str:
    if eff <= 0.01:
        return 'Very small'
    elif eff <= 0.2:
        return 'Small'
    elif eff <= 0.5:
        return 'Medium'
    elif eff <= 0.8:
        return 'Large'
    elif eff <= 1.2:
        return 'Very large'
    else:
        return 'Huge'


def compute_stats(df_gen: pd.DataFrame, df_fuz: pd.DataFrame, output_file: str, avg_output_file: str, stat_csv: str):
    combined_df = pd.concat([df_gen, df_fuz], keys=["genetic", "fuzzer"]).reset_index()
    combined_df.columns = ['source', *combined_df.columns[1:]]
    del combined_df[combined_df.columns[1]]

    plt.figure(figsize=(18, 8))
    sns.set(style="whitegrid")
    sns.boxplot(data=combined_df, x="file", y="score", hue="source")
    plt.yticks(range(0, 101, 10))
    plt.savefig(output_file)

    plt.figure(figsize=(18, 8))
    df_avg = combined_df.groupby(['file', 'source']).mean().reset_index()
    sns.set(style="whitegrid")
    sns.barplot(data=df_avg, x="file", y="score", hue="source")
    plt.yticks(range(0, 101, 10))
    plt.savefig(avg_output_file)

    df_avg = df_avg.pivot(index='file', columns='source', values='score').rename_axis(None, axis=1)
    df_avg['cohen-d'] = [math.nan] * len(df_avg.index)
    df_avg['interpretation'] = [math.nan] * len(df_avg.index)
    df_avg['wilcoxon'] = [math.nan] * len(df_avg.index)

    for f in combined_df['file'].drop_duplicates():
        list_gen = df_gen.loc[(df_gen.file == f), 'score'].tolist()
        list_fuz = df_fuz.loc[(df_fuz.file == f), 'score'].tolist()

        df_avg.loc[f, 'cohen-d'] = cohen_d(list_gen, list_fuz)
        df_avg.loc[f, 'interpretation'] = effect_size(df_avg.loc[f, 'cohen-d'])
        df_avg.loc[f, 'wilcoxon'] = wilcoxon(list_gen, list_fuz, zero_method='zsplit').pvalue

    df_avg.to_csv(stat_csv)


def run_mutpy(test_path: str, source_path: str) -> float:
    output = subprocess.check_output(
        [sys.executable, MUT_PY_PATH, '-t', source_path, '-u', test_path]).decode('utf-8')
    score = re.search('Mutation score \\[.*]: (\\d+\\.\\d+)%', output).group(1)
    return float(score)


def mutate_suite(out_file: str, in_test_dir: str, to_test: List[str]):
    scores: List[Dict[str, any]] = []

    if os.path.isfile(out_file):  # do not re-generate if file exists
        return pd.read_csv(out_file, index_col=0)

    for filename in tqdm(to_test, desc=f"mut.py [{os.path.basename(out_file)}]"):
        source_path = os.path.join(IN_SOURCE_DIR, f"{filename}.py")
        test_path = os.path.join(in_test_dir, f"test_{filename}.py")
        scores.append({
            'file': filename,
            'score': run_mutpy(test_path, source_path)
        })

    df = pd.DataFrame.from_records(scores)
    df.to_csv(out_file)
    return df


def main():
    files = [os.path.splitext(f) for f in os.listdir(IN_SOURCE_DIR)]
    to_test = [file[0] for file in files if file[1] == ".py"]
    to_test = [e for t in to_test for e in ([t] * REPS)]

    df_gen = mutate_suite(os.path.join(OUT_DIR, 'mutation_results_genetic.csv'), IN_TEST_DIR, to_test)
    df_fuz = mutate_suite(os.path.join(OUT_DIR, 'mutation_results_fuzzer.csv'), IN_FUZZER_TEST_DIR, to_test)

    compute_stats(df_gen, df_fuz,
                  os.path.join(OUT_DIR, "mutation_scores.png"),
                  os.path.join(OUT_DIR, "mutation_scores_mean.png"),
                  os.path.join(OUT_DIR, "stats.csv"))


if __name__ == "__main__":
    main()