# -*- coding: utf-8 -*-
"""ESMF regridding module

This module provides fast ESMF-based 2-D regridders for the National Water
Model v2.0 forcing engine.

Authors: Kent Sparrow and Greg Fall
Date: May 10, 2018

United States Department of Commerce
National Oceanic and Atmospheric Administration
National Weather Service
Office of Water Prediction
"""

from __future__ import print_function
import sys
import time
from netCDF4 import Dataset
from scipy.sparse import csr_matrix
import numpy as np
import errMod

def with_weights(statusMeta,source_grid, weight_file_path, no_data_value):

    """ESMF "offline" regridder.
    Perform sparse matrix regridding using a sparse matrix of regridding
    weights generated by ESMF_RegridWeightGen or its equivalent.

    Args:
        source_grid (float): A 2-D grid of data to regrid.
        weight_file_path (string): A NetCDF file of regridding weights.
        no_data_value (float): The missing or no-data value for input and
                               output.

    Returns:
        The result of the regridding, a 2-D grid, or None if the regrid
        fails.

    Authors: Kent Sparrow and Greg Fall
    Date: May 10, 2018

    United States Department of Commerce
    National Oceanic and Atmospheric Administration
    National Weather Service
    Office of Water Prediction

    """

    # Establish whether to measure wall times of code sections.
    measure_wall_times = False

    if measure_wall_times:
        time_start = time.time()

    # Establish whether to use scipy sparse matrix capability, which will
    # execute MUCH faster than an explicit sparse matrix regrid.
    use_scipy_sparse = True

    if len(source_grid.shape) != 2:
        statusMeta.errMsg = "Input grid must be a 2-D array."
        try:
            errMod.logErr(statusMeta)
        except:
            errMod.errOutScreen(statusMeta)
        raise Exception()

    source_grid_rows, source_grid_cols = source_grid.shape

    row = None
    col = None
    sparse_matrix_weight = None
    dest_grid_dims = None
    dest_grid_rows = None
    dest_grid_cols = None

    # Replace masked values with the no-data value. If this is not done, then
    # the regrid will return a generic no-data value of 9999 wherever
    # a masked value of source_grid is encountered in the csr_matrix.dot
    # process. Maybe there is 
    if isinstance(source_grid, np.ma.MaskedArray):
        source_grid = source_grid.filled(no_data_value)
        statusMeta.statusMsg = "Replaced mask in source grid with no-data values."
        try:
            errMod.logMsg(statusMeta)
        except:
            errMod.errOutScreen(statusMeta)

    # Check the structure of the weight file.
    master_dict = {}
    try:
        with Dataset(weight_file_path, 'r') as nc_weight:

            # Get the names and sizes of the n_a, n_b, and n_s dimensions.
            for name, dim in nc_weight.dimensions.items():
                master_dict[name] = len(dim)
            if "n_a" not in master_dict:
                statusMeta.errMsg = "Missing 'n_a' dimension."
                try:
                    errMod.logErr(statusMeta)
                except:
                    errMod.errOutScreen(statusMeta)
                raise Exception()
            if "n_b" not in master_dict:
                statusMeta.errMsg = "Missing 'n_b' dimension"
                try:
                    errMod.logErr(statusMeta)
                except:
                    errMod.errOutScreen(statusMeta)
                raise Exception()
            if "n_s" not in master_dict:
                statusMeta.errMsg = "Missing 'n_s' dimension"
                try:
                    errMod.logErr(statusMeta)
                except:
                    errMod.errOutScreen(statusMeta)
                raise Exception()
            try:
                if master_dict['src_grid_rank'] != 2:
                    statusMeta.errMsg = "Incorrect source grid rank " + str(master_dict['src_grid_rank'])
                    try:
                        errMod.logErr(statusMeta)
                    except:
                        errMod.errOutScreen(statusMeta)
                    raise Exception()
            except KeyError: # GF What about other kinds of exceptions?
                statusMeta.errMsg = "missing source grid rank"
                try:
                    errMod.logErr(statusMeta)
                except:
                    errMod.errOutScreen(statusMeta)
                raise Exception()
            try:
                if master_dict['dst_grid_rank'] != 2:
                    statusMeta.errMsg = "Incorrect destination grid rank " + str(master_dict['dst_grid_rank'])
                    try:
                        errMod.logErr(statusMeta)
                    except:
                        errMod.errOutScreen(statusMeta)
                    raise Exception()
            except KeyError:
                statusMeta.errMsg = "Missing destination grid rank."
                try:
                    errMod.logErr(statusMeta)
                except:
                    errMod.errOutScreen(statusMeta)
                raise Exception()

            # Check source grid dimensions against input data
            for nc_var_name, nc_var in nc_weight.variables.items():

                # Verify source grid information from weight file against
                # source_grid.
                if nc_var_name == "src_grid_dims":
                    src_grid_dims_size = nc_var.shape
                    if len(src_grid_dims_size) != 1:
                        statusMeta.errMsg = "src_grid_dims must be a 1-D array."
                        try:
                            errMod.logErr(statusMeta)
                        except:
                            errMod.errOutScreen(statusMeta)
                        raise Exception()
                    if src_grid_dims_size[0] != 2:
                        statusMeta.errMsg = "src_grid_dims must have two elements."
                        try:
                            errMod.logErr(statusMeta)
                        except:
                            errMod.errOutScreen(statusMeta)
                        raise Exception()
                    if nc_var[0] != source_grid_cols:
                        statusMeta.errMsg = "src_grid_dims_lists " + str(nc_var[0]) + " columns; " + \
                                            str(source_grid_cols) + " expected"
                        try:
                            errMod.logErr(statusMeta)
                        except:
                            errMod.errOutScreen(statusMeta)
                        raise Exception()
                    if nc_var[1] != source_grid_rows:
                        statusMeta.errMsg = "src_grid_dims_lists " + str(nc_var[1]) + " rows; " + \
                                            str(source_grid_rows) + " expected"
                        try:
                            errMod.logErr(statusMeta)
                        except:
                            errMod.errOutScreen(statusMeta)
                        raise Exception()
                    if master_dict["n_a"] != \
                       (source_grid_rows * source_grid_cols):
                        statusMeta.errMsg = "n_a does not match source grid size."
                        try:
                            errMod.logErr(statusMeta)
                        except:
                            errMod.errOutScreen(statusMeta)
                        raise Exception()

                # Sanity check destination grid and "n_b" dimensions.
                if nc_var_name == 'dst_grid_dims':
                    dest_grid_dims_size = nc_var.shape
                    if len(dest_grid_dims_size) != 1:
                        statusMeta.errMsg = "dst_grid_dims must be a 1-D array."
                        try:
                            errMod.logErr(statusMeta)
                        except:
                            errMod.errOutScreen(statusMeta)
                        raise Exception()
                    if dest_grid_dims_size[0] != 2:
                        statusMeta.errMsg = "dst_grid_dims must have two elements."
                        try:
                            errMod.logErr(statusMeta)
                        except:
                            errMod.errOutScreen(statusMeta)
                        raise Exception()
                    dest_grid_dims = nc_weight.variables['dst_grid_dims'][:]
                    dest_grid_cols = dest_grid_dims[0]
                    dest_grid_rows = dest_grid_dims[1]
                    if master_dict["n_b"] != (dest_grid_rows * dest_grid_cols):
                        statusMeta.errMsg = "n_b does not match destination grid"
                        try:
                            errMod.logErr(statusMeta)
                        except:
                            errMod.errOutScreen(statusMeta)
                        raise Exception()

                # Sanity check and read weight ("S") variable.
                if nc_var_name == 'S':
                    if len(nc_var.shape) != 1:
                        statusMeta.errMsg = "S Must be a 1-D array."
                        try:
                            errMod.logErr(statusMeta)
                        except:
                            errMod.errOutScreen(statusMeta)
                        raise Exception()
                    if len(nc_var) != master_dict["n_s"]:
                        statusMeta.errMsg = "S must use the 'n_s' dimension"
                        try:
                            errMod.logErr(statusMeta)
                        except:
                            errMod.errOutScreen(statusMeta)
                        raise Exception()
                    sparse_matrix_weight = nc_weight.variables['S'][:]

                # Sanity check and read destination index ("row") variable.
                if nc_var_name == 'row': # destination index
                    if len(nc_var.shape) != 1:
                        statusMeta.errMsg = "row must be a 1-D array."
                        try:
                            errMod.logErr(statusMeta)
                        except:
                            errMod.errOutScreen(statusMeta)
                        raise Exception()
                    if len(nc_var) != master_dict["n_s"]:
                        statusMeta.errMsg = "row must use the 'n_s' dimension."
                        try:
                            errMod.logErr(statusMeta)
                        except:
                            errMod.errOutScreen(statusMeta)
                        raise Exception()
                    row = nc_weight.variables['row'][:]
                    if np.min(row) < np.int(1):
                        statusMeta.errMsg = "Non-positive values found in row variable."
                        try:
                            errMod.logErr(statusMeta)
                        except:
                            errMod.errOutScreen(statusMeta)
                        raise Exception()
                    if np.max(row) > master_dict["n_b"]:
                        statusMeta.errMsg = "Out of bounds values found in 'row' variable."
                        try:
                            errMod.logErr(statusMeta)
                        except:
                            errMod.errOutScreen(statusMeta)
                        raise Exception()

                # Sanity check and read source index ("col") variable.
                if nc_var_name == 'col': # source index
                    if len(nc_var.shape) != 1:
                        statusMeta.errMsg = "col must be a 1-D array."
                        try:
                            errMod.logErr(statusMeta)
                        except:
                            errMod.errOutScreen(statusMeta)
                        raise Exception()
                    if len(nc_var) != master_dict["n_s"]:
                        statusMeta.errMsg = "col must use the 'n_s' dimension."
                        try:
                            errMod.logErr(statusMeta)
                        except:
                            errMod.errOutScreen(statusMeta)
                        raise Exception()
                    col = nc_weight.variables['col'][:]
                    if np.min(col) < np.int(1):
                        statusMeta.errMsg = "Non-positive values found in col variable."
                        try:
                            errMod.logErr(statusMeta)
                        except:
                            errMod.errOutScreen(statusMeta)
                        raise Exception()
                    if np.max(col) > master_dict["n_a"]:
                        statusMeta.errMsg = "Out of bounds values in 'col' variable."
                        try:
                            errMod.logErr(statusMeta)
                        except:
                            errMod.errOutScreen(statusMeta)
                        raise Exception()

            # Verify the "ESMF_regrid_method" is recognized and supported.
            try:
                regrid_method = nc_weight.ESMF_regrid_method
            except RuntimeError:
                statusMeta.errMsg = "Missing ESMF_regrid_method attribute."
                try:
                    errMod.logErr(statusMeta)
                except:
                    errMod.errOutScreen(statusMeta)
                raise Exception()

            if regrid_method.lower() == "first-order conservative":
                # Get the frac_b variable
                try:
                    frac_b = nc_weight.variables['frac_b'][:]
                except RuntimeError:
                    statusMeta.errMsg = "Missing frac_b variable."
                    try:
                        errMod.logErr(statusMeta)
                    except:
                        errMod.errOutScreen(statusMeta)
                    raise Exception()
                if len(frac_b.shape) != 1:
                    statusMeta.errMsg = "frac_b must be a 1-D array."
                    try:
                        errMod.logErr(statusMeta)
                    except:
                        errMod.errOutScreen(statusMeta)
                    raise Exception()
                if len(frac_b) != master_dict["n_b"]:
                    statusMeta.errMsg = "frac_b must use the \"n_b\" dimension."
                    try:
                        errMod.logErr(statusMeta)
                    except:
                        errMod.errOutScreen(statusMeta)
                    raise Exception()
            elif regrid_method.lower() != "bilinear":
                statusMeta.errMsg = "Unrecognized ESMF_regrid_method of \"%s\"." % regrid_method
                try:
                    errMod.logErr(statusMeta)
                except:
                    errMod.errOutScreen(statusMeta)
                raise Exception()

    except (IOError, OSError):
        statusMeta.errMsg = "Weight file " + weight_file_path + " not found or wrong format."
        try:
            errMod.logErr(statusMeta)
        except:
            errMod.errOutScreen(statusMeta)
        raise Exception()

    # Initialize output grid.
    dest_grid = np.full((dest_grid_rows, dest_grid_cols), no_data_value)

    # Adjust row and col to be zero-based.
    row = row - 1
    col = col - 1

    if measure_wall_times:
        time_finish = time.time()
        statusMeta.statusMsg = "opened weight file and checked dimensions in %s seconds" % (time_finish - time_start)
        try:
            errMod.logMsg(statusMeta)
        except:
            errMod.errOutScreen(statusMeta)
        time_start = time.time()

    # Get unique values of row to initialize destination grid
    row_uniq = np.unique(row)
    uniq_dest_col = row_uniq % dest_grid_cols
    uniq_dest_row = row_uniq / dest_grid_cols
    uniq_dest_col = uniq_dest_col.astype(int)
    uniq_dest_row = uniq_dest_row.astype(int)

    if measure_wall_times:
        time_finish = time.time()
        statusMeta.statusMsg = "calculated unique values of row in %s seconds" % (time_finish - time_start)
        try:
            errMod.logMsg(statusMeta)
        except:
            errMod.errOutScreen(statusMeta)
        time_start = time.time()

    dest_grid[uniq_dest_row, uniq_dest_col] = 0.0
    if measure_wall_times:
        time_finish = time.time()
        statusMeta.statusMsg = "initialized output grid in %s seconds" % (time_finish - time_start)
        try:
            errMod.logMsg(statusMeta)
        except:
            errMod.errOutScreen(statusMeta)

    if regrid_method.lower() == "first-order conservative":

        # Verify frac_b are nonzero for all output indices listed in "row"

        if measure_wall_times:
            time_start = time.time()

        frac_b_tmp = frac_b[row_uniq]
        countTmp = np.where(frac_b_tmp == 0)
        if countTmp[0].shape[0] > 0:
            statusMeta.errMsg = "Unexpected zero value in frac_b variable."
            try:
                errMod.logErr(statusMeta)
            except:
                errMod.errOutScreen(statusMeta)
            raise Exception()
        frac_b_tmp = None
        countTmp = None

        if measure_wall_times:
            time_finish = time.time()
            statusMeta.statusMsg = "checked frac_b consistency in %s seconds" % (time_finish - time_start)
            try:
                errMod.logMsg(statusMeta)
            except:
                errMod.errOutScreen(statusMeta)
       

    # Calculate row and column indices for sparse matrix

    if measure_wall_times:
        time_start = time.time()

    if measure_wall_times:
        time_finish = time.time()
        statusMeta.statusMsg = "calculated source and destination row/column values in %s seconds" % (time_finish - time_start)
        try:
            errMod.logMsg(statusMeta)
        except:
            errMod.errOutScreen(statusMeta)

    # Identify elements of the destination grid NOT listed in row.
    #dest_cell_not_in_row = np.in1d(np.arange(master_dict["n_b"]), row,
    #                               invert=True)

    if use_scipy_sparse:

        # Perform sparse matrix multiplication using scipy.sparse

        if measure_wall_times:
            time_start = time.time()

        wgt = csr_matrix((sparse_matrix_weight, (row, col)),
                         shape=(master_dict["n_b"], master_dict["n_a"]))

        # Generate 1-D version of source grid.
        source_grid_1d = np.reshape(source_grid, master_dict["n_a"])

        # Perform sparse matrix multiply.
        dest_grid_1d = wgt.dot(source_grid_1d)

        # Generate a "no-data-value flag". Any nonzero result means a
        # no_data_value in the source grid is associated with a nonzero weight
        # (sparse_matrix_weight).
        dest_grid_1d_ndv = wgt.dot(source_grid_1d == no_data_value)

        # Apply no-data values identified above to the destination grid.
        dest_grid_1d_ndv_ind = np.where(dest_grid_1d_ndv > 0.0)
        if len(dest_grid_1d_ndv_ind[0]) > 0:
            dest_grid_1d[dest_grid_1d_ndv_ind] = no_data_value
            #print("Adjusted %s destination cells to ndv based on "
            #      "source data" % len(dest_grid_1d_ndv_ind[0]))
            statusMeta.statusMsg = "Adjusted %s destination cells to ndv based on source data" % len(dest_grid_1d_ndv_ind[0])
            try:
                errMod.logMsg(statusMeta)
            except:
                errMod.errOutScreen(statusMeta)

        # Identify destination grid cells that were initialized to
        # no_data_value.
        dest_grid_init_ndv = np.where(dest_grid == no_data_value)

        # Convert destination grid from 1-D to 2-D.
        dest_grid = np.reshape(dest_grid_1d, (dest_grid_rows, dest_grid_cols))

        # Apply no-data values from dest_grid initialization.
        dest_grid[dest_grid_init_ndv] = no_data_value

        if measure_wall_times:
            time_finish = time.time()
            statusMeta.statusMsg = "sparse matrix regrid using scipy.sparse took %s seconds" % (time_finish - time_start)
            try:
                errMod.logMsg(statusMeta)
            except:
                errMod.errOutScreen(statusMeta)

    else:

        # Perform sparse matrix multiplication with source no_data_value /
        # nonzero weight check. This is MUCH slower than the use_scipy_sparse
        # option.

        # Calculate 2-D indices corresponding to row and col.
        dest_col = row % dest_grid_cols
        dest_row = row / dest_grid_cols
        source_col = col % source_grid_cols
        source_row = col / source_grid_cols

        if measure_wall_times:
            time_start = time.time()

        for i in range(len(sparse_matrix_weight)):
            if dest_grid[dest_row[i], dest_col[i]] == no_data_value:
                continue
            source_grid_val = source_grid[source_row[i], source_col[i]]
            weight = sparse_matrix_weight[i]
            if (source_grid_val == no_data_value) and (weight > 0.0):
                dest_grid[dest_row[i], dest_col[i]] = no_data_value
                continue
            dest_grid[dest_row[i], dest_col[i]] += weight * source_grid_val

        if measure_wall_times:
            time_finish = time.time()
            statusMeta.statusMsg = "sparse matrix regrid took %s seconds" % (time_finish - time_start)
            try:
                errMod.logMsg(statusMeta)
            except:
                errMod.errOutScreen(statusMeta)

    if regrid_method.lower() == "first-order conservative":

        # Adjust dest_grid values with frac_b.

        if measure_wall_times:
            time_start = time.time()

        frac_b_grid = frac_b.reshape(dest_grid_rows, dest_grid_cols)
        ind_tmp = np.where((dest_grid != no_data_value) & (frac_b_grid != 0.0))
        dest_grid[ind_tmp] = dest_grid[ind_tmp] / frac_b_grid[ind_tmp]
        ind_tmp = None
        frac_b_grid = None

        if measure_wall_times:
            time_finish = time.time()
            statusMeta.statusMsg = "final adjustment for conservative regrid took %s seconds" % (time_finish - time_start)
            try:
                errMod.logMsg(statusMeta)
            except:
                errMod.errOutScreen(statusMeta)

    return dest_grid

offline = with_weights # Alterative name for ESMF_regrid.with_weights