# Program Name: postprocess_reservoir_output.py
# Author: Katelyn FitzGerald and Mehdi Rezaeianzadeh
# Date: 09/15/2020
# This program will
# 1) Masks out the Great Lakes water level and outflow 
# 2) Removes all mention of USACE and RFC sites from the reservoir type, data file source and assimilated value
# Usage: python postprocess_reservoir_output.py $reservoirFile

#!/usr/bin/env python 
#!/usr/bin/env python
import os
import sys
import numpy as np
from netCDF4 import Dataset

fileName = sys.argv[1]

ds = Dataset(fileName,"r+")

ds['outflow'][np.isin(ds['feature_id'][:],[4800002,4800004,4800005,4800006,4800007])]=-9999.

ds['water_sfc_elev'].missing_value = -9999.
del ds['water_sfc_elev'].comment
ds['water_sfc_elev'][np.isin(ds['feature_id'][:],[4800002,4800004,4800005,4800006,4800007])]=-9999.

ds['outflow'][np.isin(ds['reservoir_type'][:],[3,4])] = -9999.

ds['reservoir_assimilated_value'].missing_value = -9999.
ds['reservoir_assimilated_value'][np.isin(ds['reservoir_type'][:],[3,4])] = -9999.

#if len(set(ds['reservoir_type'][:])) != 1:
#  ds['reservoir_assimilated_source_file'][np.isin(ds['reservoir_type'][:],[3,4])] = ""


ds['reservoir_type'].missing_value = -9999
ds['reservoir_type'][np.isin(ds['reservoir_type'][:],[3,4])] = -9999
ds['reservoir_type'].flag_values = np.array([1, 2]).astype(np.int32)
ds['reservoir_type'].flag_meanings = "Level_pool USGS-persistence"


ds.close()