#!/usr/bin/env python

import sys
import numpy as np
import argparse
from desitarget.gfa import select_gfas
from desitarget.geomask import is_in_gal_box, is_in_box
from desitarget import io
from desitarget.io import check_both_set
from time import time
t0 = time()

from desiutil.log import get_logger
log = get_logger()

import multiprocessing
nproc = multiprocessing.cpu_count() // 2
nside = io.desitarget_nside()   # ADM default HEALPix Nside used throughout desitarget

from argparse import ArgumentParser
ap = ArgumentParser(description='Generates a file of GFA (Guide/Focus/Alignment) targets via matching to Gaia')
ap.add_argument("surveydir",
                help="Base directory for a Legacy Surveys Data Release (e.g. '/global/project/projectdirs/cosmo/data/legacysurvey/dr6/' at NERSC)")
ap.add_argument("dest",
                help="Output GFA targets directory (the file name is built on-the-fly from other inputs)")
ap.add_argument("--gaiadr",
                help="Gaia Data Release on which to base GFAs. Options are dr2 or edr3. Defaults to dr2.",
                default="dr2")
ap.add_argument('-m', '--maglim', type=float,
                help="Magnitude limit on GFA targets in Gaia G-band (defaults to [21])",
                default=21)
ap.add_argument('-n', "--numproc", type=int,
                help="number of concurrent processes to use (defaults to [{}]). Note that if numproc > 4, I/O limited parts of the code revert to numproc=4".format(nproc),
                default=nproc)
ap.add_argument('-s2', "--surveydir2",
                help='Additional Legacy Surveys Data Release directory (useful for combining, e.g., DR8 into one file of GFAs)',
                default=None)
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('-dec', "--mindec", type=float,
                help="Minimum declination to include in output file for NON-LEGACY-SURVEYS sources (degrees; defaults to [-90])",
                default=-90.)
ap.add_argument('-b', "--mingalb", type=float,
                help="Closest latitude to Galactic plane to output for NON-LEGACY-SURVEYS sources (e.g. send 10 to limit to areas beyond -10o <= b < 10o; defaults to [0])",
                default=0.)
ap.add_argument("--nourat", action='store_true',
                help="If sent, then DO NOT add URAT proper motions for Gaia sources that are missing measurable PMs")

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 "maglim", "mindec", "mingalb", "nourat", "gaiadr":
    if isinstance(nsdict[nskey], bool):
        if nsdict[nskey]:
            extra += " --{}".format(nskey)
    else:
        extra += " --{} {}".format(nskey, nsdict[nskey])

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

if ns.bundlefiles is None:
    log.info('running on {} processors...t = {:.1f}s'.format(ns.numproc, time()-t0))
    # ADM formally writing pixelized files requires both the nside
    # ADM and the list of healpixels to be set.
    check_both_set(ns.healpixels, ns.nside)

# 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(',')]

gfas = select_gfas(infiles, maglim=ns.maglim, numproc=ns.numproc, nside=ns.nside,
                   pixlist=pixlist, bundlefiles=ns.bundlefiles, extra=extra,
                   mindec=ns.mindec, mingalb=ns.mingalb, addurat=not(ns.nourat),
                   dr=ns.gaiadr)

# ADM only proceed if we're not writing a slurm script.
if ns.bundlefiles is None:
    # ADM extra header keywords for the output fits file.
    extra = {k: v for k, v in zip(
        ["maglim", "mindec", "mingalb", "nourat", "gaiadr"],
        [ns.maglim, ns.mindec, ns.mingalb, ns.nourat, ns.gaiadr])}

    log.info('Writing GFAs to file...t = {:.1f} mins'.format((time()-t0)/60.))
    ngfas, outfile = io.write_gfas(ns.dest, gfas, indir=ns.surveydir,
                                   indir2=ns.surveydir2, nside=nside,
                                   nsidefile=ns.nside, hpxlist=pixlist,
                                   extra=extra)
    log.info('{} GFAs written to {}...t = {:.1f} mins'.format(
        ngfas, outfile, (time()-t0)/60.))
