import math
from math import *
from java.awt.Color import GREEN
from java.awt.Color import BLUE

band="blue"
dir='/Users/balog/Work/HERSCHEL/DPWorkshop2013/'
#get the PSF file
bm = fitsReader(file = \
'/Users/balog/Work/HERSCHEL//PSF/blu_20_vesta_od160_OA+63_recentered.fits')
src=fitsReader(dir+'AlphaTau.blue.fits')

pabm = 0.0
obsid = src.getMeta()["obsid"].value
pa    = src.getMeta()["posAngle"].value 
speed = src.getMeta()["mapScanSpeed"].value 
sa    = src.getMeta()["mapScanAngle"].value

#rotate the PSF to the source
angle=pa-pabm
bmr=rotate(image=bm, angle=angle, subsampleBits=32,\
  interpolation=rotate.INTERP_BICUBIC) 
# optional information
bmr.setDescription("beam rotated for "+str(obsid)+", "+band) 
Display(bm)
Display(bmr)
print "WCS of rotated beam:",bmr.wcs

#fixing the WCS after rotating
wcs=bm.getWcs()
bmr.setWcs(wcs)
bmr.setDescription("beam rotated for "+str(obsid)+", "+band\
  +" with original wcs")
print "WCS of rotated beam with old WCS imposed:",bmr.wcs

#matching the WCS of the beam and the source
xpixstep=bmr["image"].meta["cdelt1"].value*3600.
ypixstep=bmr["image"].meta["cdelt2"].value*3600.
if (xpixstep<0): xpixstep=-1*xpixstep
if (ypixstep<0): ypixstep=-1*ypixstep

#RA=bmr.meta["raNominal"].value
#DEC=bmr.meta["decNominal"].value
RA=bmr["image"].meta["crval1"].value
DEC=bmr["image"].meta["crval2"].value
cxpix=bmr.wcs.getPixelCoordinates(RA,DEC)[1] 
cypix=bmr.wcs.getPixelCoordinates(RA,DEC)[0]

nxpix=bmr["image"].meta["naxis1"].value
nypix=bmr["image"].meta["naxis2"].value
cxpix=nxpix/2.
cypix=nypix/2.

#finding the exact center of the beam
maxshift = 50. 
print "   Fitting rotated beam ..."
for boxsize in range(5,60):
  try:
    minX=cxpix-boxsize/2.
    minY=cypix-boxsize/2.
    sfit   = sourceFitting(elongated=True,slope=False,image=bmr,\
            minX=minX,minY=minY,width=boxsize,height=boxsize)
    print "      boxsize "+str(boxsize)+"...success!"   
    poo=1 # seems to be necessary to make this stop when it has success
  except:
    print "      boxsize "+str(boxsize)+"...failed"
    cxpixfit=-1  # or some other value to indicate failure
    cypixfit=-1  # eg maybe the centre of the map, ... 
    pixfit_RA=-1 
    pixfit_Dec=-1
    poo=0
  else:
    # upon success, grab the results of the sourceFitting
    cxpixfit   = sfit.getCenterX()
    cypixfit   = sfit.getCenterY()
    pixfit_RA  = sfit.getCenterRA()
    pixfit_Dec = sfit.getCenterDec()
    # very occasionally sourceFitting gives nonsence results, so:
    if ((ABS(cxpix-cxpixfit) > maxshift) or (ABS(cypix-cypixfit) > maxshift)):
      print "         ...however, source found too far from original coords"
      print"             so setting to -1,-1" 
      print "            orig coords: "+str(cxpix)+", "+str(cypix)+\
            " found coords: "+str(bm_cxpixfit)+", "+str(bm_cypixfit) 
      cxpixfit=-1 # or some other value to indicate failure
      cypixfit=-1 # eg maybe the centre of the map, ... 
      pixfit_RA=-1
      pixfit_Dec=-1
      poo=0
    else:
      bm_cxpixfit   = sfit.getCenterX()
      bm_cypixfit   = sfit.getCenterY()
      bm_pixfit_RA  = sfit.getCenterRA()
      bm_pixfit_Dec = sfit.getCenterDec()
      sigma_x    = sfit.getSigmaXPixels()
      sigma_y    = sfit.getSigmaYPixels()
      bm_fwhm_x     = round(abs(2.*SQRT(2.*LOG(2))*sigma_x*xpixstep),1)
      bm_fwhm_y     = round(abs(2.*SQRT(2.*LOG(2))*sigma_y*ypixstep),1)
      print "Beam details:"
      print "   beam position: pixel:",bm_cxpixfit, bm_cypixfit
      print "                  coordinates:",bm_pixfit_RA, bm_pixfit_Dec
      print "   beam fwhm x,y:",bm_fwhm_x,bm_fwhm_y
      poo=1
    if (poo == 1): break
    pass

#same for the source
xpixstep=src["image"].meta["cdelt1"].value*3600. # pixel size in arcsec
ypixstep=src["image"].meta["cdelt2"].value*3600.
if (xpixstep<0): xpixstep=-1*xpixstep # avoid negative values 
if (ypixstep<0): ypixstep=-1*ypixstep
# starting RA, Dec values in decimal degrees...from the meta data
# or you can work it out yourself
# decimal values are necessary
RA=src.meta["raNominal"].value
DEC=src.meta["decNominal"].value
# ...and in pixel values
cxpix=src.wcs.getPixelCoordinates(RA,DEC)[1] 
cypix=src.wcs.getPixelCoordinates(RA,DEC)[0]
maxshift = 10. 
print "   Fitting astronomical source ..."
for boxsize in range(5,60):
  try:
    minX=cxpix-boxsize/2.
    minY=cypix-boxsize/2.
    sfit   = sourceFitting(elongated=True,slope=False,image=src,\
            minX=minX,minY=minY,width=boxsize,height=boxsize)
    print "      boxsize "+str(boxsize)+"...success!"   
    poo=1 # seems to be necessary to make this stop when it has success
  except:
    print "      boxsize "+str(boxsize)+"...failed"
    cxpixfit=-1 # or some other value to indicate failure
    cypixfit=-1 # eg maybe the centre of the map, ... 
    pixfit_RA=-1 
    pixfit_Dec=-1
    poo=0
  else:
    # upon success, grab the results of the sourceFitting
    cxpixfit   = sfit.getCenterX()
    cypixfit   = sfit.getCenterY()
    pixfit_RA  = sfit.getCenterRA()
    pixfit_Dec = sfit.getCenterDec()
    # very occasionally sourceFitting gives nonsence results, so:
    if ((ABS(cxpix-cxpixfit) > maxshift) or (ABS(cypix-cypixfit) > maxshift)):
      print "         ...however, source found too far from original coords,"
      print "         so setting to -1,-1" 
      print "            orig coords: "+str(cxpix)+","+str(cypix)+\
            " found coords: "+str(bm_cxpixfit)+", "+str(bm_cypixfit) 
      cxpixfit=-1 # or some other value to indicate failure
      cypixfit=-1 # eg maybe the centre of the map, ...
      pixfit_RA=-1
      pixfit_Dec=-1
      poo=0
    else:
      src_cxpixfit   = sfit.getCenterX()
      src_cypixfit   = sfit.getCenterY()
      src_pixfit_RA  = sfit.getCenterRA()
      src_pixfit_Dec = sfit.getCenterDec()
      sigma_x    = sfit.getSigmaXPixels()
      sigma_y    = sfit.getSigmaYPixels()
      src_fwhm_x     = round(abs(2.*SQRT(2.*LOG(2))*sigma_x*xpixstep),1)
      src_fwhm_y     = round(abs(2.*SQRT(2.*LOG(2))*sigma_y*ypixstep),1)
      print "Source details:"
      print "   source position: pixel:",src_cxpixfit, src_cypixfit
      print "                  coordinates:",src_pixfit_RA, src_pixfit_Dec
      print "   source fwhm x,y:",src_fwhm_x,src_fwhm_y
      poo=1
    if (poo == 1): break
    pass

#match the coordinates of the beam and the source
bmr.getWcs().setCrval1(src_pixfit_RA) 
bmr.getWcs().setCrval2(src_pixfit_Dec)
bmr.getWcs().setCrpix1(bm_cxpixfit+1.0) 
bmr.getWcs().setCrpix2(bm_cypixfit+1.0)
print "WCS of rotated beam with old WCS imposed and shifted to the astro source position:"\
  ,bmr.wcs
d=Display(src)
d.setTitle(str(obsid)+","+band)
d.setCutLevels(99.0)
d.setTitle(str(obsid)+","+band+",rotated by "+str(angle))
# red circle where the astro source was found to be
d.addCircle(src_cypixfit,src_cxpixfit,2,2,java.awt.Color.red)
# If you want to add contours of the beam and the astro source to the greyscale map"
# Beam:
# ->must be as many as there are contour levels
cs=[BLUE,BLUE,BLUE,BLUE,BLUE,BLUE,BLUE,BLUE,BLUE,BLUE] 
# ->change the min and max according to the flux range in your map
contoursb = automaticContour(image=bmr, levels=10, min=0.001, max=max(bmr.image)/2, distribution=0,\
   colors=cs)
d.addWcsImageContour(contoursb)
# Astro source:
#cs=[GREEN,GREEN,GREEN,GREEN,GREEN]
#contourss = automaticContour(image=src, levels=5, min=0.002, max=0.004, \
#  distribution=0, colors=cs)
#d.addWcsImageContour(contourss)
d.setZoomFactor(16)
#EEF curves
rskyin_src  = 60
rskyout_src = 75
rapers=Double1d.range(41)*2 
rapers=rapers[1:]
fluxes=Double2d(len(rapers),2)
#photometry of the source
cnt=0
for raper in (rapers):
  apphot = fixedSkyAperturePhotometry(image=src,centroid=False,\
           centerX=round(src_cxpixfit,3),centerY=round(src_cypixfit,3),\
           radiusArcsec=raper,sky=0)
  fluxes[cnt,0]=apphot["Results table"]["Total flux"].data[1]
  cnt+=1
cnt=0
#photometry of the PSF
for raper in (rapers):
  apphot = fixedSkyAperturePhotometry(image=bmr,centroid=False,\
           centerX=round(bm_cxpixfit,1),centerY=round(bm_cypixfit,1),\
           radiusArcsec=raper,sky=0)
  fluxes[cnt,1]=apphot["Results table"]["Total flux"].data[1] 
  cnt+=1

fluxes[:,0]=fluxes[:,0]/fluxes[0,0] 
fluxes[:,1]=fluxes[:,1]/fluxes[0,1] 
idx=rapers.where((rapers>rskyin_src)&(rapers<rskyout_src))
fluxes[:,0]=fluxes[:,0]/MEAN(fluxes[idx,0]) 
idx=rapers.where((rapers>rskyin_bm)&(rapers<rskyout_bm))
fluxes[:,1]=fluxes[:,1]/MEAN(fluxes[idx,1]) 
#plot the results
p=PlotXY(titleText="My EEF for obsid "+str(obsid)+", "+band)
p.addLayer(LayerXY(rapers,fluxes[:,0],line=1,color=java.awt.Color.black))
p.addLayer(LayerXY(rapers,fluxes[:,1],line=3,color=java.awt.Color.red))
p[1].setName("beam")
p[0].setName("astro source")
p.xaxis.title.text="arcsec"
p.yaxis.title.text="flux"
p.getLegend().setVisible(True)