#!/usr/bin/env python

from __future__ import print_function, division

import sys
import numpy as np
import fitsio
import os

from desitarget import io
from desitarget.cmx.cmx_cuts import select_targets
from desitarget.targets import decode_targetid

from time import time
start = time()

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

import multiprocessing
nproc = multiprocessing.cpu_count() // 2
nside = io.desitarget_nside()

from desiutil.log import get_logger
log = get_logger()

from argparse import ArgumentParser
ap = ArgumentParser(description='Generates DESI commissioning (cmx) target bits from Legacy Surveys sweeps or tractor files')
ap.add_argument("sweepdir",
                help="Tractor/sweeps file or root directory with tractor/sweeps files")
ap.add_argument("dest", 
                help="Output target selection directory (the file name is built on-the-fly from other inputs)")
ap.add_argument('-s2', "--sweepdir2",
                help='Additional Tractor/sweeps file or directory (useful for combining, e.g., DR8 into one file of targets)',
                default=None)
ap.add_argument("--numproc", type=int,
                help='number of concurrent processes to use [{}]'.format(nproc),
                default=nproc)
ap.add_argument('--nside', type=int,
                help="Process GFAs in HEALPixels at this resolution (defaults to None). See also the 'healpixels' input flag",
                default=None)
ap.add_argument('--healpixels',
                help="HEALPixels corresponding to `nside` (e.g. '6,21,57'). Only process files that touch these pixels and return targets within these pixels",
                default=None)
ap.add_argument("--bundlefiles", type=int,
                help="(overrides all options but `surveydir`) print slurm script to parallelize by sending (any) integer. This is an integer rather than boolean for historical reasons",
                default=None)
ap.add_argument("--cmxdir", type=str,
                help="Pass to set the CMX_DIR environment variable directly in the code",
                default=None)
ap.add_argument("--noresolve", action='store_true',
                help="Do NOT resolve into northern targets in northern regions and southern targets in southern regions")
ap.add_argument("--nobackup", action='store_true',
                help="Do NOT run the Gaia-only backup targets (which require the GAIA_DIR environment variable to be set).")

ns = ap.parse_args()

# ADM build the list of command line arguments as
# ADM bundlefiles potentially needs to know about them.
extra = " --numproc {}".format(ns.numproc)
nsdict = vars(ns)
for nskey in "noresolve", "cmxdir":
    if isinstance(nsdict[nskey], bool):
        if nsdict[nskey]:
            extra += " --{}".format(nskey)
    else:
        if nsdict[nskey] is not None:
            extra += " --{} {}".format(nskey, nsdict[nskey])

infiles = io.list_sweepfiles(ns.sweepdir)
if ns.sweepdir2 is not None:
    infiles2 = io.list_sweepfiles(ns.sweepdir2)
    infiles += infiles2
if len(infiles) == 0:
    infiles = io.list_tractorfiles(ns.sweepdir)
    if ns.sweepdir2 is not None:
        infiles2 = io.list_tractorfiles(ns.sweepdir2)
        infiles += infiles2
if len(infiles) == 0:
    log.critical('no sweep or tractor files found')
    sys.exit(1)

# ADM Only coded for objects with Gaia matches
# ADM (e.g. DR6 or above). Fail for earlier Data Releases.
# ADM Guard against a single file being passed.
fns = infiles
if isinstance(infiles, str):
    fns = [infiles]
# ADM check the first and last file, in case --sweepdir2 was sent.
for fn in fns[0], fns[-1]:
    data = fitsio.read(fn, columns=["RELEASE", "PMRA"])
    if np.any(data["RELEASE"] < 6000):
        log.critical('Commissioning cuts only coded for DR6 or above')
        raise ValueError
    if (np.max(data['PMRA']) == 0.) & np.any(data["RELEASE"] < 7000):
        d = "/project/projectdirs/desi/target/gaia_dr2_match_dr6"
        log.info("Zero objects have a proper motion.")
        log.critical(
            "Did you mean to send the Gaia-matched sweeps in, e.g., {}?"
            .format(d)
        )
        raise IOError

if ns.bundlefiles is None:
    log.info("running on {} processors".format(ns.numproc))

# ADM parse the list of HEALPixels in which to run.
pixlist = ns.healpixels
if pixlist is not None:
    pixlist = [int(pix) for pix in pixlist.split(',')]

targets = select_targets(infiles, numproc=ns.numproc,
                         resolvetargs=not(ns.noresolve), backup=not(ns.nobackup),
                         nside=ns.nside, pixlist=pixlist, extra=extra,
                         bundlefiles=ns.bundlefiles, cmxdir=ns.cmxdir)

if ns.bundlefiles is None:
    # ADM write out two sets of targets. The gaia-only and the not gaia-only.
    _, _, _, _, _, gaiadr = decode_targetid(targets["TARGETID"])
    isgaia = gaiadr > 0
    for supp, ii in zip([False, True], [~isgaia, isgaia]):
        ntargs, filename = io.write_targets(
            ns.dest, targets[ii], indir=ns.sweepdir,
            indir2=ns.sweepdir2, survey="cmx", nside=nside,
            resolve=not(ns.noresolve), supp=supp,
            nsidefile=ns.nside, hpxlist=pixlist)
        log.info('{} targets written to {}...t={:.1f}s'
                 .format(ntargs, filename, time()-start))
