Simulation.py

'''
Copyleft Oct 10, 2015 Arya Iranmehr, PhD Student, Bafna's Lab, UC San Diego,  Email: airanmehr@gmail.com
'''

from __future__ import division

import numpy as np;
import pandas as pd;
np.set_printoptions(linewidth=140, precision=5, suppress=True)
import subprocess, uuid, os,sys
import pylab as plt
import UTILS.Util as utl

stdout_old=sys.stdout;sys.stdout=open('/dev/null','w');import simuPOP as sim;sys.stdout=stdout_old # to avoid simuPop welcome message!


def sig(x): return 1./(1+np.exp(-x));
def logit(p): return (np.inf if p==1  else  np.log(p/(1.-p)))




a='';


def fff(msg):
    global a
    a += msg


class MSMS:
    @staticmethod
    def Simulate(n=200, mu=2*1e-9, L=50000, Ne=1e6,r=1e-9,verbose=False,seed=None,intPos=False):
        L=int(L)
        a= MSMS.Song(F=n, mu=mu, L=L, Ne=Ne, r=r,verbose=verbose,seed=seed)
        c=pd.Series(a.columns)
        if c.round().value_counts().max()==1:
            a.columns=c.round().astype(int)
        elif c.astype(int).value_counts().max()==1:
            a.columns = c.astype(int)
        if intPos:
            a.columns=map(int,np.sort(np.random.choice(L, a.shape[1], replace=False)))
        return a

    @staticmethod
    def Song(F=200, mu=2*1e-9, L=50000, Ne=1e6,r=4e-9, uid=None, theta=None, msmsFile=None, dir=None,verbose=False,seed=None):
        """
        Everything is exactly the sam
        """
#         print 'mu: {} r:{} NE:{} ,theta={} '.format(mu,r,Ne,4*Ne*mu*L), theta
        if msmsFile is not None:
            pop=MSMS.load(filename=msmsFile)[0]
        else:
            if theta:
                pop=MSMS.MSMS(n=F, numReps=1, theta=theta, rho=2*Ne*(L-1)*r, L=L, Ne=Ne, uid=uid, dir=dir,verbose=verbose,seed=seed)[0]
            else:
                pop=MSMS.MSMS(n=F, numReps=1, theta=2*Ne*mu*L, rho=2*Ne*(L-1)*r, L=L, Ne=Ne, uid=uid, dir=dir,verbose=verbose,seed=seed)[0]
        pop.r=r
        pop.Ne=Ne
        pop.L=L
        return pop
    
    @staticmethod
    def MSMS(n, numReps, theta, rho, L, Ne=None,uid=None,oneMutationEvery=None, dir=dir,verbose=False,seed=None):
        """
        Returns a list of dataframe for each replicate
        """
        if dir is None:
            dir= utl.PATH.simout;dir+= 'msms/';
        os.system('mkdir -p ' +dir)
        if oneMutationEvery is not None:
            nSS=L/oneMutationEvery
            theta=nSS/sum(1./np.arange(1,n))
        if uid is None:
            uid=str(uuid.uuid4())
        unique_filename = dir+uid+'.msms'
        if seed is None:
            seed=''
        else:
            seed=' -seed {} '.format(seed)
        cmd="java -jar -Xmx2g ~/bin/msms/lib/msms.jar -ms {} {} -t {:.0f} -r {:.0f} {:.0f} -oFP 0.000000000000E00 {} > {}".format(n, numReps, theta, rho, L, seed,unique_filename)
        if verbose:
            print cmd
        subprocess.call(cmd,shell=True)
        return MSMS.load(unique_filename)

    @staticmethod
    def getSeed(filename):
        file=open(filename);cmd=np.array(file.readline().strip().split(' '));seed=file.readline().strip()
        return seed

    @staticmethod
    def load(filename):
        n, R, L, posUnderSelection = MSMS.getParams(open(filename).readline())
        lines=np.array(map(str.strip,open(filename).readlines()) )
        posIdx= np.where(map(lambda x: x[:len('positions:')]=='positions:',lines))[0]
        try:
            theta = lines[np.where(map(lambda x: 'ThetaW Estimate Summaray:' in x, lines))[0][0]].split(':')[1].strip()
        except:
            theta = None
        POS=[map(lambda x: (float(x)*L), lines[ii].split()[1:]) for ii in posIdx]
        dfs=[pd.DataFrame(map(list ,lines[i +1 +range(n)]),columns=pos ) for i,pos in zip(posIdx,POS)]
        for df in dfs:
            df[df!='0']=1
            df[df=='0']=0
            df.L = L
            if posUnderSelection is not None:
                df.posUnderSelection = posUnderSelection * L
            if theta is not None:
                df.stat = pd.Series(theta.split(), index=['W', 'Pi', 'D']).astype(float)
        return  dfs

    @staticmethod
    def getParams(line):
        """
        Args:
            params: takes the first line of msmsm file

        Returns:
            n,R,L: number of individuals in the sample, the number of the replicates, genome length

        """
        params=np.array(line.strip().split(' '))
        offset=np.where(map(lambda x: 'ms'in x,  params))[0][0]
        if params[offset+1] == '-N':
            i=3
        else:
            i=1
        posUnderSelection = None
        if '-Sp' in params:  posUnderSelection = float(params[np.where(params == '-Sp')[0][0] + 1])
        return int(params[offset + i]), int(params[offset + i + 1]), int(
                params[np.where(params == '-r')[0][0] + 2]), posUnderSelection

    @staticmethod
    def fixDuplicatePositions(pos,L):
        pos=pd.Series(range(len(pos)),index=pos)
        posHits=pos.index.value_counts()
        invalidPOS=posHits[posHits>1]
        if not invalidPOS.shape[0]:
            return pos.index.values
        for invalidPos in invalidPOS.index:
            mini=pos.loc[invalidPos].min()
            maxi=pos.loc[invalidPos].max()
            lowerBound=pos[pos==mini-1].index.max()
            upperBound=pos[pos==maxi+1].index.min(); 
            if maxi==pos.shape[0]-1: upperBound=L
            if mini==0: lowerBound=0
            validRange=np.arange((upperBound-lowerBound)/2) # only second and third quartiles,
            offset=validRange+validRange.shape[0]/2 # first qunatulw
            newPos=pos.index.values;
            newPos[mini:maxi+1]=np.sort(np.random.choice(offset,pos.loc[invalidPos].shape[0],replace=False))+lowerBound
            pos.index=newPos
        assert pos.index.value_counts().max()==1
        return pos.index.values

    @staticmethod
    def Selection(msms, Ne, n, numReplicates, theta, rho, window_size, s, origin_count, posUnderSelection, gens, path):
        seed = ''
        for ii, gen in enumerate(gens):
            fname = path + '{}.msms'.format(int(gen))
            if (not ii) and s != 0:
                # while (nu0 < 0.95) or (nu0 > 0.99):
                cmd = "{} -N {} -ms {} {} -t {} -r {} {:.0f} -SAA {} -SaA {} -SI {} 1 {} -Sp {} -oOC -Smark -oFP 0.000000000000E00 {}  -SForceKeep -SFC   -oTW  >{}".format(
                        msms, Ne, n, numReplicates, theta, rho, window_size, 2 * Ne * s, Ne * s, gen / (4. * Ne),
                                                                             origin_count / Ne,
                        posUnderSelection, ('-seed {}'.format(seed), '')[seed is ''], fname)
                os.system(cmd)
            else:
                cmd = "{} -N {} -ms {} {} -t {} -r {} {:.0f} -SAA {} -SaA {} -SI {} 1 {} -Sp {} -oOC -Smark -oFP 0.000000000000E00 {}  -SFC -SForceKeep  -oTW  >{}".format(
                        msms, Ne, n, numReplicates, theta, rho, window_size, 2 * Ne * s, Ne * s, gen / (4. * Ne),
                                                                             origin_count / Ne,
                        posUnderSelection, ('-seed {}'.format(seed), '')[seed is ''], fname)
                os.system(cmd)
            if not ii: seed = MSMS.getSeed(fname)

    @staticmethod
    def SelectionFinale(msms, Ne, n, numReplicates, theta, rho, window_size, s, origin_count, posUnderSelection, gens,
                            path):
        seed = ''
        nu0 = 0
        for ii, gen in enumerate(gens):
            fname = path + '{}.msms'.format(int(gen))
            if (not ii) and s != 0:
                while (nu0 < 0.9):
                    cmd = "{} -N {} -ms {} {} -t {} -r {} {:.0f} -SAA {} -SaA {} -SI {} 1 {} -Sp {} -oOC -Smark -oFP 0.000000000000E00 {}  -SForceKeep -SFC   -oTW  >{}".format(
                            msms, Ne, n, numReplicates, theta, rho, window_size, 2 * Ne * s, Ne * s, gen / (4. * Ne),
                                                                                 origin_count / Ne,
                            posUnderSelection, ('-seed {}'.format(seed), '')[seed is ''], fname)
                    os.system(cmd)

                    nu0 = MSMS.load(fname)[0].mean(0).loc[25000]
            else:
                cmd = "{} -N {} -ms {} {} -t {} -r {} {:.0f} -SAA {} -SaA {} -SI {} 1 {} -Sp {} -oOC -Smark -oFP 0.000000000000E00 {}  -SFC -SForceKeep  -oTW  >{}".format(
                        msms, Ne, n, numReplicates, theta, rho, window_size, 2 * Ne * s, Ne * s, gen / (4. * Ne),
                                                                             origin_count / Ne,
                        posUnderSelection, ('-seed {}'.format(seed), '')[seed is ''], fname)
                os.system(cmd)
            if not ii: seed = MSMS.getSeed(fname)

    @staticmethod
    def SelectionNu(msms, Ne, n, numReplicates, theta, rho, window_size, s, posUnderSelection, nu, path=None):
        seed = ''
        if path is None: path = '~/tmp.msms'
        fname = path + '{}.msms'.format(nu)
        cmd = "{} -N {} -ms {} {} -t {} -r {} {:.0f} -SAA {} -SaA {} -SF 0 {}  -Sp {} -oOC -Smark -oFP 0.000000000000E00 {}  -SFC   -oTW  >{}".format(
                msms, Ne, n, numReplicates, theta, rho, window_size, 2 * Ne * s, Ne * s, nu, posUnderSelection,
                ('-seed {}'.format(seed), '')[seed is ''], fname)
        print cmd
        os.system(cmd)
        return MSMS.load(fname)

    @staticmethod
    def SelectionNuForward(msms, Ne, n, numReplicates, theta, rho, window_size, s, origin_count, posUnderSelection,
                               gens, path):
        nu0 = 0
        for ii, gen in enumerate(gens):
            fname = path + '{}.msms'.format(gen)
            if (not ii) and s != 0:
                while (nu0 < 0.95) or (nu0 > 0.99):
                    cmd = "{} -N {} -ms {} {} -t {} -r {} {:.0f} -SAA {} -SaA {} -SI {} 1 {} -Sp {} -oOC -Smark -oFP 0.000000000000E00 {}  -SFC   -oTW  >{}".format(
                            msms, Ne, n, numReplicates, theta, rho, window_size, 2 * Ne * s, Ne * s, gen / (4. * Ne),
                                                                                 origin_count / Ne,
                            posUnderSelection, ('-seed {}'.format(seed), '')[seed is ''], fname)
                    os.system(cmd)
                    nu0 = MSMS.load(fname)[0].mean(0).loc[25000]
            print nu0, gen, cmd
            if not ii: seed = MSMS.getSeed(fname)


class Simulation:
    @staticmethod
    def setSeed(seed):
        if seed is None: return
        sim.setRNG('rand', seed + 1);
        np.random.seed(seed)



    @staticmethod
    def load(ExperimentName, s=0.1, L=50000, experimentID=0, nu0=0.005, isFolded=False, All=False, startGeneration=0,
             maxGeneration=50, numReplicates=3, numSamples=5, step=10, replicates=None, coverage=np.inf):
        path='{}{}/simpop/'.format(utl.PATH.simout, ExperimentName) + Simulation.getSimulationName(s=s, L=L, experimentID=experimentID, initialCarrierFreq=nu0, isFolded=isFolded) + '.pkl'
        sim= pd.read_pickle(path)
        sim.savedPath=path
        if replicates is not None:          sim.setReplicates(sorted(replicates))
        elif numReplicates is not None:     sim.setReplicates(range(numReplicates))

        if coverage != np.inf:
            sim.Xi = sim.X
            sim.X = sim.C.loc[coverage] / sim.D.loc[coverage].astype(float)
            sim.X = np.array(map(lambda x: utl.roundto(x, 5), sim.X.reshape(-1) * 1e4)).reshape(sim.X.shape) / 1e4
            sim.CD=sim.getCD(coverage)
            sim.CD.columns.names=['REP','GEN','READ']

        if not All: sim.setSamplingTimes(maxGeneration=min(maxGeneration,sim.getGenerationTimes()[-1]),numSamples=numSamples,step=step,startGeneration=startGeneration)
        return sim

    @staticmethod
    def getSimulationName(s,L,experimentID,initialCarrierFreq,isFolded,msms=False):
        if msms:
            return 'L{:.0f}K.{:04.0f}'.format(L/1000,experimentID)
        if s:
            return 'Nu{:E}.s{:E}.L{:.0f}K.{:04.0f}{}'.format(np.round(float(initialCarrierFreq), 3), s, L / 1000,
                                                             experimentID, ('', '.Folded')[isFolded])
        else:
            return 'Nu{:E}.s{:E}.L{:.0f}K.{:04.0f}{}'.format(0, s * 100, L / 1000, experimentID,
                                                             ('', '.Folded')[isFolded])

    def setReplicates(self,replicates):
        self.numReplicates=len(replicates)
        self.X=self.X[:,:,replicates]
        self.C = self.C.apply(lambda x: x[:, :, replicates])
        self.D = self.D.apply(lambda x: x[:, :, replicates])

    def __init__(self, outpath=utl.PATH.simout, N=1000, generationStep=10, maxGeneration=None,
                 s=0.05, r=4e-9, Ne=1e6, mu=2e-9, F=200, h=0.5, L=50000, startGeneration=0, numReplicates=3, H0=None,
                 foldInitialAFs=False, save=True, foutName=None,
                 doForwardSimulationNow=True, experimentID=-1,
                 msmsFile=None,initialCarrierFreq=0, ExperimentName=None, simulateNeutrallyFor=0,
                 initialNeutralGenerations=0, ignoreInitialNeutralGenerations=True,
                 makeSureSelectedSiteDontGetLost=True, onlyKeep=None, verbose=0, sampingTimes=None, minIncrease=0,
                 model=None,initDiploidPop=None,posUnderSelection=-1,haplotypes=False,seed=None,recombinator=None
                 ):
        """
        A General Simulation Class; with params
        H0: Dataframe F x m for F individuals and m segregation sites ;  Initial Haplotypes; dataframe with columns as positions 
        """
        self.recombinator=recombinator
        if seed is not None:
            Simulation.setSeed(seed)
        self.s = s;
        self.r = r;
        self.Ne = Ne;
        self.mu = mu;
        self.F = F;
        self.h = h;
        self.L = int(L);
        self.startGeneration = startGeneration;
        self.numReplicates = numReplicates;
        self.posUnderSelection = -1
        self.initDiploidPop = initDiploidPop
        self.initialCarrierFreq= initialCarrierFreq if initialCarrierFreq else 1./self.F
        if foutName is not None:
            self.uid=foutName
            self.uidMSMS=None
        elif experimentID>=0:
            self.uid=Simulation.getSimulationName(self.s, self.L, self.experimentID, initialCarrierFreq=self.initialCarrierFreq, isFolded=self.foldInitialAFs)
            self.uidMSMS=Simulation.getSimulationName(self.s, self.L, self.experimentID, initialCarrierFreq=self.initialCarrierFreq, isFolded=self.foldInitialAFs,msms=True)
        else:
            self.uid=str(uuid.uuid4())
            self.uidMSMS=self.uid


        if H0 is None:
            self.simulateH0()
            H0=self.H0
        else:
            self.setH0(H0);

        if posUnderSelection >= 0:
            if self.positions is None:
                self.positions=map(int, self.initDiploidPop.lociPos())
            self.set_posUnderSelection(posUnderSelection)
        assert ExperimentName != None
        self.save=save
        self.model=model
        self.minIncrease = minIncrease
        self.samplingTimes=sampingTimes
        self.initialNeutralGenerations=initialNeutralGenerations
        self.onlyKeep=onlyKeep
        self.makeSureSelectedSiteDontGetLost=makeSureSelectedSiteDontGetLost
        self.ignoreInitialNeutralGenerations=ignoreInitialNeutralGenerations
        self.msmsFile=msmsFile;self.outpath=outpath; self.outpath=outpath ; self.N=N; self.generationStep=generationStep; self.maxGeneration= maxGeneration;
        self.foldInitialAFs=foldInitialAFs;self.doForwardSimulationNow=doForwardSimulationNow;self.experimentID=experimentID
        self.simulateNeutrallyFor=simulateNeutrallyFor

        self.setH0(H0);
        if not os.path.exists(self.outpath) : os.makedirs(self.outpath)
        self.outpath+=ExperimentName
        if not os.path.exists(self.outpath) : os.makedirs(self.outpath)
        self.outpathmsms=self.outpath+'/msms/';self.outpath+='/simpop/'
        if not os.path.exists(self.outpath) : os.makedirs(self.outpath)
        if not os.path.exists(self.outpathmsms) : os.makedirs(self.outpathmsms)
        if self.maxGeneration is None: self.maxGeneration=Simulation.getFixationTime(self.s, Ne=self.F, roundto10=True)
        self.theta=2*self.Ne*self.mu*self.L
        self.pops=[]

        if self.model is None:
            import simuPOP.demography as dmg
            self.model=dmg.LinearGrowthModel(T=self.maxGeneration, N0=self.N, NT=self.N)

        if self.doForwardSimulationNow:
            self.forwardSimulation()

    @staticmethod
    def simulateSingleLoci(nu0=0.005, T=100, s=0.1, N=1000,verbose=True,h=0.5,seed=None):
        if verbose:
            print '.',
        step = 1
        Simulation.setSeed(seed)
        pop = sim.Population(size=N, ploidy=2, loci=[1],infoFields=['fitness']);sim.initGenotype(pop, prop=[1-nu0,nu0]);simulator = sim.Simulator(pop.clone(), rep=1);
        # sim.stat(pop, alleleFreq=[0]);        print pop.dvars().alleleFreq[0][1]
        global a;a = "0;;{}\n".format(nu0)
        simulator.evolve(initOps=[sim.InitSex()],
                         preOps=sim.MapSelector(loci=0, fitness={(0, 0): 1, (0, 1): 1 + s *h, (1, 1): 1 + s}),
                         matingScheme=sim.RandomMating(), postOps=[sim.Stat(alleleFreq=[0], step=step),
                                                                   sim.PyEval("'%d;;' % (gen+1)", reps=0, step=step,
                                                                              output=fff), sim.PyEval(
                    r"'{}\n'.format(map(lambda x: round(x[1],5),alleleFreq.values())[0])", step=step, output=fff)],
                         gen=T)
        return pd.DataFrame(zip(*map(lambda x: x.split(';;'), a.strip().split('\n')))).T.set_index(0)[1].astype(float)



    def createInitialDiploidPopulation(self):
        """
        initHaps : np 2D array which m x nSS where m i number of individual haps and nSS is number of SS
        return a homozygote diploid population which every haplotype is copied n times
        """
        if self.initDiploidPop is not None: return self.initDiploidPop
        assert int(2*self.N/self.F)==2*self.N/float(self.F)  # N should be  a multiplier of F
        nSS=self.H0.shape[1];n=int(self.N/self.F)
        try:
            pop = sim.Population(size=self.N, ploidy=2, loci=nSS,lociPos=list(self.positions), infoFields='fitness')
        except:
            import traceback
            print(traceback.format_exc())
            print list(self.positions), nSS,n,self.H0.shape[0]
            exit()
        assert (self.N % self.H0.shape[0]) ==0
        H= [[list(h.values),list(h.values)] for _ in range(n) for _,h in self.H0.iterrows()]
        for (i,h) in zip(pop.individuals(),H): # for each indv assing first and second chromosome
            i.setGenotype(h[0],0 );i.setGenotype(h[1],1 ) #homozygote population of diploid
        # sim.stat(pop, alleleFreq=range(nSS));print np.array([pop.dvars().alleleFreq[x][1] for x in range(nSS)])
        return pop

    @staticmethod
    def getGT(pop, i=None, pos=None):
        if i == None and pos == None:
            df = pd.concat([pd.DataFrame([list(i.genotype(0)) for i in pop.individuals()]),
                            pd.DataFrame([list(i.genotype(1)) for i in pop.individuals()])],
                           keys=[0, 1]).sort_index().reorder_levels([1, 0]).sort_index()
            df.columns = map(int, pop.lociPos())
            return df
        i = np.where(np.array(pop.lociPos()).astype(int) == pos)[0][0]
        a, b = [], []
        for ind in pop.individuals():
            a += [ind.genotype(0)[i]]
            b += [ind.genotype(1)[i]]
        return pd.concat([pd.Series(a), pd.Series(b)], keys=[0, 1]).reorder_levels([1, 0]).sort_index()
    @staticmethod
    def createDiploidPopulationFromDataFrame(df):
        """
        initHaps : np 2D array which m x nSS where m i number of individual haps and nSS is number of SS
        return a homozygote diploid population which every haplotype is copied n times
        """
        pop = sim.Population(size=df.shape[0]/2, ploidy=2, loci=df.shape[1], lociPos=list(df.columns), infoFields='fitness')
        for j,i in enumerate(pop.individuals()): # for each indv assing first and second chromosome
            i.setGenotype(df.loc[j].loc[0].tolist(),0 );i.setGenotype(df.loc[j].loc[1].tolist(),1 )
        return pop

    @staticmethod
    def _simualtePop(pop, s=0, h=0.5, r=2e-8, siteUnderSelection=0,gen=1,recombinator=None,seed=None):
        "Gets population and returns population"
        Simulation.setSeed(seed)
        simulator = sim.Simulator(pop.clone(), rep=1)
        if recombinator is None:recombinator=sim.Recombinator(intensity=r)
        simulator.evolve(
            initOps=[sim.InitSex()],
            preOps=sim.MapSelector(loci=siteUnderSelection, fitness={(0, 0): 1, (0, 1): 1 + s * h, (1, 1): 1 + s}),
            matingScheme=sim.RandomMating(ops=recombinator),
            gen=gen)
        return simulator.population(0).clone()

    @staticmethod
    def _simualte(pop,s,h,r,siteUnderSelection,positions,startGeneration,generationStep,maxGeneration,model=None,makeSureSelectedSiteDontGetLost=True):
        "Gets population and returns Dataframe, Static method"
        N = int(pop.popSize())
        if model is None:
            import simuPOP.demography as dmg
            model = dmg.LinearGrowthModel(T=maxGeneration, N0=N, NT=N)
        simulator = sim.Simulator(pop.clone(), rep=1)
        global a;a = ""
        pops=[]
        step=1# this is slow but safe, dont change it
        simulator.evolve(
            initOps=[sim.InitSex()],
            preOps=sim.MapSelector(loci=siteUnderSelection, fitness={(0, 0): 1, (0, 1): 1 + s * h, (1, 1): 1 + s}),
            matingScheme=sim.RandomMating(ops=sim.Recombinator(intensity=r),subPopSize=model),
            postOps=[sim.Stat(alleleFreq=range(int(pop.numLoci()[0])), step=step), sim.PyEval("'Gen %4d;;' % (gen+1)", reps=0,step= step, output=fff), sim.PyEval(r"'{},'.format(map(lambda x: round(x[1],5),alleleFreq.values()))", step=step, output=fff),sim.PyOutput('\n', reps=-1, step=step, output=fff)],
            gen = maxGeneration)
        # idx=np.arange(self.generationStep-1,self.maxGeneration,self.generationStep)+self.initialNeutralGenerations
        print a
        _,data=zip(*map(lambda x: x.split(';;'),a.strip().split('\n')))
        data=np.array(map(eval,data))[:,0,:]
        print data
        # if data[-1, self.siteUnderSelection] >= self.initialCarrierFreq + self.minIncrease or self.s == 0 or not self.makeSureSelectedSiteDontGetLost:
        if data[-1, siteUnderSelection] or s == 0 or not makeSureSelectedSiteDontGetLost:
            try:
                pops+=[simulator.extract(0) ]
            except:
                print 'Error'
            return data[int(startGeneration/generationStep):,:]
        else:
            return Simulation._simualte()


    def simualte(self):
        "Gets population and returns Dataframe, Class method"
        import simuPOP.demography as dmg
        # model=dmg.ExponentialGrowthModel(T=50, N0=1000, NT=200)
        simulator = sim.Simulator(self.initDiploidPop.clone(), rep=1)
        # sim.dump(self.initDiploidPop)
        global a;a = ""
        if self.recombinator is None:
            self.recombinator=sim.Recombinator(intensity=self.r)
        step=1# this is slow but safe, dont change it
        simulator.evolve(
            initOps=[sim.InitSex()],
            preOps=sim.MapSelector(loci=self.siteUnderSelection, fitness={(0,0):1, (0,1):1+self.s*self.h, (1,1):1+self.s}),
            matingScheme=sim.RandomMating(ops=self.recombinator,subPopSize=self.model),

            postOps=[sim.Stat(alleleFreq=range(len(self.positions)), step=step),
                     sim.PyEval("'Gen %4d;;' % (gen+1)", reps=0,step= step, output=fff), sim.PyEval(r"'{},'.format(map(lambda x: round(x[1],5),alleleFreq.values()))", step=step, output=fff),sim.PyOutput('\n', reps=-1, step=step, output=fff)],
            gen = self.maxGeneration)
        # idx=np.arange(self.generationStep-1,self.maxGeneration,self.generationStep)+self.initialNeutralGenerations
        _,data=zip(*map(lambda x: x.split(';;'),a.strip().split('\n')))
        data=np.array(map(eval,data))[:,0,:]
        # if data[-1, self.siteUnderSelection] >= self.initialCarrierFreq + self.minIncrease or self.s == 0 or not self.makeSureSelectedSiteDontGetLost:
        if data[-1, self.siteUnderSelection] or self.s == 0 or not self.makeSureSelectedSiteDontGetLost:
            try:
                self.pops+=[simulator.extract(0) ]
            except:
                print 'Error'
            return data[int(self.startGeneration/self.generationStep):,:]
        else:
            # print pd.Series(data[:, self.siteUnderSelection])
            return self.simualte()
    
    def simulateH0(self):
        self.H0=MSMS.Song(F=self.F, L=self.L, Ne=self.Ne, r=self.r, mu=self.mu,uid=self.uidMSMS)
        
        
    def set_siteUnderSelection(self,x):
        self.siteUnderSelection=x
        self.posUnderSelection=self.positions[self.siteUnderSelection]
    
    def set_posUnderSelection(self,x):
        self.posUnderSelection=x
        self.siteUnderSelection=np.where(self.positions==self.posUnderSelection)[0][0]
    
    def setH0(self,H0):
        self.H0=H0
        self.positions=self.H0.columns.values
        self.F=self.H0.shape[0]

    def set_BeneficialLoci(self,selectionOnRandomSite=False,siteUnderSelection=None,posUnderSelection =None):
        if selectionOnRandomSite:
            self.set_siteUnderSelection(np.random.randint(0,self.H0.shape[1]))
        elif siteUnderSelection is not None:
            self.set_siteUnderSelection(siteUnderSelection)
        elif posUnderSelection is not None:
            self.set_siteUnderSelection(posUnderSelection)
        else:
            if not self.s:
                self.set_siteUnderSelection(self.X0.argmax())
            else:
                sites=np.sort(np.where(self.X0== self.initialCarrierFreq)[0]);
                if not len(sites):
                    sites=np.sort(np.where(( self.X0 <= self.initialCarrierFreq +0.025) & ( self.X0 >= self.initialCarrierFreq -0.025) ) [0]);
                    if not len(sites):
                        print 'Try again. No site at freq ',self.initialCarrierFreq, self.uid; return
                self.set_siteUnderSelection(sites[np.random.randint(0,len(sites))])

    def createInitHaps(self):
        assignPositions=True
        if self.H0 is None:
            H0 = MSMS.Song(F=self.F, L=self.L, Ne=self.Ne, r=self.r, mu=self.mu, uid=self.uidMSMS,
                           msmsFile=self.msmsFile, dir=self.outpathmsms)
        else:
            H0 = self.H0
            assignPositions=False
        if self.foldInitialAFs:
            idx = H0.mean(0) > 0.5
            H0.iloc[:, idx.values] = 1 - H0.iloc[:, idx.values]
        self.setH0(H0)
        if assignPositions:
            self.positions_msms = self.H0.columns.values.copy(True)
            self.positions = sorted(np.random.choice(self.L, self.H0.shape[1], replace=False))
        self.H0 = pd.DataFrame(self.H0.values, columns=self.positions)
        self.X0 = self.H0.mean(0).values

    def forwardSimulation(self):
        """
        returns np 3D array T x nSS x R which T=|{t_1,t_2,..}| (nnumber of times), nSS is number of SS , and R is the number of replicates  
        """
        import numpy as np
        # df = pd.DataFrame([list(i.genotype(j)) for j in range(2) for i in self.initDiploidPop.individuals()])
        if self.posUnderSelection<0 and self.initDiploidPop is None:
            self.createInitHaps()
            self.set_BeneficialLoci()
            self.initDiploidPop=self.createInitialDiploidPopulation()
        elif self.initDiploidPop is None:
            self.createInitHaps()
            self.initDiploidPop = self.createInitialDiploidPopulation()
            # self.X0=self.H0.mean().values
        else:
            self.X0=Simulation.getGT(self.initDiploidPop).mean().values

        # df = pd.DataFrame([list(i.genotype(j)) for j in range(2) for i in self.initDiploidPop.individuals()])
        # print pd.concat([df.mean(),self.H0.mean().reset_index(drop=True)],1)
        self.X=np.array([self.simualte() for _ in range(self.numReplicates)]).swapaxes(0, 2).swapaxes(0, 1)
        self.X=np.append(np.tile(self.X0[:,None],(1,self.X.shape[2]))[None,:,:],self.X,axis=0)

        self.sampleDepths()
        if self.save:
            pd.to_pickle(self,self.outpath+self.uid+'.pkl')
        # self.createDF()


     
    def getGenerationTimes(self,step=None,includeZeroGeneration=True):
        if step is None: step=self.generationStep
        times= np.arange(0,self.maxGeneration-self.startGeneration+1,step)
        if includeZeroGeneration:
            return times
        else:
            return times[1:]
    
    def getTrueGenerationTimes(self,step=None,includeZeroGeneration=True):
        if step is None: step=self.generationStep
        times= np.arange(self.startGeneration,self.maxGeneration+1,step)
        if includeZeroGeneration:
            return times
        else:
            return times[1:]

    
    @staticmethod
    def getFixationTime(s,Ne=200,roundto10=True):
        if s==0: s=0.01
        t=-4*int(logit(1./Ne)/s)
        if roundto10:
            return (t//10 +1)*10
        else:
            return t
         
    @staticmethod
    def sampleInitSamplingTime(s,Ne=200,phase=0,samplingWindow=50,startOfEpoch=False):
        fix=Simulation.getFixationTime(s, Ne=Ne)
        if phase==0:    lower,upper=(0, fix-samplingWindow)
        if phase==1:    lower,upper=(0, fix/3-samplingWindow)
        if phase==2:    lower,upper=(fix/3, 2*fix/3-samplingWindow)
        if phase==3:    lower,upper=(2*fix/3, fix-samplingWindow)
        if startOfEpoch:
            rnd=lower
        else:
            rnd=np.random.randint(lower,max(lower,upper)+1)
        return int(rnd)//10 *10
    @staticmethod
    def sampleStartTimesforAlls(samplingWindow=50):
        S=[0.1, 0.05, 0.02, 0.01,0]
        for phase in [1,2,3]:
            pd.DataFrame([[Simulation.sampleInitSamplingTime(s, phase=phase, samplingWindow=samplingWindow, startOfEpoch=True) for _ in range(100)] for s in S], index=S).T.to_pickle('/home/arya/out/startSamplingTimes.phase{}.sampleWin{}.pkl'.format(phase, samplingWindow))
        
    def setSamplingTimes(self,maxGeneration=None,numSamples=5,step=None,startGeneration=None):
            GT=pd.Series(range(len(self.getTrueGenerationTimes(includeZeroGeneration=True))),index=self.getTrueGenerationTimes(includeZeroGeneration=True))
            if startGeneration is not None:   self.startGeneration=startGeneration
            if maxGeneration is not None:   self.maxGeneration = maxGeneration
            if step is not None:self.generationStep=step
            else:   self.generationStep=(self.maxGeneration-self.startGeneration)/numSamples
            i = GT.loc[self.getTrueGenerationTimes(includeZeroGeneration=True)[:self.X.shape[0]]].values
            self.X = self.X[i, :, :]
            self.C = self.C.apply(lambda x: x[i, :, :])
            self.D = self.D.apply(lambda x: x[i, :, :])
            self.X0=self.X[0,:,0]

    @staticmethod
    def getSamplingTimeBasedOnFreq(sim,phase,samplingWin=50):
        carrier_freq=[0.1,0.5,0.9][phase-1]
        a= np.where(sim.X[:,sim.siteUnderSelection,:].mean(1)>carrier_freq)[0]
        ft=sim.getTrueGenerationTimes().max()
        if len(a):
            t= sim.getTrueGenerationTimes()[np.where(sim.X[:,sim.siteUnderSelection,:].mean(1)>carrier_freq)[0].min()]
        else:
            t=sim.getTrueGenerationTimes().max()
        return min(t,ft-samplingWin)

    @staticmethod
    def Load(s=0.1, experimentID=0, nu0=0.005, numReplicates=3, step=10, ModelName='TimeSeries', samplingWindow=50,
             L=50000, depthRate=30):
        if not s: nu0=0.005
        sim = Simulation.load(s=s, experimentID=experimentID % 100, nu0=nu0, numReplicates=numReplicates, step=step,
                              ExperimentName=ModelName, All=True, L=L, replicates=range(numReplicates),
                              coverage=depthRate)
        sim.experimentID=experimentID
        startGen=0
        sim.setSamplingTimes(maxGeneration=min(startGen+samplingWindow,sim.getTrueGenerationTimes()[-1]),step=step,startGeneration=startGen)
        sim.createDF()
        return sim

    def getHardSweepMutations(self):
        MAF=1./self.H0.shape[0]
        dups=self.H0[self.H0.duplicated()]
        x0=pd.Series(self.X0, index=self.positions)
        hard=[]
        for _,dup in dups.iterrows():
            numDup=self.H0.apply(lambda x:(x==dup).all(),axis=1).sum()
            hard=np.append(hard, (dup*x0==numDup*MAF).replace({False:None}).dropna().index.values)
        hard=np.sort(np.append(hard,(x0==MAF).replace({False:None}).dropna().index.values).astype(int))
        return hard

    @property
    def df(self):
        reps=range(self.numReplicates)
        self.df=pd.concat([pd.DataFrame(self.X[:,:,r],columns=self.positions,index=pd.MultiIndex.from_product([[r],range(self.X.shape[0])],names=['REP','TIME'])).T for r in reps],axis=1)
        if self.numReplicates==1:
            self.df=self.df[0]
        return self.df
    def computeCDi(self, EE, depthRate):
        E = EE.loc[depthRate]
        index = pd.Series(range(E.shape[0]), E.index)
        C = pd.concat([pd.DataFrame(self.C.loc[depthRate][:, :, r], columns=self.H0.columns,
                                    index=pd.MultiIndex.from_product([[r], self.getTrueGenerationTimes()],
                                                                     names=['REP', 'GEN'])).T for r in
                       range(self.numReplicates)], axis=1)
        D = pd.concat([pd.DataFrame(self.D.loc[depthRate][:, :, r], columns=self.H0.columns,
                                    index=pd.MultiIndex.from_product([[r], self.getTrueGenerationTimes()],
                                                                     names=['REP', 'GEN'])).T for r in
                       range(self.numReplicates)], axis=1)
        self.cd = pd.concat([pd.Series(zip(C[i], D[i])) for i in C.columns], axis=1)
        self.cd.columns = C.columns;
        self.cd.index = C.index
        self.cdi = self.cd.applymap(lambda x: index.loc[x])

    def sampleDepths(self,depths = [30, 100, 300]):
        self.D = pd.Series(None, index=depths)
        self.C = pd.Series(None, index=depths)
        for depthRate in depths:
            self.D.loc[depthRate] = np.random.poisson(depthRate,
                                                      self.X.shape[0] * self.X.shape[1] * self.X.shape[2]).reshape(
                self.X.shape).astype(object)
            self.C.loc[depthRate] = np.array([np.random.binomial(d, x) for x, d in
                                              zip(self.X.reshape(-1), self.D.loc[depthRate].reshape(-1))]).reshape(
                    self.X.shape).astype(object)

    @staticmethod
    def sampleDepthX(X,cov):
        D= np.random.poisson(cov,X.size)
        C= np.array([np.random.binomial(d, x) for x, d in zip(X, D)])
        return C,D
    @staticmethod
    def sampleDepthXSeries(X,cov):
        C,D=Simulation.sampleDepthX(X.values,cov)
        a=pd.DataFrame([C,D],columns=X.index,index=['C','D']).T
        return a

    @staticmethod
    def computeCDdf(a, E):
        index = pd.Series(range(E.shape[0]), E.index)
        def f(x):
            try:
                return index.loc[x]
            except:
                return -1
        z=a.groupby(level=[0,1],axis=1).apply(lambda x: x.apply(lambda y:(y.iloc[0],y.iloc[1]),1)).applymap(f)
        return z[(z<0).sum(1)==0]
    def getCD(self,coverage):
        T=self.getTrueGenerationTimes()
        Ti=T
        if T[-1]!=self.C[coverage].shape[0]-1: Ti=range(self.C[coverage].shape[0])
        C=pd.concat([pd.DataFrame(self.C[coverage][Ti,:,i],columns=self.positions,index=T).T for i in range(self.numReplicates)],1,keys=range(self.C[coverage].shape[2]))
        D=pd.concat([pd.DataFrame(self.D[coverage][Ti,:,i],columns=self.positions,index=T).T for i in range(self.numReplicates)],1,keys=range(self.C[coverage].shape[2]))
        CD=pd.concat([C,D],1,keys=['C','D']).reorder_levels([1,2,0],1).sort_index(1)
        CD.columns.names=['REP','GEN','READ']
        return CD

    @staticmethod
    def Recombinator(rate, loci):
        """
        Recombination at loci, after variant index. Loci can take value in [0, NumSNPs-1]
        Args:
            rate: recombination rate
            loci: index of the loci in which rec is is being performed
        Returns: recombinator which is an argument of Simulation, _simulation2 and evolve. It can be list of loci
        """
        if not isinstance(loci, list):
            loci = [loci]
        return sim.Recombinator(intensity=rate, loci=loci)


class POP:
    @staticmethod
    def createISOGenicDiploidPopulation(df):
        """
        initHaps : np 2D array which m x nSS where m i number of individual haps and nSS is number of SS
        return a homozygote diploid population which every haplotype is copied n times
        """
        pop = sim.Population(size=df.shape[0], ploidy=2, loci=df.shape[1], lociPos=list(df.columns),
                             infoFields='fitness')
        for (i, (_, h)) in zip(pop.individuals(), df.iterrows()):
            i.setGenotype(h.tolist(), 0);
            i.setGenotype(h.tolist(), 1)
        return pop

    @staticmethod
    def toDF(pop):
        x = pd.concat(map(pd.DataFrame, [map(list, [i.genotype(0), i.genotype(1)]) for i in pop.allIndividuals()]),
                      keys=range(pop.popSize()))
        x.columns = list(pop.lociPos())
        return x

    @staticmethod
    def freq(pop):
        sim.stat(pop, alleleFreq=range(pop.numLoci()[0]), vars=['alleleFreq'])
        return pd.Series(pd.DataFrame(pop.vars()['alleleFreq']).loc[1].reindex().values,map(int,pop.lociPos())).fillna(0)

    @staticmethod
    def Haplotypes(pop,counts=False,unique=True):
        if isinstance(pop,sim.Population):
            a=POP.toDF(pop)
        else:
            a=pop

        H=a.reset_index(drop=True)
        H.columns=map(int,H.columns)
        b=H.loc[H.sum(1).sort_values().index].astype(str).apply(lambda x: ''.join(x), 1).reset_index(drop=True)
        if counts:
            return b.value_counts().sort_index()
        else:
            if unique:
                b=b.drop_duplicates()
            return b.loc[b.sort_values().index].reset_index(drop=True)


    @staticmethod
    def establish(H, ba, k=5):
        N = H.shape[0]
        car = H[H[ba] == 1]
        n = car.shape[0]
        return pd.concat([car.iloc[np.random.choice(n, k)], H.iloc[np.random.choice(N, N - k)]]).reset_index(drop=True)

class Drift:
    @staticmethod
    def nextGeneration(N,x):
        return (np.random.random(N)<=x).mean()

    @staticmethod
    def sampleReads(D,x):
        return [Drift.sampleReadsDerived(D,x),D]

    @staticmethod
    def sampleReadsDerived(D,x):
        return (np.random.random(D)<=x).sum()

    @staticmethod
    def simulateAF(N,x,T):
        Xt=[]
        for i in range(1, T[-1]+1):
            x=Drift.nextGeneration(N,x)
            if i in T:Xt.append(x)
        return Xt

    @staticmethod
    def simulatePoolCD(N,n,cd):
        x=cd[0].C/float(cd[0].D)
        D=cd.xs('D',level=1)
        Xt=[]
        for i in range(1, D.index[-1]+1):
            x=Drift.nextGeneration(N,x)
            if i in D.index:
                y=Drift.nextGeneration(n,x)
                Xt.append(Drift.sampleReads(D[i], y))
        return pd.DataFrame([[cd[0].C,cd[0].D]]+Xt,index=D.index,columns=['C','D'])

    @staticmethod
    def simulatePoolDerivd(N,n,cd):
        x=cd[0].C/float(cd[0].D)
        D=cd.xs('D',level=1)
        Xt=[]
        for i in range(1, D.index[-1]+1):
            x=Drift.nextGeneration(N,x)
            if i in D.index:
                Xt+=[Drift.sampleReadsDerived(D[i], Drift.nextGeneration(n,x))]
        return [cd[0].C]+Xt

    @staticmethod
    def simulatePools(N,cd,M):
        return pd.concat([Drift.simulatePool(N,cd) for _ in range(M)],keys=range(M))


    @staticmethod
    def simulateAFs(N,x,T,M):
        return pd.DataFrame([Drift.simulateAF(N,x,T) for _ in range(M)],columns=T)