def test_specify_args(self):
     y = (0.70,  -0.16,  0.77, -1.37, -1.99,  1.35, 0.08,
          0.02,  -1.48, -0.08,  0.34,  0.03, -0.42, 0.87,
          -1.36,  1.43,  0.80, -0.48, -1.61, -1.27)
     code = """
         data {
             real y[20];
         }
         parameters {
             real mu;
             real<lower=0> sigma;
         }
         model {
             y ~ normal(mu, sigma);
         }"""
     stepsize0 = 0.15
     sf = stan(model_code=code, data=dict(y=y), iter=200,
               control=dict(adapt_engaged=False, stepsize=stepsize0))
     self.assertEqual(sf.get_sampler_params()[0]['stepsize__'][0], stepsize0)
     sf2 = stan(fit=sf, iter=20, algorithm='HMC', data=dict(y=y),
                control=dict(adapt_engaged=False, stepsize=stepsize0))
     self.assertEqual(sf2.get_sampler_params()[0]['stepsize__'][0], stepsize0)
     sf3 = stan(fit=sf, iter=1, data=list(y=y), init=0, chains=1)
     i_u = sf3.unconstrain_pars(sf3.get_inits[0])
     self.assertEqual(i_u, [0, 0])

 def test_user_init(self):
     model_code = self.model_code
     data = self.data
     beta = np.ones((data['K'], data['D']))
     fit1 = pystan.stan(model_code=model_code, iter=10, chains=1, seed=2,
                        data=data, init=[dict(beta=beta)], warmup=0)
     np.testing.assert_equal(fit1.get_inits()[0]['beta'], beta)
     beta = 5 * np.ones((data['K'], data['D']))
     fit2 = pystan.stan(model_code=model_code, iter=10, chains=1, seed=2,
                        data=data, init=[dict(beta=beta)], warmup=0)
     np.testing.assert_equal(fit2.get_inits()[0]['beta'], beta)

    def _initialize(self, xs):
        print("The following message exists as Stan instantiates the model.")
        if hasattr(self, 'file'):
            self.model = pystan.stan(file=self.file,
                                     data=xs, iter=1, chains=1)
        else:
            self.model = pystan.stan(model_code=self.model_code,
                                     data=xs, iter=1, chains=1)

        self.num_vars = sum([sum(dim) if sum(dim) != 0 else 1
                             for dim in self.model.par_dims])
        self.flag_init = True

    def setUpClass(cls):

        np.random.seed(1)

        n = 10000
        p = 3

        cls.beta_true = beta_true = (1, 3, 5)
        X = np.random.normal(size=(n, p))
        X = (X - np.mean(X, axis=0)) / np.std(X, ddof=1, axis=0, keepdims=True)
        y = np.dot(X, beta_true) + np.random.normal(size=n)

        model_code = """
        data {
            int<lower=0> N;
            int<lower=0> p;
            matrix[N,p] x;
            vector[N] y;
        }
        parameters {
            vector[p] beta;
            real<lower=0> sigma;
        }
        model {
            y ~ normal(x * beta, sigma);
        }
        """

        data = {'N': n, 'p': p, 'x': X, 'y': y}

        cls.fit = pystan.stan(model_code=model_code, data=data, iter=500)

 def test_stan_file(self):
     schools_code = self.schools_code
     schools_dat = self.schools_dat
     with tempfile.NamedTemporaryFile(delete=False) as f:
         f.write(schools_code.encode('utf-8'))
     fit = pystan.stan(file=f.name, data=schools_dat, iter=1000, chains=4)
     validate_data(fit)

def fit():
    data = make_data()
    model = pystan.stan(model_code=stan_code, data=data, 
                        iter=1000, chains=5)
    #model.plot()
    #plt.show()
    return model

    def test_stan_args_basic(self):
        y = np.array([0.70,  -0.16,  0.77, -1.37, -1.99,  1.35, 0.08,
                      0.02,  -1.48, -0.08,  0.34,  0.03, -0.42, 0.87,
                      -1.36,  1.43,  0.80, -0.48, -1.61, -1.27])

        code = '''
        data {
            real y[20];
        }
        parameters {
            real mu;
            real<lower=0> sigma;
        }
        model {
            y ~ normal(mu, sigma);
        }'''

        sf = pystan.stan(model_code=code, iter=10, thin=3, data=dict(y=y))
        args = sf.stan_args[0]
        self.assertEqual(args['iter'], 10)
        self.assertEqual(args['thin'], 3)
        self.assertEqual(args['init'], b'random')

        sampling = args['ctrl']['sampling']
        self.assertEqual(sampling['adapt_engaged'], True)
        self.assertEqual(sampling['adapt_window'], 25)
        self.assertEqual(sampling['adapt_init_buffer'], 75)
        self.assertEqual(sampling['adapt_gamma'], 0.05)
        self.assertEqual(sampling['adapt_delta'], 0.8)
        self.assertEqual(sampling['adapt_kappa'], 0.75)
        self.assertEqual(sampling['adapt_t0'], 10)

def fit(model_code, *args, **kwargs):
    """
    Fit a Stan model. Caches the compiled model.

    *args and **kwargs are passed to the pystan.stan function.

    Arguments you most likely want to pass: data, init, iter, chains.        

    Unlike pystan.stan, if the n_jobs kwarg is not specified, it defaults to
    -1.

    Parameters
    -------------------
    model_code : string
        Stan model


    Returns
    -------------------
    pystan StanFit4Model instance : the fit model
    """
    kwargs = dict(kwargs)
    kwargs['model_code'] = model_code
    if 'n_jobs' not in kwargs:
        kwargs['n_jobs'] = -1
    if model_code in FIT_CACHE:
        print("Reusing model.")
        kwargs['fit'] = FIT_CACHE[model_code]
    else:
        print("NOT reusing model.")
    start = time.time()
    FIT_CACHE[model_code] = pystan.stan(*args, **kwargs)
    print("Ran in %0.3f sec." % (time.time() - start))
    return FIT_CACHE[model_code]

def run_stan(stan_data, target):

    project_dir = os.path.join(os.environ['HOME'], 'Projects', 'Kaggle', 'otto')
    data_dir = os.path.join(project_dir, 'data')

    iguesses = get_initial_values(stan_data['counts'])

    stan_file = os.path.join(project_dir, 'stan', 'single_component.stan')
    fit = pystan.stan(file=stan_file, model_name='single_component_' + target, data=stan_data,
                      chains=4, iter=400, warmup=200, init=iguesses)

    # dump the MCMC samples to an HDF file
    samples = fit.extract()

    cnames = ['concentration', 'negbin_nfailures']
    nbins = stan_data['counts'].shape[1]
    cnames += ['bin_prob_' + str(i + 1) for i in range(nbins)]

    raw_samples = np.column_stack((samples['concentration'], samples['negbin_nfailures'], samples['bin_probs']))
    samples = pd.DataFrame(data=raw_samples, columns=cnames)
    samples.to_hdf(os.path.join(data_dir, 'single_component_' + target + '_samples.h5'), 'df')

    # dump the stan results to a text file
    with open(os.path.join(data_dir, 'single_component_' + target + '_stan_summary.txt'), 'w') as f:
        print >> f, fit

    # make plots of the stan results
    plot_dir = os.path.join(project_dir, 'plots')
    fit.plot()
    plt.savefig(os.path.join(plot_dir, 'single_component_' + target + '_trace.png'))

    return

def main():
  # data
  J = 8
  data_y = np.array([28, 8, -3, 7, -1, 1, 18, 12])
  data_sigma = np.array([15, 10, 16, 11, 9, 11, 10, 18])

  standata = dict(J=J, y=data_y, sigma=data_sigma)
  fit = pystan.stan('eight_schools.stan', data=standata, iter=100000)
  print(fit)

def gamma_tutorial():
    '''
    Basic stan gamma GLM with log link function
    Taken from http://seananderson.ca/2014/04/08/gamma-glms.html
    :return:
    '''
    # define stan data
    N = 100
    x = np.random.uniform(-1, 1, N)
    a = 0.5
    b = 1.2
    y_true = np.exp(a + b * x)
    shape = 10.0

    # Passing a vector for scale generates one sample per value in vector; |y_true| = 100 => |y| = 100
    # note that scale is generated in odd fashion (w/ shape parameter), hence odd shape parameter in stan model
    y = np.random.gamma(shape, scale=y_true/shape)

    # Now we put data into a dictionary so we can give to stan
    this_data = {
        'N': N,
        'x': x,
        'y': y
    }

    # define stan model
    stan_code = """
    data {
      int<lower=0> N;
      vector[N] x;
      vector[N] y;
    }
    parameters {
      real a;
      real b;
      real<lower=0.001,upper=100> shape;
    }
    model {
      a ~ normal(0,100);
      b ~ normal(0,100);
      for(i in 1:N)
        y[i] ~ gamma(shape, (shape / exp(a + b * x[i])));
    }
    """

    # fit model
    fit = pystan.stan(model_code = stan_code, data=this_data, iter=1000, chains=4, thin=3)

    '''
    Note that the following two statements are equivalent:
    y ~ normal(alpha + beta * x, sigma)
    for (n in 1:N):
        y[n] ~ normal(alpha + beta * x[n], sigma)
    ...meaning that stan automatically vectorizes
    '''

    """

def main():
    # Load the data for the toy model
    ids, x, y, xerr, yerr, pxy = np.loadtxt("../data/hogg-toy-model.data", unpack=True)

    with open("models/mixture-model-with-uncertainties.stan", "r") as fp:
        model_code = fp.read()

    # Fit the model
    fit = pystan.stan(model_code=model_code, iter=10000, chains=8,
        data={
            "x_measured": x, "x_uncertainty": xerr,
            "y_measured": y, "y_uncertainty": yerr,
            "N": len(x)
        })

    print(fit)
    fit.traceplot()

    samples = fit.extract(permuted=True)
    parameters = pd.DataFrame({"m": samples["m"], "b": samples["b"]})

    # Predictive model
    pred_x = np.arange(0, 300)
    model = lambda theta: pd.Series({"fitted": theta[0] + theta[1] * pred_x})

    median_parameters = parameters.median()

    yhat = model(median_parameters)

    # get the predicted values for each chain
    chain_predictions = parameters.apply(model, axis=1)

    fig = plt.figure()
    ax = fig.add_subplot(111)

    num_chains = 50
    indices = np.random.choice(300, num_chains)

    for i, index in enumerate(indices):
        ax.plot(pred_x, chain_predictions.iloc[index, 0], color="lightgrey")

    #  data
    ax.errorbar(x, y, xerr=xerr, yerr=yerr, fmt=None, facecolor="k", ecolor="k", zorder=10)
    ax.plot(x, y, 'ko', zorder=10)

    # fitted values
    ax.plot(pred_x, yhat["fitted"], "k", lw=2)

    # supplementals
    ax.set_xlim(0, 300)
    ax.set_ylim(0, 750)

    ax.set_xlabel('X')
    ax.set_ylabel('Y')

    plt.show()

def gamma_glm(data):
    '''
    Gamma GLM for use in bayesian class project
    input: data - dictionary of data for stan model
    :return:
    '''

    # if you want to make this better/run more expts, dynamically generate some of this model.
    # presently, assume flat priors for all parameters
    stan_data = data[0]

    stan_model = """
    data {
        int<lower=0> N;              // number of samples
        int<lower=0> N_test;
        int<lower=0> K;              // length of each X vector
        matrix[N,K] x;
        vector[N] y;

        matrix[N_test,K] x_test;    // samples for testing
    }

    parameters {
        real<lower=.001> beta_0;
        vector<lower=.001>[K] beta;
        real<lower=.001> phi;
        real<lower=.001> alpha;
    }

    model {
        alpha ~ normal(20,10);
        for (k in 1:K)
            beta[k] ~ normal(0,1000);
        beta_0 ~ normal(10000,1000);
        phi ~ normal(0,2);
        for (k in 1:K)
            y ~ gamma(alpha, inv(x[k] * beta + beta_0) + phi);       // inverse-link
    }

    generated quantities {                              // predictions!
        vector[N_test] y_test;
        for (n in 1:N_test)
            y_test[n] <- gamma_rng(alpha, inv(x_test[n] * beta + beta_0) + phi);
    }
    """

    # fit model
    fit = pystan.stan(model_code=stan_model, data=stan_data, iter=2000, chains=4, thin=1)

    print "Model:"

    print fit

    eval_acc(fit, data)

    def test_user_initfun_chainid(self):
        model_code = self.model_code
        data = self.data

        def make_inits(chain_id):
            return dict(mu=chain_id)

        chain_id = [9, 10, 11, 12]
        fit1 = pystan.stan(model_code=model_code, iter=10, chains=4, seed=2,
                           data=data, init=make_inits, warmup=0, chain_id=chain_id)
        for i, inits in zip(chain_id, fit1.get_inits()):
            self.assertEqual(inits['mu'], i)

def get_basic_model_fit(prior_params, num_iters, num_chains, init_f, seed, the_type, data):
    """
    convention is to return by permuted and unpermuted traces
    """
    import problem_props.constants as constants
    import pystan
    stan_file = constants.basic_model_stan_file
    stan_data = data_to_basic_model_pystan_input(data)
    other_inputs = {\
        'n_types': max(stan_data['types']),\
            'n_props': len(data),\
            'n_events': len(stan_data['types']),\
            'the_type': the_type + 1\
            }
    all_stan_data = dict(prior_params.items() + stan_data.items() + other_inputs.items())
    if init_f == None:
        fit = pystan.stan(file = stan_file, data = all_stan_data, seed = seed, iter = num_iters, chains = num_chains, verbose = True)
    else:
        init_d = [init_f(data, i) for i in range(num_chains)]
        fit = pystan.stan(file = stan_file, data = all_stan_data, seed = seed, iter = num_iters, chains = num_chains, verbose = True, init = init_d)
    return fit.extract(permuted=True), fit.extract(permuted=False)

def test_array_param():
    """
    Make sure shapes are getting unraveled correctly. Mixing up row-major and
    column-major data is a potential issue.
    """
    model_code = """
    data {
      int<lower=2> K;
    }
    parameters {
      real beta[K,1,2];
    }
    model {
      for (k in 1:K)
        beta[k,1,1] ~ normal(0,1);
      for (k in 1:K)
        beta[k,1,2] ~ normal(100,1);
    }"""
    fit = stan(model_code=model_code, data=dict(K=4))

    # extract, permuted
    beta = fit.extract()['beta']
    assert beta.shape == (4000, 4, 1, 2)
    beta_mean = np.mean(beta, axis=0)
    assert beta_mean.shape == (4, 1, 2)
    assert np.all(beta_mean[:, 0, 0] < 4)
    assert np.all(beta_mean[:, 0, 1] > 100 - 4)

    # extract, permuted=False
    extracted = fit.extract(permuted=False)
    assert extracted.shape == (1000, 4, 9)
    # in theory 0:4 should be
    # 'beta[0,0,0]'
    # 'beta[1,0,0]'
    # 'beta[2,0,0]'
    # 'beta[3,0,0]'
    #
    # and 4:8 should be
    # 'beta[0,0,1]'
    # 'beta[1,0,1]'
    # 'beta[2,0,1]'
    # 'beta[3,0,1]'
    assert np.all(np.mean(extracted[:, :, 0:4], axis=(0, 1)) < 4)
    assert np.all(np.mean(extracted[:, :, 4:8], axis=(0, 1)) > 100 - 4)
    assert np.all(extracted[:, :, 8] < 0)  # lp__

    # optimizing
    sm = fit.stanmodel
    op = sm.optimizing(data=dict(K=4))
    beta = op['beta']
    assert beta.shape == (4, 1, 2)
    assert np.all(beta[:, 0, 0] < 4)
    assert np.all(beta[:, 0, 1] > 100 - 4)

    def test_user_initfun(self):
        model_code = self.model_code
        data = self.data

        beta = np.ones((data['K'], data['D']))

        def make_inits(chain_id):
            return dict(beta=beta * chain_id)

        fit1 = pystan.stan(model_code=model_code, iter=10, chains=4, seed=2,
                           data=data, init=make_inits, warmup=0)
        for i, inits in enumerate(fit1.get_inits()):
            np.testing.assert_equal(beta * i, inits['beta'])

def get_gaussian_subspace_model_traces(num_iters, num_chains, seed, init_d, n_clusters, data, hypers):
    """
    hypers should be in dictionary form, read to pass to pystan
    parallel version would map this function partialed with only seed undetermined to a list of seeds, then reduce the results
    returns list of dictionary of param_name:traces
    """
    
    import crime_pattern.constants as constants

    stan_file = '%s/%s' % (constants.stan_folder, 'gaussian_subspaces.stan')
    N, n_dim = len(data), len(iter(data).next())
    import pystan
    pystan_data = hypers.copy()
    pystan_data['d'] = data
    pystan_data['N'] = N
    pystan_data['n_dim'] = n_dim
    pystan_data['n_clusters'] = n_clusters

    if init_d == None:
        fit = pystan.stan(file=stan_file, data=pystan_data, seed = seed, iter=num_iters, chains=num_chains, verbose=True)
    else:
        fit = pystan.stan(file=stan_file, data=pystan_data, seed = seed, iter=num_iters, chains=num_chains, verbose=True, init=[init_d for i in xrange(num_chains)])
    return fit.extract()

def basic_linear(data):
    '''
    Use as baseline for stan implementation
    :return:
    '''

    # Model for matrix linear regression with D = {X, Y}, X is an n*k matrix, y is n*1 values
    # explicit priors are specified in the model section, otherwise a uniform prior is assumed
    # use generated quantities to predict N new values:
    stan_data = data[0]

    stan_model = """
    data {
        int<lower=0> N;
        int<lower=0> N_test;
        int<lower=0> K;
        matrix[N,K] x;
        vector[N] y;

        matrix[N_test,K] x_test;
    }

    parameters {
        real alpha;
        vector[K] beta;
        real <lower=0> sigma;
    }

    model {
        y ~ normal(alpha + x * beta, sigma);
    }
    generated quantities {
        vector[N_test] y_test;
        for (n in 1:N_test)
            y_test[n] <- normal_rng(x_test[n] * beta + alpha, sigma);   // NOTE: generated quantities do NOT vectorize
    }

    """

    fit = pystan.stan(model_code=stan_model, data=stan_data, iter=1000, chains=4, thin=1)

    print 'Model:'
    print fit

    # fit.traceplot()
    # py.show()
    eval_acc(fit, data)

 def test_user_init_unspecified(self):
     model_code = """
     data {
       real x;
     }
     parameters {
       real mu;
       real<lower=0> sigma;
     }
     model {
       x ~ normal(mu, sigma);
     }
     """
     data = self.data
     # NOTE: we are only specifying 'mu' and not 'sigma' (partial inits)
     fit = pystan.stan(model_code=model_code, iter=10, chains=1, seed=2, data=data, init=[dict(mu=4)], warmup=0)
     self.assertIsNotNone(fit)