import healpy as hp
import numpy as np
import cosmotool as ct
import argparse
import h5py as h5
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()

parser=argparse.ArgumentParser(description="Generate Skymaps from CIC maps")
parser.add_argument('--boxsize', type=float, required=True)
parser.add_argument('--Nside', type=int, default=128)
parser.add_argument('--base_h5', type=str, required=True)
parser.add_argument('--base_fig', type=str, required=True)
parser.add_argument('--start', type=int, required=True)
parser.add_argument('--end', type=int, required=True)
parser.add_argument('--step', type=int, required=True)
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('--iid', type=int, default=0)
parser.add_argument('--ksz_map', type=str, required=True)
args = parser.parse_args()

L = args.boxsize
Nside = args.Nside

def build_unit_vectors(N):
  ii = np.arange(N)*L/N - 0.5*L
  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 build_radial_v(v):
  u = build_unit_vectors(N)
  vr = v[0] * u[0]
  vr += v[1] * u[1]
  vr += v[2] * u[2]
  return vr

def generate_from_catalog(vfield,Boxsize):
  import progressbar as pbar
  
  cat = np.load("2m++.npy")

  profiler = KSZ_Isothermal(2.37)

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

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

  ksz_template = np.zeros(12*Nside**2, dtype=np.float64)
  ksz_mask = np.zeros(12**Nside**2, dtype=np.uint8)

  pb = pbar.ProgressBar(maxval = cat.size, widgets=[pb.Bar(), pb.ETA()]).start()

  for k,i in np.ndenumerate(cat):
    pb.update(k)
    l,b=i['gal_long'],i['gal_lat']

    l *= deg2rad
    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))

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

    vr = interp3d(DA * x0, DA * y0, DA * z0, vfield, Boxsize)

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

    cos_theta = x0*xp1+y0*yp1+z0*zp1
    cos_theta /= np.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] += vr * m
    ksz_mask[idx_masked] = 0

  pb.finish()

  return ksz_template, ksz_mask

for i in xrange(args.start,args.end,args.step):
  ff=plt.figure(1)
  plt.clf()
  v=[]
  with h5.File(args.base_h5 % i, mode="r") as f:
    p = wrapper_impulsion(f)
    v.append(p[i] / f['density'][:])

  # Rescale by Omega_b / Omega_m
  vr = build_radial_v(v)
  vr *= -ksz_normalization*rho_mean_matter*1e6
  del v

  proj = generate_from_catalog(vfield,Boxsize)

  hp.write_map(args.ksz_map % i, proj)

  hp.mollview(proj, fig=1, coord='CG', cmap=plt.cm.coolwarm, title='Sample %d' % i, min=args.minval, 
              max=args.maxval)

  ff.savefig(args.base_fig % i)