#!/usr/bin/env python

from __future__ import print_function, division

import os, sys
import numpy as np
import fitsio
from time import time

from desitarget.QA import make_qa_page, _parse_tcnames, read_data
from desitarget.io import read_targets_header

from desiutil.log import get_logger
log = get_logger()

import multiprocessing
# ADM the rather large denominator guards against memory issues.
nproc = multiprocessing.cpu_count() // 8

from argparse import ArgumentParser
ap = ArgumentParser()
ap.add_argument("src", 
                help="Input target file or directory of HEALPixel-split target files (e.g. /project/projectdirs/desi/target/catalogs/targets-dr6-0.20.0.fits")
ap.add_argument("dest", 
                help="Output directory to make QA webpage hierarchy (e.g. /project/projectdirs/desi/www/users/USERNAME/ will appear at http://portal.nersc.gov/project/desi/users/USERNAME)")
ap.add_argument('-m', "--mocks", action='store_true', 
                help="Input target file is a set of mock targets, so we can make extra, special QA plots")
ap.add_argument('-n', "--noplots", action='store_true', 
                help="Do not produce any new plots, only the body of the QA webpage")
ap.add_argument('--nside', default=None, type=int , 
                help="Specify the healpixel resolution of the sky maps")
ap.add_argument('-w','--weightmapfile', default=None, 
                help="Location (e.g. /project/projectdirs/desi/target/catalogs/dr6/pixweight-dr6-0.20.2.fits) of the imaging HEALPix map file for the Data Release that corresponds to these targets (e.g. as made by desitarget.imagefootprint.pixweight()")
ap.add_argument('-t','--tcnames', default=None, 
                help="Names of specific target classes for which to make the QA pages. (e.g. QSO,LRG,ALL)")
ap.add_argument('--nosystematics', action='store_true', 
                help="Don't make systematics plots on the front page (perhaps because the weightmapfile was not passed)")
ap.add_argument("--numproc", type=int,
                help='number of concurrent processes to use when generating plots [defaults to {}]'.format(nproc),
                default=nproc)
ap.add_argument("--downsample", type=int,
                help='downsample the targets by (roughly) this value, as a speed-up.',
                default=None)
ap.add_argument("--prepare", action='store_true',
                help='write a file to dest that is the equivalent of what is read from src (and just contains the columns needed for QA). This file can then be used in place of src in a subsequent run as a speed-up. Most effective is to run with --prepare set on a login node and then use the written file on compute nodes. Overrides all inputs except src, dest and downsample.')

start = time()
ns = ap.parse_args()

# ADM fail if systematics were requested but no systematics map was passed.
if ns.weightmapfile is None and ns.nosystematics is False and not ns.prepare:
    log.error("The weightmap file was not passed so systematics cannot be tracked. Try again sending --nosystematics.")
    raise IOError

# ADM check that the passed pixweight file corresponds to the passed target file.
if ns.weightmapfile is not None and not ns.mocks:
    dsurv = read_targets_header(ns.src)["SURVEY"]
    try:
        psurv = fitsio.read_header(ns.weightmapfile, "PIXWEIGHTS")["SURVEY"]
        if dsurv != psurv and 'cmx' not in dsurv:
            msg = 'target file is type {} but pixweight file is type {}!!!'.format(dsurv, psurv)
            log.critical(msg)
            raise ValueError(msg)
    # ADM we'll flag an error for early versions of files, which is fine.
    except KeyError:
        pass

tcnames = ns.tcnames
if tcnames is not None:
    tcnames = _parse_tcnames(tcstring=ns.tcnames, add_all=True)

makeplots = not(ns.noplots)
systematics = not(ns.nosystematics)

if ns.nside:
    import healpy as hp
    max_bin_area = hp.nside2pixarea(ns.nside, degrees=True)
else:
    max_bin_area = 1.0

if ns.prepare:
    fn = os.path.join(os.getenv("SCRATCH"), ns.dest)
    objs, _, _, hdr = read_data(ns.src, header=True, downsample=ns.downsample)
    fitsio.write(fn+'.tmp', objs, extname='TARGETS', header=hdr, clobber=True)
    os.rename(fn+'.tmp', fn)
    log.info('{} targets written to {}'.format(len(objs), fn))
else:
    log.info("generating plots on {} processors".format(ns.numproc))
    make_qa_page(ns.src, qadir=ns.dest, numproc=ns.numproc,
                 mocks=ns.mocks, makeplots=makeplots, max_bin_area=max_bin_area,
                 imaging_map_file=ns.weightmapfile, tcnames=tcnames,
                 systematics=systematics, downsample=ns.downsample)
    log.info('Targeting QA pages and plots written to {}...t={:.1f}s'.format(
        ns.dest, time()-start))
    log.info('Mocks option was set to {}'.format(ns.mocks))
    log.info('No plots option was set to {}'.format(ns.noplots))
    log.info('No systematics option was set to {}'.format(ns.nosystematics))
