#!/usr/bin/env python
import sys

from netCDF4 import Dataset
import numpy as np


def change_qpe_mpe(mpe_file, ext_ana_file):
    try:
        f2 = Dataset(ext_ana_file, 'a')
    except IOError:
        print("Can not create output file {}. Check disk space and permissions".format(ext_ana_file), file=sys.stderr)
        sys.exit(22)

    try:
        f1 = Dataset(mpe_file, 'r')
    except IOError:
        print("Input file {} not found or can not be accessed".format(mpe_file), file=sys.stderr)
        sys.exit(42)

    f1_rainrate = f1.variables["RAINRATE"]
    dims = f1_rainrate.dimensions
    ndims = len(dims)

    if ndims == 2:
        precip_rate = f1_rainrate[::-1, :].astype(float)
    else:
        precip_rate = f1_rainrate[0, ::-1, :].astype(float)

    f2_rainrate = f2.variables["RAINRATE"]
    precip_rate1 = f2_rainrate[0, :, :]

    precip_rate = np.where(precip_rate.mask, precip_rate1, precip_rate)

    #print(precip_rate)
    #print(f2_rainrate.dimensions)
    #print(str(len(f2.dimensions['time'])) + " x " + str(len(f2.dimensions['y'])) + " x " + str(len(f2.dimensions['x'])))
    #print(precip_rate.shape)
    #print(f2_rainrate[:].shape)
    #print(f2_rainrate[0,:].shape)

    f2_num_dim = len(f2_rainrate[:].shape)
    if f2_num_dim == 2:
        f2_rainrate[:] = precip_rate
    elif f2_num_dim == 3:
        f2_rainrate[0,:] = precip_rate

    f2.close()
    f1.close()

    print("End of precipitation replacement for analysis.")


if __name__ == '__main__':
    if len(sys.argv) != 3:
        print("usage: python3 change_qpe_mpe.py <mpeFile.nc> <extAnAFile.nc>", file=sys.stderr)
        exit(1)
    else:
        change_qpe_mpe(sys.argv[1], sys.argv[2])