%pylab inline

from __future__ import division

from qinfer.abstract_model import Model, Simulatable
from qinfer.distributions import UniformDistribution
from qinfer.smc import SMCUpdater
import scipy.stats

# We start off by making a new class that inherits from Model.
class MultiCosModel(Model):
    
    # We need to specify how many model parameters this model has.
    @property
    def n_modelparams(self):
        return 2
    
    # The number of outcomes is always 2, so we indicate that it's constant
    # and define the n_outcomes method to return that constant.
    @property
    def is_n_outcomes_constant(self):
        return True
    def n_outcomes(self, expparams):
        return 2
    
    # Next, we denote that the experiment parameters are represented by a
    # single field of two floats.
    @property
    def expparams_dtype(self):
        return [('ts', '2float')]
    
    # Valid models are defined to lie in the range [0, 1].
    def are_models_valid(self, modelparams):
        return np.all(np.logical_and(modelparams > 0, modelparams <= 1), axis=1)
    
    # Finally, the likelihood function itself.
    def likelihood(self, outcomes, modelparams, expparams):
        # We first call the superclass method, which basically
        # just makes sure that call count diagnostics are properly
        # logged.
        super(MultiCosModel, self).likelihood(outcomes, modelparams, expparams)
        
        # Next, since we have a two-outcome model, everything is defined by
        # Pr(0 | modelparams; expparams), so we find the probability of 0
        # for each model and each experiment.
        #
        # We do so by taking a product along the modelparam index (len 2,
        # indicating omega_1 or omega_2), then squaring the result.
        pr0 = np.prod(
            np.cos(
                # shape (n_models, 1, 2)
                modelparams[:, np.newaxis, :] *
                # shape (n_experiments, 2)
                expparams['ts']
            ), # <- broadcasts to shape (n_models, n_experiments, 2).
            axis=2 # <- product over the final index (len 2)
        ) ** 2 # square each element
        
        # Now we use pr0_to_likelihood_array to turn this two index array
        # above into the form expected by SMCUpdater and other consumers
        # of likelihood().
        return Model.pr0_to_likelihood_array(outcomes, pr0)

m = MultiCosModel()

prior = UniformDistribution([[0, 1], [0, 1]])

updater = SMCUpdater(m, 2500, prior)

true = prior.sample()
for idx_exp in xrange(50):
    # Use the variance in the posterior to set the time we evolve for.
    # Note that this is a slight generalization of the
    # particle guess heuristic introduced in [].
    #
    # We slow it down by 30% to ensure unimodality, since we
    # don't have an inversion argument here.
    expparams = np.array([
        # We want one element...
        (
             # ...containing an array.
             0.7 / (2 * np.pi * np.diag(scipy.linalg.sqrtm(updater.est_covariance_mtx()))),
        )
    ], dtype=m.expparams_dtype)
    
    o = m.simulate_experiment(true, expparams)
    
    updater.update(o, expparams)

updater.plot_posterior_contour(res1=60, res2=60)
plot(true[0, 0], true[0, 1], 'r.', markersize=12)

w1s, w2s, Z = updater.posterior_mesh(res1=60, res2=60)
plt.imshow(Z, origin='lower', extent=[np.min(w1s), np.max(w1s), np.min(w2s), np.max(w2s)])
colorbar()
plot(true[0, 0], true[0, 1], 'r.', markersize=12)