#!python

import sys
import fitsio

from desitarget.io import desitarget_nside
from desitarget.streams import io, cuts

from time import time
start = time()

from desiutil.log import get_logger
log = get_logger()

from argparse import ArgumentParser
ap = ArgumentParser(description=("Generates DESI target bits from Legacy Surveys"
                                 " sweep files for the MWS 1B targets.")
                    )
ap.add_argument("sweepdir",
                help=("Root directory of LS sweeps for a given data release for "
                      "ONE of EITHER north or south.")
                )
ap.add_argument("dest",
                help=("Output target selection directory (the file name is built"
                      " on-the-fly from other inputs).")
                )
ap.add_argument('-t','--targnames_in', default=None,
                help=("Comma-separated names of DESI 1B target classes to run "
                      " (e.g. x,y,z). Default is to run all MWS targets.")
                )
ap.add_argument("--readpertarg", action="store_true",
                help=("Default is to read targets by looping over all sweeps. "
                      "Pass to instead loop over DESI MWS 1B target classes."
                      )
                )
ap.add_argument("--donotaddnors", action="store_true",
                help=("Both the south/north LS files are read by default. Pass "
                      "this to read only the specific files in passed sweepdir.")
                )
ap.add_argument("--donotreadcache", action="store_true",
                help=("Default is to read from previously cached data files "
                      "where possible. Pass this to start from scratch (and "
                      "overwrite caches).")
                )
ap.add_argument("--numproc", type=int, default=1,
                help=("Number of parallel processes to use. numproc of 16 is a "
                "good balance between speed and I/O. Defaults to 1 for serial.")
                )
ap.add_argument("--mindec", type=float, default=-20,
                help=("Hard data limit. Targets south of this are not returned. "
                      "Defaults to -20.")
                )
ap.add_argument("--nside", type=int, default=None,
                help=("(NESTED) HEALPixel nside used with pixint. Only used if "
                "readpertarg is False.")
                )
ap.add_argument("--pixint", type=int, default=None,
                help=("Read and cache targets in (NESTED) HEALpixels at nside. "
                "For parallelizing. Only used if readpertarg is False.")
                )


ns = ap.parse_args()

# ADM build the list of command line arguments to
# ADM write to the output header.
hdr = fitsio.FITSHDR()
nsdict = vars(ns)
for k in nsdict:
    # ADM we set the pixel information separately.
    if k not in ["nside", "pixint", "targnames_in"]:
        hdr[k.upper()] = nsdict[k]
# ADM also explicitly add the command used, just in case.
hdr["CMDLINE"] = ' '.join(sys.argv)

# ADM change input comma-separated string of target class names to list.
tn = ns.targnames_in
hdr["TARGSRUN"] = "ALL"
if tn is not None:
    tn = tn.split(',')
    # ADM add the input target classes to the header to keep track.
    hdr["TARGSRUN"] = tn

# ADM flip the sense of some inputs.
addnors = not(ns.donotaddnors)
readcache = not(ns.donotreadcache)

targets = cuts.select_targets(
    ns.sweepdir, targnames_in=tn, readpertarg=ns.readpertarg,
    addnors=addnors, readcache=readcache, numproc=ns.numproc, mindec=ns.mindec,
    nside=ns.nside, pixint=ns.pixint
)

# ADM note that we hard code for BRIGHT, which is hidden in write_targets.
ntargs, outfile = io.write_targets(ns.dest, targets, hdr, nside=ns.nside,
                                   pixint=ns.pixint, nsidecol=desitarget_nside())

log.info(f"{ntargs} targets written to {outfile}...t={time()-start:.1f}s")
