#!/usr/bin/env python
import os
import sys

from netCDF4 import Dataset


def nc_subset(in_file, out_file, ens_member):
    # Remove the output file, if it exists
    if os.path.exists(out_file):
        os.remove(out_file)

    # Open the input file:
    nc_in = Dataset(in_file, 'r')

    # Create the output file
    try:
        nc_out = Dataset(out_file, 'w', format="NETCDF4")
    except IOError as e:
        print("Could not create output file {}: {}".format(out_file, e.args[1]), 
              file=sys.stderr)
        exit(2)

    # Copy the dimensions
    for name, dimension in nc_in.dimensions.items():
        nc_out.createDimension(name, len(dimension) if not dimension.isunlimited() else None)

    # Copy the variables
    for name, variable in nc_in.variables.items(): 
        if name in ['time', 'reference_time', 'x', 'y', 'T2D', 'RAINRATE']:
            nc_out.createVariable(name, variable.datatype, variable.dimensions)
            nc_out[name].setncatts(nc_in[name].__dict__)	
            nc_out[name][:] = nc_in[name][:]
        elif name in ['crs']:
            nc_out.createVariable(name, 'S1')
            nc_out[name].setncatts(nc_in[name].__dict__)
        

    # Copy the global attributes
    nc_out.model_output_valid_time = nc_in.model_output_valid_time
    nc_out.model_initialization_time = nc_in.model_initialization_time
    nc_out.model_output_type = nc_in.model_output_type
    nc_out.model_configuration = nc_in.model_configuration
    nc_out.model_total_valid_times = nc_in.model_total_valid_times
    nc_out.NWM_version_number = nc_in.NWM_version_number
    nc_out.setncattr('ensemble_member_number', ens_member)

    nc_out.close()
    nc_in.close()


if __name__ == '__main__':
    if len(sys.argv) != 4:
        print("usage: python get_nc.py input_file output_file ens_member", file=sys.stderr)
        sys.exit(1)
    
    if os.path.exists(sys.argv[1]):
        nc_subset(sys.argv[1], sys.argv[2], sys.argv[3])
    else:
        print("Input file not found: {}".format(sys.argv[1]), file=sys.stderr)
        sys.exit(1)