import pymc, seaborn as sns

import gzip

import commons

import sys

import matplotlib

import matplotlib.ticker

import matplotlib.pyplot as plt

import pickle

import random

import time

from pymc import Exponential, deterministic, Poisson, Uniform, Normal, observed, MCMC, Matplot, stochastic

from argparse import ArgumentParser

from operator import itemgetter

import PyGammaCombo

tstr = PyGammaCombo.TString

import numpy as np

def parse_dv(args, bins):

return names, lower_bounds, upper_bounds, None

ap = ArgumentParser("Using Monte-Carlo Markov-Chains for estimating LHC fits.")

# Default constants

ap.add_argument('-n', action='store', default=1000000, type=int, help='Number of events')

ap.add_argument('-b', action='store', default=5000, type=int, help='Number of burnout')

ap.add_argument('-c', action='store', required=True, type=int, help='Number of the combination')

ap.add_argument('-bins', action='store_true', help='If set, then output as histrograms')

ap.add_argument('-i', action='store_true', help='Print information about combiner')

ap.add_argument('-plot', action='store_true', help='Plot bars and graphs')

ap.add_argument('-u', action='store_true', help='Switch combination function to uniform')

ap.add_argument('-vars', nargs='+', required=True, help='Variables')

args = vars(ap.parse_args())

print(' '.join(sys.argv))

# Constants

N = args['n']

combination = args['c']

print_information = args['i']

is_mcmc = not args['u']

bins = args['bins']

plot_graph = args['plot']

burnout = args['b']

desired_variables, lower_bounds, upper_bounds, edges = parse_dv(args['vars'], bins)

sns.set_style("whitegrid")

matplotlib.rcParams.update({'font.size': 10})

pgc_utils = PyGammaCombo.gammacombo_utils

cout = PyGammaCombo.gammacombo_utils.getCout()

extract = PyGammaCombo.gammacombo_utils.extractFromRooArgSet

toRooRealVar = PyGammaCombo.gammacombo_utils.toRooRealVar

if __name__ == "__main__":

start_time = time.time()

gce = PyGammaCombo.gammacombo_utils.getGammaComboEngine("")

cmb = gce.getCombiner(combination)

cmb.combine()

parameters = cmb.getParameters()

parameter_names = list(cmb.getParameterNames())

pdf = cmb.getPdf()

if print_information:

print("Required parameters for combination:")

for p in parameter_names:

param = extract(parameters, p)

param.Print()

weights=None

if is_mcmc:

# Declare vars

var_dict = {}

for i, p in enumerate(desired_variables):

var_dict[p] = Uniform(p, doc="{}".format(p), lower=lower_bounds[i], upper=upper_bounds[i])

# Dynamically create PyMC sampling function

stochastic_args = ','.join(["{}=var_dict['{}']".format(k, k) for k in var_dict.keys()])

exec("@stochastic\n"

"def combi({}, value=0):\n"

"\tfor p in desired_variables:\n"

"\t\tparameters.setRealValue(p, var_dict[p])\n"

"\treturn pdf.getLogVal()\n".format(stochastic_args))

# Define and start sampling

mcmc = MCMC([combi] + list(var_dict.values()),

db='pickle',

dbmode='w',

dbname='mcmc-{}_combiner.pickle'.format(combination))

mcmc.sample(iter=N, burn=burnout, thin=1)

# Output

data = {v: mcmc.trace(v)[:] for v in desired_variables}

if bins:

with gzip.open('output/bins.dat.gz', 'w') as file:

data_to_save = {}

for v in data:

data_to_save[v] = np.histogram(data[v], bins=edges[v])

pickle.dump({'data': data_to_save}, file, protocol=2)

else:

with gzip.open('output/raw.dat.gz', 'w') as file:

pickle.dump({'data': data}, file, protocol=2)

else: #not MCMC

var_dict = {}

weights = []

data = {v: [] for v in desired_variables}

for i in range(N):

for i, p in enumerate(desired_variables):

rval = np.random.uniform(lower_bounds[i], upper_bounds[i])

data[p].append(rval)

parameters.setRealValue(p, rval)

weights.append(pdf.getVal())

if bins:

with gzip.open('output/bins.dat.gz', 'w') as file:

data_to_save = {}

for v in data:

data_to_save[v] = np.histogram(data[v], bins=edges[v], weights=weights)

pickle.dump({'data': data_to_save}, file, protocol=2)

else:

with gzip.open('output/raw.dat.gz', 'w') as file:

pickle.dump({'data': data, 'weights': weights}, file, protocol=2)

if plot_graph:

for v in desired_variables:

commons.plot(np.array(data[v]), combination, v, weights=weights)

print("Total time: {}".format(time.time() - start_time))