import random
import numpy as np

from labirinth_ai.Models.Genotype import NeatLike


def fib(n):
    if n == 0:
        return [1]
    elif n < 0:
        return [0]
    else:
        return [fib(n - 1)[0] + fib(n - 2)[0]] + fib(n - 1)


class Population:
    def __init__(self, subject_class, world, subject_number):
        self.subjects = []
        self.world = world
        for _ in range(subject_number):
            px, py = self.world.generate_free_coordinates()
            self.subjects.append(subject_class(px, py, genotype_class=NeatLike))
        self.subject_number = subject_number
        self.subject_class = subject_class

    def select(self):
        ranked = list(self.subjects)
        ranked.sort(key=lambda subject: subject.accumulated_rewards, reverse=True)

        return ranked[:int(self.subject_number / 2)]

    @classmethod
    def scatter(cls, n, buckets):
        out = np.zeros(buckets)
        if n == 0:
            return out

        fib_number = 0
        fibs = fib(fib_number)
        while np.sum(fibs) <= n and len(fibs) <= buckets:
            fib_number += 1
            fibs = fib(fib_number)
        fib_number -= 1
        fibs = fib(fib_number)

        for bucket in range(buckets):
            if bucket < len(fibs):
                out[bucket] += fibs[bucket]
            else:
                break

        return out + cls.scatter(n - np.sum(fibs), buckets)

    def evolve(self):
        # get updated weights from the models
        for subject in self.subjects:
            subject.model.update_genes_with_weights()

        # crossbreed the current pop
        best_subjects = self.select()
        distribution = list(self.scatter(self.subject_number - int(self.subject_number / 2), int(self.subject_number / 2)))

        new_subjects = list(best_subjects)
        for index, offspring_num in enumerate(distribution):
            for _ in range(int(offspring_num)):
                parent_1 = best_subjects[index]
                parent_2 = best_subjects[random.randint(index + 1, len(best_subjects) - 1)]

                new_genes = parent_1.model.genes.cross(parent_2.model.genes)

                # position doesn't matter, since mutation will set it
                new_subject = self.subject_class(0, 0, new_genes)
                new_subject.history = parent_1.history
                new_subject.samples = parent_1.samples + parent_2.samples
                new_subjects.append(new_subject)

        assert len(new_subjects) == self.subject_number, 'All generations should have constant size!'

        # mutate the pop
        mutated_subjects = []
        innovation_num = max(map(lambda subject: max(map(lambda connection: connection.innvovation_num,
                                                         subject.model.genes.connections
                                                         )
                                                     )
                             , new_subjects))
        for subject in new_subjects:
            subject.accumulated_rewards = 0

            innovation_num = subject.model.genes.mutate(innovation_num)

            px, py = self.world.generate_free_coordinates()
            new_subject = self.subject_class(px, py, subject.model.genes)
            new_subject.history = subject.history
            new_subject.samples = subject.samples
            mutated_subjects.append(new_subject)

        self.subjects = mutated_subjects