"""
Sparse matrix multiplication (SMM) using scipy.sparse library.
"""

import xarray as xr
import scipy.sparse as sps
import warnings


def read_weights(filename, N_in, N_out):
    '''
    Read regridding weights into a scipy sparse COO matrix.

    Parameters
    ----------
    filename : str
        Offline weight file generated by ESMPy.

    N_in, N_out : integers
        ``(N_out, N_in)`` will be the shape of the returning sparse matrix.
        They are the total number of grid boxes in input and output grids::

              N_in = Nx_in * Ny_in
              N_out = Nx_out * Ny_out

        We need them because the shape cannot always be infered from the
        largest column and row indices, due to unmapped grid boxes.

    Returns
    -------
    A : scipy sparse COO matrix.

    '''
    ds_w = xr.open_dataset(filename)

    col = ds_w['col'].values - 1  # Python starts with 0
    row = ds_w['row'].values - 1
    S = ds_w['S'].values

    A = sps.coo_matrix((S, (row, col)), shape=[N_out, N_in])
    return A


def apply_weights(A, indata, Ny_out, Nx_out):
    '''
    Apply regridding weights to data.

    Parameters
    ----------
    A : scipy sparse COO matrix

    indata : numpy array of shape ``(..., Nlat, Nlon)`` or ``(..., Ny, Nx)``.
        Should be C-ordered. Will be then tranposed to F-ordered.

    Ny_out, Nx_out : integers
        Output data shape for unflatten operation.
        For rectilinear grid, it is just ``(Nlat_out, Nlon_out)``.

    Returns
    -------
    outdata : numpy array of shape ``(..., Ny_out, Nx_out)``.
        Extra dimensions are the same as `indata`.
        If input data is C-ordered, output will also be C-ordered.
    '''

    assert Nx_out * Ny_out == A.shape[0], (
        "Nx_out * Ny_out should equal to A.shape[0]")

    # COO matrix is fast with F-ordered array but slow with C-array, so we
    # take in a C-ordered and then transpose)
    # (CSR or CRS matrix is fast with C-ordered array but slow with F-array)
    if not indata.flags['C_CONTIGUOUS']:
        warnings.warn("Input array is not C_CONTIGUOUS. "
                      "Will affect performance.")

    # get input shape information
    s = indata.shape
    Ny_in, Nx_in = (s[-2], s[-1])
    N_extra_list = s[0:-2]

    # use flattened array for dot operation
    indata_flat = indata.reshape(-1, Ny_in*Nx_in)
    outdata_flat = A.dot(indata_flat.T).T

    # unflattened output array
    outdata = outdata_flat.reshape([*N_extra_list, Ny_out, Nx_out])
    return outdata
