import matplotlib
matplotlib.use('Agg')
import numpy as np
import pylab as plt
import os
from astrometry.util.file import *

suffix = '-default-nw064-bw008'
fnpat = 'mcmc%s-%%03i.pickle' % (suffix)
burn = 0

params = []
lnp = []
step = []
steps = np.arange(burn, 2000+1, 10)
for k in steps:
	fn = fnpat % k
	print fn
	if not os.path.exists(fn):
		break
	(k,pos,lnprob,state) = unpickle_from_file(fn)
	params.append(pos)
	lnp.append(lnprob)
	step.append(k)

step = np.array(step)
params = np.array(params)
lnp = np.array(lnp)
print 'params', params.shape
# (nsteps, nwalkers, nparams)
print 'lnp', lnp.shape

pgood = params[:,:,6]
pexif = params[:,:,7]
lncent = params[:,:,8]

(ns,nw) = pgood.shape
print 'steps', ns, 'walkers', nw

def xax():
	plt.xlabel('MCMC iteration')
	plt.xticks(np.arange(500,2000+1, 500))
	plt.xlim(-10,2010)

def plotparam(x, A):
	plt.plot(x, A, 'k-', alpha=0.05)
	y = np.mean(A, axis=1)
	yerr=np.std(A, axis=1)
	sub = 5
	x = np.append(x[::sub], x[-1])
	y = np.append(y[::sub], y[-1])
	yerr = np.append(yerr[::sub], yerr[-1])
	print 'x', x
	print 'y', y
	print 'yerr', yerr
	plt.errorbar(x, y, yerr, color='k')


#dpi=300
dpi=100
plt.figure(figsize=(4,4), dpi=dpi)
spargs = dict(bottom=0.1, left=0.2, right=0.95, top=0.9)
matplotlib.rc('font', family='computer modern roman')
matplotlib.rc('font', size=14)
matplotlib.rc('text', usetex=True)
matplotlib.rc('savefig', dpi=dpi)

plt.clf()
plt.subplots_adjust(**spargs)
plotparam(step, pgood)
xax()
plt.ylabel(r'$p_{\mathrm{good}}$')
plt.savefig('pgood.png')
plt.savefig('pgood.pdf')

plt.clf()
plt.subplots_adjust(**spargs)
plotparam(step, pexif)
xax()
plt.ylabel(r'$p_\mathrm{EXIF}$')
plt.savefig('pexif.png')
plt.savefig('pexif.pdf')

plt.clf()
plt.subplots_adjust(**spargs)
plotparam(step, np.exp(lncent))
xax()
plt.ylabel(r'$\eta$')
plt.savefig('eta.png')
plt.savefig('eta.pdf')

mx = np.mean(lnp, axis=1)[-1]
plt.clf()
plt.subplots_adjust(**spargs)
plotparam(step, lnp - mx)
xax()
plt.ylabel('ln posterior probability (relative)')
plt.ylim(-50, 10)
plt.savefig('lnp.png')
plt.savefig('lnp.pdf')