#!/usr/bin/env python

from __future__ import print_function

import sys
import argparse
import netCDF4 as nc
import numpy as np

"""
What this script does:

Make the OASIS masks.nc file. The masks.nc needs the ice T grid mask and the atmosphere T, U and V grid masks. The name of the mask variables depend on what is in the OASIS namcouple file. 

How to run this script: 

"""

def make_ice_mask(input_file, output_file, mask_name):

    ice_mask = input_file.variables['kmt'][:]
    ny, nx = ice_mask.shape

    output_file.createDimension('ny0', ny)
    output_file.createDimension('nx0', nx)

    msk = output_file.createVariable(mask_name, 'int32', dimensions=('ny0', 'nx0'))
    msk.units = '0/1:o/l'
    msk.title = 'Ice grid T-cell land-sea mask.'
    # Flip the mask. 
    msk[np.where(ice_mask > 0)] = 0
    msk[np.where(ice_mask < 1)] = 1


def make_atm_masks(input_file, output_file,
                   t_mask_name, u_mask_name, v_mask_name):

    atm_frac = input_file.variables['lsm'][:]
    ny, nx = atm_frac.shape[2:]

    # Convert fraction to mask. 
    atm_msk = atm_frac
    atm_msk[atm_msk < 1.0] = 0

    output_file.createDimension('ny1', ny)
    output_file.createDimension('nx1', nx)
    output_file.createDimension('ny2', ny - 1)
    output_file.createDimension('nx2', nx)

    t_msk = output_file.createVariable(t_mask_name, 'int32', dimensions=('ny1', 'nx1'))
    t_msk.units = '0/1:o/l'
    t_msk.title = 'Atm grid T-cell land-sea mask.'
    t_msk[:] = atm_msk[:]

    u_msk = output_file.createVariable(u_mask_name, 'int32', dimensions=('ny1', 'nx1'))
    u_msk.units = '0/1:o/l'
    u_msk.title = 'Atm grid U-cell land-sea mask.'
    # No masking needed on this grid. 
    u_msk[:] = 0

    v_msk = output_file.createVariable(v_mask_name, 'int32', dimensions=('ny2', 'nx2'))
    v_msk.units = '0/1:o/l'
    v_msk.title = 'Atm grid V-cell land-sea mask.'
    # No masking needed on this grid. 
    v_msk[:] = 0


def main():

    parser = argparse.ArgumentParser()
    parser.add_argument('input_ice_mask', help='The input ice mask.')
    parser.add_argument('input_atm_mask', help='The input atmosphere mask.')
    parser.add_argument('output_file', help='Name of the output file.')
    parser.add_argument('--ice_t_mask_name', default='cice.msk',
                        help='Name of the ice mask variable.')
    parser.add_argument('--atm_t_mask_name', default='um1t.msk',
                        help='Name of the atm t mask variable.')
    parser.add_argument('--atm_u_mask_var_name', default='um1u.msk',
                        help='Name of the atm u mask variable.')
    parser.add_argument('--atm_v_mask_var_name', default='um1v.msk',
                        help='Name of the atm u mask variable.')

    args = parser.parse_args()

    # Make the output file. 
    with nc.Dataset(args.output_file, 'w') as f_w:
        with nc.Dataset(args.input_ice_mask, 'r') as f_r:
            make_ice_mask(f_r, f_w, args.ice_t_mask_name)
        
        with nc.Dataset(args.input_atm_mask, 'r') as f_r:
            make_atm_masks(f_r, f_w, args.atm_t_mask_name,
                           args.atm_u_mask_var_name, args.atm_v_mask_var_name)

        # Update the history. 
        f_w.history = ' '.join(sys.argv)


if __name__ == "__main__":
    sys.exit(main())
