import random

import argparse
import os
import sys
import numpy as np

from evolalg.base.lambda_step import LambdaStep
from evolalg.experiment import Experiment
from evolalg.fitness.fitness_step import FitnessStep
from evolalg.mutation_cross.frams_cross_and_mutate import FramsCrossAndMutate
from evolalg.population.frams_population import FramsPopulation
from evolalg.repair.multistep import MultistepRepair
from evolalg.repair.remove.field import FieldRemove
from evolalg.selection.tournament import TournamentSelection
from evolalg.statistics.halloffame_stats import HallOfFameStatistics
from evolalg.statistics.statistics_deap import StatisticsDeap
from evolalg.base.union_step import UnionStep
from evolalg.utils.population_save import PopulationSave
from evolalg.utils.stable_generation import StableGeneration
from FramsticksLib import FramsticksLib


def ensureDir(string):
    if os.path.isdir(string):
        return string
    else:
        raise NotADirectoryError(string)


def parseArguments():
    parser = argparse.ArgumentParser(
        description='Run this program with "python -u %s" if you want to disable buffering of its output.' % sys.argv[
            0])
    parser.add_argument('-path', type=ensureDir, required=True, help='Path to Framsticks without trailing slash.')
    # parser.add_argument('-opt', required=True,
    #                     help='optimization criteria : vertpos, velocity, distance, vertvel, lifespan, numjoints, numparts, numneurons, numconnections. Single or multiple criteria.')
    parser.add_argument('-lib', required=False, help="Filename of .so or .dll with framsticks library")
    parser.add_argument('-genformat', required=False, default="1",
                        help='Genetic format for the demo run, for example 4, 9, or B. If not given, f1 is assumed.')

    parser.add_argument("-popsize", type=int, default=50, help="Size of population, default 50.")
    return parser.parse_args()


def extract_fitness(ind):
    return ind.fitness


def print_population_count(pop):
    print("Current:",len(pop))
    return pop # Each step must return a population


def random_remove(pop):
    return [_ for _ in pop if random.randint(0, 1) == 0]


def main():
    parsed_args = parseArguments()
    frams_lib = FramsticksLib(parsed_args.path, parsed_args.lib, "eval-movement.sim")

    hall_of_fame = HallOfFameStatistics(100, "fitness")
    statistics_union = UnionStep([
        hall_of_fame,
        StatisticsDeap([
            ("avg", np.mean),
            ("stddev", np.std),
            ("min", np.min),
            ("max", np.max),
            ("count", len)
        ], extract_fitness)
    ])

    fitness_remove = UnionStep([
        FitnessStep(frams_lib, fields={"velocity": "fitness", "data->recording": "recording"},
                    fields_defaults={"velocity": None, "data->recording": None}),
        FieldRemove("recording", None),
        print_population_count,  # Stages can be also any Callable
    ])

    selection = TournamentSelection(5, copy=True)
    new_generation_steps = [
        FramsCrossAndMutate(frams_lib, cross_prob=0.2, mutate_prob=0.9),
        fitness_remove
    ]

    generation_modifications = [
        statistics_union  # Or niching, novelty
    ]

    init_stages = [FramsPopulation(frams_lib, parsed_args.genformat, parsed_args.popsize),
                   fitness_remove,  # It is possible to create smaller population
                   statistics_union]

    end_steps = [PopulationSave("halloffame.gen", provider=hall_of_fame.halloffame, fields={"genotype": "genotype",
                                                                                           "fitness": "fitness",
                                                                                           "custom": "recording"})]

    experiment = Experiment(init_population=init_stages,
                            selection=selection,
                            new_generation_steps=new_generation_steps,
                            generation_modification=generation_modifications,
                            end_steps=end_steps,
                            population_size=parsed_args.popsize
                            )
    experiment.init()
    experiment.run(3)
    for ind in hall_of_fame.halloffame:
        print("%g\t%s" % (ind.fitness, ind.genotype))


if __name__ == '__main__':
    main()