#!/usr/bin/env python

import os, sys
import numpy as np

from desitarget import io
from desitarget.brightmask import make_bright_star_mask, get_mask_dir
from desitarget.tychomatch import get_tycho_nside, get_tycho_dir
from desitarget.uratmatch import get_urat_dir
from desitarget.gaiamatch import get_gaia_dir

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

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

from desiutil.log import get_logger
log = get_logger()

from argparse import ArgumentParser
ap = ArgumentParser()
ap.add_argument("--maskdir",
                help='Output directory to which to write mask [defaults to the $MASK_DIR environment variable]')
ap.add_argument("--maglim", type=float,
                help="Magnitude limit for building the mask [defaults to 12]",
                default=12.)
ap.add_argument("--matchrad", type=float,
                help="Tycho stars are discarded if they match a Gaia star at this radius",
                default=1.)
ap.add_argument("--nside", type=int,
                help="Write output files in HEALPixels at this nside [{}]".format(nside))
ap.add_argument("--numproc", type=int,
                help="number of concurrent processes to use [{}]".format(nproc),
                default=nproc)
ap.add_argument("--maskepoch", type=float,
                help="Stars are shifted (by just proper motion) to this epoch to make the mask [2023.0]",
                default=2023.0)
ap.add_argument("--gaiaepoch", type=float,
                help="Epoch of Gaia observations [2015.5]",
                default=2015.5)

ns = ap.parse_args()

# ADM if the MASK_DIR directory was passed, set it...
maskdir = ns.maskdir
if maskdir is None:
    # ADM ...otherwise retrieve it from the environment variable.
    maskdir = get_mask_dir()

mask = make_bright_star_mask(
    maglim=ns.maglim, matchrad=ns.matchrad, numproc=ns.numproc,
    maskepoch=ns.maskepoch, gaiaepoch=ns.gaiaepoch)

# ADM extra header keywords for the output fits file.
extra = {k: v for k, v in zip(
    ["matchrad", "gaiaepoc", "gaiadir", "tychodir", "uratdir"],
    [ns.matchrad, ns.gaiaepoch, get_gaia_dir(), get_tycho_dir(), get_urat_dir()])}

# ADM write the mask file(s).
nmasks, mxdir = io.write_masks(maskdir, mask, maglim=ns.maglim,
                              maskepoch=ns.maskepoch, nside=nside, extra=extra)

log.info('wrote {} total masks to files in {}'.format(nmasks, mxdir))
