"""Interface to the `emcee` ensemble Monte Carlo sampler.

See http://dan.iel.fm./emcee for details on `emcee`.
"""


import numpy as np
import flavio

try:
    import emcee
except:
    pass

class emceeScan(object):
    """Ensemble Monte Carlo sampler using the `emcee` package.

    Methods:

    - run: run the sampler
    - result: get a flat array of the sampled points of all walkers
    - save_result: save the result to a `.npy` file

    Important attributes:

    - mc: the `emcee.EnsembleSampler` instance

    """

    def __init__(self, fit, nwalkers=None, **kwargs):
        """Initialize the emceeScan.

        Parameters:

        - fit: an instance of `flavio.statistics.fits.BayesianFit`
        - nwalkers (optional): number of walkers. Defaults to ten times the number
          of dimensions.

        Additional keyword argumements will be passed to
        `emcee.EnsembleSampler`.
        """

        assert isinstance(fit, flavio.statistics.fits.BayesianFit), "emcee fit object must be an instance of BayesianFit"
        self.fit = fit

        self.dimension = len(fit.get_random_start)
        if nwalkers is None:
            self.nwalkers = self.dimension * 10
        else:
            self.nwalkers = nwalkers
        def get_random_good():
            # iterate until a random point with finite probability is found
            good = False
            i = 0
            while not good:
                x = fit.get_random_start
                good = np.isfinite(fit.log_target(x))
                i += 1
                if(i == 1000):
                    raise ValueError("Could not find enough starting values with finite probability. "
                                     " Try reducing the starting priors for the Wilson coefficients.")
            return x
        self.start = [get_random_good() for i in range(self.nwalkers)]


        self.mc = emcee.EnsembleSampler(nwalkers=self.nwalkers,
                                        dim=self.dimension,
                                        lnpostfn=self.fit.log_target,
                                        **kwargs)

    def run(self, steps, burnin=1000):
        """Run the sampler.

        Parameters:

        - steps: number of steps per walker
        - burnin (optional): number of steps for burn-in (samples will not be
          retained); defaults to 1000

        Note that the total number of samples will be `steps * nwalkers`!
        """
        pos = self.start
        if burnin > 0:
            pos, prob, state = self.mc.run_mcmc(self.start, burnin)
        self.mc.reset()
        self.mc.run_mcmc(pos, steps)

    @property
    def result(self):
        """Return a flat array of the samples."""
        return self.mc.flatchain[:]

    def save_result(self, file):
        """Save the samples to a `.npy` file."""
        res = self.result
        np.save(file, res)

@@ -1249,52 +1249,341 @@
     
-def get_random_good():%0A            # iterate
+self.start = None%0A        self.mc = emcee.EnsembleSampler(nwalkers=self.nwalkers,%0A                                        dim=self.dimension,%0A                                        lnpostfn=self.fit.log_target,%0A                                        **kwargs)%0A%0A    def _get_random_good():%0A        %22%22%22Generate random starting points
  unt
@@ -1587,23 +1587,16 @@
 until a 
-random 
 point wi
@@ -1596,32 +1596,40 @@
 oint with finite
+%0A       
  probability is 
@@ -1633,26 +1633,29 @@
 is found
+.%22%22%22
 %0A
- 
         
    good 
@@ -1642,27 +1642,24 @@
 %22%22%22%0A        
-   
 good = False
@@ -1667,20 +1667,16 @@
         
-    
 i = 0%0A  
@@ -1681,20 +1681,16 @@
         
-   
- 
 while no
@@ -1697,20 +1697,16 @@
 t good:%0A
-    
         
@@ -1742,28 +1742,24 @@
             
- 
-   
 good = np.is
@@ -1800,23 +1800,15 @@
     
-  
-  
 i += 1%0A
-    
     
@@ -1826,28 +1826,24 @@
 i == 1000):%0A
-    
             
@@ -1962,20 +1962,16 @@
         
-    
 %22 Try re
@@ -2032,28 +2032,24 @@
 .%22)%0A        
-    
 return x%0A   
@@ -2049,324 +2049,180 @@
 n x%0A
+%0A
     
-    self.start = %5Bget_random_good() for i in range(self.nwalkers)%5D%0A%0A%0A        self.mc = emcee.EnsembleSampler(nwalkers=self.nwalkers,%0A                                        dim=self.dimension,%0A                                        lnpostfn=self.fit.log_target,%0A                                        **kwarg
+def initialize_sampler(self):%0A        %22%22%22Initialize the emcee.EnsembleSampler instance.%22%22%22%0A        self.start = %5Bself._get_random_good() for i in range(self.nwalker
 s)
+%5D
 %0A%0A  
@@ -2551,24 +2551,93 @@
         %22%22%22%0A
+        if self.start is None:%0A            self.initialize_sampler()%0A
         pos 

Move chain initialazation out of __init__