import numpy as np
import matplotlib.pyplot as pl
import scipy.optimize as op
import pdb
%matplotlib inline

class adaptiveSMC(object):
    def __init__(self):
        self.M = 50  # Number of pseudo observations for each particle
        self.N = 1000  # Number of particles
        self.alpha = 0.95 # Percentage of particles surviving
        self.theta = np.random.uniform(low=-10.0, high=10.0, size=self.N)  # Initial parameters from the prior        
        
        self.pseudoObs = self.simulateFromModel(self.theta)  # Simulate particles from the prior
        
        self.W = np.ones(self.N) / self.N
        
        self.epsilon = 100
        self.epsilonOld = 0
        self.epsilonTarget = 0.025
        self.ESSThreshold = self.N / 2
        
        self.epsilonStore = []
        self.ESSStore = []
        
    def simulateFromModel(self, pars):
        pseudoObs = np.zeros((self.M,self.N))
        for i in range(self.M):
            u = np.random.rand(self.N)
            a = np.where(u < 0.5)
            b = np.where(u >= 0.5)
            pseudoObs[i,a] = pars[a] + np.random.rand(a[0].size)
            pseudoObs[i,b] = pars[b] + np.random.rand(b[0].size)
        return pseudoObs
    
    def aliveParticles(self, epsilon, x):
        npa = np.sum(np.abs(x) < epsilon, axis=0)
        return np.sum(npa>0) / self.N
    
    def sample(self):
#        while (self.epsilon > self.epsilonTarget):

# Find next level
        self.epsilonOld = self.epsilon
        refLevel = self.alpha * self.aliveParticles(self.epsilonOld, self.pseudoObs)
        self.epsilon = op.bisect(lambda eps: self.aliveParticles(eps, self.pseudoObs)-refLevel, 0, self.epsilonOld)
        if (self.epsilon < self.epsilonTarget):
            self.epsilon = self.epsilonTarget
        
        self.ESS = 1.0 / np.sum(self.W**2)
        self.epsilonStore.append(self.epsilon)
        self.ESSStore.append(self.ESS)
        
# Compute associated weights
        d = np.abs(self.pseudoObs)
        npaOld = np.sum(d < self.epsilonOld, axis=0)
        npa = np.sum(d < self.epsilon, axis=0)
        a = np.where(npaOld > 0)
        b = np.where(npaOld == 0)
        self.W[a] = self.W[a] * npa[a] / npaOld[a]
        self.W[b] = np.zeros(b[0].size)
        self.W /= np.sum(self.W)
        
# Resample if necessary
        if (np.sum(self.W**2) * self.ESSThreshol > 1):
            
        

smc = adaptiveSMC()

smc.sample()

10.0076821269

a[0].size

2