#%% Define model
import numpy,math
import matplotlib.pyplot as plt
import random as random

import theano.tensor as t
import theano

random.seed(1)  # set random seed

# copy-pasted function from the specific model used in the source and adapted with as_op
@theano.compile.ops.as_op(itypes=[t.iscalar, t.dscalar, t.fscalar, t.fscalar],otypes=[t.dvector])
def sampleFromSalpeter(N, alpha, M_min, M_max):
    log_M_Min = math.log(M_min)
    log_M_Max = math.log(M_max)
    maxlik = math.pow(M_min, 1.0 - alpha)
    Masses = []
    while (len(Masses) < N):
        logM = random.uniform(log_M_Min,log_M_Max)
        M    = math.exp(logM)
        likelihood = math.pow(M, 1.0 - alpha)
        u = random.uniform(0.0,maxlik)
        if (u < likelihood):
            Masses.append(M)
    return Masses

# SAME function as above, used to make test data (so no Theano here)
def sampleFromSalpeterDATA(N, alpha, M_min, M_max):
    log_M_Min = math.log(M_min)
    log_M_Max = math.log(M_max)
    maxlik = math.pow(M_min, 1.0 - alpha)
    Masses = []
    while (len(Masses) < N):
        logM = random.uniform(log_M_Min,log_M_Max)
        M    = math.exp(logM)
        likelihood = math.pow(M, 1.0 - alpha)
        u = random.uniform(0.0,maxlik)
        if (u < likelihood):
            Masses.append(M)
    return Masses

# Generate toy data.
N      = 1000000  # Draw 1 Million stellar masses.
alpha  = 2.35
M_min  = 1.0
M_max  = 100.0
Masses = sampleFromSalpeterDATA(N, alpha, M_min, M_max)

#%% Estimation process
import pymc3 as pm
basic_model = pm.Model()
with basic_model:

    # Priors for unknown model parameters
    alpha2 = pm.Normal('alpha2', mu=3, sd=10)

    N2=t.constant(1000000)
    M_min2  = t.constant(1.0)
    M_max2  = t.constant(100.0)

    # Expected value of outcome
    m =  sampleFromSalpeter(N2, alpha2, M_min2, M_max2)

    # Likelihood (sampling distribution) of observations
    Y_obs = pm.Normal('Y_obs', mu=m, sd=10, observed=Masses)

#%% Sample
with basic_model:
    step = pm.Metropolis()
    trace = pm.sample(10000, step=step)

#%% Plot posteriors
_ = pm.traceplot(trace)
pm.summary(trace)