#!/usr/bin/env python
"""
A contrived toy model to see if I can get this thing to work

"""
import pylab
import numpy as np
import os
import sys
import logging
import pymc
from pymc.Matplot import plot as mcplot
import scipy.integrate

N = 1000
true_alpha = 2
# Simulated flux densities
s_arr = pymc.rpareto(true_alpha, 1, size=N)


def make_model(dbname, noise=0.0):
    alpha = pymc.Uniform('alpha', 1, 3)
    if noise is None:
        flux_meas = pymc.Pareto('flux', alpha, 1, value=s_arr, observed=True)
        vars = [flux_meas]
    else:
        flux = pymc.Pareto('flux', alpha, 1, size=N)
        flux_meas = pymc.Normal('flux_meas', mu=flux, tau=noise**-2, observed=True,
                             value=pymc.rnormal(s_arr, tau=noise**-2, size=N))

        vars = [flux, flux_meas]

    vars.append(alpha)
    pylab.figure(30)
    pylab.hist(np.log10(flux_meas.value[flux_meas.value > 0]), log=True, histtype='step', label=dbname)
    pylab.legend(frameon=False)

    model =  pymc.MCMC(vars, db='pickle', dbname=dbname)
    return model



def runmodel(name, noise, extralbl=None):
    if extralbl is None:
        extralbl = ''

    dbname='toymodel5.%s.%s.mcmc' % (name, extralbl)
    M1 = make_model(dbname, noise)

    print "Model", name, ":",
    print ', '.join([v.__name__ for v in M1.variables])
    for v in M1.variables:
        print " Name:", v
        print "    parents:", v.parents
        print "    children:", v.children


    if os.path.exists(dbname):
        db =  pymc.database.pickle.load(dbname)
        db.connect_model(M1)
        dat = db
    else:
        M1.sample(5000, 1000, 5)
        mcplot(M1)
        M1.db.close()
        dat = M1
        M1.summary()

    al = dat.trace('alpha')[:]
    pylab.figure(1)
    pylab.hist(al, histtype='step', normed=True, label='%s %0.3f' % (name, al.std()), bins=20)
    pylab.legend(frameon=False)

    print name, "alpha std",  al.std()
    return M1


def _main():
    from optparse import OptionParser
    parser = OptionParser()
    parser.set_usage('%prog [options] args')
    parser.set_description('Script description')
    parser.add_option('-v', '--verbose', dest='verbose', action='store_true', help='Be verbose [Default %default]')
    parser.set_defaults(verbose=False)
    (values, args) = parser.parse_args()
    if values.verbose:
        logging.basicConfig(level=logging.DEBUG)
    else:
        logging.basicConfig(level=logging.INFO)


    if len(args) >= 1:
        extralbl = args[0]
    else:
        extralbl = None

    runmodel('nonoise', None, extralbl)
    runmodel('lllnoise', 0.001, extralbl)
    runmodel('lllnoise2', 0.001, extralbl)
    runmodel('lllnoise3', 0.001, extralbl)
    runmodel('llnoise', 0.01, extralbl)
    runmodel('lnoise', 0.1, extralbl)
    runmodel('noise', 1, extralbl)
#    runmodel('vnoise', 10, extralbl)
#    runmodel('vvnoise', 100, extralbl)
#    runmodel('vvcnoise', 1000, extralbl)



    pylab.show()


if __name__ == '__main__':
    _main()