"""
A model for the disasters data with a changepoint, with GOF diagnostics added

changepoint ~ U(0,110)
early_mean ~ Exp(1.)
late_mean ~ Exp(1.)
disasters[t] ~ Po(early_mean if t <= switchpoint, late_mean otherwise)

"""

import pymc as pm
from numpy import array, concatenate, ones
from numpy.random import randint

__all__ = [
    'disasters_array',
    'switchpoint',
    'early_mean',
    'late_mean',
    'disasters']

disasters_array = array([4, 5, 4, 0, 1, 4, 3, 4, 0, 6, 3, 3, 4, 0, 2, 6,
                         3, 3, 5, 4, 5, 3, 1, 4, 4, 1, 5, 5, 3, 4, 2, 5,
                         2, 2, 3, 4, 2, 1, 3, 2, 2, 1, 1, 1, 1, 3, 0, 0,
                         1, 0, 1, 1, 0, 0, 3, 1, 0, 3, 2, 2, 0, 1, 1, 1,
                         0, 1, 0, 1, 0, 0, 0, 2, 1, 0, 0, 0, 1, 1, 0, 2,
                         3, 3, 1, 1, 2, 1, 1, 1, 1, 2, 4, 2, 0, 0, 1, 4,
                         0, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 1, 0, 1])

n = len(disasters_array)

# Define data and stochastics

switchpoint = pm.DiscreteUniform('switchpoint', lower=0, upper=110)
early_mean = pm.Exponential('early_mean', beta=1.)
late_mean = pm.Exponential('late_mean', beta=1.)


@pm.stochastic(observed=True, dtype=int)
def disasters(value=disasters_array,
              early_mean=early_mean,
              late_mean=late_mean,
              switchpoint=switchpoint):
    """Annual occurences of coal mining disasters."""
    return pm.poisson_like(
        value[:switchpoint], early_mean) + pm.poisson_like(value[switchpoint:], late_mean)


@pm.deterministic
def disasters_sim(early_mean=early_mean,
                  late_mean=late_mean,
                  switchpoint=switchpoint):
    """Coal mining disasters sampled from the posterior predictive distribution"""
    return concatenate((pm.rpoisson(early_mean, size=switchpoint), pm.rpoisson(
        late_mean, size=n - switchpoint)))


@pm.deterministic
def expected_values(early_mean=early_mean,
                    late_mean=late_mean,
                    switchpoint=switchpoint):
    """Discrepancy measure for GOF using the Freeman-Tukey statistic"""

    # Sample size
    n = len(disasters_array)
    # Expected values
    return concatenate(
        (ones(switchpoint) * early_mean, ones(n - switchpoint) * late_mean))


if __name__ == '__main__':
    vars = [
        switchpoint,
        early_mean,
        late_mean,
        disasters,
        disasters_sim,
        expected_values]
    # Instiatiate model
    M = pm.MCMC(vars)
    # Sample
    M.sample(10000, burn=5000, verbose=2)
    # Calculate discrepancy function
    D = pm.diagnostics.discrepancy(
        disasters_array,
        disasters_sim,
        expected_values)
    # Plot GOF graphics
    pm.Matplot.discrepancy_plot(D, 'D')
    pm.Matplot.gof_plot(disasters_sim, disasters_array, 'disasters')
