#!/usr/bin/env python

# Packages we're going to use
import math
import numpy as np
import pandas as pd
from itertools import chain

# Abstract base class for population and pedigree
from collections import MutableMapping

# Other pydigree objects
from pydigree.individual import Individual
from pydigree.common import *
from pydigree.recombination import recombine
from pydigree.individualcontainer import IndividualContainer
from pydigree.genotypes import ChromosomeSet

missing_genotype = (0, 0)


def is_missing_genotype(g):
    return g == missing_genotype


# Population growth models.
# These are easy enough to supply your own if you wanted
def exponential_growth(p, r, t):
    """
    Models exponential growth over discrete generations.
    
    :param p: initial population
    :param r: growth rate
    :param t: number of generations
    :type p: numeric
    :type t: numeric
    :type r: numeric

    :returns: population size at time t
    :rtype: numeric
    """
    return p * math.exp(r * t)


def logistic_growth(p, r, k, t):
    """
    Models logistic growth over discrete generations.
    
    :param p: initial population
    :param r: growth rate
    :param k: final population
    :param t: number of generations

    :returns: population size at time t
    :rtype: numeric
    """
    return (p * k) / (p + (k - p) * math.exp(-r * t))


# Classes
class Population(IndividualContainer):
    # Methods for mapping types

    def __init__(self, intial_pop_size=0, name=None):
        self.chromosomes = ChromosomeSet()
        self.pool = None
        self.population = {}
        self.n0 = intial_pop_size
        self.name = name

    def __hash__(self):
        return id(self)

    def __getitem__(self, key):
        return self.population[key]

    def __contains__(self, item):
        return item in self.population.values()

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

    def __setitem__(self, key, value):
        self.population[key] = value

    def __delitem__(self, key):
        del self.population[key]

    def size(self):
        """ Returns the number of individuals in the population. """
        return len(self.population)

    def remove_ancestry(self):
        """ Makes every individual in the population a founder """
        for x in self.individuals:
            x.remove_ancestry()

    # Adding and removing people
    #
    #
    def register_individual(self, ind):
        ''' Adds an individual to the population '''
        if ind.label in self.population:
            raise ValueError('ID %s already in population!' % ind.label)
        self.population[ind.label] = ind

    def remove_individual(self, ind):
        ''' Removes an individual from the population '''
        del self[ind.label]

    def add_founders(self, n):
        for i in range(n):
            self.founder_individual(register=True)

    def update(self, other):
        '''
        Merges two datasets (i.e. performs Individual.update for each individual in the pedigree)

        Assumes unique individual IDs
        
        :param other: New data to merge in
        :type other: Population
        
        :return: void
        '''
        self.chromosomes = other.chromosomes
        self.clear_genotypes()

        selfids = {x.label for x in self.individuals}
        otherids = {x.label for x in other.individuals}
        overlap = set.intersection(selfids, otherids)

        if not overlap:
            return

        for x in overlap:
            self.population[x].update(other._getindividual(x))

    def _getindividual(self, label):
        return self[label]

    @property
    def individuals(self):
        ''' Returns a list of individuals in the population '''
        return [x for x in self.population.values()]


    # Chromosome functions
    #
    #
    def add_chromosome(self, chrom):
        """ Adds a chromosome to the population """
        self.chromosomes.add_chromosome(chrom)

    def chromosome_count(self):
        """ Returns the number of chromosomes in this population """
        return len(self.chromosomes)

    # Random mating
    #
    #
    def mate(self, ind1, ind2, id, sex=None):
        """
        Creates an individual as the child of two specificied individual
        objects and randomly chooses a sex.

        :param ind1: The first parent
        :param ind2: The second parent
        :type ind1: Individual
        :type ind2: Individual
        :param id: ID label for the child
        :param sex: Sex of child, randomly chosen if not specified
        :type sex: {0,1}
        :return: An individual with ind1 and ind2 as parents
        :rtype: Individual
        """
        if sex is None:
            sex = np.random.choice([0, 1])
        child = Individual(self, id, ind1, ind2, sex)
        return child

    def advance_generation(self, gensize):
        '''
        Simulates a generation of random mating.

        :param gensize: The size of the new generation  
        :type gensize: numeric
        '''

        males = self.males()
        females = self.females()

        fathers = np.random.randint(0, len(males), gensize)
        mothers = np.random.randint(0, len(females), gensize)
        sexes = np.random.randint(0, 2, gensize)

        newpop = {}
        for indidx in range(gensize):
            dadidx, momidx, sex = fathers[indidx], mothers[indidx], sexes[indidx]
            newind = self.mate(males[dadidx], females[momidx], indidx, sex=sex)
            newpop[newind.label] = newind

        self.population = newpop

    def founder_individual(self, register=True, sex=None):
        "Creates a new founder individual and adds to the population"

        if sex is not None:
            sex = sex.lower()
        sexd = {'m': 0, 'f': 1, None: np.random.choice([0, 1])}
        i = Individual(self, self.size(), None, None, sexd[sex])
        if register:
            self.register_individual(i)
        return i

    # Genotype Functions
    #
    #
    def get_founder_genotypes(self):
        ''' 
        Gives genotypes to each founder in the population with chromosomes 
        from the chromosome pool. If there is no pool, genotypes are generated
        under linkage equilibrium
        '''
        for ind in self.individuals:
            if not self.pool:
                genotypes = self.get_linkage_equilibrium_genotypes()
            else: 
                genotypes = self.pool.get_genotype_set()

            ind.genotypes = genotypes

    def get_genotypes(self):
        ''' 
        Causes each Individual object in the pedigree to request genotypes
        from its parents
        '''
        for x in self.individuals:
            x.get_genotypes()

    def get_linkage_equilibrium_genotypes(self):
        '''
        Returns a set of genotypes for an individual in linkage equilibrium
        '''
        return [[c.linkageequilibrium_chromosome(),
                 c.linkageequilibrium_chromosome()]
                for c in self.chromosomes]

    def ld(self, locusa, locusb, allelea=None, alleleb=None, method='r2'):
        # """
        # Returns a measure of linkage disquilibirum between alleles at two
        # loci, for methods in {D, Dprime, r2}. If alleles aren't specified,
        # the major alleles at each locus are chosen.

        # LD metrics are functions of D, the measure of deviance from
        # independence in sampling alleles. D = x11 - p1q1, where x11 is the
        # frequency of allele 1 at locus p occuring with allele 1 at locus q.

        # Lewontin's D' is a normalization of D, since D is dependent on
        # frequency.

        # D' = D / Dmax, where:
        # Dmax = min(p1*q1, p2*q2) if D > 0
        # Dmax = min(p2*q1,p1*q2) if D < 0

        # r2 is a third metric, equal to
        # D**2 / (p1 * p2 * q1 * q2)
        # r2 ranges from 0 to 1

        # Alleles are in linkage equilibrium if any of these metrics are 0.

        # Arguments
        # -----
        # locusa:  The first locus
        # locusb:  The second locus
        # allelea: The allele at the first locus
        #          (the major allele if not specified)
        # alleleb: The allele at the second locus
        #          (the major allele if not specified)
        # method:  The metric of linkage disequilibrium
        #          (r2 [default], Dprime, or D)

        # Returns: a double
        # """
        raise NotImplementedError('LD requires phase known data')
        method = method.lower()
        if method not in set(['r2', 'dprime', 'd']):
            raise ValueError("LD metric must be one of r2, Dprime or D")
        if not allelea:
            allelea = self.major_allele(locusa)
        if not alleleb:
            alleleb = self.major_allele(locusb)
        pop = [x for x in self.individuals if
               (not is_missing_genotype(x.get_genotype(locusa))) and
               (not is_missing_genotype(x.get_genotype(locusb)))]
        # Allele 1
        p1 = self.allele_frequency(locusa, allelea)
        p2 = 1 - p1
        # Allele 2
        q1 = self.allele_frequency(locusa, allelea)
        q2 = 1 - q1
        x = sum(1 for x in pop
                if x.has_allele(locusa, allelea) and
                x.has_allele(locusb, alleleb))
        D = x - p1 * q1
        # All of these metrics are normalizations of D by dividing it by
        # some expression. If D==0 all of them will equal 0.
        if method == 'd':
            return D
        elif method == 'dprime':
            if D > 0:
                Dmax = min(p1 * q1, p2 * q2)
            else:
                Dmax = min(p2 * q1, p1 * q2)
            return D / Dmax
        elif method == 'r2':
            return (D ** 2) / (p1 * p2 * q1 * q2)


