import random
from typing import Tuple

import numpy as np
import numpy.typing as npt

from common import plot_plan, read_data


class GeneticTSP:

    def __init__(self, population: int, crossover_prob: float, mutation_prob: float, data):
        self.generation: int = 0
        self.population: int = population
        self.data: npt.NDArray = data
        self.genes: int = len(data)
        self.crossover_prob: float = crossover_prob
        self.mutation_prob: float = mutation_prob
        self.generate_first_generation()

    def generate_first_generation(self):
        self.candidates: npt.NDArray = np.array([np.random.permutation(np.arange(self.genes)) for _ in range(self.population)])


    def get_distance(self, candidate):
        return sum([self.data[candidate[i - 1], candidate[i]] for i in range(self.genes)])


    def fitness(self):
        distances = np.array([ self.get_distance(candidate) for candidate in self.candidates ])
        max_distance = max(distances)
        fitness = max_distance - distances
        fitness_sum = sum(fitness)
        self.fitness_probs: npt.NDArray = fitness / fitness_sum


    def crossover(self, parent1: npt.NDArray, parent2: npt.NDArray) -> Tuple[npt.NDArray, npt.NDArray]:
        if self.crossover_prob < np.random.random():
            return (parent1, parent2)

        cut: int = np.random.randint(0, self.genes)

        offspring1 = parent1[:cut]
        offspring2 = parent2[:cut]

        offspring1 = np.concatenate((offspring1, np.array([gene for gene in parent2 if gene not in offspring1])))
        offspring2 = np.concatenate((offspring2, np.array([gene for gene in parent1 if gene not in offspring2])))

        return (offspring1, offspring2)


    def mutate(self, offspring):
        if self.mutation_prob < np.random.random():
            return

        pos1: int = np.random.randint(0, self.genes)
        pos2: int = np.random.randint(0, self.genes)

        offspring[pos1], offspring[pos2] = offspring[pos2], offspring[pos1]

    def select_individual(self):
        choice = np.random.choice(self.population, 1, p=self.fitness_probs)[0]
        return self.candidates[choice]


    def generate_next_generation(self):
        new_generation = []
        self.fitness()
        
        for _ in range(0,self.population,2):
            offspring1, offspring2 = self.crossover(self.select_individual(), self.select_individual())
            self.mutate(offspring1)
            self.mutate(offspring2)
            new_generation.append(offspring1)
            new_generation.append(offspring2)

        self.candidates = np.array(new_generation)


    def run(self, generations = 10):
        for _ in range(generations):
            self.generate_next_generation()

    def get_candidates(self):
        return self.candidates


if __name__ == "__main__":
    cities, data = read_data("./european_cities.csv")

    np.random.seed(1987)
    gen = GeneticTSP(500, .8, .4, data[:10,:10])

    original_cands = gen.get_candidates()
    res = [ (gen.get_distance(cand), cand) for cand in original_cands ]
    res.sort(key=lambda i: i[0])

    print(f"Original population")
    print(f"Distance: {res[0][0]}")
    plot_plan([ cities[i] for i in res[-1][1] ])

    gen.run(500)

    arr = gen.get_candidates()
    res = [ (gen.get_distance(cand), cand) for cand in arr ]
    res.sort(key=lambda i: i[0])

    print(f"Improved population")
    print(f"Distance: {res[0][0]}")
    plot_plan([ cities[i] for i in res[0][1] ])