#!python

"""
Create a nested directory tree with file symlinks to "copy" one prod to another

Stephen Bailey
March 2021
"""

import os, sys, glob
import shutil
from collections import Counter
import optparse

import numpy as np
from astropy.table import Table

from desiutil.log import get_logger
from desispec.workflow.tableio import load_table, write_table

parser = optparse.OptionParser(usage = "%prog [options] indir outdir")
parser.add_option("--explist", help="file with NIGHT EXPID to copy")
parser.add_option("-n", "--night", default='20??????',
        help="YEARMMDD to copy (can include glob wildcards)")
parser.add_option("-t", "--tiles",
        help="Commas separated list of tiles to copy")
parser.add_option("--fullcopy", action="store_true",
        help="Copy files instead of linking")
parser.add_option("--abspath", action="store_true",
        help="Link to absolute path instead of relative path")
parser.add_option("--filter-exptables", action="store_true",
        help="Filter exposure_tables to ignore entries not in input explist")

opts, args = parser.parse_args()
inroot = os.path.abspath(args[0])
outroot = os.path.abspath(args[1])

log = get_logger()

#- Get list of NIGHT EXPID to copy
if opts.explist is not None:
    explist = Table.read(opts.explist, format='ascii')
else:
    rows = list()
    for dirname in sorted(glob.glob(f'{inroot}/exposures/{opts.night}/????????')):
        nightdir, expid = os.path.split(dirname)
        try:
            night = int(os.path.basename(nightdir))
            expid = int(expid)
            rows.append((night, expid))
        except ValueError:
            pass

    explist = Table(rows=rows, names=('NIGHT', 'EXPID'))

if opts.tiles is not None:
    tiles = np.array([int(t) for t in opts.tiles.split(',')])

    keep = np.zeros(len(explist), dtype=bool)
    for night in np.unique(explist['NIGHT']):
        yearmm = night//100
        exptabfile = f'{inroot}/exposure_tables/{yearmm}/exposure_table_{night}.csv'
        if not os.path.exists(exptabfile):
            # commissioning data without exposure tables isn't supported
            # by copyprod, but if later nights are missing exptab then error
            if night > 20201119:
                log.error(f'Missing exposure table for {night}; stopping')
                sys.exit(1)
            else:
                log.warning(f'No exposure table for CMX {night}; skipping')
                continue

        exptab = Table.read(exptabfile)
        ii = np.isin(exptab['TILEID'], tiles)
        keep |= np.isin(explist['EXPID'], exptab['EXPID'][ii])

    if np.sum(keep) > 0:
        explist = explist[keep]
    else:
        raise ValueError(f'Tile {tileid} not found on nights {np.unique(explist["NIGHT"])}')

num_nights = len(np.unique(explist['NIGHT']))
num_expids = len(np.unique(explist['EXPID']))
assert num_expids == len(explist)
log.info(f'Linking {num_expids} exposures on {num_nights} nights')

#- What type of copy or link are we making?
if opts.fullcopy:
    link = shutil.copy2
    opts.abspath = True
    log.debug('Performing full copy of selected NIGHT/EXPID')
else:
    link = os.symlink
    log.debug('Creating symlinks')

inroot = os.path.abspath(inroot)
outroot = os.path.abspath(outroot)
log.info(f'Creating links in {outroot} to files in {inroot}')

if not os.path.exists(outroot):
    os.makedirs(outroot)

def link_dirfiles(indir, outdir, abspath=False):
    """
    Create relative links in outdir to all files in indir
    """
    os.makedirs(outdir, exist_ok=True)
    if abspath:
        srcpath = indir
    else:
        srcpath = os.path.relpath(indir, outdir)

    for infile in sorted(glob.glob(f'{indir}/*.*')):
        basename = os.path.basename(infile)
        src = os.path.join(srcpath, basename)
        dst = os.path.join(outdir, basename)

        # changed or broken links need to remove dst first
        if os.path.islink(dst):
            orig_src = os.readlink(dst)
            if orig_src != src or not os.path.exists(dst):
                log.warning(f'Replacing {dst} -> {orig_src} with {src}')
                os.remove(dst)

        if not os.path.exists(dst):
            link(src, dst)
        elif not os.path.islink(dst):
            log.warning(f'Not replacing non-link {dst}')


#- calibnight: link all files for requested nights
for night in np.unique(explist['NIGHT']):
    log.info(f'calibnight/{night}')
    indir = f'{inroot}/calibnight/{night}'
    outdir = f'{outroot}/calibnight/{night}'
    link_dirfiles(indir, outdir, opts.abspath)

#- preproc, exposures: link all files for requested night/expid
for subdir in ['preproc', 'exposures']:
    for night, expid in explist['NIGHT', 'EXPID']:
        log.info(f'{subdir}/{night}/{expid:08d}')
        indir = f'{inroot}/{subdir}/{night}/{expid:08d}'
        outdir = f'{outroot}/{subdir}/{night}/{expid:08d}'
        link_dirfiles(indir, outdir, opts.abspath)

#- exposure tables: trim to requested exposures
for night in np.unique(explist['NIGHT']):
    yearmm = night//100
    infile = f'{inroot}/exposure_tables/{yearmm}/exposure_table_{night}.csv'
    outfile = f'{outroot}/exposure_tables/{yearmm}/exposure_table_{night}.csv'
    os.makedirs(os.path.dirname(outfile), exist_ok=True)
    if opts.filter_exptables:
        exptab = load_table(infile, tabletype='exptable')
        ignore = np.isin(exptab['EXPID'], explist['EXPID'], invert=True)
        if np.any(ignore):
            msg = f'Night {night} ignoring '
            msg += ', '.join([f'{n} {obstype}' for obstype, n in Counter(exptab['OBSTYPE'][ignore]).items()])
            msg += ' exposures not included in explist'
            log.warning(msg)

        exptab['LASTSTEP'][ignore] = 'ignore'
        write_table(exptab, outfile, tabletype='exptable')
    else:
        log.info(f'Night {night} copying exposure tables unchanged')
        shutil.copy2(infile, outfile)


#- processing tables: trim to requested exposures.
#- Disabled for now while we figure out exactly what we want.
# for night in np.unique(explist['NIGHT']):
#     night_expids = explist['EXPID'][explist['NIGHT'] == night]
#     tmp = glob.glob(f'{inroot}/processing_tables/processing_table_*-{night}.csv')
#     if len(tmp) == 1:
#         infile = tmp[0]
#     elif len(tmp) == 0:
#         log.error(f'Unable to find processing table for {night}')
#         continue
#     elif len(tmp) == 0:
#         log.error(f'Multiple processing tables for {night} ?!?')
#         continue
# 
#     outfile = '{}/processing_tables/{}'.format(
#             outroot, os.path.basename(infile))
#     proctab = load_table(infile, tabletype='proctable')
# 
#     #- Convert string "expid1|expid2" entries into keep or not
#     keep = np.zeros(len(proctab), dtype=bool)
#     for i, expids in enumerate(proctab['EXPID']):
#         if np.any(np.isin(expids, night_expids)):
#             keep[i] = True
# 
#     proctab = proctab[keep]
# 
#     os.makedirs(os.path.dirname(outfile), exist_ok=True)
#     write_table(proctab, outfile, tabletype='proctable')

log.info(f'Production copied into {outroot}')

