#!/usr/bin/env python
# encoding: utf-8
#
# @Author: Tom Dwelly
# @Date: Jan-2020
# @Filename: build_sdssv_plateholes
# @License: BSD 3-Clause
# @Copyright: Tom Dwelly


from __future__ import (absolute_import, division,
                        print_function, unicode_literals)

import os
import subprocess

import numpy as np
from astropy.table import Table


# the base directory where we can find the platelist product
# platelist_dir = os.path.expanduser('~/scratch/platelist')
# platelist_dir = (
#   '/uufs/chpc.utah.edu/common/home/sdss05/software/svn.sdss.org/data/sdss/platelist/trunk' )
# TODO it would be better to read this from the $PLATELIST_DIR environment variable
try:
    platelist_dir = os.environ['PLATELIST_DIR']
except:
    raise Exception("You must first set the PLATELIST_DIR environment variable.\n"
                    "For example, by running this command (on the Utah system):\n"
                    " #> module load platelist")

platelist_pattern = 'plates/015?XX/015???/plateHoles-015???.par'

platelist_expr = os.path.join(platelist_dir,
                              platelist_pattern)

# the following should be adjusted by the user to indicate where outputs should be written:
outdir = '.'
# outdir = os.path.expanduser('~/scratch/test_yanny_scraper2')

outstem = 'plateholes_'
outfmt = 'fits'
out = os.path.join(outdir, outstem)

# Now, the lists of exception plates
# derived from here:
# https://wiki.sdss.org/pages/viewpage.action?pageId=69043025
# updated on 19th Feb 2021 (added MWM plate 15188 to exclusion list)

# list of plates that were designed but never drilled:
phantom_plates = [
    15140,  # eFEDS-not-drilled
    15142,  # Do not Mark! (probably not pluggable)
    15145,  #
    15148,  #
    15151,  #
    15154,  #
    15157,  #
    15160,  #
    15163,  #

    15097,  # MWM Plates that will never be drilled
    15098,  #
    15099,  #
    15104,  #
    15117,  #
    15118,  #
    15123,  #
    15124,  #
    15126,  #
    15128,  #
    15130,  #
    15131,  #
    15132,  #
    15134,  #
    15135,  #
    15210,  #
    15211,  #
    15212,  #
    15213,  #
    15214,  #
    15215,  #
    15216,  #
]

# list of plates that were drilled but never shipped, or shipped then immediately rejected
rejected_plates = []

# list of plates that have no hope of observation during SDSS-V 2020/21 plate operations
hopeless_plates = [
    15188,
    15204,   # Do Not Mark! (they have set already
    15205,   # and won't be observed during plate program)
    15206,   #
    15207,   #
    15208,   #
    15209,   #

    15096,   # Should be retired do not mark
    15136,   #
]

# run the main yanny scraper code to build the
# fits versions of the yanny plateHoles files

subprocess.run(['yanny_scraper',
                '-i',
                platelist_expr,
                '-o',
                out,
                '-f',
                outfmt])


# open up the 'meta' file produced by the above
# - i.e. a table with one row per known plate design
# we can predict the output filename pretty easily
output_meta = out + 'meta.' + outfmt

try:
    t = Table.read(output_meta)
except:
    raise Exception(f"Unable to open {output_meta} as a Table")

print(f"Flagging all known non-drilled/rejected/hopeless plates in {output_meta}")

# add a new column to represent if the plate has been drilled/not rejected etc.

# now remove those plates that didn't get drilled
# set to true for all initially
phantom = np.isin(t['plateId'], phantom_plates)
rejected = np.isin(t['plateId'], rejected_plates)
hopeless = np.isin(t['plateId'], hopeless_plates)

valid = np.where(phantom | rejected | hopeless, False, True)

nvalid = np.count_nonzero(valid)
print(f"After filtering {nvalid}/{len(t)} plates are marked as 'VALID'=True")

t['ISVALID'] = valid

t.write(output_meta, format=outfmt, overwrite=True)

# done
