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, FITN = [1,1,1], 
           NOOFFARMPOP=10,NOOFGEN=10,WANTEDSTARTFST=0.042):

  FST=[]
  STARTFST=[]

  for iter in range(0,noofiter):
  
    #create population:
    pop=sim.Population(size=[POPSIZE0]*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 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 populations:
    pop.resize([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]}),
      matingScheme=sim.RandomMating(),
      gen=NOOFGEN,
      postOps=sim.PyOperator(func=setSex))

    #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()
    
  return [FST]

#subprocedure ends here

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

for popsize in range(10,40,10):
  fst0=OneRun(noofiter=NOOFITER,POPSIZE=popsize)[0] 
  fst1=OneRun(noofiter=NOOFITER,POPSIZE=popsize,FITN=FITN1)[0]

  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)
  print popsize
  outfile.write (str(popsize) + '\t' + str(power) + '\n' )

outfile.close()