import cPickle,numpy,pyfits as py

import pymc

from pylens import *

from imageSim import SBModels,convolve

import indexTricks as iT

from SampleOpt import AMAOpt

import pylab as pl

import numpy as np

import myEmcee_blobs as myEmcee

#import myEmcee

from matplotlib.colors import LogNorm

from scipy import optimize

from scipy.interpolate import RectBivariateSpline

# likelihood seems even slower to increase than before. Debug this at some point - it should work?

X = 'xytest'

print X

# plot things

def NotPlicely(image,im,sigma):

#pl.savefig('/data/ljo31/Lens/TeXstuff/plotrun'+str(X)+'.png')

img1 = py.open('/data/ljo31/Lens/J0901/F606W_sci_cutout.fits')[0].data.copy()

sig1 = py.open('/data/ljo31/Lens/J0901/F606W_noisemap.fits')[0].data.copy()

psf1 = py.open('/data/ljo31/Lens/J0901/F606W_psf2.fits')[0].data.copy()

psf1 = psf1[5:-6,5:-6]

psf1 = psf1/np.sum(psf1)

img2 = py.open('/data/ljo31/Lens/J0901/F814W_sci_cutout.fits')[0].data.copy()

sig2 = py.open('/data/ljo31/Lens/J0901/F814W_noisemap.fits')[0].data.copy()

psf2 = py.open('/data/ljo31/Lens/J0901/F814W_psf2.fits')[0].data.copy()

psf2 = psf2[4:-6,3:-6]

psf2 = psf2/np.sum(psf2)

## identify pixels of interest- can just do this from one lensed image I think, as they should all end up in the same place.

px1,py1 = 28, 30 # pixel indices of brightest peak

px2,py2 = 52,40 # upper SB peak # or could make this more sohpisticated by subtracting galaxy model and detecting brightest pixels left over!

guiFile = '/data/ljo31/Lens/J0901/gui38_7'

print guiFile

imgs = [img1,img2]

sigs = [sig1,sig2]

psfs = [psf1,psf2]

PSFs = []

OVRS = 3

yc,xc = iT.overSample(img1.shape,OVRS)

yo,xo = iT.overSample(img1.shape,1)

#mask = np.ones(img1.shape)

mask = py.open('/data/ljo31/Lens/J0901/mask.fits')[0].data.copy()

tck = RectBivariateSpline(xo[0],yo[:,0],mask)

mask2 = tck.ev(xc,yc)

mask2[mask2<0.5] = 0

mask2[mask2>0.5] = 1

mask2 = mask2.T

mask2 = mask2==1

mask = mask==1

xpeak,ypeak,imgpeak = np.load('/data/ljo31/Lens/J0901/peakpixels.npy').T

for i in range(len(imgs)):

psf = psfs[i]

image = imgs[i]

psf /= psf.sum()

psf = convolve.convolve(image,psf)[1]

PSFs.append(psf)

G,L,S,offsets,shear = numpy.load(guiFile)

pars = []

cov = []

### first parameters need to be the offsets

xoffset = offsets[0][3]

yoffset = offsets[1][3]

pars.append(pymc.Uniform('xoffset',-5.,5.,value=xoffset))

pars.append(pymc.Uniform('yoffset',-5.,5.,value=yoffset))

cov += [0.4,0.4]

gals = []

for name in G.keys():

s = G[name]

p = {}

if name == 'Galaxy 1':

for key in 'x','y','q','pa','re','n':

lo,hi,val = s[key]['lower'],s[key]['upper'],s[key]['value']

pars.append(pymc.Uniform('%s %s'%(name,key),lo,hi,value=val))

p[key] = pars[-1]

cov.append(s[key]['sdev']*10)

elif name == 'Galaxy 2':

for key in 'x','y','q','pa','re','n':

if s[key]['type']=='constant':

p[key] = s[key]['value']

else:

lo,hi,val = s[key]['lower'],s[key]['upper'],s[key]['value']

if key == 're':

pars.append(pymc.Uniform('%s %s'%(name,key),lo,hi,value=val))

else:

pars.append(pymc.Uniform('%s %s'%(name,key),lo,hi,value=val))

p[key] = pars[-1]

cov.append(s[key]['sdev']*1)

#for key in 'x','y':

# p[key] = gals[0].pars[key]

gals.append(SBModels.Sersic(name,p))

lenses = []

for name in L.keys():

s = L[name]

p = {}

for key in 'x','y','q','pa','b','eta':

lo,hi,val = s[key]['lower'],s[key]['upper'],s[key]['value']

pars.append(pymc.Uniform('%s %s'%(name,key),lo,hi,value=val))

cov.append(s[key]['sdev']*1)

p[key] = pars[-1]

lenses.append(MassModels.PowerLaw(name,p))

p = {}

p['x'] = lenses[0].pars['x']

p['y'] = lenses[0].pars['y']

pars.append(pymc.Uniform('extShear',-0.3,0.3,value=shear[0]['b']['value']))

cov.append(1)

p['b'] = pars[-1]

pars.append(pymc.Uniform('extShear PA',-180.,180,value=shear[0]['pa']['value']))

cov.append(100.)

p['pa'] = pars[-1]

lenses.append(MassModels.ExtShear('shear',p))

xind,yind,imgind = np.load('/data/ljo31/Lens/J0901/peakpixels.npy').T

xll,yll = pylens.getDeflections(lenses,[xind,yind])

xpeak = np.sum(xll*imgind)/np.sum(imgind)

ypeak = np.sum(yll*imgind)/np.sum(imgind)

l = {'x':xpeak, 'y':ypeak}

dics = []

srcs = []

for name in S.keys():

s = S[name]

p = {}

if name == 'Source 2':

print name

for key in 'q','re','n','pa':

lo,hi,val = s[key]['lower'],s[key]['upper'],s[key]['value']

pars.append(pymc.Uniform('%s %s'%(name,key),lo,hi,value=val))

p[key] = pars[-1]

if key == 'pa':

cov.append(s[key]['sdev']*100)

elif key == 're':

cov.append(s[key]['sdev']*10)

else:

cov.append(s[key]['sdev']*10)

for key in 'x','y': # subtract lens potition - to be added back on later in each likelihood iteration!

lo,hi,val = s[key]['lower'],s[key]['upper'],s[key]['value']

#print key, '= ', val

lo,hi = lo - l[key], hi - l[key]

val = val - l[key]

pars.append(pymc.Uniform('%s %s'%(name,key),lo-1 ,hi+1,value=val )) # the parameter is the offset between the source centre and peak source brightness pixel position in the source plane...!

p[key] = pars[-1] # the unsubtracted value - this is what we want to keep in the dictionary

print 'here',p[key]

cov.append(s[key]['sdev'])

elif name == 'Source 1':

print name

for key in 'q','re','n','pa':

lo,hi,val = s[key]['lower'],s[key]['upper'],s[key]['value']

pars.append(pymc.Uniform('%s %s'%(name,key),lo,hi,value=val))

p[key] = pars[-1]

if key == 'pa':

cov.append(s[key]['sdev']*100)

else:

cov.append(s[key]['sdev']*10)

for key in 'x','y':

p[key] = srcs[0].pars[key]

dics.append(p)

srcs.append(SBModels.Sersic(name,p))

print srcs[0].pars['x'].value

npars = []

for i in range(len(npars)):

pars[i].value = npars[i]

@pymc.deterministic

def logP(value=0.,p=pars):

return lp #,models

@pymc.observed

def likelihood(value=0.,lp=logP):

return lp #[0]

def resid(p):

lp = -2*logP.value

return self.imgs[0].ravel()*0 + lp

optCov = None

if optCov is None:

optCov = numpy.array(cov)

print len(cov), len(pars)

S = myEmcee.PTEmcee(pars+[likelihood],cov=optCov,nthreads=5,nwalkers=80,ntemps=4)

S.sample(1000)

outFile = '/data/ljo31/Lens/J0901/emcee'+str(X)

f = open(outFile,'wb')

cPickle.dump(S.result(),f,2)

f.close()

result = S.result()

lp = result[0]

trace = numpy.array(result[1])

a1,a2,a3 = numpy.unravel_index(lp.argmax(),lp.shape)

for i in range(len(pars)):

pars[i].value = trace[a1,a2,a3,i]

print "%18s %8.3f"%(pars[i].__name__,pars[i].value)

jj=0

for jj in range(12):

S = myEmcee.PTEmcee(pars+[likelihood],cov=optCov,nthreads=5,nwalkers=80,ntemps=4,initialPars=trace[a1])

S.sample(1000)

outFile = '/data/ljo31/Lens/J0901/emcee'+str(X)

f = open(outFile,'wb')

cPickle.dump(S.result(),f,2)

f.close()

result = S.result()

lp = result[0]

trace = numpy.array(result[1])

a1,a2,a3 = numpy.unravel_index(lp.argmax(),lp.shape)

for i in range(len(pars)):

pars[i].value = trace[a1,a2,a3,i]

print jj

jj+=1

## now we need to interpret these resultaeten

logp,coeffs,dic,vals = result

ii = np.where(logp==np.amax(logp))

#coeff = coeffs[ii][0]

colours = ['F555W', 'F814W']

#mods = S.blobs

models = []

for i in range(len(imgs)):

#mod = mods[i]

#models.append(mod[a1,a2,a3])

if i == 0:

dx,dy = 0,0

else:

dx = pars[0].value

dy = pars[1].value

xp,yp = xc+dx,yc+dy

xop,yop = xo+dy,yo+dy

image = imgs[i]

sigma = sigs[i]

psf = PSFs[i]

imin,sigin,xin,yin = image.flatten(), sigma.flatten(),xp.flatten(),yp.flatten()

n = 0

model = np.empty(((len(gals) + len(srcs)),imin.size))

for gal in gals:

gal.setPars()

tmp = xc*0.

tmp = gal.pixeval(xp,yp,1./OVRS,csub=1) # evaulate on the oversampled grid. OVRS = number of new pixels per old pixel.

tmp = iT.resamp(tmp,OVRS,True) # convert it back to original size

tmp = convolve.convolve(tmp,psf,False)[0]

model[n] = tmp.ravel()

n +=1

for lens in lenses:

lens.setPars()

x0,y0 = pylens.lens_images(lenses,srcs,[xp,yp],1./OVRS,getPix=True)

for src in srcs:

src.setPars()

tmp = xc*0.

tmp = src.pixeval(x0,y0,1./OVRS,csub=1)

tmp = iT.resamp(tmp,OVRS,True)

tmp = convolve.convolve(tmp,psf,False)[0]

model[n] = tmp.ravel()

n +=1

rhs = (imin/sigin) # data

op = (model/sigin).T # model matrix

fit, chi = optimize.nnls(op,rhs)

components = (model.T*fit).T.reshape((n,image.shape[0],image.shape[1]))

model = components.sum(0)

models.append(model)

NotPlicely(image,model,sigma)

pl.suptitle(str(colours[i]))

pl.show()

pl.figure()

for i in range(4):

pl.plot(lp[:,i,:])

pl.show()