#!/usr/bin/python
# line_mcmc -- MCMC errors for slope and intercept of straight line
# line_mcmc  datafile t1init t2init step1 step2 nchain nthin outfile
#   datafile = file with columns x y error
#   t1init = initial guess for intercept
#   t2init = initial guess for slope
#   step1 = step size for intercept
#   step2 = step size for slope
#   nchain = desired number of steps in chain
#   nthin = thinning factor, nthin=1 for no thinning
#   outfile = name of output file
#
# The underlying model is y = t1 + t2*x.
# The code reads data and errors (assumed uncorrelated) from datafile
# and does an MCMC chain for t1 and t2 using these data values.

import numpy as np
from numpy import matrix
from math import *
import sys

datafile=sys.argv[1]
t1init=float(sys.argv[2])
t2init=float(sys.argv[3])
step1=float(sys.argv[4])
step2=float(sys.argv[5])
nchain=int(sys.argv[6])
nthin=int(sys.argv[7])
outfile=sys.argv[8]

# reads columns 1, 2, and 3 into arrays x, y, and errors
x,y,errors=np.loadtxt(datafile,unpack=True)

def prob(x,cinv,prefac):
    """
    Return multivariate Gaussian probability 
    x = vector of data values  (matrix)
    cinv = inverse covariance matrix 
    prefac = prefactor for properly normalized Gaussian (float),
             taken as input so that it isn't computed every call
             should be [(2\pi)^{M/2} \sqrt{det(C)}]^{-1}
    """
    arg=float(x*cinv*x.T)
    return (prefac*exp(-arg/2.))


### MCMC chain for parameters ###

cov=np.diag(errors)
cinv=np.linalg.inv(cov)
prefac=1./(2*pi*sqrt(np.linalg.det(cov)))

t1=t1init
t2=t2init
deltay = y - (t1 + t2*x)
dym=matrix(deltay)
p1=prob(dym,cinv,prefac)        # probability at starting point

chain=np.zeros((nchain,3))    # store (theta1,theta2,p) as elements of chain
chain[0][0]=t1
chain[0][1]=t2
chain[0][2]=p1

ichain=1
naccept=0
while (ichain<nchain):
    t1trial=t1+step1*np.random.normal()
    t2trial=t2+step2*np.random.normal()
    deltay = y - (t1trial + t2trial*x)
    dym=matrix(deltay)
    p2=prob(dym,cinv,prefac)            # compute probability
    if ((p2>p1 or p1==0) or np.random.random()<p2/p1):    # accept step?
        chain[ichain][0]=t1trial
        chain[ichain][1]=t2trial
        chain[ichain][2]=p2
        t1=t1trial
        t2=t2trial
        p1=p2
        naccept+=1
    else:
        chain[ichain][0]=t1
        chain[ichain][1]=t2
        chain[ichain][2]=p2
    ichain+=1

xaccept=float(naccept)/float(ichain)
print '%.4f of trials accepted' % (xaccept)
np.savetxt(outfile,chain[::nthin],'%11.8f')