#!/usr/bin/env python

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

from desitarget import io
from desitarget.randoms import supplement_randoms
import fitsio

#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(description='Make a random catalog with "zeros" for pixel-level quantities in missing (i.e. outside-of-the-footprint) bricks')
ap.add_argument("randomcat",
                help='An existing "inside-of-the-footprint" random catalog. Must contain at least the column "BRICKNAME" and the header card "SEED" (e.g /project/projectdirs/desi/target/catalogs/randoms-dr4-0.20.0.fits). If SEED does not exist in the header, the seed defaults to 1 in desitarget.randoms.supplement_randoms()"')
ap.add_argument("dest",
                help="Output directory for random catalog (the file name is built on-the-fly from other inputs)")
ap.add_argument("--density", type=int,
                help='Number of points per sq. deg. to generate (defaults to the value of DENSITY in the header of the "randomcat" input file)',
                default=None)
ap.add_argument("--seed", type=int,
                help='Force a particular seed to be used instead of reading from the header card "SEED"',
                default=None)
ap.add_argument("--numproc", type=int,
                help='number of concurrent processes to use [{}]'.format(nproc),
                default=nproc)
ap.add_argument("--dustdir",
                help="Directory of SFD dust maps (defaults to the equivalent of $DUST_DIR+'/maps')",
                default=None)

ns = ap.parse_args()

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

log.info('running on {} processors...t = {:.1f}s'.format(ns.numproc, time()-start))

# ADM check that the passed file looks like a random catalog.
try:
    ranhead = fitsio.read_header(ns.randomcat, "RANDOMS")
except OSError:
    msg = 'Passed input file does not appear to be a random catalog: {}'.format(ns.randomcat)
    log.critical(msg)
    sys.exit(1)

# ADM determine the seed used for the existing random catalog.
try:
    seed = ranhead["SEED"]
    log.info('read original seed of {} from {}'.format(seed, ns.randomcat))
except KeyError:
    seed = 1
    log.info('defaulting to an original SEED of {}'.format(seed))
if ns.seed is not None:
    # ADM if we enforced a new seed, store the original seed
    # ADM but update to use the new seed.
    origseed = seed
    writeseed = ns.seed
    batch = 20
    if writeseed > batch-1:
        log.error("Only set to batch 20 catalogs: Seed must be in range {}-{}"
                  .format(0, writeseed-1))
        raise IOError
    seed = origseed*20 + writeseed
    log.info('actually using a seed of {}*{}+{}={}'.format(
        origseed, batch, writeseed, seed))

# ADM grab the density used for the existing random catalog, if needed.
if ns.density is None:
    try:
        density = ranhead["DENSITY"]
        log.info('read DENSITY of {} from {}'.format(density, ns.randomcat))
    except KeyError:
        msg = 'No DENSITY keyword in header of: {}, so --density must be passed'.format(ns.randomcat)
        log.critical(msg)
        sys.exit(1)
else:
    density = ns.density

# ADM find the bricknames covered by the existing random catalog.
try:
    donebns = fitsio.read(ns.randomcat, columns="BRICKNAME")
except ValueError:
    msg = 'No BRICKNAME column in passed input file: {}'.format(ns.randomcat)
    log.critical(msg)
    sys.exit(1)

# ADM make the array of "zerod" bricks.
randoms = supplement_randoms(donebns, density=density, numproc=ns.numproc,
                             dustdir=ns.dustdir, seed=seed)

# ADM extra header keywords for the output fits file.
extra = {k: v for k, v in zip(["density", "seed"],
                              [density, writeseed])}
# ADM if we enforced a new seed, also store the original seed
# ADM that was used to make the original file of randoms.
if ns.seed is not None:
    extra["origseed"] = origseed

# ADM write out the supplemental random catalog.
nrands, outfile = io.write_randoms(ns.dest, randoms, indir=ns.randomcat,
                                   supp=True, extra=extra)

log.info('wrote file of {} randoms to {}...t = {:.1f}s'
         .format(nrands, outfile, time()-start))
