import simuPOP as sim
import random
import numpy
from scipy import stats


NOOFITER = 1000
ALPHA = 0.05  #alpha for declaring null hypothesis (no selection) not valid
FITN1 = [1,0.8,0.6] #fitness under selection

def setSex(pop):
    for idx,ind in enumerate(pop.individuals()):
        ind.setSex(sim.MALE if idx % 2 == 0 else sim.FEMALE)
    return True

#start of main subprocedure:
def OneRun(noofiter, POPSIZE=50,POPSIZE0=50, POPSIZE2=500, FITN = [1,1,1], NOOFWILDPOP=10,
           NOOFFARMPOP=10,NOOFGEN=10,WANTEDSTARTFST=0.042):

  FST=[]
  STARTFST=[]

  for iter in range(0,noofiter):
  
    #create population:
    pop=sim.Population(size=[POPSIZE0]*(NOOFWILDPOP+NOOFFARMPOP),ploidy=2,loci=1,infoFields=['fitness'])
    #assign sex:
    sim.initSex(pop, sex=[sim.MALE,sim.FEMALE]) 

    #assign random start freq. in source population:
    ok=False
    while not ok:  #if ok then start-fst has been successfully generated 
      startfreq=random.random()
      sim.initGenotype(pop,freq=[startfreq,1-startfreq])
      fst=0
      k=0
      while fst<(WANTEDSTARTFST + 0.0)/100 and k<500:
        k=k+1
        pop.evolve(
          matingScheme=sim.RandomMating(),
          gen=1,
          postOps=sim.PyOperator(func=setSex))
        sim.stat(pop,structure=0,vars=['F_st'])
        fst=pop.dvars().F_st
      if k<500: ok=True 

    STARTFST.append(fst)

    #resize wild populations (so that they become large and experience little drift):
    pop.resize([POPSIZE2]*NOOFWILDPOP+ [POPSIZE]*NOOFFARMPOP, propagate=True)
  
    #apply selection
    pop.evolve(
      preOps=sim.MapSelector(loci=0, fitness={(0,0):FITN[0], (0,1):FITN[1], (1,0):FITN[1],(1,1):FITN[2]},
        subPops=range(NOOFWILDPOP,NOOFWILDPOP+NOOFFARMPOP)),
      matingScheme=sim.RandomMating(),
      gen=NOOFGEN,
      postOps=sim.PyOperator(func=setSex))

    #resize wild populations again:
    pop.resize([POPSIZE]*(NOOFWILDPOP+NOOFFARMPOP),propagate=True)

    #merge populations:
    pop.mergeSubPops(subPops=range(0,NOOFWILDPOP),name='wild')
    pop.mergeSubPops(subPops=range(1,len(pop.subPopNames())),name='farm')

    #calculate Fst:
    sim.stat(pop,structure=0,vars=['F_st'])
    FST.append(pop.dvars().F_st)
    #has now appended one Fst (one iteration)
 
  FST.sort()
  FST.reverse()
  
  return [FST, STARTFST]

#subprocedure ends here

outfile=open('fst.txt','w')

for popsize in range(10,20,10):
  [fst0, startfst0] = OneRun(noofiter=NOOFITER,POPSIZE=popsize) 
  [fst1, startfst1] = OneRun(noofiter=NOOFITER,POPSIZE=popsize,FITN=FITN1)  
  x_alpha= fst0[int(round(NOOFITER*ALPHA))]  
  n=0
  for fst in fst1:
    if fst>x_alpha:
      n = n + 1
  power = (n + 0.0)/(NOOFITER + 0.0)
  outfile.write (str(popsize) + '\t' + str(power) + '\t' + str(numpy.mean(startfst0)) + '\t' + str(numpy.mean(startfst1)) + '\n')

outfile.close()