%matplotlib inline
import pymc as pm
import numpy as np
from pymc.examples import disaster_model

switchpoint = pm.DiscreteUniform('switchpoint', lower=0, upper=110)

early_mean = pm.Exponential('early_mean', beta=1., value=4)
late_mean = pm.Exponential('late_mean', beta=1., value=3)

@pm.stochastic
def switchpoint(value=1900, t_l=1851, t_h=1962):
    """The switchpoint for the rate of disaster occurrence."""
    if value > t_h or value < t_l:
        # Invalid values
        return -np.inf
    else:
        # Uniform log-likelihood
        return -np.log(t_h - t_l + 1)

def switchpoint_logp(value, t_l, t_h):
    if value > t_h or value < t_l:
        return -np.inf
    else:
        return -np.log(t_h - t_l + 1)

def switchpoint_rand(t_l, t_h):
    return np.round( (t_l - t_h) * np.random.random() ) + t_l

switchpoint = pm.Stochastic( logp = switchpoint_logp,
                doc = 'The switchpoint for the rate of disaster occurrence.',
                name = 'switchpoint',
                parents = {'t_l': 1851, 't_h': 1962},
                random = switchpoint_rand,
                trace = True,
                value = 1900,
                dtype=int,
                rseed = 1.,
                observed = False,
                cache_depth = 2,
                plot=True,
                verbose = 0)

from scipy.stats.distributions import poisson

@pm.observed
def likelihood(value=[1, 2, 1, 5], parameter=3):
    return poisson.logpmf(value, parameter).sum()

disasters = pm.Poisson('disasters', mu=2, 
                       value=disaster_model.disasters_array, 
                       observed=True)

@pm.deterministic
def rate(s=switchpoint, e=early_mean, l=late_mean):
    ''' Concatenate Poisson means '''
    out = np.empty(len(disaster_model.disasters_array))
    out[:s] = e
    out[s:] = l
    return out

from pymc import Lambda
x = pm.MvNormal('x', np.ones(3), np.eye(3))
y = pm.MvNormal('y', np.ones(3), np.eye(3))
sum_xy = Lambda('sum_xy', lambda x=x, y=y: x+y)

print(x[0])

print(x[0]+y[2])

def rate_eval(switchpoint=switchpoint, early_mean=early_mean, late_mean=late_mean):
    value = np.zeros(111)
    value[:switchpoint] = early_mean
    value[switchpoint:] = late_mean
    return value

rate = pm.Deterministic(eval = rate_eval,
                  name = 'rate',
                  parents = {'switchpoint': switchpoint, 
                          'early_mean': early_mean, 
                          'late_mean': late_mean},
                  doc = 'The rate of disaster occurrence.',
                  trace = True,
                  verbose = 0,
                  dtype=float,
                  plot=False,
                  cache_depth = 2)

N = 10
x_0 = pm.Normal('x_0', mu=0, tau=1)

x = np.empty(N, dtype=object)
x[0] = x_0

for i in range(1, N):

    x[i] = pm.Normal('x_%i' % i, mu=x[i-1], tau=1)

@pm.observed
def y(value=1, mu=x, tau=100):
    return pm.normal_like(value, np.sum(mu**2), tau)

x = [pm.Normal('x_%i' % i, mu=x[i-1], tau=1) for i in range(1,N)]

@pm.potential
def rate_constraint(l1=early_mean, l2=late_mean):
    if np.abs(l2 - l1) > 1:
        return -np.inf
    return 0

def rate_constraint_logp(l1=early_mean, l2=late_mean):
    if np.abs(l2 - l1) > 1:
        return -np.inf
    return 0

rate_constraint = pm.Potential(logp = rate_constraint_logp,
                    name = 'rate_constraint',
                    parents = {'l1': early_mean, 'l2': late_mean},
                    doc = 'Constraint on rate differences',
                    verbose = 0,
                    cache_depth = 2)

# Log dose in each group
log_dose = [-.86, -.3, -.05, .73]

# Sample size in each group
n = 5

# Outcomes
deaths = [0, 1, 3, 5]

from pymc import invlogit

def bioassay():
    
    a = Normal('a', 0, 0.01, value=0)
    b = Normal('b', 0, 0.01, value=0)
    
    p = Lambda('p', lambda a=a, b=b: invlogit(a + b*log_dose))
    
    y = Binomial('y', n, p=p, value=deaths, observed=True)
    
    return(locals())

from pymc.examples import gelman_bioassay
M = pm.MAP(gelman_bioassay)
M.fit()

M.alpha.value

M.beta.value

M.AIC

M.BIC

N = pm.NormApprox(gelman_bioassay)
N.fit()

N.mu[N.alpha]

N.C[N.alpha, N.beta]

import matplotlib.pyplot as plt

N.sample(100)
plt.hist(N.trace('alpha')[:])

M = pm.MCMC(gelman_bioassay)

M.use_step_method(pm.Metropolis, M.alpha, proposal_sd=1., proposal_distribution='Normal')

from pymc.examples import disaster_model_linear
M = pm.MCMC(disaster_model_linear)
M.use_step_method(pm.AdaptiveMetropolis, M.params_of_mean)

M = pm.MCMC(gelman_bioassay)
M.sample(10000, burn=5000)

pm.Matplot.plot(M.LD50)

from IPython.core.display import HTML
def css_styling():
    styles = open("styles/custom.css", "r").read()
    return HTML(styles)
css_styling()