import pyfits
import scipy.fftpack
import scipy.signal
import scipy
import pylab as pb
import numpy as np

def nolabels():
    pb.gca().xaxis.set_ticklabels([])
    pb.gca().yaxis.set_ticklabels([])

def build2dgauss(l,fwhm):
  x = np.empty([l,l])
  y = np.empty([l,l])
  for i in range(len(x)):
       x[i]=np.arange(0,l,1)
       y[i]=np.ones(len(x))*i

  print "shape(x) = ", np.shape(x)
  x0 = y0 = l/2+1
  g = np.exp(-4*np.log(2)*((x-x0)**2+(y-y0)**2)/fwhm**2)
  print "shape(g) = ", np.shape(g)
  return(g)

pb.clf()
files = ['uid___A002_X2d0501_X1a2-DV03-H-USB.fits',
         'uid___A002_X2d0501_X1a2-DA42-H-USB.fits',
         'uid___A002_X2d0501_X1a2-PM04-H-USB.fits']
i = 1
shrink = 1
pb.rc('text',usetex=True)
for ap in [0,1]:
  for f in range(len(files)):
    adesc=pb.subplot(3,3,i)
    im = pyfits.open(files[f])
    data = np.array(im[0].data[0])
    xy = np.transpose(data)[ap]
    if (ap==1):
         xy *= 180/np.pi
         pb.imshow(xy,vmin=-20,vmax=20)
    else:
         xy = (xy/np.max(xy))**2
         pb.imshow(xy)
    nolabels()
    if (ap == 0):
      pb.title(files[f].split('-')[1])
    if (i==1):
      pb.ylabel('(Amplitude)$^2$',size=14)
    if (i==7):
      pb.ylabel('Phase (deg)',size=14)
    pb.colorbar(shrink=shrink)
    if (ap==0):
      adesc=pb.subplot(3,3,i+3)
      pb.imshow(10*np.log10(xy),vmin=-20,vmax=0)
      nolabels()
      pb.colorbar(shrink=shrink)
      if (i==1):
        pb.ylabel('Log[(Amplitude)$^2$]',size=14)
    i += 1
  i += 3
pb.savefig('almabeams2d.png')
pb.show()