""" Classes for Population and Individual objects. """

import numpy as np
from scipy.stats import norm
from collections.abc import MutableSequence
from .settings import DOMINANCE, RECOMBINATION_RATE, LOCUS_NUM, FITNESS_FUNCTION

class Population(MutableSequence):
    """ Creates instance of Population object. individuals and generation values
    are required to initialize a Population object. """
    def __init__(self, individuals):
        self.individuals = individuals # Main list
        self._check_individuals_type()

    @property
    def size(self):
        return len(self.individuals)

    @property
    def fitness_distribution(self):
        return np.array([ind.fitness for ind in self.individuals])

    @property
    def phenotype_distribution(self):
        return np.array([ind.phenotype for ind in self.individuals])

    def make_next_generation(self):
        next_generation = []

        # Iterate for number of children individuals
        # Here we assumes population size is constant over time.
        for _ in range(len(self.individuals)):
            parent1, parent2 = self.choose_parents()
            child = Mating(parent1, parent2).reproduce_a_child()

            next_generation.append(child)

        return Population(next_generation)

    def choose_parents(self, test_neutral=False):
        # Choose without considering fitness. This if for test.
        if test_neutral:
            ixs = np.random.randint(len(self.individuals), size=2)
            return self.individuals[ixs[0]], self.individuals[ixs[1]]

        # Choose two parents randomly but weighted by their fitness_distribution
        else:
            parent1 = self.weighted_random_choice(
                self.fitness_distribution, self.individuals)
            parent2 = self.weighted_random_choice(
                self.fitness_distribution, self.individuals)

            return parent1, parent2

    @staticmethod
    def weighted_random_choice(weights, values):
        rnd = np.random.rand() * sum(weights)
        for w, v in zip(weights, values):
            rnd -= w
            if rnd < 0:
                return v

    def _check_individuals_type(self):
        assert isinstance(self.individuals, list)

    def __len__(self):
        return len(self.individuals)

    def __delitem__(self, index):
        self.individuals.__delitem__(index)

    def insert(self, index, value):
        self.individuals.insert(index, value)

    def __setitem__(self, index, value):
        self.individuals.__setitem__(index, value)

    def __getitem__(self, index):
        return self.individuals.__getitem__(index)

    def __repr__(self):
        return '< {name}: {size} individuals exist. >'\
            .format(name=type(self).__name__, 
                    size=self.size)

class Individual(object):
    """ Creates instance of Individual object which represents an individual of 
    diploid organism. genotype is required to initialize an Individual object. 
    """
    def __init__(self, genotype):
        genotype = np.array(genotype)
        self.genotype = genotype
        self.phenotype = self.genotype_to_phenotype()
        self.fitness = self.phenotype_to_fitness()

    def genotype_to_phenotype(self, mode='mean'):
        _locus_trait_contribution = lambda locus: DOMINANCE \
            if len(set(locus)) > 1 else np.mean(locus)

        if mode == 'mean':
            return np.apply_along_axis(
                _locus_trait_contribution, 0, self.genotype).mean()
        else:
            raise Exception('Other trait value calculator is not supported yet. '\
                'Please input "mean".')

    def phenotype_to_fitness(self):
        if FITNESS_FUNCTION == 'norm':
            fitness_function = norm.pdf

        return fitness_function(self.phenotype)

    def produce_gamete(self):
        """ Return gamete genotype. One of two chromosomes is drawn randomly. """
        chr_index = np.random.randint(2)

        return self.genotype[chr_index]

    def __repr__(self):
        return str(self.genotype)

class Mating(object):
    def __init__(self, parent1, parent2):
        assert type(parent1).__name__ == 'Individual'
        assert type(parent2).__name__ == 'Individual'
        self.parent1 = parent1
        self.parent2 = parent2

    def get_gametes(self):
        gamete1 = self.parent1.produce_gamete()
        gamete2 = self.parent2.produce_gamete()

        return gamete1, gamete2

    @staticmethod
    def recombination(gamete1, gamete2):
        if np.random.rand() <= RECOMBINATION_RATE:
            recomb_pos = np.random.randint(low=1, high=LOCUS_NUM)
        
        # Return input without changing anything
        else:
            return gamete1, gamete2

        rec_gamete1 = np.concatenate(
            [gamete1[:recomb_pos], gamete2[recomb_pos:]])
        rec_gamete2 = np.concatenate(
            [gamete2[:recomb_pos], gamete1[recomb_pos:]])

        return rec_gamete1, rec_gamete2

    def reproduce_a_child(self):
        gamete1, gamete2 = self.get_gametes()

        child_genotype = self.recombination(gamete1, gamete2)

        return Individual(child_genotype)

