import Utils.Util as utl
import Utils.Simulation as sim
import numpy as np
import CLEAR.Libs.Markov as mkv
import simuPOP.demography as dmg
import pandas as pd
import Utils.Plots as pplt
import pylab as plt
def getPopSize(model):
    model.plot(utl.home+'aa.png')
    return pd.Series([model.init_size[0]]+(model.pop_regions['0'][3::4][:-1].tolist()))
# filename = '{}{}/msms/L{:E}.{:E}.msms'.format(utl.simoutpath, 'TimeSeries', 50000,0)
# a=sim.Simulation(s=0, experimentID=0, numReplicates=3, initialCarrierFreq=0.005, save=False,L=5000000,   maxGeneration=50, ExperimentName='TimeSeries2', generationStep=1)
# b=sim.Simulation(s=0.1, experimentID=0, numReplicates=3, initialCarrierFreq=0.1, save=False,L=5000000,   maxGeneration=50, ExperimentName='TimeSeries2', generationStep=1)
# c=sim.Simulation(s=0.1, experimentID=0, numReplicates=3, initialCarrierFreq=0.005, save=False,L=5000000,   maxGeneration=50, ExperimentName='TimeSeries2', generationStep=1)
expansion=dmg.ExponentialGrowthModel(T=80, N0=100, NT=10000)
model=dmg.LinearGrowthModel(T=80, N0=100, NT=100)
bottleneck=dmg.ExponentialGrowthModel(T=20, N0=10000, NT=500)
model=dmg.MultiStageModel([expansion,bottleneck])
B=getPopSize(model)
B.plot()
fname='/media/arya/d4565cf2-d44a-4b67-bf97-226a486c01681/Data/SimulationOutFiles/TimeSeries2/msms/bd5fba32-fbdc-4056-bd6f-441828938b01.msms'
a=sim.Simulation(s=0,F=100, save=True,L=2e7,  Ne=2e5, numReplicates=5,maxGeneration=100, ExperimentName='TimeSeries2', generationStep=1,model=model,msmsFile=fname)
a.sampleDepths([100])
CD=a.getCD(100)
a.model
CD.to_pickle()
rangeN=[200]+range(0,2001,500)[1:]+range(3000,10001,1000)
cd=CD.loc[:,pd.IndexSlice[:,[80,81]]]
cd
likes=mkv.estimateN(cd,rangeN=rangeN,Nt=True)

likes.plot()
    ,likes.loc[80]
NN.loc[80]

NN=getPopSize(a.model)
    .plot();likes.idxmax().plot()
a.X.shape
a=sim.Simulation(s=0,F=100, save=False,L=2e7,  Ne=2e5, maxGeneration=100, ExperimentName='TimeSeries2', generationStep=1,msmsFile=fname)
a

cc=sim.Simulation(s=0.5, experimentID=0, numReplicates=3, initialCarrierFreq=0.005, save=False,L=5000000,   maxGeneration=50, ExperimentName='TimeSeries2', generationStep=1)
# pd.to_pickle(a,utl.home+'5M.pkl')
# pd.to_pickle(b,utl.home+'5M.sel.pkl')
# pd.to_pickle(c,utl.home+'5M.sel.hard.strong.pkl')
# pd.to_pickle(cc,utl.home+'5M.sel.hard.strongs5.pkl')
c=pd.read_pickle('/home/arya/5M.sel.pkl')
d=pd.read_pickle('/home/arya/5M.sel.hard.pkl')
cd=c.getCD(100).loc[:,pd.IndexSlice[:,range(0,51,10)]]

f=lambda x: (x.alt-x.null)*x.s.apply(np.sign)
import time
start=time.time();mkv.HMM(CD=cd,gridH=[0.5],N=300,n=200,saveCDE=False,loadCDE=False,verbose=1);print time.time()-start
reload(mkv)
mkv.HMM(CD=cd.iloc[:10],gridH=[0.5],N=300,n=200,saveCDE=False,loadCDE=False,verbose=1).fit(False)
aa

a=pd.concat([f(mkv.HMM(CD=cd,gridH=[0.5],N=300,n=200,saveCDE=False,loadCDE=False,verbose=1).fit(False)[0.5])],keys=[1]);a.index.names=['CHROM','POS']
b=a.copy(True)
c.posUnderSelection
pplt.Manhattan(utl.scanGenome(a,winSize=1000000).abs());plt.axvline(c.posUnderSelection)
pplt.Manhattan(b);plt.axvline(d.posUnderSelection)

A=[]
for gen in range(10,51,10):
    CD=cd.loc[:,pd.IndexSlice[:,(gen-10,gen)]]
    CD=CD[CD.groupby(level=0,axis=1).apply(lambda x: utl.polymorphic(x,minAF=1e-4)).sum(1)==3]
    I=np.random.choice(CD.shape[0],5000)
    print CD.shape,gen
    A+=[mkv.HMM(CD=CD.iloc[I],gridH=[0.5],saveCDE=False,loadCDE=False,verbose=0).fitN2(rangeN=range(200,1001,100),n=200).rename(gen)]
    print gen,A[-1].idxmax()

A=pd.concat(A,1)
A.idxmax().plot();B.plot()
# drift=pd.concat(A,1)