import mwa_hyperbeam,os,astropy,numpy as np,datetime,skyfield.api as si,matplotlib.pyplot as plt,paircarsdata,warnings,erfa
from optparse import OptionParser
from astropy.io import fits
import astropy.wcs as pywcs
from astropy.time import Time
from numpy.linalg import inv
from astropy.coordinates import EarthLocation,SkyCoord,AltAz
from casatasks import *
from scipy.interpolate import RegularGridInterpolator

def main():
	usage= 'Correct images for MWA primary beam response'
	parser = OptionParser(usage=usage)
	parser.add_option('--MWA_PB_path',dest="MWA_PB_path",default=None,help="MWA primary beam file path",metavar="Directory path")
	parser.add_option('--imagename',dest="imagename",default=None,help="Name of the image file",metavar="String")
	parser.add_option('--metafits',dest="metafits",default=None,help="Name of the metafits file",metavar="String")
	parser.add_option('--gridpoint',dest="gridpoint",default=-1,help="MWA sweet spot pointing number",metavar="Integer")
	parser.add_option('--num_threads',dest="nthreads",default=1,help="Number of threads to use for parallel processing",metavar="Integer")
	parser.add_option('--image_conv',dest="conv",default='iau',help="Image convension",metavar="String")
	parser.add_option('--outfile',dest="outfile",default='MWAPB_cor',help="Output file basename",metavar="String")		
	parser.add_option('--restore_correction',dest="restore",default=False,help="Correct for primary beam or restore the correction",metavar="Boolean")		
	parser.add_option('--verbose',dest="verbose",default=False,help="Verbose output or not",metavar="Boolean")		
	(options, args) = parser.parse_args()
	mwapb_cor(options.imagename,options.outfile,metafits=options.metafits,verbose=eval(str(options.verbose)),gridpoint=int(options.gridpoint),nthreads=int(options.nthreads),\
				conv=str(options.conv),restore=eval(str(options.restore)),MWA_PB_path=str(options.MWA_PB_path))

datadir = os.path.dirname(paircarsdata.__file__)
datafile=datadir+'/mwa_full_embedded_element_pattern.h5'
haslam_map_file=datadir+'/haslam_map.fits'
sweet_spot_file=datadir+'/MWA_sweet_spots.npy'
os.environ['MWA_BEAM_FILE']=datafile
MWALON=116.67
MWALAT=-26.7
MWAALT=377.8
MWAPOS = EarthLocation.from_geodetic(lon="116:40:14.93", lat="-26:42:11.95", height=377.8)
######################################################################
def altaz_to_parallactic_angle(alt,az,LAT):
	'''
	Function to calculate parallactic angle given Altitude, Azimuith of the source and Latitude on the observatory
	Parameters
	----------
	alt : float
		Altitude in degree
	az : float
		Azimuith in degree
	lat : float
		Latitude in degree
	Returns
	-------
	float
		Parallactic angle in radian
	''' 
	alt=np.deg2rad(alt)
	az=np.deg2rad(az)
	lat=np.deg2rad(LAT)
	p=-np.arctan2(np.sin(az)*np.cos(lat),np.cos(alt)*np.sin(lat) - np.sin(alt)*np.cos(lat)*np.cos(az))
	return p

def get_azza_from_fits(filename,metafits='',verbose=False):
	'''
	Get azimuith and zenith angle arrays from fits file
	Parameters
	----------
	filename : str
		Name of the fits file
	metafits : str
		Metafits file name
	verbose : bool
		Verbose output or not
	Returns
	-------
	dict
		{'za_rad': theta,'astro_az_rad': phi}
	'''
	try:
		f=fits.open(filename)
	except IOError as err:
		print('Unable to open %s for reading\n%s', filename, err)
		return None

	h = f[0].header
	f.close()

	wcs = pywcs.WCS(h)
	naxes = h['NAXIS']

	if 'HPX' in h['CTYPE1']:
		print('Cannot deal with HPX coordinates')
		return None

	if 'RA' not in h['CTYPE1']:
		print('Coordinate 1 should be RA')
		return None
	if 'DEC' not in h['CTYPE2']:
		print('Coordinate 1 should be DEC')
		return None
	
	# construct the basic arrays
	x = np.arange(1, h['NAXIS1'] + 1)
	y = np.arange(1, h['NAXIS2'] + 1)
	# assume we want the first frequency
	# if we have a cube this will have to change
	ff = 1
	Y, X = np.meshgrid(y, x)

	Xflat = X.flatten()
	Yflat = Y.flatten()
	FF = ff * np.ones(Xflat.shape)

	if naxes >= 4:
		Tostack = [Xflat, Yflat, FF]
		for i in range(3, naxes):
			Tostack.append(np.ones(Xflat.shape))
	else:
		Tostack = [Xflat, Yflat]
	pixcrd = np.vstack(Tostack).transpose()

	try:
		# Convert pixel coordinates to world coordinates
		# The second argument is "origin" -- in this case we're declaring we
		# have 1-based (Fortran-like) coordinates.
		sky = wcs.wcs_pix2world(pixcrd, 1)
	except Exception as e:
		print('Problem converting to WCS: %s' % e)
		return None

	# extract the important pieces
	ra = sky[:, 0]
	dec = sky[:, 1]

	# and make them back into arrays
	RA = ra.reshape(X.shape)
	Dec = dec.reshape(Y.shape)

	# get the date so we can convert to Az,El
	if os.path.exists(metafits)==False:
		try:
			d = h['DATE-OBS']
		except Exception:
			print('Unable to read observation date DATE-OBS from %s' % filename)
			return None
	else:
		metadata=fits.getheader(metafits)
		d=metadata['DATE-OBS']

	if '.' in d:
		d = d.split('.')[0]
	dt = datetime.datetime.strptime(d, '%Y-%m-%dT%H:%M:%S')
	mwatime = Time(dt)
	if verbose:
		print('Image time = '+str(mwatime.iso)+', GPS time = '+str(mwatime.gps)+'\n')

	source = SkyCoord(ra=RA, dec=Dec, frame='icrs', unit=(astropy.units.deg, astropy.units.deg))
	source.location=MWAPOS
	source.obstime = mwatime

	source_altaz = source.transform_to('altaz')
	Alt, Az = source_altaz.alt.deg, source_altaz.az.deg  # Transform to Topocentric Alt/Az at the current epoch

	# go from altitude to zenith angle
	theta = (90 - Alt) * np.pi / 180
	phi = Az * np.pi / 180

    # BUGFIX by MS 2016-04-20 - described in details in odt document. The essence is that python reads images transposed with respect to cross-diagonal
    #  however later Tim calculates map of ra,dec then az,alt using WCS parameters which results in having coordinates for non-transposed image.
    # I verified pixel values vs. RA,DEC in the map in ds9 and the code and calling .transpose() here fixes the problem.
    # Also makes images look correct, otherwise there were "leakage-sources" on the right (east) of the beam-corrected image !
    #    return {'za_rad':theta,'astro_az_rad':phi,'frequencies':frequencies,'ra_now':RAnow,'dec_now':Decnow}
    #     return {'za_rad':theta,'astro_az_rad':phi,'frequencies':frequencies}
	return {'za_rad': theta.transpose(),'astro_az_rad': phi.transpose()}

def get_IQUV(filename,stokesaxis=4):
	"""Get IQUV from a CASA image file exported to fits."""
	data = fits.getdata(filename)
	stokes = {}
	if stokesaxis==3:
		stokes['I'] = data[0, 0, :, :]
		stokes['Q'] = data[0, 1, :, :]
		stokes['U'] = data[0, 2, :, :]
		stokes['V'] = data[0, 3, :, :]
	else:
		stokes['I'] = data[0, 0, :, :]
		stokes['Q'] = data[1, 0, :, :]
		stokes['U'] = data[2, 0, :, :]
		stokes['V'] = data[3, 0, :, :]
	return stokes

def get_inst_pols(stokes):
	"""Return instrumental polaristations matrix (Vij)"""
	XX = stokes['I'] + stokes['Q']
	XY = stokes['U'] + stokes['V'] * 1j
	YX = stokes['U'] - stokes['V'] * 1j
	YY = stokes['I'] - stokes['Q']
	Vij = np.array([[XX, XY], [YX, YY]])
	Vij=np.swapaxes(np.swapaxes(Vij,0,2),1,3)
	return Vij

def B2IQUV(B):
	"""Convert sky brightness matrix to I, Q, U, V"""
	B11 = B[:,:,0, 0]
	B12 = B[:,:,0, 1]
	B21 = B[:,:,1, 0]
	B22 = B[:,:,1, 1]

	stokes = {}
	stokes['I'] = (B11 + B22) / 2.
	stokes['Q'] = (B11 - B22) / 2.
	stokes['U'] = (B12 + B21) / 2.
	stokes['V'] = 1j * (B21 - B12) / 2.
	return stokes

def mwapb_cor(imagename,outfile,metafits='',verbose=False,gridpoint=-1,nthreads=1,conv='iau',restore=False,MWA_PB_path='',output_stokes=''):
	'''
	Correct FITS or CASA image for MWA primary beam
	Parameters
	----------
	imagename : str
		Name of the input file
	outfile : str
		Basename of the outputfile
	metafits : str
		Metafits file path
	verbose : bool
		Verbose output or not
	gridpoint : int
		MWA gridpoint number (default : -1, provide if you do not have metafits file)
	nthreads : int
		Number of cpu threads use for parallel computing
	conv : str
		Image coordinate convention (default : \'iau\') (options : \'iau\' and \'mwa\')
	restore : bool
		Whether correct for MWA primary beam or restore the correction
	MWA_PB_path : str
		MWA primary beam data file path (default : In built datafile from P-AIRCARS)
	output_stokes : str
		Output Stokes planes ('I' or 'IQUV', default : input image stokes)
	Returns
	-------
	str
		Output image name
	'''
	if MWA_PB_path!='' and os.path.exists(str(MWA_PB_path)+'/mwa_full_embedded_element_pattern.h5'):
		os.environ['MWA_BEAM_FILE']=str(MWA_PB_path)+'/mwa_full_embedded_element_pattern.h5'
	beam=mwa_hyperbeam.FEEBeam()
	outfile=os.path.basename(outfile)
	if imagename[-1]=='/':
		imagename=imagename[:-1]
	try:
		image_header=fits.getheader(imagename)
		naxes=int(image_header['NAXIS'])
		stokesaxis=4
		if naxes >= 4:
			if 'STOKES' not in image_header['CTYPE4']:
				stokesaxis=3
				if 'STOKES' not in image_header['CTYPE3']:
					print('Coordinate 3 or 4 should be STOKES')
					return 1
		else:
			print ('Number of axes is not 4.\n')
			return 1
		if stokesaxis==4:
			if int(image_header['NAXIS4'])!=4:
				if int(image_header['CRVAL4'])!=1:
					print ('1.Stokes axes are either not \'I\' or \'IQUV\'.\n')
					return 1
				else:
					data=fits.getdata(imagename)
					header=fits.getheader(imagename)
					header['NAXIS4']=4.
					header['CRVAL4']=1.
					header['CDELT4']=1.
					data=np.repeat(data,4,0)
					data[1,...]*=0
					data[2,...]*=0
					data[3,...]*=0
					fits.writeto(imagename+'.allStokes',data=data,header=header,overwrite=True)
					fitsfile=imagename+'.allStokes'
					org_stokes='I'
			else:
				org_stokes='IQUV'
				fitsfile=imagename
			imagetype='fits'
		else:
			if int(image_header['NAXIS3'])!=4:
				if int(image_header['CRVAL3'])!=1:
					print ('Stokes axes are either not \'I\' or \'IQUV\'.\n')
					return 1
				else:
					data=fits.getdata(imagename)
					header=fits.getheader(imagename)
					header['NAXIS3']=4.
					header['CRVAL3']=1.
					header['CDELT3']=1.
					data=np.repeat(data,4,1)
					data[:,1,...]*=0
					data[:,2,...]*=0
					data[:,3,...]*=0
					fits.writeto(imagename+'.allStokes',data=data,header=header,overwrite=True)
					fitsfile=imagename+'.allStokes'
					org_stokes='I'
			else:
				org_stokes='IQUV'
			imagetype='fits'
	except Exception as e:
		try:
			if imagename[-1]=='/':
				imagename=imagename[:-1]
			if os.path.exists(imagename+'.fits'):
				os.system('rm -rf '+imagename+'.fits')
			exportfits(imagename=imagename,fitsimage=imagename+'.fits',stokeslast=False)
			fitsfile=imagename+'.fits'
			image_header=fits.getheader(fitsfile)
			naxes=int(image_header['NAXIS'])
			stokesaxis=4
			if naxes >= 4:
				if 'STOKES' not in image_header['CTYPE4']:
					stokesaxis=3
					if 'STOKES' not in image_header['CTYPE3']:
						print('Coordinate 3 or 4 should be STOKES')
						return 1
			else:
				print ('Number of axes is not 4.\n')
				return 1
			if stokesaxis==4:
				if int(image_header['NAXIS4'])!=4:
					if int(image_header['CRVAL4'])!=1:
						print ('1.Stokes axes are either not \'I\' or \'IQUV\'.\n')
						return 1
					else:
						data=fits.getdata(fitsfile)
						header=fits.getheader(fitsfile)
						header['NAXIS4']=4.
						header['CRVAL4']=1.
						header['CDELT4']=1.
						data=np.repeat(data,4,0)
						data[1,...]*=0
						data[2,...]*=0
						data[3,...]*=0
						fits.writeto(fitsfile+'.allStokes',data=data,header=header,overwrite=True)
						fitsfile=fitsfile+'.allStokes'
						org_stokes='I'
				else:
					org_stokes='IQUV'
					fitsfile=fitsfile
				imagetype='CASA'
			else:
				if int(image_header['NAXIS3'])!=4:
					if int(image_header['CRVAL3'])!=1:
						print ('Stokes axes are either not \'I\' or \'IQUV\'.\n')
						return 1
					else:
						data=fits.getdata(fitsfile)
						header=fits.getheader(fitsfile)
						header['NAXIS3']=4.
						header['CRVAL3']=1.
						header['CDELT3']=1.
						data=np.repeat(data,4,1)
						data[:,1,...]*=0
						data[:,2,...]*=0
						data[:,3,...]*=0
						fits.writeto(fitsfile+'.allStokes',data=data,header=header,overwrite=True)
						fitsfile=fitsfile+'.allStokes'
						org_stokes='I'
				else:
					org_stokes='IQUV'
			imagetype='CASA'
		except Exception as e:
			print ('Image is not in fits or CASA image format. Error occured : '+str(e)+'\n')
			return 1
	if restore==False:
		print ('Correcting image : '+os.path.basename(imagename)+' for MWA primary beam response......\n')
	else:
		print ('Undo the correction image : '+os.path.basename(imagename)+' for MWA primary beam response......\n')
	naxes=image_header['NAXIS']
	freqaxis=3
	stokesaxis=4
	if naxes >= 4:
		if 'FREQ' not in image_header['CTYPE3']:
			freqaxis=4
			stokesaxis=3
			if 'FREQ' not in image_header['CTYPE4']:
				print('Coordinate 3 or 4 should be FREQ')
				return 1
		if freqaxis==3:
			freq = image_header['CRVAL3']  # read number of frequency channels
		else:
			freq = image_header['CRVAL4']
	else:
		print ('Number of axes is not 4.\n')
		return 1
	if metafits==None or os.path.isfile(metafits)==False:
		if gridpoint==-1:
			print ('Either provide correct metafits file or grid point number.\n')
			return 1
		else:
			sweet_spots=np.load(sweet_spot_file,allow_pickle=True).all()
			delay=sweet_spots[int(gridpoint)][-1]
	else:
		metadata=fits.getheader(metafits)
		gridpoint=metadata['GRIDNUM']
		sweet_spots=np.load(sweet_spot_file,allow_pickle=True).all()
		delay=sweet_spots[int(gridpoint)][-1]
	if nthreads<=0:
		nthreads=1
	os.environ['RAYON_NUM_THREADS']=str(nthreads)
	if verbose:
		print ('Calculating azimuith and zenith angles...\n')
	alt_az_array=get_azza_from_fits(fitsfile,metafits=metafits,verbose=verbose)
	if verbose:
		print ('Calculating Jones arrays....\n')
	'''jones_array=beam.calc_jones_array(alt_az_array['astro_az_rad'].flatten(),alt_az_array['za_rad'].flatten(), freq, delay, [1]*16, True)
	jones_array=jones_array.reshape(jones_array.shape[0],2,2) # Array reshaping to (xpix*ypix,2,2)
	pa_angles=erfa.hd2pa(alt_az_array['astro_az_rad'].flatten(),np.deg2rad(90-np.rad2deg(alt_az_array['za_rad']).flatten()),np.deg2rad(MWALAT))
	#altaz_to_parallactic_angle(90-np.rad2deg(alt_az_array['za_rad']).flatten(),np.rad2deg(alt_az_array['astro_az_rad']).flatten(),MWALAT)
	jones_pa=np.array([np.matrix([[np.cos(parang),np.sin(parang)],[-np.sin(parang),np.cos(parang)]]) for parang in pa_angles])

	if conv=='iau':
		print ('Converting beam jones convention to IAU...\n')
		jones_array=np.matmul(np.array([[0,-1],[-1,0]]),jones_array)
	jones_array=np.matmul(jones_pa,jones_array)'''
	jones_array=beam.calc_jones_array(alt_az_array['astro_az_rad'].flatten(),alt_az_array['za_rad'].flatten(), freq, delay, [1]*16, True,np.deg2rad(MWALAT),True)
	jones_array=jones_array.reshape(jones_array.shape[0],2,2)
	
	imagedata=fits.getdata(fitsfile)
	if stokesaxis==4:
		imagedata_swaped=np.swapaxes(np.swapaxes(imagedata,0,2),1,3)
	else:
		imagedata_swaped=np.swapaxes(np.swapaxes(np.swapaxes(imagedata,0,2),1,3),2,3)
	if stokesaxis==3:
		imagedata_reshaped=imagedata_swaped[:,:,:,0].reshape(imagedata.shape[-1]*imagedata.shape[-1],2,2)
	else:
		imagedata_reshaped=imagedata_swaped[:,:,:,0].reshape(imagedata.shape[-1]*imagedata.shape[-2],2,2)

	stokes=get_IQUV(fitsfile,stokesaxis=stokesaxis)
	Vij=get_inst_pols(stokes)
	Vij_reshaped=Vij.reshape(Vij.shape[0]*Vij.shape[1],2,2)
	jones_array_H=np.transpose(jones_array.conj(),axes=((0,2,1)))
	if eval(str(restore))==False:
		Vij_corrected=np.matmul(inv(jones_array),np.matmul(Vij_reshaped,inv(jones_array_H)))
	else:
		Vij_corrected=np.matmul(jones_array,np.matmul(Vij_reshaped,jones_array_H))
	Vij_reshaped=Vij_corrected.reshape(Vij.shape)
	B=B2IQUV(Vij_reshaped)
	if stokesaxis==3:
		imagedata[0,0,:,:]=np.real(B['I'])
		imagedata[0,1,:,:]=np.real(B['Q'])
		imagedata[0,2,:,:]=np.real(B['U'])
		imagedata[0,3,:,:]=np.real(B['V'])
	else:
		imagedata[0,0,:,:]=np.real(B['I'])
		imagedata[1,0,:,:]=np.real(B['Q'])
		imagedata[2,0,:,:]=np.real(B['U'])
		imagedata[3,0,:,:]=np.real(B['V'])

	output_fitsfile=os.path.dirname(os.path.abspath(imagename))+'/'+outfile+'.fits'
	if os.path.exists(output_fitsfile):
		os.system('rm -rf '+output_fitsfile)
	header_keys=image_header.keys()
	if 'BMAJ' not in header_keys or 'BMIN' not in header_keys or 'BPA' not in header_keys:
		try:
			f=fits.open(fitsfile)[1].data[0]
			maj=f[0]
			minor=f[1]
			pa=f[2]
			image_header['BMAJ']=maj
			image_header['BMIN']=minor
			image_header['BPA']=pa
			has_beam=True
		except:
			has_beam=False
	else:
		maj=image_header['BMAJ']*3600
		minor=image_header['BMIN']*3600
		pa=image_header['BPA']
		has_beam=True
	fits.writeto(output_fitsfile,data=imagedata,header=image_header,overwrite=True)
	if imagetype=='CASA':
		image_extension=imagename.split('.')[-1]
		if os.path.exists(os.path.dirname(os.path.abspath(imagename))+'/'+outfile+'.'+image_extension):
			os.system('rm -rf '+os.path.dirname(os.path.abspath(imagename))+'/'+outfile+'.'+image_extension)
		if has_beam:
			importfits(fitsimage=output_fitsfile,imagename=os.path.dirname(os.path.abspath(imagename))+'/'+outfile+'.'+image_extension,\
						beam=[str(maj)+'arcsec',str(minor)+'arcsec',str(pa)+'deg'])
		else:
			importfits(fitsimage=output_fitsfile,imagename=os.path.dirname(os.path.abspath(imagename))+'/'+outfile+'.'+image_extension)
		os.system('rm -rf '+output_fitsfile)
		if (output_stokes=='' and org_stokes=='I') or output_stokes=='I':
			if os.path.exists(os.path.dirname(os.path.abspath(imagename))+'/'+outfile+'.'+image_extension+'.StokesI'):
				os.system('rm -rf '+os.path.dirname(os.path.abspath(imagename))+'/'+outfile+'.'+image_extension+'.StokesI')
			imsubimage(imagename=os.path.dirname(os.path.abspath(imagename))+'/'+outfile+'.'+image_extension,outfile=os.path.dirname(os.path.abspath(imagename))+\
						'/'+outfile+'.'+image_extension+'.StokesI',stokes='I')
			if os.path.exists(os.path.dirname(os.path.abspath(imagename))+'/'+outfile+'.'+image_extension):
				os.system('rm -rf '+os.path.dirname(os.path.abspath(imagename))+'/'+outfile+'.'+image_extension)
			os.system('mv '+os.path.dirname(os.path.abspath(imagename))+'/'+outfile+'.'+image_extension+'.StokesI '+os.path.dirname(os.path.abspath(imagename))+'/'+\
						outfile+'.'+image_extension)
		os.system('rm -rf '+output_fitsfile+' '+fitsfile+' casa*log')
		print ('Output image written to : '+os.path.dirname(os.path.abspath(imagename))+'/'+outfile+'.'+image_extension+'\n')
		return os.path.dirname(os.path.abspath(imagename))+'/'+outfile+'.'+image_extension
	else:
		if (output_stokes=='' and org_stokes=='I') or output_stokes=='I':
			data=fits.getdata(output_fitsfile)
			fits.writeto(output_fitsfile,data=data[0,...],header=image_header,overwrite=True)
		print ('Output image written to : '+output_fitsfile+'\n')
		os.system('rm -rf casa*log')
		return output_fitsfile

def get_pb_radec(ra,dec,freq,obstime='',metafits=None,gridpoint=-1,MWA_PB_path=''):
	'''
	Function to get MWA primary beam at specific RA, DEC
	Parameters
	----------
	ra : str 
		RA either in degree or 'hh:mm:ss' or '%fh%fm%fs' format
	dec : str 
		DEC either in degree or 'dd:mm:ss' or '%fd%fm%fs'format
	freq : float
		Frequency in MHz
	obstime : str
		Time of the observation in 'yyyy-mm-dd hh:mm:ss' format
	metafits : str
		MWA metafits file
	gridpoint : int
		MWA sweetspot number
	MWA_PB_path : str
		MWA primary beam file path
	Returns
	-------
	int
		Success message (0 or 1)
	np.array
		Jones matrix
	float
		Stokes I beam value
	float
		XX power beam value
	float
		YY power beam value
	'''
	if MWA_PB_path!='' and os.path.exists(str(MWA_PB_path)+'/mwa_full_embedded_element_pattern.h5'):
		os.environ['MWA_BEAM_FILE']=str(MWA_PB_path)+'/mwa_full_embedded_element_pattern.h5'
	beam=mwa_hyperbeam.FEEBeam()
	if obstime=='' and metafits==None:
		print ('No time of observation is provided.\n')
		return 1,np.array([]),None
	else:
		metadata=fits.getheader(metafits)
		obstime=metadata['DATE-OBS']	
	observing_time=Time(obstime)  
	aa=AltAz(location=MWAPOS,obstime=observing_time)
	try:
		ra=float(ra)
		dec=float(dec)
		coord=SkyCoord(ra,dec,frame='icrs',unit='deg')
	except:
		try:
			coord=SkyCoord(ra,dec)
		except:
			coord=SkyCoord(ra,dec,unit=(u.hourangle,u.deg))
	altaz_object=coord.transform_to(aa)
	alt=altaz_object.alt.degree
	az=altaz_object.az.degree
	if metafits==None or os.path.isfile(metafits)==False:
		if gridpoint==-1:
			print ('Either provide correct metafits file or grid point number.\n')
			return 1,np.array([])
		else:
			sweet_spots=np.load(sweet_spot_file,allow_pickle=True).all()
			delay=sweet_spots[int(gridpoint)][-1],None
	else:
		metadata=fits.getheader(metafits)
		gridpoint=metadata['GRIDNUM']
		sweet_spots=np.load(sweet_spot_file,allow_pickle=True).all()
		delay=sweet_spots[int(gridpoint)][-1]
	jones_array=beam.calc_jones(np.deg2rad(az),np.deg2rad(90-alt), freq*10**6, delay, [1]*16, True,np.deg2rad(MWALAT),True)
	stokesI_beam=((np.abs(jones_array[0]))**2+(np.abs(jones_array[1]))**2+(np.abs(jones_array[2]))**2+(np.abs(jones_array[3]))**2)/2
	power_beam_array_XX=np.real(jones_array[0]*jones_array[0].conjugate()+jones_array[1]*jones_array[1].conjugate())
	power_beam_array_YY=np.real(jones_array[2]*jones_array[2].conjugate()+jones_array[3]*jones_array[3].conjugate())
	return 0,jones_array,float(stokesI_beam),power_beam_array_XX,power_beam_array_YY

def makeAZZA_dOMEGA(npix=256, projection='SIN'):
	'''
	Make azimuth and zenith angle arrays for a square image of side npix
	Parameters
	----------
	npix : int
		Number of pixels of the grid
	projection : str
		SIN or ZEA
	Returns
	-------
	list
		Azimuth angles in radians
	list
		Zenith angle in radians
	int
		Total number of pixels above the horizon
	float
		Differential solid angle (dOMEGA)	
	'''
	# build az and za arrays
	# use linspace to ensure we go to horizon on all sides
	z = np.linspace(-npix / 2.0, npix / 2.0, num=npix, dtype=np.float32)
	x = np.empty((npix, npix), dtype=np.float32)
	y = np.empty((npix, npix), dtype=np.float32)
	dOMEGA = np.empty((npix, npix), dtype=np.float32)
	for i in range(npix):
		y[i, 0:] = z
		x[0:, i] = z
	d = np.sqrt(x * x + y * y) / (npix / 2)
	# only select pixels above horizon
	t = (d <= 1.0)
	n_total = t.sum()
	dOMEGA.fill(np.pi * 2.00 / n_total)
	za = np.zeros((npix, npix), dtype=np.float32) * np.NaN
	if projection == 'SIN':
		za[t] = np.arcsin(d[t])
		dOMEGA = np.cos(za) * np.pi * 2.00 / n_total
	elif projection == 'ZEA':
		d = d * 2 ** 0.5  # ZEA requires R to extend beyond 1.
		za[t] = 2 * np.arcsin(d[t] / 2.0)
	else:
		e = 'Projection %s not found' % projection
		print(e)
		raise ValueError(e)
	az = np.arctan2(y, x)
	az = az + np.pi  # 0 to 2pi
	az = 2 * np.pi - az  # Change to clockwise from top (when origin is in top-left)
	return az, za, n_total, dOMEGA


def get_Haslam(freq, scaling=-2.55):
	'''
	Get the Haslam 408 MHz all sky map extrapolated to desired frequency.
	Parameters
	----------
	freq : float
		Frequency in MHz
	scaling : float
		Power law index for extrapolation (default : -2.55)
	Returns
	-------
	dict
		A python dictionary 
			i) Haslam map at given frequency in 10xK unit
			ii) RA in degree
			iii) DEC in degree
	'''
	radio_image_touse=haslam_map_file
	if not os.path.exists(radio_image_touse):
		print("Could not find 408 MHz image: %s\n" % (radio_image_touse))
		return None
	try:
		f=fits.open(radio_image_touse)
	except Exception:
		print("Error opening 408 MHz image: %s\n" % (radio_image_touse))
		return None
	skymap = f[0].data[0] / 10.0  # Haslam map is in 10xK
	skymap = skymap * (freq / 408.0) ** scaling  # Scale to frequency
	RA_1D = (f[0].header.get('CRVAL1') +
		     (np.arange(1, skymap.shape[1] + 1) - f[0].header.get('CRPIX1')) * f[0].header.get('CDELT1'))  # /15.0
	dec_1D = (f[0].header.get('CRVAL2') +
		      (np.arange(1, skymap.shape[0] + 1) - f[0].header.get('CRPIX2')) * f[0].header.get('CDELT2'))
	return {'skymap': skymap, 'RA': RA_1D, 'dec': dec_1D}  # RA, dec in degs

def map_sky(skymap,RA,dec,az_grid,za_grid,obstime=''):
	'''
	Reprojects Haslam map onto an input az, ZA grid.
	Parameters
	----------
	skymap : np.array
		Haslam map in RA DEC
	RA : np.array
		1D range of RAs (deg)
	dec : np.array 
		1D range of decs (deg)
	az_grid : np.array 
		Grid of azes onto which we map sky
	za_grid : np.array 
		Grid of ZAs onto which we map sky
	obstime : str
		Time of the observation in 'yyyy-mm-dd hh:mm:ss' format
	'''
	# Get az, ZA grid transformed to equatorial coords
	grid2eq = horz2eq(az_grid, za_grid, obstime)
	my_interp_fn = RegularGridInterpolator((dec, RA[::-1]), skymap[:, ::-1], fill_value=np.nan)
	# interpolate map onto az,ZA grid
	# Convert to RA=-180 - 180 format (same as Haslam)
	# We do it this way so RA values are always increasing for RegularGridInterpolator
	grid2eq['RA'][grid2eq['RA'] > 180] = grid2eq['RA'][grid2eq['RA'] > 180] - 360
	my_map = my_interp_fn(np.dstack([grid2eq['DEC'], grid2eq['RA']]))
	return my_map

def horz2eq(az,ZA,obstime):
	'''
	Convert from horizontal (az, ZA) to equatorial (RA, dec)
	Parameters
	----------
	az : np.array 
		Grid of azes onto which we map sky
	za : np.array 
		Grid of ZAs onto which we map sky
	obstime : str
		Time of the observation in 'yyyy-mm-dd hh:mm:ss' format
	Returns
	-------
	dict
		A python dictionary {'RA' : degress, 'DEC' : degrees} 
	'''
	MWA_TOPO = si.Topos(longitude=(116, 40, 14.93), latitude=(-26, 42, 11.95), elevation_m=377.8)
	skyfield_loader = si.Loader(datadir, verbose=False, expire=True)
	PLANETS = skyfield_loader('de421.bsp')
	TIMESCALE = skyfield_loader.timescale(builtin=True)
	S_MWAPOS = PLANETS['earth'] + MWA_TOPO
	observing_time=Time(obstime)
	observer=S_MWAPOS.at(TIMESCALE.from_astropy(observing_time))
	coords=observer.from_altaz(alt_degrees=(90 - ZA), az_degrees=az, distance=si.Distance(au=9e90))
	ra_a, dec_a, _ = coords.radec()
	return {'RA':ra_a._degrees, 'DEC':dec_a.degrees}

def make_primarybeammap(freq=0,metafits='',gridpoint=-1,obstime='',resolution=1,MWA_PB_path='',nthreads=1,verbose=False):
	'''
	Parameters
	----------
	freq : float
		Frequency in MHz
	metafits : str
		Metafits file
	gridpoint : int
		MWA sweetspot number
	obstime : str
		Time of the observation in 'yyyy-mm-dd hh:mm:ss' format (If not given automatically obtain from metafits file)
	resolution : float
		Beam resolution in degree (default : 1deg)
	MWA_PB_path : str
		MWA primary beam file path
	nthreads : int
		Number of cpu threads use for parallel computing
	verbose : bool
		Verbose output or not
	Returns
	-------
	float 
		Sum of full Beam*Sky (XX) 
	float
		Sum of full Beam (XX)
	float
		Antenna temperature (XX)
	float
		Total beam area (XX)
	float 
		Sum of full Beam*Sky (YY) 
	float
		Sum of full Beam (YY)
	float
		Antenna temperature (YY)
	float
		Total beam area (YY)
	'''
	warnings.filterwarnings("ignore")
	if verbose:
		print ('Estimating beam averaged sky temperatures at frequency : '+str(freq)+' MHz.\n')
	if MWA_PB_path!='' and os.path.exists(str(MWA_PB_path)+'/mwa_full_embedded_element_pattern.h5'):
		os.environ['MWA_BEAM_FILE']=str(MWA_PB_path)+'/mwa_full_embedded_element_pattern.h5'
	n_pix=int(360/resolution)
	az_grid, za_grid, n_total, dOMEGA=makeAZZA_dOMEGA(n_pix,'ZEA')
	az_grid = az_grid * 180 / np.pi
	za_grid = za_grid * 180 / np.pi
	alt_grid = 90 - (za_grid)
	# first go from altitude to zenith angle
	theta = (90 - alt_grid) * np.pi / 180
	phi = az_grid * np.pi / 180
	beam=mwa_hyperbeam.FEEBeam()
	if metafits==None or os.path.isfile(metafits)==False:
		if gridpoint==-1:
			print ('Either provide correct metafits file or grid point number.\n')
			return 1
		else:
			sweet_spots=np.load(sweet_spot_file,allow_pickle=True).all()
			delay=sweet_spots[int(gridpoint)][-1]
		if obstime=='':
			print ('Either provide correct metafits file or observation time.\n')
			return 1
	else:
		metadata=fits.getheader(metafits)
		gridpoint=metadata['GRIDNUM']
		sweet_spots=np.load(sweet_spot_file,allow_pickle=True).all()
		delay=sweet_spots[int(gridpoint)][-1]
		obstime=metadata['DATE-OBS']	
	if nthreads<=0:
		nthreads=1
	os.environ['RAYON_NUM_THREADS']=str(nthreads)
	jones_array=beam.calc_jones_array(phi.flatten(),theta.flatten(),freq*10**6,delay,[1]*16,True,np.deg2rad(MWALAT),True)
	power_beam_array={}
	power_beam_array['XX']=np.real(jones_array[:,0]*jones_array[:,0].conjugate()+jones_array[:,1]*jones_array[:,1].conjugate()).reshape(az_grid.shape)
	power_beam_array['YY']=np.real(jones_array[:,2]*jones_array[:,2].conjugate()+jones_array[:,3]*jones_array[:,3].conjugate()).reshape(az_grid.shape)	
	# Get Haslam and interpolate onto grid
	haslam_map = get_Haslam(freq)
	mask = np.isnan(za_grid)
	za_grid[np.isnan(za_grid)] = 90.0  # Replace nans as they break the interpolation
	sky_grid = map_sky(haslam_map['skymap'], haslam_map['RA'], haslam_map['dec'],az_grid, za_grid,obstime=obstime)
	sky_grid[mask] = np.nan  # Remask beyond the horizon
	beamsky_sum_XX = 0
	beam_sum_XX = 0
	Tant_XX = 0
	beam_dOMEGA_sum_XX = 0
	beamsky_sum_YY = 0
	beam_sum_YY = 0
	Tant_YY = 0
	beam_dOMEGA_sum_YY = 0
	pols = ['XX', 'YY']
	for pol in pols:
		# Get gridded sky
		beam = power_beam_array[pol]
		beamsky = beam * sky_grid
		beam_dOMEGA = beam * dOMEGA
		beamsky_sum=np.nansum(beamsky)
		beam_sum=np.nansum(beam)
		beam_dOMEGA_sum=np.nansum(beam_dOMEGA)
		Tant=np.nansum(beamsky)/np.nansum(beam)
		if pol == 'XX':
			beamsky_sum_XX = beamsky_sum
			beam_sum_XX = beam_sum
			Tant_XX = Tant
			beam_dOMEGA_sum_XX = beam_dOMEGA_sum
		if pol == 'YY':
			beamsky_sum_YY = beamsky_sum
			beam_sum_YY = beam_sum
			Tant_YY = Tant
			beam_dOMEGA_sum_YY = beam_dOMEGA_sum
	return beamsky_sum_XX,beam_sum_XX,Tant_XX,beam_dOMEGA_sum_XX,beamsky_sum_YY,beam_sum_YY,Tant_YY,beam_dOMEGA_sum_YY

if __name__=='__main__':
	main()


