#!/usr/bin/env python
'''
Masks the 12US2 portion of a 36US3 gridded file
'''

import sys
import numpy as np
import fauxioapi as io

def get_mask_range(grid, mask_grid):
    '''
    Calc range on output grid to mask
    '''
    range_dict = {}
    # Output grid column and row mask start
    range_dict['row_start'] = (mask_grid.YORIG - grid.YORIG)/grid.YCELL
    range_dict['col_start'] = (mask_grid.XORIG - grid.XORIG)/grid.XCELL
    for key in ('row_start','col_start'):
        if range_dict[key] < 0:
            range_dict[key] = 0
    # Output grid column and row mask end
    range_dict['row_end'] = ((mask_grid.YORIG + mask_grid.YCELL * mask_grid.NROWS) - grid.YORIG)/grid.YCELL
    range_dict['col_end'] = ((mask_grid.XORIG + mask_grid.XCELL * mask_grid.NCOLS) - grid.XORIG)/grid.XCELL
    if range_dict['row_end'] > grid.NROWS:
        range_dict['row_end'] = grid.NROWS
    if range_dict['col_end'] > grid.NCOLS:
        range_dict['col_end'] = grid.NCOLS
    for key in ('row_start','col_start','row_end','col_end'):
        if range_dict[key] == int(range_dict[key]):
            range_dict[key] = int(range_dict[key])
        else:
            raise ValueError('Mask not completely nested at %s: %s' %(key, range_dict[key]))
    return (slice(range_dict['row_start'],range_dict['row_end']), 
              slice(range_dict['col_start'],range_dict['col_end']))

big_grid = '36US3'
small_grid = '12US2'
griddesc = '/work/EMIS/em_v7.2/ge_dat/gridding/griddesc_lambertonly_18jan2019_v7.txt'
infile_name = sys.argv[1]
outfile_name = sys.argv[2]
grid = io.Grid(big_grid, griddesc)
mask_grid = io.Grid(small_grid, griddesc)
row_slice, col_slice = get_mask_range(grid, mask_grid)
print(infile_name,outfile_name)
with io.IODataset(infile_name) as infile, io.IODataset(outfile_name, 'w', diskless=True, persist=True) as outfile:
    outfile.set_dimensions(LAY=infile.NLAYS, ROW=grid.NROWS, COL=grid.NCOLS, VAR=infile.NVARS)
    var_list = [var for var in infile.variables if var not in ('TFLAG','ACRESBURNED','HFLUX','NH3_FERT')]
    for var_name in var_list:
        print(var_name)
        in_var = infile.variables[var_name]
        out_arr = in_var[:]
        out_arr[:,:,row_slice,col_slice] = 0
        out_var = outfile.create_variable(var_name, 'REAL', ('TSTEP','LAY','ROW','COL'))
        out_var.units = in_var.units
        out_var[:] = out_arr
        outfile.sync()
    outfile.set_attributes(infile.SDATE, grid)
    outfile.write_TFLAG()


