import time

import numpy as np
import matplotlib.pyplot as plt
import pymc3 as pm
import theano.tensor as tt
import pystan


# CONFIG
# ----------------------------------------------------------------------
R = 10                          # number of runs
K = 5                           # predictor dimension
N = 10000                       # number of samples

# Data generation
# ----------------------------------------------------------------------
def scale(w0, w1):
    scale = np.sqrt(w0**2 + np.sum(w1**2))
    return w0/scale, w1/scale


X = np.random.rand(N, K)        # predictors
w1 = np.random.rand(K)          # weights
w0 = -np.median(np.dot(X, w1))  # 50% of elements for each class
w0, w1 = scale(w0, w1)          # unit length (euclidean)
y = w0 +  np.dot(X, w1) > 0     # labels


def stan_model(n_samples=1000, n_chains=8, verbose=False):
    logit_code = r"""
    functions {
        real sigmoid(real x) {
            return 1.0/(1.0 + exp(-x));
        }
    }
    data {
        int<lower=0> N;             // number of samples
        int<lower=0> K;             // number of predictors
        real X[N, K];               // sample predictors
        int<lower=0,upper=1> y[N];  // labels
    }

    parameters {
        real w0;
        real w1[K];
    }

    transformed parameters {
        real p[N];
        real x;
        for (i in 1:N) {
            x = w0;
            for (j in 1:K) {
                x = x + w1[j]*X[i, j];
            }
            p[i] = sigmoid(x);
        }
    }

    model {
        w0 ~ normal(0, 20);
        w1 ~ normal(0, 20);
        y ~ bernoulli(p);
    }
    """

    logit_data = {
        'N': len(y),
        'K': K,
        'X': X,
        'y': y.astype(int)
    }
    fit = pystan.stan(
        model_code=logit_code,
        data=logit_data,
        iter=n_samples,
        chains=8,
        verbose=verbose)
    trace = fit.extract()
    return trace


def pymc_model(n_samples=8000, verbose=False):
    with pm.Model():
        w0_pred = pm.Normal('w0', 0, 20)
        w1_pred = pm.Normal('w1', 0, 20, shape=K)
        y_pred = pm.Bernoulli('y',
                              p=tt.nnet.sigmoid(w0_pred + tt.dot(X, w1_pred)),
                              observed=y.astype(int))
        trace = pm.sample(n_samples, progressbar=verbose)
        return trace


def get_weights(trace, discard=0):
    if discard == 'half':
        discard = len(trace['w0'])/2
    return scale(
        np.mean(trace['w0'][discard:], axis=0),
        np.mean(trace['w1'][discard:], axis=0))

w_true = np.hstack((w0, w1))
w_stan = np.empty((R, K + 1))
w_pymc = np.empty((R, K + 1))
for i in range(R):
    trace_stan = stan_model(100, 8)
    w_stan[i, :] = np.hstack(get_weights(trace_stan))

    trace_pymc = pymc_model(800)
    w_pymc[i, :] = np.hstack(get_weights(trace_pymc, discard='half'))

def avg_error(w_true, w_pred):
    return np.mean(np.sum((w_pred - w_true)**2, axis=1), axis=0)

print('Avg error (stan):', avg_error(w_true, w_stan))
print('Avg error (pymc):', avg_error(w_true, w_pymc))

# Avg error (stan): 4.60143514284e-06
# Avg error (pymc): 0.00473672889666