import  matplotlib.pyplot as plt
%pylab inline
%matplotlib inline

import numpy as np

# Doris, can you grab the Gaussian coefficients from the Hogg and Lang paper (download source and grab with python from the tex, or by hand..) and code them up in the deVauc function?

def N(xgrid,ygrid,m,Vinv):
  """
  2d normal (Gaussian) distribution.
  #http://en.wikipedia.org/wiki/Gaussian_function#Two-dimensional_Gaussian_function
  xgrid,ygrid is the output of a meshgrid.
  #m is the mean (x_0,y_0)
  Vinv is a 2x2 np.array, you already took the inverse!
  """
  assert Vinv[0,1] == Vinv[0,1]
  exparg = (xgrid - m[0])*Vinv[0,0]*(xgrid - m[0]) + \
           2*(xgrid - m[0])*Vinv[1,0]*(ygrid -m[1]) + \
           (ygrid - m[1])*Vinv[1,1]*(ygrid - m[1])
  
  return (np.linalg.det(Vinv))**0.5/(2.*np.pi)*np.exp(-0.5*exparg)

def N2diso(r,sig2inv):
  exparg = r**2*sig2inv
  return (0.5/np.pi)*sig2inv*np.exp(-0.5*exparg)
  
def deVauc1d(r,Ie,re,aopt=1):
  """
  Input r in units of the half-light radius.
  aopt = 0: straight-up deVauc formula.
  aopt = 1: SDSS approximation ('luv') Eqn 17-18 of Hogg and Lang.
  aopt = 2: Hogg and Lang approximation to luv (sum over 10 2d isotropic Gaussians).
  """
  xi = np.fabs(r/re)
  if aopt == 0:
    return Ie*np.exp(-7.66924944*((xi)**(0.25)-1))
  if aopt == 1:
    result = np.zeros(len(r))
    result = Ie*np.exp(-7.66925*(((xi)**2 + 0.0004)**(0.125)-1))
    xx = np.where(xi > 8.)[0]
    if len(xx) > 0:
      result[xx] = 0.
    xx = np.where((xi >= 7.) & (xi < 8.))[0]
    if len(xx) > 0.:
      result[xx] = result[xx] * (1.-(xi[xx]-7)**2)**2
    return result
  if aopt == 2:
    ## hard code Hogg/Lang coefficients.
    mluv10={0.01468:0.01190**-2, \
            0.09627:0.02210**-2, \
            0.28454:0.03995**-2, \
            0.63005:0.07117**-2, \
            1.19909:0.12586**-2, \
            2.03195:0.22240**-2, \
            3.07255:0.39593**-2, \
            4.10682:0.71922**-2, \
            4.83948:1.37549**-2, \
            4.94943:3.13117**-2}
            
    #mdev10={0.00139:0.00087, \
    #        0.00941:0.00296, \
    #        0.04441:0.00792, \
    #        0.16162:0.01902, \
    #        0.48121:0.04289, \
    #        1.20357:0.09351, \
    #        2.54182:0.20168, \
    #        4.46441:0.44126, \
    #        6.22820:1.01833, \
    #        6.15393:2.74555}
            
    result = np.zeros(len(r))
    xx = np.where(xi > 8.)[0]
    for kk, vv in mluv10.iteritems():
      result += kk*N2diso(xi,vv)
    if len(xx) > 0:
      result[xx] = 0.
    result = Ie*result
    return result
    
    
    

def intT1d(r,pofr):
  """
  Trapezoid rule for integration.
  """
  return 0.5*((r[1:]-r[:-1])*(pofr[1:]+pofr[:-1])).sum()

## Plot of the deVauc and LUV profiles.  Can you add the plot of the Gaussian approximation?

r=10**np.arange(-3,2,0.001)
print len(r)

print 'this shows that parameter dependence works!'

plt.figure()
Ieval = 1.; reval = 1.
plt.loglog(r,deVauc1d(r,Ieval,reval,0),'b')
plt.loglog(r,deVauc1d(r,Ieval,reval,1),'g')
plt.loglog(r,deVauc1d(r,Ieval,reval,2),'r')

plt.figure()
Ieval = 2.; reval = 1.
plt.loglog(r,deVauc1d(r,Ieval,reval,0),'b')
plt.loglog(r,deVauc1d(r,Ieval,reval,1),'g')
plt.loglog(r,deVauc1d(r,Ieval,reval,2),'r')

plt.figure()
Ieval = 1.; reval = 2.
plt.loglog(r,deVauc1d(r,Ieval,reval,0),'b')
plt.loglog(r,deVauc1d(r,Ieval,reval,1),'g')
plt.loglog(r,deVauc1d(r,Ieval,reval,2),'r') 
    


r=np.arange(100)
r = np.array([float(i) for i in r])
print r
plt.ylim(0,2)
plt.plot(r,deVauc1d(r,0.5,0.5,2)) 



result = N2diso(np.arange(1000),10)
# plt.plot(result)
# plt.figure()
plt.loglog(result)

result = N2diso(np.arange(100),0.001)
# plt.plot(result)
# plt.figure()
plt.loglog(result)
plt.xlim(0,10**2)
plt.ylim(10**-8,10**-3)

result = N2diso(np.arange(1000),0.001)
plt.plot(result)

x=r
y=deVauc1d(r,1.,1.,0)*2*np.pi*x
y2=deVauc1d(r,1.,1.,1)*2*np.pi*x
y3=deVauc1d(r,1.,1.,2)*2*np.pi*x
ihalf = np.where(r > 1.)[0].min()
print 'half light radius correct?', x[ihalf], intT1d(x[:ihalf],y[:ihalf])/intT1d(x,y)
print 'half light approx correct for LUV profile?',intT1d(x[:ihalf],y2[:ihalf])/intT1d(x,y2)
print 'half light approx correct for Hogg profile?',intT1d(x[:ihalf],y3[:ihalf])/intT1d(x,y3)

print 'note that the profiles are not normalized to have unit flux, but to be 1 at the half light radius!'
print 'therefore the total integrated flux is ~21, agrees with Hogg Sum over a_luv'
intT1d(x,y),intT1d(x,y2),intT1d(x,y3)

## Plot of the deVauc and LUV profiles.  Can you add the plot of the Gaussian approximation?

r=10**np.arange(-3,2,0.001)
print len(r)
# plt.loglog(r,deVauc(r,1.,1.,0),'b')
plt.loglog(r,deVauc(r,1.,1.,0),'b')
plt.loglog(r,deVauc(r,1.,1.,1),'g')
# plt.loglog(r,deVauc(r,1.,1.,2),'r')

x=r
y=deVauc(r,1.,1.,0)*2*np.pi*x
y2=deVauc(r,1.,1.,1)*2*np.pi*x
ihalf = np.where(r > 1.)[0].min()
print 'half light radius correct?', x[ihalf], intT1d(x[:ihalf],y[:ihalf])/intT1d(x,y)
print 'half light approx correct for LUV profile?',intT1d(x[:ihalf],y2[:ihalf])/intT1d(x,y2)

V = np.array([[80,0],[0,10]])
# V = np.array([[80,0],[0,80]]) #circular profile
Vinv = np.linalg.inv(V)
if 0==1:
  print V
  print Vinv
  print V*Vinv

## set up the model using the geometry of the image files.
# https://www.sdss3.org/instruments/camera.php#Filters
nx = 2048
ny = 1361
psize = 0.396 # arcseconds per pixel.

# coordinates of center of pixels
x = (np.arange(0,nx,1)+0.5)*psize
y = (np.arange(0,ny,1)+0.5)*psize

xM, yM = np.meshgrid(x,y,indexing='xy')

# 2048*1361

print x[-10:]
print y[-10:]

m = np.array([100,300])
## test with a fat Gaussian that spans many pixels.
I = N(xM,yM,m,Vinv)
# check -- is the Gaussian profile normalized?
parea = psize**2
print 'should sum to 1.  Error:',(I[:,:].flatten()).sum()*parea-1

plt.matshow(I,origin='lower',extent = [x.min(),x.max(),y.min(),y.max()],interpolation='nearest')

mysize=1
xx = np.where((np.fabs(xM - m[0]) < V[0,0]**0.5*mysize) & (np.fabs(yM - m[1]) < V[1,1]**0.5*mysize))
extent = max(xx[0].max()-xx[0].min(),xx[1].max()-xx[1].min())
print extent
xcen = int(xx[0].mean())
ycen = int(xx[1].mean())

## zoom
#plt.matshow(I[xx[0].min():xx[0].max()+1,xx[1].min():xx[1].max()+1],origin='lower',extent = [x.min(),x.max(),y.min(),y.max()])
## zoom, keeping aspect ratio.
#myextent = xM[
plt.matshow(I[xcen-extent/2:xcen+extent/2,ycen-extent/2:ycen+extent/2],origin='lower',interpolation='nearest')#extent = [xM[xx[0]].min(),xM[xx[0]].max(),yM[xx[1]].min(),yM[xx[1]].max()])

mdev6 = {0.01308: 0.00263, 0.12425: 0.01202,0.63551: 0.04031,2.22560: 0.12128,5.63989: 0.36229,9.81523: 1.23604}
# check sum as 18.454 as in Table 1 
sum(mdev6.keys())

r = np.linspace(0,8)
sig = 2
mu=0
N= 1./(sig*np.sqrt(2*np.pi))*np.exp(-(r-mu)**2/(2*sig**2))

plt.plot(N)

def N_1d(r,mu,sig):
    # sig = std_dev = sqrt(v) (just the scalar given in the table, nothing needs to be done to it)
    mu=0
    return 1./(sig*np.sqrt(2*np.pi))*np.exp(-(r-mu)**2/(2*sig**2))

def Iconv_1d(r, mdev):
#     I_conv = numpy.zeros((50,))
    I_conv = np.zeros_like(r)
    for a in mdev.keys():
        sig =mdev[a]
#         r = np.linspace(0,8)
        I_conv += a*N_1d(r,0,sig)
    return I_conv

# plt.plot(Iconv_1d(mdev6))
r = np.linspace(0,8)
print len(r)
# plt.plot(r,Iconv_1d(r,mdev6))
plt.loglog(r,Iconv_1d(r,mdev6))

# plt.plot(Iconv_1d(mdev6))
r = np.linspace(0,10,100000)
print len(r)
# plt.plot(r,Iconv_1d(r,mdev6))
plt.loglog(r,Iconv_1d(r,mdev6))

r=10**np.arange(0,1.676,0.001)
plt.loglog(r,Iconv_1d(r,mdev6))

mdev8 = {0.00262:0.00113, 0.02500:0.00475,0.13413:0.01462,0.51326:0.03930,1.52005:0.09926,3.56204:0.24699,6.44845:0.63883,8.10105:1.92560}
mdev10={0.00139:0.00087, 0.00941:0.00296,0.04441:0.00792,0.16162:0.01902,0.48121:0.04289,1.20357:0.09351,2.54182:0.20168,4.46441:0.44126,6.22820:1.01833,6.15393:2.74555}

# r=10**np.arange(-3,2,0.001)
r=10**np.linspace(-3,1,5000)
plt.loglog(r,deVauc(r,1.,1.,0),'b')
plt.loglog(r,deVauc(r,1.,1.,1),'g')
plt.loglog(r,Iconv_1d(r,mdev6),'r')

# r=10**np.arange(-3,2,0.001)
r=10**np.linspace(-3,1,5000)
# print len(r)
# plt.loglog(r,deVauc(r,1.,1.,0),'b')
plt.loglog(r,deVauc(r,1.,1.,0),'b')
plt.loglog(r,deVauc(r,1.,1.,1),'g')
# plt.loglog(r,Iconv_1d(r,mdev6),'r')
# plt.loglog(r,Iconv_1d(r,mdev8),'c')
plt.loglog(r,Iconv_1d(r,mdev10),'m')

r=10**np.linspace(-3,1,5000)
plt.loglog(r,Iconv_1d(r,mdev6),'r')
plt.loglog(r,Iconv_1d(r,mdev8),'c')
plt.loglog(r,Iconv_1d(r,mdev10),'m')

# V = np.array([[0.,0.],[0.,0.]])
def Iconv(mdev):
    I_conv = numpy.zeros((1361, 2048))
    for a in mdev.keys():
        v=mdev[a]
        m = np.array([100,300])
        v_inv = np.linalg.inv(np.array([[v**2,0],[0,v**2]]))
        I_conv += a*N(xM,yM,m,v_inv)
    return I_conv

def plotI(data):
    mysize=1
    V = np.array([[5.,0.],[0., 5.]]) 
    xx = np.where((np.fabs(xM - m[0]) < V[0,0]**0.5*mysize) & (np.fabs(yM - m[1]) < V[1,1]**0.5*mysize))
    extent = max(xx[0].max()-xx[0].min(),xx[1].max()-xx[1].min())
    xcen = int(xx[0].mean())
    ycen = int(xx[1].mean())
    plt.matshow(data[xcen-extent/2:xcen+extent/2,ycen-extent/2:ycen+extent/2],origin='lower',interpolation='nearest')#extent = [xM[xx[0]].min(),xM[xx[0]].max(),yM[xx[1]].min(),yM[xx[1]].max()])

I_conv= Iconv(mdev6)
plotI(I_conv)

mdev8 = {0.00262:0.00113, 0.02500:0.00475,0.13413:0.01462,0.51326:0.03930,1.52005:0.09926,3.56204:0.24699,6.44845:0.63883,8.10105:1.92560}
I_conv8= Iconv(mdev8)
plotI(I_conv8)

mdev10={0.00139:0.00087, 0.00941:0.00296,0.04441:0.00792,0.16162:0.01902,0.48121:0.04289,1.20357:0.09351,2.54182:0.20168,4.46441:0.44126,6.22820:1.01833,6.15393:2.74555}
I_conv10= Iconv(mdev10)
plotI(I_conv10)

# Look very simmilar but not identical by element-wise comparison
print np.array_equal(I_conv,I_conv8) 
print np.array_equal(I_conv8,I_conv10) 

I_conv.shape

m = np.array([100,300])
m[0]

n=0
rlist = []
Iconvlst = []
m = np.array([100,300]) #center
# too large can not run on ipynb
for i in arange(I_conv.shape[0]-1):
    for j in arange(I_conv.shape[1]-1):
#         print(i,j)
        r = np.sqrt(abs(i-m[0])**2+abs(j-m[1])**2) #treating center as origin
#         print I_conv[i][j]
        rlist.append(r)
        Iconvlst.append(I_conv[i][j])

# print rlist
# print Iconvlst

%matplotlib inline

len(rlist)

len(Iconvlst)

r=10**np.arange(-3,2,0.001)
print len(r)
plt.figure()
# plt.loglog(I_conv,'b')
# plt.show()
plt.loglog(rlist,Iconvlst)

np.nonzero(I_conv.flatten())
#truncating all the zero values so that this can be plotted on loglog scale, otherwise -inf error
I_conv = [I_conv.flatten()[i] for i in np.nonzero(I_conv.flatten())] 
# np.where(I_conv>0)

print(I_conv)

a=plt.plot(I_conv,'o'); #semicolon supress matplotlib output (direct output to vars)

# plt.loglog(I_conv,'o') # can not loglog since negative exp
# a = plt.plot(log(I_conv),'o')#Stilll a delta function

## work in progress -- want to finish Doris?  We want to evaluate the model (deVauc convolved with Gaussian PSF) on a fine grid, and then sum up the fine pixels that contribute to 
## the pixels of SDSS camera size.  Given some reasonable parameters for galaxy size (re = 1.0, 1.5, 2.0) and PSF values (take some percentiles from the galaxy files I give you), figure out
## how fine the fine grid needs to be.

## truncate the fit 5*sigma away from the center.
fitsigma = 5.
tgal = 1.5 ## rough typical size of a galaxy, in arcseconds.
nsub = 10 ## how fine should we subsample the model?  Let's test this empirically for realistic values of p.
pfine = psize/float(nsub)

nbig = fitsigma*tgal/psize # number of native size pixels.
nsmall = nbig*nsub # number of small pixels we sum over to get model in the big pixel.

## make mbig the location of the center of the source within the central pixel.
mbig = np.array([0.2,0.3])
assert (np.fabs(mbig) < 0.99).all()
xbig = 

xf = np.arange(0,nsmall+0.5,1)
yf = np.arange(0,nsmall+0.5,1)
assert len(xf) == nsmall
xbig = np.arange(

## width of Gaussian is tg
mg = np.array([0.,0.])
Vg = np.array([[tgal**2,0],[0,tgal**2]])
Vginv = np.linalg.inv(Vg)

# coordinates of center of pixels
xf = (np.arange(-fitsigma*tgal/pfine+mg[0]+0.5,fitsigma*tgal/pfine+mg[0]-0.4,1))*pfine
yf = (np.arange(-fitsigma*tgal/pfine+mg[1]+0.5,fitsigma*tgal/pfine+mg[1]-0.4,1))*pfine

xMf, yMf = np.meshgrid(xf,yf,indexing='xy')
If = N(xMf,yMf,mg,Vginv)

plt.imshow(If,origin='lower',interpolation='nearest')