import math
import random
from model.base_model import Model
from random import randint, sample

class Gene:  # City
    def __init__(self, name, lat, lng):
        self.name = name
        self.lat = lat
        self.lng = lng

    def get_distance_to(self, dest):
        return math.sqrt( (self.lng - dest.lng) ** 2 + (self.lat - dest.lat) ** 2 )

class Individual:  # Route: possible solution to TSP
    def __init__(self, genes):
        assert(len(genes) > 3)
        self.genes = genes
        self.__reset_params()

    def swap(self, gene_1, gene_2):
        self.genes[0]
        a, b = self.genes.index(gene_1), self.genes.index(gene_2)
        self.genes[b], self.genes[a] = self.genes[a], self.genes[b]
        self.__reset_params()

    def add(self, gene):
        self.genes.append(gene)
        self.__reset_params()

    @property
    def fitness(self):
        if self.__fitness == 0:
            self.__fitness = 1 / self.travel_cost  # Normalize travel cost
        return self.__fitness

    @property
    def travel_cost(self):  # Get total travelling cost
        if self.__travel_cost == 0:
            for i in range(len(self.genes)):
                origin = self.genes[i]
                if i == len(self.genes) - 1:
                    dest = self.genes[0]
                else:
                    dest = self.genes[i+1]

                self.__travel_cost += origin.get_distance_to(dest)

        return self.__travel_cost

    def __reset_params(self):
        self.__travel_cost = 0
        self.__fitness = 0

class Population:  # Population of individuals
    def __init__(self, individuals):
        self.individuals = individuals

    @staticmethod
    def gen_individuals(sz, genes):
        individuals = []
        for _ in range(sz):
            individuals.append(Individual(sample(genes, len(genes))))
        return Population(individuals)

    def add(self, route):
        self.individuals.append(route)

    def rmv(self, route):
        self.individuals.remove(route)

    def get_fittest(self):
        fittest = self.individuals[0]
        for route in self.individuals:
            if route.fitness > fittest.fitness:
                fittest = route

        return fittest

class MyGAModel(Model):
    def __init__(self):
        super().__init__()
        self.iteration = 0

    def init(self, nodes):
        super().init(nodes)

    def evolve(self, pop, tourn_size, mut_rate):
        new_generation = Population([])
        pop_size = len(pop.individuals)
        elitism_num = pop_size // 4

        # Elitism
        for _ in range(elitism_num):
            fittest = pop.get_fittest()
            new_generation.add(fittest)
            pop.rmv(fittest)

        # Crossover
        for _ in range(elitism_num, pop_size):
            parent_1 = self.selection(new_generation, tourn_size)
            parent_2 = self.selection(new_generation, tourn_size)
            child = self.crossover(parent_1, parent_2)
            new_generation.add(child)

        # Mutation
        for i in range(elitism_num, pop_size):
            self.mutate(new_generation.individuals[i], mut_rate)

        return new_generation

    def crossover(self, parent_1, parent_2):
        def fill_with_parent1_genes(child, parent, genes_n):
            start_at = randint(0, len(parent.genes)-genes_n-1)
            finish_at = start_at + genes_n
            for i in range(start_at, finish_at):
                child.genes[i] = parent_1.genes[i]

        def fill_with_parent2_genes(child, parent):
            j = 0
            for i in range(0, len(parent.genes)):
                if child.genes[i] == None:
                    while parent.genes[j] in child.genes:
                        j += 1
                    child.genes[i] = parent.genes[j]
                    j += 1

        genes_n = len(parent_1.genes)
        child = Individual([None for _ in range(genes_n)])
        fill_with_parent1_genes(child, parent_1, genes_n // 2)
        fill_with_parent2_genes(child, parent_2)

        return child

    def mutate(self, individual, rate):
        for _ in range(len(individual.genes)):
            if random.random() < rate:
                sel_genes = sample(individual.genes, 2)
                individual.swap(sel_genes[0], sel_genes[1])

    def selection(self, population, competitors_n):
        return Population(sample(population.individuals, competitors_n)).get_fittest()

    def fit(self, max_it=100):
        """
        Execute simulated annealing algorithm.
        """
        pop_size = 1000
        mut_rate = 0.9
        tourn_size = 100

        genes = [Gene(num, city[0], city[1]) for num, city in enumerate(self.coords)]
        self.genes = genes

        population = Population.gen_individuals(pop_size, genes)
        

        for it in range(0, max_it):
            mut_rate = mut_rate * 0.95
            if mut_rate < 0.05:
                mut_rate = 0.05
            population = self.evolve(population, tourn_size, mut_rate)
            cost = population.get_fittest().travel_cost

            it += 1
            self.fitness_list.append(cost)
            # print("[step] ", it, " [mut] ", mut_rate, " [best] ", self.fitness_list[self.fitness_list.index(min(self.fitness_list))])

        self.best_solution = [gene.name for gene in population.get_fittest().genes]

        return self.best_solution, self.fitness_list