#!/usr/bin/env python

"""
Example program to add a variable containing all zeroes to an
xarray Dataset object.  Useful when you have to insert new
species into GEOS-Chem diagnostic or restart files.

Calling sequence:
-----------------
add_blank_var.py varname infile outfile
"""

import glob
from os.path import join, basename
from sys import argv
import numpy as np
import xarray as xr
from gcpy.util import create_blank_dataarray


def add_blank_var_to_ncfile(
        varname,
        infile,
        outfile,
        fill_value=np.float32(0.0),
        varattrs=None
):
    """
    Adds a variable containing all zeroes to a netCDF file

    Args:
    -----
    varname : str
        Name of the variable to add.
    infile: str
        Name of the input netCDF file (will not be overwritten).
    outfile: str
        Name of the output netCDF file containing varname.

    Keyword Args:
    -------------
    varattrs: dict
        A dictionary containing netCDF variable attributes to be written
        to outfile.
    """
    with xr.set_options(keep_attrs=True):

        ds = xr.open_dataset(infile)
        
        if varattrs is None:
            varattrs = ds.attrs

        da = create_blank_dataarray(
            varname,
            ds.sizes,
            ds.coords,
            varattrs,
            fill_value=fill_value,
            fill_type=np.float32
        )

        ds = xr.merge([ds, da])

        ds.to_netcdf(outfile)


if __name__ == '__main__':

    # Get lists of files
    gcc_files = glob.glob("../v2021-09/GEOSChem.Restart*")
    gchp_files = glob.glob("../v2021-09/GCHP.Restart*")
    
    # Name of the blank varible to add (EDIT AS NEEDED)
    # NOTE: For GCHP, the prefix must be "SPC_" instead of "SpeciesRst_"

    # Variable attributes (EDIT AS NEEDED)
    varattrs = {
        "MW_g"            : "146.98",
        "long_name"       : "Dummy species to track production rate of RO2",
        "units"           : "mol mol-1 dry",
        "Is_Gas"          : "true",
        "averaging_method": "instantaneous",
        "_FillValue"      : np.float(-1.0e-32)
    }

    # Add blank variable to GEOS-Chem Classic restarts
    for infile in gcc_files:
        varname = "SpeciesRst_PRO2"
        outfile = basename(infile)
        print(f"Creating {outfile}")       
        add_blank_var_to_ncfile(
            varname,
            infile,
            outfile,
            varattrs=varattrs
        )

    # Add blank variable to GCHP restarts
    varattrs = {
        "standard_name" : "Dummy species to track production rate of RO2",
        "long_name"     : "Dummy species to track production rate of RO2",
        "units"         : "mol mol-1 dry",
    }

    for infile in gchp_files:
        varname = "SPC_PRO2"
        outfile = basename(infile).replace("GCHP", "GEOSChem")
        print(f"Creating {outfile}") 
        add_blank_var_to_ncfile(
            varname,
            infile,
            outfile,
            varattrs=varattrs,
        )