import numpy as np
from pixell import enmap, utils

shape, wcs = enmap.geometry(pos=np.array([[1,-1],[-1,1]])*2*utils.degree, res=0.1*utils.arcmin)
flat = 2

def sim_beam_avg(shape, wcs, flat=2, bsize=2*utils.fwhm*utils.arcmin, lknee=3000, alpha=-4):
	lmap = enmap.lmap(shape, wcs)
	# Set up the beam and inv N for a horizontal and vertical scan.
	# The beam is elongated by flattening in the scan direction
	Bs, Ns = enmap.zeros((2,2)+shape, wcs)
	for i in range(2):
		# y-aligned scans for i==0, otherwise x-aligned
		l1, l2 = (lmap if i == 0 else lmap[::-1])
		# beam should be physically long in l1-direction
		Bs[i] = np.exp(-0.5*(l1**2*flat + l2**2/flat)*bsize**2)
		with utils.nowarn():
			# scan in l1-direction
			Ns[i] = np.maximum(1e-10,(1 + (l1/lknee)**alpha)**-1)
	# Plain average beam
	B_avg    = np.mean(Bs,0)
	# Weighted average beam
	B_wavg = np.sum(Bs*Ns,0)/np.sum(Ns,0)
	return B_avg, B_wavg

B_par_avg, B_par_wavg = sim_beam_avg(shape, wcs, flat=flat)
B_ort_avg, B_ort_wavg = sim_beam_avg(shape, wcs, flat=1/flat)

enmap.write_map("B_par_wavg.fits", B_par_wavg.lform())
enmap.write_map("B_par_avg.fits",  B_par_avg.lform())
enmap.write_map("B_ort_wavg.fits", B_ort_wavg.lform())
enmap.write_map("B_ort_avg.fits",  B_ort_avg.lform())

# Real-space stuff too
enmap.write_map("B_par_wavg_real.fits", enmap.fftshift(enmap.ifft(B_par_wavg+0j).real))
enmap.write_map("B_par_avg_real.fits",  enmap.fftshift(enmap.ifft(B_par_avg+0j).real))
enmap.write_map("B_ort_wavg_real.fits", enmap.fftshift(enmap.ifft(B_ort_wavg+0j).real))
enmap.write_map("B_ort_avg_real.fits",  enmap.fftshift(enmap.ifft(B_ort_avg+0j).real))