import numpy as np
import glob
#import astropy.io.fits as fits
import matplotlib.pyplot as plt

#takes the cubes from the pipeline and smooths and regrids, recreating the .com files
# in the case that one or more planes has a bad beam and messed up the original ones.

imlist=glob.glob('*IQU.iter3.image.pbcor.tt0.subim.fits')

bmmeanlist=[]
bmarlist=[]
bmajlist=[]
bminlist=[]
freqlist=[]
bpalist=[]

for im in imlist:
    hdrdict=imhead(im)
    bmaj=hdrdict['perplanebeams']['beams']['*0']['*0']['major']['value']
    bmajlist.append(bmaj)
    bmin=hdrdict['perplanebeams']['beams']['*0']['*0']['minor']['value']
    bminlist.append(bmin)
    bpa=hdrdict['perplanebeams']['beams']['*0']['*0']['positionangle']['value']
    freqlist.append(hdrdict['refval'][2])
    #print(im,bmaj,bmin,bpa)
    bav=np.sqrt(bmaj*bmin)
    bar=bmaj/bmin
    bmmeanlist.append(bav)
    bmarlist.append(bar)
    bpalist.append(bpa)

#make initial guess at lowest ferquency beam
# check that each higher frequency beam is within the lowest frequency beam
# if not, increase major and minor axes by 10% ieteratively until it is

freqorder=np.argsort(np.array(freqlist))
bmajarr=np.array(bmajlist)
bminarr=np.array(bminlist)
bpaarr=np.array(bpalist)
nfreq=len(freqlist)

print(bmajarr[freqorder])
#print(bminarr[freqorder])
#print(bpaarr[freqorder])

maxindx=np.argmax(bmajarr)
majbm=np.max(bmajarr)
minbm=bminarr[maxindx]
medbpa=bpaarr[maxindx]

print('First pass beam parameters ',majbm,minbm,medbpa)

medbparad=medbpa*np.pi/180.

#parameterize best ellipse, see if any points from the other ellipses are outside it
t=np.arange(100)/100*2*np.pi - np.pi

xo=majbm/2.*np.cos(t)*np.cos(medbparad+np.pi/2)-minbm/2.*np.sin(t)*np.sin(medbparad+np.pi/2)
yo=majbm/2.*np.sin(t)*np.cos(medbparad+np.pi/2)+minbm/2.*np.cos(t)*np.sin(medbparad+np.pi/2)
#print(xo)
#print(yo)
#plt.plot(xo,yo)
disto=np.sqrt(xo**2+yo**2)

distdiffarr=np.zeros((100,nfreq))

for i in range(nfreq):
    x=bmajarr[i]/2.*np.cos(t)*np.cos(medbparad+np.pi/2)-bminarr[i]/2.*np.sin(t)*np.sin(medbparad+np.pi/2)
    y=bmajarr[i]/2.*np.sin(t)*np.cos(medbparad+np.pi/2)+bminarr[i]/2.*np.cos(t)*np.sin(medbparad+np.pi/2)
    #xdiff=abs(x)-abs(xo)
    #ydiff=abs(y)-abs(yo)
    dist=np.sqrt(x**2+y**2)
    distdiff=dist-disto
    distdiffarr[:,i]=distdiff
    plt.plot(t,distdiff)
    #print(xdiff)
    #print(ydiff)

bad= distdiffarr>0.0001
print(bad)
check=np.any(bad)
print('Is beam expansion needed?',check)
# if any points do lie outside the ellipse, then try bumping ellipse size by 10% in maj & min
if check:
    majbm=majbm*1.1
    minbm=minbm*1.1
    print('New bmaj, bmin = ',majbm,minbm)

#we're going to assume that worked, now apply beam and coalign images

for im in imlist:
    imroot=im.split('.spw')[0]
    spw=im.split('.')[8]
    print(imroot)
    imV=imroot+'.'+spw+'.V.iter3.image.pbcor.tt0.subim.fits'
    outsmIQU=imroot+'.'+spw+'.IQU.iter3.image.pbcor.tt0.subim.sm'
    outregridIQU=imroot+'.'+spw+'.IQU.iter3.image.pbcor.tt0.subim.com'
    print(outsmIQU)
    outsmV=imroot+'.'+spw+'.V.iter3.image.pbcor.tt0.subim.sm'
    outregridV=imroot+'.'+spw+'.V.iter3.image.pbcor.tt0.subim.com'
    print(outsmV)
    # get beam info to see if any channels don't need convolution (typically spw2)
    hdrdict=imhead(im)    
    bmaj=hdrdict['perplanebeams']['beams']['*0']['*0']['major']['value']
    bmin=hdrdict['perplanebeams']['beams']['*0']['*0']['minor']['value']
    #figure out original pointing
    hist=imhistory(im)
    sub='Uncorrected CRVAL'
    crvlist=[s for s in hist if sub in s]
    crv1=crvlist[0]
    crv2=crvlist[1]
    origcrv1=float(crv1.split(' ')[3])*np.pi/180.
    origcrv2=float(crv2.split(' ')[3])*np.pi/180.
    templateIQU=imregrid(im,template='get')
    templateV=imregrid(imV,template='get')
    print(templateIQU['csys']['direction0']['crval'])
    templateIQU['csys']['direction0']['crval']=[origcrv1,origcrv2]
    templateV['csys']['direction0']['crval']=[origcrv1,origcrv2]
    print(origcrv1,origcrv2)
    if (bmaj < majbm/1.02): #if beam maj axis is <2% smaller than consensus beam, smooth
        imsmooth(imagename=im,major=majbm,minor=minbm,pa=medbpa,targetres=True,outfile=outsmIQU,overwrite=True)
        imsmooth(imagename=imV,major=majbm,minor=minbm,pa=medbpa,targetres=True,outfile=outsmV,overwrite=True)
    else:
        imsmooth(imagename=im,kernel='common',targetres=True,outfile=outsmIQU,overwrite=True)
        os.system('cp -r '+imV+' '+outsmV)
        print('Not smoothing '+spw)
    imregrid(outsmIQU,output=outregridIQU, template=templateIQU)
    imregrid(outsmV,output=outregridV, template=templateV)
    exportfits(outregridIQU,outregridIQU+'.fits',overwrite=True)
    exportfits(outregridV,outregridV+'.fits',overwrite=True)


"""
    outsm=imroot+'_fixed_sm'
    imsmooth(imagename=im,major=majbm,minor=minbm,pa=medbpa,targetres=True,outfile=outsm,overwrite=True)
    rmsname=imroot+'.image.pbcor.tt0.rms.subim.fits'
    importfits(rmsname,imroot+'_rms',overwrite=True)   
    template=imregrid(imroot+'_rms',template='get')
    outregrid=imroot+'_fixed_sm_regridded'
    fitsname=outregrid+'.fits'
    imregrid(outsm,output=outregrid, template=template)
    exportfits(outregrid,fitsname)
"""