import numexpr as ne
import healpy as hp
import numpy as np
import cosmotool as ct
import argparse
import h5py as h5
import matplotlib
matplotlib.use('Agg')
from matplotlib import pyplot as plt
import ksz
from ksz.constants import *
from cosmotool import interp3d

def wrapper_impulsion(f):

  class _Wrapper(object):
    def __init__(self):
      pass
      
    def __getitem__(self,direction):
    
      if 'velocity' in f:
        return f['velocity'][:,:,:,direction]
      
      n = "p%d" % direction
      return f[n]

  return _Wrapper()


def build_unit_vectors(N):
  ii = np.arange(N,dtype=np.float64)/N - 0.5
  d = np.sqrt(ii[:,None,None]**2 + ii[None,:,None]**2 + ii[None,None,:]**2)
  d[N/2,N/2,N/2] = 1
  ux = ii[:,None,None] / d
  uy = ii[None,:,None] / d
  uz = ii[None,None,:] / d

  return ux,uy,uz

def generate_from_catalog(dmin,dmax,Nside,do_random=False):
  import progressbar as pbar
  
  cat = np.load("2m++.npy")

  cat['distance'] = cat['best_velcmb']
  cat = cat[np.where((cat['distance']>100*dmin)*(cat['distance']<dmax*100))]

  deg2rad = np.pi/180
  Npix = 12*Nside**2
  xp,yp,zp = hp.pix2vec(Nside, np.arange(Npix))
  N2 = np.sqrt(xp**2+yp**2+zp**2)

  ksz_template = np.zeros(Npix, dtype=np.float64)
  ksz_mask = np.ones(Npix, dtype=np.uint8)

  for i in pbar.ProgressBar(maxval = cat.size, widgets=[pbar.Bar(), pbar.ETA()])(cat):
    if do_random:
      l = np.random.rand()*360
      b = np.arcsin(2*np.random.rand()-1)*90/np.pi
    else:
      l,b=i['gal_long'],i['gal_lat']

    l=ne.evaluate('l*deg2rad')
    b=ne.evaluate('b*deg2rad')

    x0 = np.cos(l)*np.cos(b)
    y0 = np.sin(l)*np.cos(b)
    z0 = np.sin(b)

    DA =i['distance']/100
    Lgal = DA**2*10**(0.4*(tmpp_cat['Msun']-i['K2MRS']+25))

    profiler = ksz.KSZ_Isothermal(Lgal, 2.37)

    idx0 = hp.query_disc(Nside, (x0,y0,z0), 3*profiler.rGalaxy/DA)

    xp1 = xp[idx0]
    yp1 = yp[idx0]
    zp1 = zp[idx0]
    N2_1 = N2[idx0]

    cos_theta = ne.evaluate('(x0*xp1+y0*yp1+z0*zp1)/(sqrt(x0**2+y0**2+z0**2)*(N2_1))')

    idx,idx_masked,m = profiler.projected_profile(cos_theta, DA)
    idx = idx0[idx]
    idx_masked = idx0[idx_masked]
    ksz_template[idx] += m
    ksz_mask[idx_masked] = 0

  ne.evaluate('ksz_template*ksz_normalization', out=ksz_template)

  return ksz_template, ksz_mask

def get_args():
  parser=argparse.ArgumentParser(description="Generate Skymaps from CIC maps")
  parser.add_argument('--Nside', type=int, default=128)
  parser.add_argument('--minval', type=float, default=-0.5)
  parser.add_argument('--maxval', type=float, default=0.5)
  parser.add_argument('--depth_min', type=float, default=10)
  parser.add_argument('--depth_max', type=float, default=60)
  parser.add_argument('--ksz_map', type=str, required=True)
  parser.add_argument('--base_fig', type=str, default="kszfig.png")
  parser.add_argument('--build_dipole', type=bool, default=False)
  parser.add_argument('--degrade', type=int, default=-1)
  parser.add_argument('--random', type=bool, default=False)
  return parser.parse_args()

def main():

  args = get_args()

  ff=plt.figure(1)
  plt.clf()
  v=[]

  print("Generating map...")

  proj,mask = generate_from_catalog(args.depth_min,args.depth_max,args.Nside,do_random=args.random)

  if args.degrade > 0:
    proj *= mask
    proj = hp.ud_grade(proj, nside_out=args.degrade)
    mask = hp.ud_grade(mask, nside_out=args.degrade)
    Nside = args.degrade
  else:
    Nside = args.Nside

  hp.write_map(args.ksz_map + ".fits", proj)
  hp.write_map(args.ksz_map + "_mask.fits", mask)

  if args.build_dipole:
    x,y,z=hp.pix2vec(Nside, np.arange(hp.nside2npix(Nside)))
    hp.write_map(args.ksz_map + "_x.fits", proj*x)
    hp.write_map(args.ksz_map + "_y.fits", proj*y)
    hp.write_map(args.ksz_map + "_z.fits", proj*z)

  hp.mollview(proj*100*1e6, fig=1, coord='GG', cmap=plt.cm.coolwarm, title='', min=args.minval, 
              max=args.maxval)

  ff.savefig(args.base_fig)
  
main()