#!python

from __future__ import print_function, division

import os, sys
import numpy as np
import fitsio
from time import time
start = time()

from desitarget.randoms import pixmap

#import warnings
#warnings.simplefilter('error')

import multiprocessing
nproc = multiprocessing.cpu_count() // 2

from desiutil.log import get_logger
log = get_logger()

from argparse import ArgumentParser
ap = ArgumentParser("Generate a map of HEALPixels with information on survey coverage, expected stellar density, and target density")
ap.add_argument("randoms",
                help='File of random points generated by, e.g., select_randoms')
ap.add_argument("targets",
                help='File of targets generated by, e.g., select_targets (should be the same Data Release as used to make the randoms')
ap.add_argument("dest",
                help='Output file name to write map of HEALPixel weights in the NESTED scheme (e.g. /project/projectdirs/desi/target/catalogs/pixweight-dr4-0.20.0.fits)')
ap.add_argument("--nside", type=int,
                help='The resolution (HEALPixel nside number) at which to build the map (defaults to 256)',
                default="256")
ap.add_argument("--gaialoc",
                help='A FITS file that already contains the Gaia stellar densities at nside to speed-up density calculations')

ns = ap.parse_args()

if not os.path.exists(ns.randoms):
    log.critical('Input directory does not exist: {}'.format(ns.randoms))
    sys.exit(1)

if not os.path.exists(ns.targets):
    log.critical('Input directory does not exist: {}'.format(ns.targets))
    sys.exit(1)

hdr = fitsio.read_header(ns.randoms, "RANDOMS")
#ADM add HEALPixel and gaialoc information to the header
hdr['GAIALOC'] = ns.gaialoc
hdr['HPXNSIDE'] = ns.nside
hdr['HPXNEST'] = True

pixmap, survey = pixmap(ns.randoms, ns.targets, hdr["DENSITY"], nside=ns.nside, gaialoc=ns.gaialoc)
hdr["SURVEY"] = survey

#ADM write out the map
log.info('Writing pixel map to {}'.format(ns.dest))

fitsio.write(ns.dest, pixmap, extname='PIXWEIGHTS', header=hdr, clobber=True)
log.info('wrote map of HEALPixel weights (in the nested scheme) to {}'.format(ns.dest))
log.info('Done...t={:.1f}s'.format(time()-start))
