def simple_hierarchical_data(n):
    """ Generate data based on the simple one-way hierarchical model
    given in section 3.1.1::

        y[i,j] | alpha[j], sigma^2 ~ N(alpha[j], sigma^2) i = 1, ..., n_j, j = 1, ..., J;
        alpha[j] | mu, tau^2 ~ N(mu, tau^2) j = 1, ..., J.

        sigma^2 ~ Inv-Chi^2(5, 20)
        mu ~ N(5, 5^2)
        tau^2 ~ Inv-Chi^2(2, 10)

    Parameters
    ----------
    n : list, len(n) = J, n[j] = num observations in group j
    """

    inv_sigma_sq = mc.rgamma(alpha=2.5, beta=50.0)
    mu = mc.rnormal(mu=5.0, tau=5.0 ** -2.0)
    inv_tau_sq = mc.rgamma(alpha=1.0, beta=10.0)

    J = len(n)
    alpha = mc.rnormal(mu=mu, tau=inv_tau_sq, size=J)
    y = [mc.rnormal(mu=alpha[j], tau=inv_sigma_sq, size=n[j]) for j in range(J)]

    mu_by_tau = mu * pl.sqrt(inv_tau_sq)
    alpha_by_sigma = alpha * pl.sqrt(inv_sigma_sq)
    alpha_bar = alpha.sum()
    alpha_bar_by_sigma = alpha_bar * pl.sqrt(inv_sigma_sq)

    return vars()

def main():
    x_t = pm.rnormal(0, 1, 200)
    x_t[0] = 0
    y_t = np.zeros(200)
    for i in range(1, 200):
        y_t[i] = pm.rnormal(y_t[i - 1], 1)

    plt.plot(y_t, label="$y_t$", lw=3)
    plt.plot(x_t, label="$x_t$", lw=3)
    plt.xlabel("time, $t$")
    plt.legend()
    plt.show()

    colors = ["#348ABD", "#A60628", "#7A68A6"]

    x = np.arange(1, 200)
    plt.bar(x, autocorr(y_t)[1:], width=1, label="$y_t$",
            edgecolor=colors[0], color=colors[0])
    plt.bar(x, autocorr(x_t)[1:], width=1, label="$x_t$",
            color=colors[1], edgecolor=colors[1])

    plt.legend(title="Autocorrelation")
    plt.ylabel("measured correlation \nbetween $y_t$ and $y_{t-k}$.")
    plt.xlabel("k (lag)")
    plt.title("Autocorrelation plot of $y_t$ and $x_t$ for differing $k$ lags.")
    plt.show()

def test_covariate_model_sim_no_hierarchy():
    # simulate normal data
    model = data.ModelData()
    model.hierarchy, model.output_template = data_simulation.small_output()

    X = mc.rnormal(0., 1.**2, size=(128,3))

    beta_true = [-.1, .1, .2]
    Y_true = pl.dot(X, beta_true)

    pi_true = pl.exp(Y_true)
    sigma_true = .01*pl.ones_like(pi_true)

    p = mc.rnormal(pi_true, 1./sigma_true**2.)

    model.input_data = pandas.DataFrame(dict(value=p, x_0=X[:,0], x_1=X[:,1], x_2=X[:,2]))
    model.input_data['area'] = 'all'
    model.input_data['sex'] = 'total'
    model.input_data['year_start'] = 2000
    model.input_data['year_end'] = 2000

    # create model and priors
    vars = {}
    vars.update(covariate_model.mean_covariate_model('test', 1, model.input_data, {}, model, 'all', 'total', 'all'))
    vars.update(rate_model.normal_model('test', vars['pi'], 0., p, sigma_true))

    # fit model
    m = mc.MCMC(vars)
    m.sample(2)

def logit_normal_draw(cf_mean, std, N, J):
    std = pl.array(std)
    if mc.__version__ == '2.0rc2': # version on Omak 
        X = [mc.invlogit(mc.rnormal(mu=cf_mean, tau=std**-2)) for n in range(N)]
        Y = pl.array(X)
    else: 
        X = mc.rnormal(mu=cf_mean, tau=std**-2, size=(N,J))
        Y = mc.invlogit(X)
    return Y

    def pred(a1=alpha1, mu_int=mu_int, tau_int=tau_int, mu_slope=mu_slope, tau_slope=tau_slope, tau_iq=tau_iq, values=(70,75,80,85)):
        """Estimate the probability of IQ<85 for different covariate values"""
        b0 = rnormal(mu_int, tau_int, size=len(phe_pred))
        a0 = rnormal(mu_slope, tau_slope, size=len(phe_pred))

        b1 = a0 + a1*crit_pred

        iq = rnormal(b0 + b1*phe_pred, tau_iq)

        return [iq<v for v in values]

 def propose(self):
     tau = 1./(self.adaptive_scale_factor * self.proposal_sd)**2
     time = pymc.rnormal(self.stochastic.value.time, tau)
     n = pymc.rnormal(len(self.stochastic.value), tau)
     if n <= 0:
         n = 0
     times = [rand.random() for _ in range(n)]
     total = float(sum(times))
     times = [item*time/total for item in times]
     events = [Event(time=item, censored=False) for item in times]
     self.stochastic.value = MultiEvent(events)

    def step(self):
        x0 = self.value[self.n]
        u = pm.rnormal(np.zeros(self.N), 1.)
        dx = np.dot(u, self.value)
 
        self.stochastic.value = x0
        logp = [self.logp_plus_loglike]
        x_prime = [x0]
 
        for direction in [-1, 1]:
            for i in xrange(25):
                delta = direction*np.exp(.1*i)*dx
                try:
                    self.stochastic.value = x0 + delta
                    logp.append(self.logp_plus_loglike)
                    x_prime.append(x0 + delta)
                except pm.ZeroProbability:
                    self.stochastic.value = x0
 
        i = pm.rcategorical(np.exp(np.array(logp) - pm.flib.logsum(logp)))
        self.value[self.n] = x_prime[i]
        self.stochastic.value = x_prime[i]
 
        if i == 0:
            self.rejected += 1
        else:
            self.accepted += 1
 
        self.n += 1
        if self.n == self.N:
            self.n = 0    

def main():
    """ Demonstrating thinning of two autocorrelated inputs (representing
        posterior probabilities). The key point is the thinned - every 2nd / 3rd
        point - functions approach zero quicker. More thinning is better (but
        expensive)
    """

    # x_t = pm.rnormal(0, 1, 200)
    # x_t[0] = 0
    y_t = np.zeros(200)
    for i in range(1, 200):
        y_t[i] = pm.rnormal(y_t[i - 1], 1)

    max_x = 200 / 3 + 1
    x = np.arange(1, max_x)

    colors = ["#348ABD", "#A60628", "#7A68A6"]
    plt.bar(x, autocorr(y_t)[1:max_x], edgecolor=colors[0],
            label="no thinning", color=colors[0], width=1)
    plt.bar(x, autocorr(y_t[::2])[1:max_x], edgecolor=colors[1],
            label="keeping every 2nd sample", color=colors[1], width=1)
    plt.bar(x, autocorr(y_t[::3])[1:max_x], width=1, edgecolor=colors[2],
            label="keeping every 3rd sample", color=colors[2])

    plt.autoscale(tight=True)
    plt.legend(title="Autocorrelation plot for $y_t$", loc="lower left")
    plt.ylabel("measured correlation \nbetween $y_t$ and $y_{t-k}$.")
    plt.xlabel("k (lag)")
    plt.title("Autocorrelation of $y_t$ (no thinning vs. thinning) \
            at differing $k$ lags.")
    plt.show()

    def step(self):
        x0 = np.copy(self.stochastic.value)
        dx = pymc.rnormal(np.zeros(np.shape(x0)), self.proposal_tau)

        logp = [self.logp_plus_loglike]
        x_prime = [x0]

        for direction in [-1, 1]:
            for i in xrange(25):
                delta = direction*np.exp(.1*i)*dx
                try:
                    self.stochastic.value = x0 + delta
                    logp.append(self.logp_plus_loglike)
                    x_prime.append(x0 + delta)
                except pymc.ZeroProbability:
                    self.stochastic.value = x0
        
        i = pymc.rcategorical(np.exp(np.array(logp) - pymc.flib.logsum(logp)))
        self.stochastic.value = x_prime[i]

        if i == 0:
            self.rejected += 1
            if self.verbose > 2:
                print self._id + ' rejecting'
        else:
            self.accepted += 1
            if self.verbose > 2:
                print self._id + ' accepting'

def test_random_effect_priors():
    model = data.ModelData()

    # set prior on sex
    parameters = dict(random_effects={'USA': dict(dist='TruncatedNormal', mu=.1, sigma=.5, lower=-10, upper=10)})


    # simulate normal data
    n = 32.
    area_list = pl.array(['all', 'USA', 'CAN'])
    area = area_list[mc.rcategorical([.3, .3, .4], n)]
    alpha_true = dict(all=0., USA=.1, CAN=-.2)
    pi_true = pl.exp([alpha_true[a] for a in area])
    sigma_true = .05
    p = mc.rnormal(pi_true, 1./sigma_true**2.)

    model.input_data = pandas.DataFrame(dict(value=p, area=area))
    model.input_data['sex'] = 'male'
    model.input_data['year_start'] = 2010
    model.input_data['year_end'] = 2010

    model.hierarchy.add_edge('all', 'USA')
    model.hierarchy.add_edge('all', 'CAN')

    # create model and priors
    vars = {}
    vars.update(covariate_model.mean_covariate_model('test', 1, model.input_data, parameters, model,
                                                     'all', 'total', 'all'))

    print vars['alpha']
    print vars['alpha'][1].parents['mu']
    assert vars['alpha'][1].parents['mu'] == .1

def plot_funnel(pi_true, delta_str):
    delta = float(delta_str)
    n = pl.exp(mc.rnormal(10, 2**-2, size=10000))
    p = pi_true*pl.ones_like(n)

    # old way:
    #delta = delta * p * n

    nb = rate_model.neg_binom_model('funnel', pi_true, delta, p, n)
    r = nb['p_pred'].value

    pl.vlines([pi_true], .1*n.min(), 10*n.max(),
              linewidth=5, linestyle='--', color='black', zorder=10)
    pl.plot(r, n, 'o', color=colors[0], ms=10,
            mew=0, alpha=.25)

    pl.semilogy(schiz['r'], schiz['n'], 's', mew=1, mec='white', ms=15,
                color=colors[1],
                label='Observed Values')

    pl.xlabel('Rate (Per 1000 PY)', size=32)
    pl.ylabel('Study Size (PY)', size=32)
    pl.axis([-.0001, .0101, 50., 15000000])
    pl.title(r'$\delta = %s$'%delta_str, size=48)
    pl.xticks([0, .005, .01], [0, 5, 10], size=30)
    pl.yticks(size=30)

def test_fixed_effect_priors():
    model = data.ModelData()

    # set prior on sex
    parameters = dict(fixed_effects={'x_sex': dict(dist='TruncatedNormal', mu=1., sigma=.5, lower=-10, upper=10)})

    # simulate normal data
    n = 32.
    sex_list = pl.array(['male', 'female', 'total'])
    sex = sex_list[mc.rcategorical([.3, .3, .4], n)]
    beta_true = dict(male=-1., total=0., female=1.)
    pi_true = pl.exp([beta_true[s] for s in sex])
    sigma_true = .05
    p = mc.rnormal(pi_true, 1./sigma_true**2.)

    model.input_data = pandas.DataFrame(dict(value=p, sex=sex))
    model.input_data['area'] = 'all'
    model.input_data['year_start'] = 2010
    model.input_data['year_start'] = 2010



    # create model and priors
    vars = {}
    vars.update(covariate_model.mean_covariate_model('test', 1, model.input_data, parameters, model,
                                                     'all', 'total', 'all'))

    print vars['beta']
    assert vars['beta'][0].parents['mu'] == 1.

 def __init__(self, stochastic, proposal_sd=None, verbose=None):
     pm.Metropolis.__init__(self, stochastic, proposal_sd=proposal_sd,
                         verbose=verbose, tally=False)
     self.proposal_tau = self.proposal_sd**-2.
     self.n = 0
     self.N = 11
     self.value = pm.rnormal(self.stochastic.value, self.proposal_tau, size=tuple([self.N] + list(self.stochastic.value.shape)))

def sim_data(N, true_cf=[[.3, .6, .1],
                           [.3, .5, .2]],
             true_std=[[.2, .05, .05], 
                       [.3, 0.1, 0.1]],
             sum_to_one=True):
    """ 
    Create an NxTxJ matrix of simulated data (T is determined by the length 
    of true_cf, J by the length of the elements of true_cf). 

    true_cf - a list of lists of true cause fractions (each must sum to one)
    true_std - a list of lists of the standard deviations corresponding to the true csmf's 
             for each time point. Can either be a list of length J inside a list of length
             1 (in this case, the same standard deviation is used for all time points) or 
             can be T lists of length J (in this case, the a separate standard deviation 
             is specified and used for each time point). 
    """

    if sum_to_one == True: 
        assert pl.allclose(pl.sum(true_cf, 1), 1), 'The sum of elements of true_cf must equal 1' 
    T = len(true_cf)
    J = len(true_cf[0])
    
    ## if only one std provided, duplicate for all time points 
    if len(true_std)==1 and len(true_cf)>1: 
        true_std = [true_std[0] for i in range(len(true_cf))]    

    ## transform the mean and std to logit space
    transformed_std = []
    for t in range(T): 
        pi_i = pl.array(true_cf[t])
        sigma_pi_i = pl.array(true_std[t])
        transformed_std.append( ((1/(pi_i*(pi_i-1)))**2 * sigma_pi_i**2)**0.5 )
            
    ## find minimum standard deviation (by cause across time) and draw from this 
    min = pl.array(transformed_std).min(0)
    common_perturbation = [pl.ones([T,J])*mc.rnormal(mu=0, tau=min**-2) for n in range(N)]
    
    ## draw from remaining variation 
    tau=pl.array(transformed_std)**2 - min**2
    tau[tau==0] = 0.000001
    additional_perturbation = [[mc.rnormal(mu=0, tau=tau[t]**-1) for t in range(T)] for n in range(N)]

    result = pl.zeros([N, T, J])
    for n in range(N):
        result[n, :, :] = [mc.invlogit(mc.logit(true_cf[t]) + common_perturbation[n][t] + additional_perturbation[n][t]) for t in range(T)]

    return result

def test_log_normal_model_sim(N=16):
    # simulate negative binomial data
    pi_true = 2.
    sigma_true = .1

    n = pl.array(pl.exp(mc.rnormal(10, 1**-2, size=N)), dtype=int)
    p = pl.exp(mc.rnormal(pl.log(pi_true), 1./(sigma_true**2 + 1./n), size=N))

    # create model and priors
    vars = dict(mu_age=mc.Uniform('mu_age', 0., 1000., value=.01),
                sigma=mc.Uniform('sigma', 0., 10000., value=1000.))
    vars['mu_interval'] = mc.Lambda('mu_interval', lambda mu=vars['mu_age']: mu*pl.ones(N))
    vars.update(rate_model.log_normal_model('sim', vars['mu_interval'], vars['sigma'], p, 1./pl.sqrt(n)))

    # fit model
    m = mc.MCMC(vars)
    m.sample(1)

def make_model(n_fmesh=11, fmesh_is_obsmesh=False):
    x = np.arange(-1., 1., .1)

    # Prior parameters of C
    nu = pm.Uniform('nu', 1., 3, value=1.5)
    phi = pm.Lognormal('phi', mu=.4, tau=1, value=1)
    theta = pm.Lognormal('theta', mu=.5, tau=1, value=1)

    # The covariance dtrm C is valued as a Covariance object.
    @pm.deterministic
    def C(eval_fun=gp.matern.euclidean,
          diff_degree=nu, amp=phi, scale=theta):
        return gp.NearlyFullRankCovariance(eval_fun, diff_degree=diff_degree, amp=amp, scale=scale)

    # Prior parameters of M
    a = pm.Normal('a', mu=1., tau=1., value=1)
    b = pm.Normal('b', mu=.5, tau=1., value=0)
    c = pm.Normal('c', mu=2., tau=1., value=0)

    # The mean M is valued as a Mean object.
    def linfun(x, a, b, c):
        return a * x ** 2 + b * x + c

    @pm.deterministic
    def M(eval_fun=linfun, a=a, b=b, c=c):
        return gp.Mean(eval_fun, a=a, b=b, c=c)

    # The actual observation locations
    actual_obs_locs = np.linspace(-.8, .8, 4)

    if fmesh_is_obsmesh:
        o = actual_obs_locs
        fmesh = o
    else:
        # The unknown observation locations
        o = pm.Normal('o', actual_obs_locs, 1000., value=actual_obs_locs)
        fmesh = np.linspace(-1, 1, n_fmesh)

    # The GP submodel
    sm = gp.GPSubmodel('sm', M, C, fmesh)

    # Observation variance
    V = pm.Lognormal('V', mu=-1, tau=1, value=.0001)
    observed_values = pm.rnormal(actual_obs_locs ** 2, 10000)

    # The data d is just array-valued. It's normally distributed about
    # GP.f(obs_x).
    d = pm.Normal(
        'd',
        mu=sm.f(o),
        tau=1. / V,
        value=observed_values,
        observed=True)

    return locals()

def simulate_age_group_data(N=50, delta_true=150, pi_true=true_rate_function):
    """ generate simulated data
    """
    # start with a simple model with N rows of data
    model = data_simulation.simple_model(N)


    # record the true age-specific rates
    model.ages = pl.arange(0, 101, 1)
    model.pi_age_true = pi_true(model.ages)


    # choose age groups randomly
    age_width = mc.runiform(1, 100, size=N)
    age_mid = mc.runiform(age_width/2, 100-age_width/2, size=N)
    age_width[:10] = 10
    age_mid[:10] = pl.arange(5, 105, 10)
    #age_width[10:20] = 10
    #age_mid[10:20] = pl.arange(5, 105, 10)

    age_start = pl.array(age_mid - age_width/2, dtype=int)
    age_end = pl.array(age_mid + age_width/2, dtype=int)

    model.input_data['age_start'] = age_start
    model.input_data['age_end'] = age_end


    # choose effective sample size uniformly at random
    n = mc.runiform(100, 10000, size=N)
    model.input_data['effective_sample_size'] = n


    # integrate true age-specific rate across age groups to find true group rate
    model.input_data['true'] = pl.nan
    model.input_data['age_weights'] = ''

    for i in range(N):
        beta = mc.rnormal(0., .025**-2)

        # TODO: clean this up, it is computing more than is necessary
        age_weights = pl.exp(beta*model.ages)
        sum_pi_wt = pl.cumsum(model.pi_age_true*age_weights)
        sum_wt = pl.cumsum(age_weights)
        p = (sum_pi_wt[age_end] - sum_pi_wt[age_start]) / (sum_wt[age_end] - sum_wt[age_start])

        model.input_data.ix[i, 'true'] = p[i]
        model.input_data.ix[i, 'age_weights'] = ';'.join(['%.4f'%w for w in age_weights[age_start[i]:(age_end[i]+1)]])

    # sample observed rate values from negative binomial distribution
    model.input_data['value'] = mc.rnegative_binomial(n*model.input_data['true'], delta_true) / n

    print model.input_data.drop(['standard_error', 'upper_ci', 'lower_ci'], axis=1)
    return model

def alpha_true_sim(model, area_list, sigma_true):
    # choose alpha^true
    alpha = dict(all=0.)
    sum_sr = 0.
    last_sr = -1
    for sr in model.hierarchy['all']:
        if sr not in area_list:
            continue

        sum_r = 0.
        last_r = -1
        for r in model.hierarchy[sr]:
            if r not in area_list:
                continue

            sum_c = 0.
            last_c = -1
            for c in model.hierarchy[r]:
                if c not in area_list:
                    continue

                alpha[c] = mc.rnormal(0., sigma_true[3]**-2.)
                sum_c += alpha[c]
                last_c = c
            if last_c >= 0:
                alpha[last_c] -= sum_c

            alpha[r] = mc.rnormal(0., sigma_true[2]**-2.)
            sum_r += alpha[r]
            last_r = r
        if last_r >= 0:
            alpha[last_r] -= sum_r

        alpha[sr] = mc.rnormal(0., sigma_true[1]**-2.)
        sum_sr += alpha[sr]
        last_sr = sr
    if last_sr >= 0:
        alpha[last_sr] -= sum_sr

    return alpha

def simulated_age_intervals(data_type, n, a, pi_age_true, sigma_true):
    # choose age intervals to measure
    age_start = np.array(mc.runiform(0, 100, n), dtype=int)
    age_start.sort()  # sort to make it easy to discard the edges when testing
    age_end = np.array(mc.runiform(age_start+1, np.minimum(age_start+10,100)), dtype=int)

    # find truth for the integral across the age intervals
    import scipy.integrate
    pi_interval_true = [scipy.integrate.trapz(pi_age_true[a_0i:(a_1i+1)]) / (a_1i - a_0i) 
                        for a_0i, a_1i in zip(age_start, age_end)]

    # generate covariates that add explained variation
    X = mc.rnormal(0., 1.**2, size=(n,3))
    beta_true = [-.1, .1, .2]
    beta_true = [0, 0, 0]
    Y_true = np.dot(X, beta_true)

    # calculate the true value of the rate in each interval
    pi_true = pi_interval_true*np.exp(Y_true)

    # simulate the noisy measurement of the rate in each interval
    p = np.maximum(0., mc.rnormal(pi_true, 1./sigma_true**2.))

    # store the simulated data in a pandas DataFrame
    data = pandas.DataFrame(dict(value=p, age_start=age_start, age_end=age_end,
                                 x_0=X[:,0], x_1=X[:,1], x_2=X[:,2]))
    data['effective_sample_size'] = np.maximum(p*(1-p)/sigma_true**2, 1.)

    data['standard_error'] = np.nan
    data['upper_ci'] = np.nan
    data['lower_ci'] = np.nan

    data['year_start'] = 2005.  # TODO: make these vary
    data['year_end'] = 2005.
    data['sex'] = 'total'
    data['area'] = 'all'
    data['data_type'] = data_type
    
    return data

def data_gen_for_rnn(samples_n=1, tau_start=75, tau_end=100, gamma=0.01, var=5):
    alpha = 1.0 / gamma
    lam = alpha
    for i in xrange(samples_n):
        con = []
        tau = pm.rdiscrete_uniform(tau_start, tau_end)
        for j in xrange(tau):
            if j == 0:
                val = round(pm.rnormal(lam, var), 2)
                con.append(val)
            elif j == 1:
                val = con[0] + pm.rnormal(0, var)
                val = round(val, 2)
                con.append(val)

            else:
                # n = len(con)
                # lam_n = float(np.array(con).sum())/n
                val = 0.7 * con[-1] + 0.3 * con[-2] + pm.rnormal(0, var)
                val = round(val, 2)
                con.append(val)
                # print val, lam_n
        yield con