import os
from typing import Callable

from deap import creator, base, tools, algorithms

import fuzzer
import instrument
from fuzzer import get_test_cases, get_test_class
import frozendict

INDMUPROB = 0.05
MUPROB = 0.1
CXPROB = 0.5
TOURNSIZE = 3
NPOP = 300
NGEN = 200
REPS = 10

to_test: str = ""

OUT_DIR = os.path.join(os.path.dirname(__file__), "tests")


def normalize(x):
    return x / (1.0 + x)


def get_test_case_generator(f_name: str, arguments: list[instrument.Arg]) -> Callable[[], list]:
    return lambda: list(list(get_test_cases(f_name, arguments, 1, enable_bar=False))[0].items())


def generate(f_name: str):
    global to_test
    to_test = f_name

    creator.create("Fitness", base.Fitness, weights=(-1.0,))
    creator.create("Individual", list, fitness=creator.Fitness)

    n_initial = 10

    args = instrument.functions[f_name]

    toolbox = base.Toolbox()
    toolbox.register("attr_test_case", get_test_case_generator(to_test, instrument.functions[to_test]))
    toolbox.register("individual",
                     tools.initRepeat,
                     creator.Individual,
                     toolbox.attr_test_case,
                     n=n_initial)
    toolbox.register("population",
                     tools.initRepeat,
                     list,
                     toolbox.individual)
    toolbox.register("evaluate", fitness)

    def mate(tc1, tc2):
        t1, t2 = frozendict.frozendict(tc1), frozendict.frozendict(tc2)

        print("ticino", tc1, tc2)
        o1, o2 = fuzzer.crossover(t1, t2, args)
        i1, i2 = creator.Individual(o1.items()), creator.Individual(o2.items())
        print("mate", i1, i2)
        return i1, i2

    def mutate(tc):
        t = frozendict.frozendict(tc)
        o = fuzzer.mutate(t, args)
        i1 = creator.Individual(o.items())
        print("mutate", i1)
        return i1,

    toolbox.register("mate", mate)
    toolbox.register("mutate", mutate)
    toolbox.register("select", tools.selTournament, tournsize=TOURNSIZE)

    coverage = []
    for i in range(REPS):
        instrument.archive_true_branches = {}
        instrument.archive_false_branches = {}
        population = toolbox.population(n=NPOP)
        print("population", population)
        algorithms.eaSimple(population, toolbox, CXPROB, MUPROB, NGEN, verbose=False)
        cov = len(instrument.archive_true_branches) + len(instrument.archive_false_branches)
        print(cov, instrument.archive_true_branches, instrument.archive_false_branches)
        coverage.append(cov)

    print(coverage)


def fitness(individuals: list[list]) -> tuple[float]:
    range_start, range_end = instrument.n_of_branches[to_test]

    # Reset any distance values from previous executions
    instrument.distances_true = {}
    instrument.distances_false = {}

    fitness = 0.0
    for individual in individuals:
        x = frozendict.frozendict(individual)

        # Run the function under test
        try:
            out = instrument.invoke(to_test, x)
        except AssertionError:
            print(to_test, x, "=", "[FAILS]")
            return 10000,

        # Sum up branch distances
        for branch in range(range_start, range_end):
            if branch in instrument.distances_true:
                if instrument.distances_true[branch] == 0 and branch not in instrument.archive_true_branches:
                    instrument.archive_true_branches[branch] = x
                if branch not in instrument.archive_true_branches:
                    fitness += normalize(instrument.distances_true[branch])
        for branch in range(range_start, range_end):
            if branch in instrument.distances_false:
                if instrument.distances_false[branch] == 0 and branch not in instrument.archive_false_branches:
                    instrument.archive_false_branches[branch] = x
                if branch not in instrument.archive_false_branches:
                    fitness += normalize(instrument.distances_false[branch])
        print(to_test, x, "=", out)

    print("fitness", fitness)
    print()
    return fitness,


def main():
    instrument.load_benchmark(save_instrumented=False)  # instrument all files in benchmark
    f_name = "railencrypt_instrumented"
    generate(f_name)
    with open(os.path.join(OUT_DIR, f_name + ".py"), "w") as f:
        cases = get_test_cases(f_name, instrument.functions[f_name], 100)
        f.write(get_test_class(f_name, cases))


if __name__ == '__main__':
    main()