import pymc
import numpy as np
import matplotlib.pyplot as plt
import scipy.stats as stats
%matplotlib inline

fig = plt.figure(1, [12, 6], dpi=300)
fig.subplots_adjust(hspace=.3)
ax1 = fig.add_subplot(2,4,1)
counts, bins, ignored = ax1.hist(data_obs, normed=False, color='black')
ax1.set_xlabel(r'$\sigma_{obs}$', fontdict={'size':18})
ax1.set_xscale('log')

ax2 = fig.add_subplot(2,4,2)
counts, bins, ignored = ax2.hist(sigma_r, normed=False, color='blue')
ax2.set_xlabel(r'$\sigma_{rock}$', fontdict={'size':18})

ax3 = fig.add_subplot(2,4,3)
counts, bins, ignored = ax3.hist(sigma_f, normed=False, color='red')
ax3.set_xlabel(r'$\sigma_{fluid}$', fontdict={'size':18})

ax4 = fig.add_subplot(2,4,4)
counts, bins, ignored = ax4.hist(sigma_m, normed=False, color='green')
ax4.set_xlabel(r'$\sigma_{mineral}$', fontdict={'size':18})

ax6 = fig.add_subplot(2,4,6)
counts, bins, ignored = ax6.hist(m_r, normed=False, color='blue')
ax6.set_xlabel(r'$m_{rock}$', fontdict={'size':18})

ax7 = fig.add_subplot(2,4,7)
counts, bins, ignored = ax7.hist(m_f, normed=False, color='red')
ax7.set_xlabel(r'$m_{fluid}$', fontdict={'size':18})

ax8 = fig.add_subplot(2,4,8)
counts, bins, ignored = ax8.hist(m_m, normed=False, color='green')
ax8.set_xlabel(r'$m_{mineral}$', fontdict={'size':18})
plt.show()

    
# prior information of unknown model parameters
phi_r = pymc.Normal('phi_r', .7, .2)
phi_f = pymc.Normal('phi_f', .5, .3)
phi_m = pymc.Normal('phi_m', .5, .3)

#m_r = pymc.Normal('m_r', .03, .015)
#m_f = pymc.Normal('m_f', 2, .3)
#m_m = pymc.Normal('m_m', 3.5, .3)

# expected outcome is governed by Archies equation of Glove 2010
@pymc.deterministic
def archies(s_r=sigma_r, s_f=sigma_f, s_m=sigma_m, 
            p_r=phi_r, p_f=phi_f, p_m=phi_m, 
            m_r=m_r, m_f=m_f, m_m=m_m):
    archie_return = (s_r*p_r**m_r)+(s_f*p_f**m_f)+(s_m*p_m**m_m)
    return archie_return


# likelihood of the observations is a normal distribution
# and is the data itself
sigma_obs = pymc.Normal('sigma_obs', mu=archies, value=data_obs, observed=True)

archie_model = pymc.Model([archies, sigma_r, sigma_f, sigma_m, 
                           phi_r, phi_f, phi_m, 
                           m_r, m_f, m_m, 
                           sigma_obs, 
                           data_obs])

model_fit = pymc.MCMC(archie_model)
model_fit.sample(iter=10000)

pr_arr = model_fit.trace('phi_r')[:]
pf_arr = model_fit.trace('phi_f')[:]
pm_arr = model_fit.trace('phi_m')[:]

pr_arr = pr_arr[np.where((pr_arr>=0)& (pr_arr<=1))]
pf_arr = pf_arr[np.where((pf_arr>=0)& (pf_arr<=1))]
pm_arr = pm_arr[np.where((pm_arr>=0)& (pm_arr<=1))]

# make sure that the rock percentage is confined to values where the 
# total proportion is equal to 1
pr_arr = pr_arr[np.where(pr_arr > 1-(stats.mode(pf_arr)[0][0]+stats.mode(pm_arr)[0][0])*1.2)]

fig_2 = plt.figure(2, [12, 6], dpi=300)
ax1 = fig_2.add_subplot(1, 3, 1)
h1 = ax1.hist(pr_arr, bins=20, color='blue')
ax1.set_xlim(0, 1)
ax1.set_xlabel(r'$\phi_{rock}$', fontdict={'size':18})

ax2 = fig_2.add_subplot(1, 3, 2)
h2 = ax2.hist(pf_arr, bins=20, color='red')
ax2.set_xlim(0, 1)
ax2.set_xlabel(r'$\phi_{fluid}$', fontdict={'size':18})


ax3 = fig_2.add_subplot(1, 3, 3)
h3 = ax3.hist(pm_arr, bins=20, color='green')
ax3.set_xlim(0, 1)
ax3.set_xlabel(r'$\phi_{mineral}$', fontdict={'size':18})


print '             Rock   Fluid    Mineral'
print 'Mean Percent   {0:.2f}   {1:.2f}   {2:.2f}'.format(pr_arr.mean(),
                                                          pf_arr.mean(),
                                                          pm_arr.mean())
print 'Median Percent {0:.2f}   {1:.2f}   {2:.2f}'.format(np.median(pr_arr),
                                                          np.median(pf_arr),
                                                          np.median(pm_arr))
print 'Mode Percent   {0:.2f}   {1:.2f}   {2:.2f}'.format(stats.mode(pr_arr)[0][0],
                                                          stats.mode(pf_arr)[0][0],
                                                          stats.mode(pm_arr)[0][0])
print 'Min Percent    {0:.2f}   {1:.2f}   {2:.2f}'.format(pr_arr.min(),
                                                          pf_arr.min(),
                                                          pm_arr.min())
print 'Max Percent    {0:.2f}   {1:.2f}   {2:.2f}'.format(pr_arr.max(),
                                                          pf_arr.max(),
                                                          pm_arr.max())
print 'STD            {0:.2f}   {1:.2f}   {2:.2f}'.format(pr_arr.std(),
                                                          pf_arr.std(),
                                                          pm_arr.std())