#!/usr/bin/env python
"""Wrapper to build HTML output from previously generated QA files.

fastspecfit-html --fastspecfile ./fastspect.fits

fastspecfit-html --fastphotfile /global/cfs/cdirs/desi/spectro/fastspecfit/everest/catalogs/fastphot-everest-sv3-dark.fits \
  --fastspecfile /global/cfs/cdirs/desi/spectro/fastspecfit/everest/catalogs/fastspec-everest-sv3-dark.fits \
  -o /global/cfs/cdirs/desi/users/ioannis/fastspecfit/everest --healpix 10016

fastspecfit-html --fastphotfile /global/cfs/cdirs/desi/spectro/fastspecfit/everest/tiles/merged/fastphot-everest-cumulative.fits \
  --fastspecfile /global/cfs/cdirs/desi/spectro/fastspecfit/everest/tiles/merged/fastspec-everest-cumulative.fits \
  -o /global/cfs/cdirs/desi/users/ioannis/fastspecfit/everest

fastspecfit-html --fastspecfile /global/cfs/cdirs/desi/spectro/fastspecfit/everest/tiles/cumulative/80607/20201219/fastspec-0-80607-thru20201219.fits \
  -o /global/cfs/cdirs/desi/users/ioannis/fastspecfit/everest

"""
import pdb # for debugging
import os, sys, time
import numpy as np

from desiutil.log import get_logger
log = get_logger()

def parse(options=None):
    """Parse input arguments.

    """
    import sys, argparse

    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    parser.add_argument('--healpix', default=None, type=str, nargs='*', help="""Generate QA for all objects
        with this healpixels (only defined for coadd-type 'healpix').""")
    parser.add_argument('--tile', default=None, type=str, nargs='*', help='Tile(s) to process.')
    parser.add_argument('--night', default=None, type=str, nargs='*', help='Night(s) to process (ignored if coadd-type is cumulative).')
    parser.add_argument('--targetids', type=str, default=None, help='Comma-separated list of target IDs to process.')
    parser.add_argument('-n', '--ntargets', type=int, help='Number of targets to process in each file.')
    parser.add_argument('--firsttarget', type=int, default=0, help='Index of first object to to process in each file (0-indexed).')
    parser.add_argument('--mp', type=int, default=1, help='Number of multiprocessing processes per MPI rank or node.')

    parser.add_argument('--outprefix', default=None, type=str, help='Optional prefix for output filename.')
    parser.add_argument('-o', '--outdir', default='.', type=str, help='Full path to desired output directory.')

    parser.add_argument('--htmlhome', type=str, default='index.html', help='Main HTML page name.')
    #parser.add_argument('--htmldir', type=str, help='Output directory for HTML files.')

    parser.add_argument('--fastphotfile', default=None, type=str, help='Full path to fastphot fitting results.')
    parser.add_argument('--fastspecfile', default=None, type=str, help='Full path to fastphot fitting results.')
    
    if options is None:
        args = parser.parse_args()
        log.info(' '.join(sys.argv))
    else:
        args = parser.parse_args(options)
        log.info('fastspecfit-html {}'.format(' '.join(options)))

    return args

def main(args=None, comm=None):
    """Make the HTML pages.

    """
    from astropy.table import Table
    from fastspecfit.io import read_fastspecfit, DESI_ROOT_NERSC, select
    from fastspecfit.html import build_htmlhome
    
    if isinstance(args, (list, tuple, type(None))):
        args = parse(args)

    #if args.htmldir is None:
    #    htmldir = os.path.join(os.getenv('FASTSPECFIT_HTML'), args.specprod)
    #else:
    #    htmldir = args.htmldir
        
    if not os.path.isdir(args.outdir):
        log.info('Creating output directory {}'.format(args.outdir))
        os.makedirs(args.outdir, exist_ok=True)

    # Read the fitting results and get all the unique targetids.
    if (args.fastphotfile is None and args.fastspecfile is None):
        log.warning('Must provide one of --fastphotfile or --fastspecfile.')
        return
    elif (args.fastphotfile is not None and args.fastspecfile is not None):
        fastspec, metaspec, specprod, coadd_type = read_fastspecfit(args.fastspecfile)
        fastphot, metaphot, _specprod, _coadd_type = read_fastspecfit(args.fastphotfile, fastphot=True)
        if (specprod != _specprod) or (coadd_type != _coadd_type):
            log.warning('Mis-matching specprod or coadd_type in fastspec vs fastphot fitting results!')
            return

        from astropy.table import join
        joinkeys = ['TARGETID']
        if 'TILEID' in metaspec.colnames:
            joinkeys += ['TILEID']
            fastspec['TILEID'] = metaspec['TILEID']
            fastphot['TILEID'] = metaphot['TILEID']

        if 'NIGHT' in metaspec.colnames:
            joinkeys += ['NIGHT']
            fastspec['NIGHT'] = metaspec['NIGHT']
            fastphot['NIGHT'] = metaphot['NIGHT']

        if coadd_type == 'healpix':
            joinkeys += ['HEALPIX']
            fastspec['HEALPIX'] = metaspec['HEALPIX']
            fastphot['HEALPIX'] = metaphot['HEALPIX']
            
        fastfit = join(fastspec, fastphot, keys=joinkeys, join_type='outer',
                       table_names=['SPEC', 'PHOT'])
        metadata = join(metaspec, metaphot, keys=joinkeys, join_type='outer',
                       table_names=['SPEC', 'PHOT'])
        
        assert(np.all(fastfit['TARGETID'] == metadata['TARGETID']))
        if 'TILEID' in metaspec.colnames:
            assert(np.all(fastfit['TILEID'] == metadata['TILEID']))

    elif args.fastphotfile is None and args.fastspecfile is not None:
        fastfit, metadata, specprod, coadd_type = read_fastspecfit(args.fastspecfile)
    elif args.fastphotfile is not None and args.fastspecfile is None:
        fastfit, metadata, specprod, coadd_type = read_fastspecfit(args.fastphotfile, fastphot=True)
    else:
        pass # should never get here
    
    # parse the targetids optional input
    if args.targetids:
        targetids = [int(x) for x in args.targetids.split(',')]
        keep = np.where(np.isin(fastfit['TARGETID'], targetids))[0]
        if len(keep) == 0:
            log.warning('No matching targetids found!')
            return
        fastfit = fastfit[keep]
        metadata = metadata[keep]

    if args.ntargets is not None:
        keep = np.arange(args.ntargets) + args.firsttarget
        log.info('Keeping {} targets.'.format(args.ntargets))
        fastfit = fastfit[keep]
        metadata = metadata[keep]

    fastfit, metadata = select(fastfit, metadata, coadd_type,
                               healpixs=args.healpix,
                               tiles=args.tile, nights=args.night)

    # read the tile info file for this production
    specproddir = os.path.join(os.environ.get('DESI_ROOT', DESI_ROOT_NERSC), 'spectro', 'redux', specprod)
    infofile = os.path.join(specproddir, 'tiles-{}.csv'.format(specprod))
    info = Table.read(infofile)
    log.info('Read {} tiles from {}'.format(len(info), infofile))
    assert(len(info) == len(np.unique(info['TILEID'])))

    # add tileid --> healpix mapping
    if coadd_type == 'healpix':
        hpxinfo = Table.read(os.path.join(specproddir, 'healpix', 'tilepix.fits'))
        metadata['TILEID'] = hpxinfo[np.isin(hpxinfo['HEALPIX'], metadata['HEALPIX'])]['TILEID']
        #info['HEALPIX'] = hpxinfo[np.isin(info['TILEID'], hpxinfo['TILEID'])]
        
    itile = np.isin(info['TILEID'], metadata['TILEID'])
    info = info[itile]
    log.info('Selecting {} tiles.'.format(len(info)))

    # Build the home (index.html) page (always, irrespective of clobber)--
    t0 = time.time()
    build_htmlhome(args.outdir, fastfit, metadata, info, htmlhome=args.htmlhome,
                   specprod=specprod, coadd_type=coadd_type, mp=args.mp)
    log.info('HTML for everything took: {:.2f} sec'.format(time.time()-t0))

if __name__ == '__main__':
    main()
    
