import pymc
import numpy as np
import pylab as pl
from measure_tau import mcmc_sampler_dict, tauoneone, tautwotwo, savepath, domillion, abundance, opr, savefig, trace_data_path
from mcmc_tools import docontours_multi, save_traces
from agpy import pymc_plotting
import pymc_tools

print "Beginning Lognormal parameter estimation using abundance=", abundance, ' opr= ', opr

mc_simple = pymc.MCMC(
    mcmc_sampler_dict(tauoneone=tauoneone, tautwotwo=tautwotwo))

graph_lognormal_simple = pymc.graph.graph(mc_simple)
graph_lognormal_simple.write_pdf(savepath + "mc_lognormal_simple_graph.pdf")
graph_lognormal_simple.write_png(savepath + "mc_lognormal_simple_graph.png")

d = mcmc_sampler_dict(
    tauoneone=tauoneone, tautwotwo=tautwotwo, truncate_at_50sigma=True)
d['b'] = pymc.Uniform(name='b', value=0.5, lower=0.3, upper=1, observed=False)


@pymc.deterministic(plot=True, trace=True)
def mach(sigma=d['sigma'], b=d['b']):
    return np.sqrt((np.exp(sigma**2) - 1) / b**2)


d['mach'] = mach
d['mach_observed'] = pymc.Normal(
    name='mach_observed', mu=mach, tau=1. / 1.5**2, value=5.1, observed=True)
    trace_data_path,
)
from mcmc_tools import docontours_multi, save_traces
from agpy import pymc_plotting
import pymc_tools
import hopkins_pdf

if "abundance" not in locals():
    from measure_tau import abundance
if "opr" not in locals():
    from measure_tau import opr

print "Beginning Hopkins parameter estimation using abundance=", abundance, " opr=", opr

# Hopkins - NO Mach number restrictions
d = mcmc_sampler_dict(tauoneone=tauoneone_hopkins, tautwotwo=tautwotwo_hopkins)
mc_hopkins_simple = pymc.MCMC(d)

graph_hopkins_simple = pymc.graph.graph(mc_hopkins_simple)
graph_hopkins_simple.write_pdf(savepath + "mc_hopkins_simple_graph.pdf")
graph_hopkins_simple.write_png(savepath + "mc_hopkins_simple_graph.png")

# Hopkins - with Mach number restrictions
d = mcmc_sampler_dict(tauoneone=tauoneone_hopkins, tautwotwo=tautwotwo_hopkins)


def Tval(sigma):
    return hopkins_pdf.T_of_sigma(sigma, logform=True)


d["Tval"] = pymc.Deterministic(name="Tval", eval=Tval, parents={"sigma": d["sigma"]}, doc="Intermittency parameter T")