import random

# We're gonna use relative selection here.
# The first value is for homozygous 0, the middle is for heterozygous, and the third is for homozygous 1
# so if s = [0.5, 1, 1], homozygous 0 individuals reproduce half as often as heterozygotes or homozygous 1, and het and hom 1 individuals reproduce equally as often (i.e. 1 is dominant)
# or if s = [0.5, 1, 1.5] hom 0 reproduce half as often as hets and one third as often as hom 1 individuals, and hets reproduce 2/3rds as often as hom 1
s = [0.5, 1, 1]

# we'll use a constant population size
popsize = 1000
# randomize the initial population to be all homozygous 0 for the "trait" under selection, and random in the non-adaptive allele
pop = {'0': [((0,0), (random.randint(0,4),random.randint(0,4))) for i in range(popsize-1)],
       'het' : [],
       '1' : [((1,1), (random.randint(0,4),random.randint(0,4)))]}

initmut = pop['1']
print('initial mutant', initmut)
# selects a type (homozygous 0, heterozygous, homozygous 1) of parent based on threshold values calulated based on s above and
# the number of each such individual present in the population to choose from
def selParentType(thresh):
    sel = random.uniform(0, thresh['end'])
    if sel <= thresh['hom0']:
        return '0'
    elif sel <= thresh['het']:
        return 'het'
    else:
        return '1'

def printohist(pop):
    otherhist = dict()
    allelecounts = [0 for i in range(5)]
    for i in range(5):
        for j in range(5):
            otherhist[tuple(sorted([i,j]))] = 0

    for t in pop:
        for i in pop[t]:
            otherhist[tuple(sorted(i[1]))] = otherhist[tuple(sorted(i[1]))]+1
            allelecounts[i[1][0]] += 1
            allelecounts[i[1][1]] += 1
    print('genotype counts')
    for k,v in otherhist.items():
        print(k,v)
    print('allele counts')
    for i,v in enumerate(allelecounts):
        print(i,v)


for i in range(100):
    if i % 10 == 0:
        print('iter {}\ndistribution of trait under selection (genotypes)\n(0,0) {}\n(0,1) {}\n(1,1) {}'.format(i,len(pop['0']),len(pop['het']),len(pop['1'])))
        print('distribution of trait *not* under selection')
        printohist(pop)

    if len(pop['0']) == popsize:
        print('oops! drift happened and the mutant allele went extinct. run again and/or increase popuation size')
        break
    # we'll fill this dictionary in a second for the new population
    newpop = {'0':[], 'het':[], '1':[]}

    # these values will be used to create intervals needed to select individuals randomly with likelihood
    # proportionate to their selective advantage to breed
    hom0sum = s[0] * len(pop['0'])
    hetsum = s[1] * len(pop['het'])
    hom1sum = s[2] * len(pop['1'])

    # by making thresholds, we can just select a random number in the range [0, endthresh] and if it is less than
    # thresh['hom0'], we randomly pick a homozygous 0 individual as a parent, if its greater than thresh['hom0'] and less than
    # thresh['het'] we choose a heterozygous individual, and if it is greater than thresh['het'] we choose a homozygous 1 individual
    thresh = dict()
    thresh['hom0'] = hom0sum
    thresh['het'] = hetsum+thresh['hom0']
    thresh['end'] = hom1sum+thresh['het']
    for j in range(popsize):
        p1 = random.choice(pop[selParentType(thresh)])
        p2 = random.choice(pop[selParentType(thresh)])

        # make a new individual by randomly chosing which allele from each parent is passed on
        newind = ((random.choice(p1[0]), random.choice(p2[0])), (random.choice(p1[1]), random.choice(p2[1])))

        if newind[0] == (0,0):
            newpop['0'].append(newind)
        elif newind[0] == (1,1):
            newpop['1'].append(newind)
        else:
            newpop['het'].append(newind)
    pop = newpop
print('initial mutant. note that the second pair of values (trait *not* under selection) are probably not over-represented in the final population')
print(initmut)