import numpy as np
import scipy.sparse as sp
from discretize.utils.matrix_utils import mkvc, sub2ind
from discretize.utils.code_utils import deprecate_function

try:
    from discretize._extensions import interputils_cython as pyx

    _interp_point_1D = pyx._interp_point_1D
    _interpmat1D = pyx._interpmat1D
    _interpmat2D = pyx._interpmat2D
    _interpmat3D = pyx._interpmat3D
    _vol_interp = pyx._tensor_volume_averaging
    _interpCython = True
except ImportError as err:
    print(err)
    import os

    # Check if being called from non-standard location (i.e. a git repository)
    # is tree_ext.cpp here? will not be in the folder if installed to site-packages...
    file_test = (
        os.path.dirname(os.path.abspath(__file__))
        + "/_extensions/interputils_cython.pyx"
    )
    if os.path.isfile(file_test):
        # Then we are being run from a repository
        print(
            """
            Unable to import interputils_cython.

            It would appear that discretize is being imported from its repository.
            If this is intentional, you need to run:

            python setup.py build_ext --inplace

            to build the cython code.
            """
        )
    _interpCython = False


def interpolation_matrix(locs, x, y=None, z=None):
    """
    Generate interpolation matrix which maps a tensor quantity to a set of locations.
    
    This function generates a sparse matrix for interpolating tensor quantities to a set of
    specified locations. It uses nD linear interpolation. The user may generate the interpolation
    matrix for tensor quantities that live on 1D, 2D or 3D tensors. This functionality is
    frequently used to interpolate quantites from cell centers or nodes to specified locations.
    
    Parameters
    ----------
    locs : numpy.ndarray
        A numpy array (n, dim) containing the locations for the interpolated values. Here *n* is
        the number of locations and *dim* is the dimension (1, 2 or 3)
    x : numpy.ndarray
        Vector defining the cell center or node locations of the tensor along the x-axis
    y : numpy.ndarray, optional
        Vector defining the cell center or node locations of the tensor along the y-axis
    z : numpy.ndarray, optional
        Vector defining the cell center or node locations of the tensor along the z-axis

    Returns
    -------
    scipy.sparse.csr_matrix
        A sparse matrix which interpolates the tensor quantity on cell centers or nodes to
        the set of specified locations. Where *n* is the number of specified locations and
        *m* is the number of tensor quantities, the output array has shape (n, m).

    Examples
    --------

    Here is a 1D example where a function evaluated on a regularly spaced grid
    is interpolated to a set of random locations. To compare the accuracy, the
    function is evaluated at the set of random locations.

    >>> from discretize.utils import interpolation_matrix
    >>> from discretize import TensorMesh
    >>> import numpy as np
    >>> import matplotlib.pyplot as plt
    >>> np.random.seed(14)
    >>> 
    >>> locs = np.random.rand(50)*0.8+0.1
    >>> x = np.linspace(0, 1, 7)
    >>> dense = np.linspace(0, 1, 200)
    >>> fun = lambda x: np.cos(2*np.pi*x)
    >>> Q = interpolation_matrix(locs, x)
    >>> 
    >>> fig1 = plt.figure(figsize=(5, 3))
    >>> ax = fig1.add_axes([0.1, 0.1, 0.8, 0.8])
    >>> ax.plot(dense, fun(dense), 'k:', lw=3)
    >>> ax.plot(x, fun(x), 'ks', markersize=8)
    >>> ax.plot(locs, Q*fun(x), 'go', markersize=4)
    >>> ax.plot(locs, fun(locs), 'rs', markersize=4)
    >>> ax.legend(
    >>>     [
    >>>         'True Function',
    >>>         'True (discrete loc.)',
    >>>         'Interpolated (computed)',
    >>>         'True (interp. loc.)'
    >>>     ],
    >>>     loc='upper center'
    >>> )
    >>> plt.show()


    Here, demonstrate a similar example on a 2D mesh using a 2D Gaussian distribution.
    We interpolate the Gaussian from the nodes to cell centers and examine the relative
    error.

    >>> hx = np.ones(10)
    >>> hy = np.ones(10)
    >>> mesh = TensorMesh([hx, hy], x0='CC')
    >>> def fun(x, y):
    >>>     return np.exp(-(x**2 + y**2)/2**2)
    >>> nodes = mesh.nodes
    >>> val_nodes = fun(nodes[:, 0], nodes[:, 1])
    >>> centers = mesh.cell_centers
    >>> val_centers = fun(centers[:, 0], centers[:, 1])
    >>> A = interpolation_matrix(
    >>>     centers, mesh.nodes_x, mesh.nodes_y
    >>> )
    >>> val_interp = A.dot(val_nodes)
    >>> fig = plt.figure(figsize=(11,3.3))
    >>> clim = (0., 1.)
    >>> ax1 = fig.add_subplot(131)
    >>> ax2 = fig.add_subplot(132)
    >>> ax3 = fig.add_subplot(133)
    >>> mesh.plot_image(val_centers, ax=ax1, clim=clim)
    >>> mesh.plot_image(val_interp, ax=ax2, clim=clim)
    >>> mesh.plot_image(val_centers-val_interp, ax=ax3, clim=clim)
    >>> ax1.set_title('Analytic at Centers')
    >>> ax2.set_title('Interpolated from Nodes')
    >>> ax3.set_title('Relative Error')
    >>> plt.show()
    

    """

    npts = locs.shape[0]
    locs = locs.astype(float)
    x = x.astype(float)
    if y is None and z is None:
        shape = [x.size]
        inds, vals = _interpmat1D(mkvc(locs), x)
    elif z is None:
        y = y.astype(float)
        shape = [x.size, y.size]
        inds, vals = _interpmat2D(locs, x, y)
    else:
        y = y.astype(float)
        z = z.astype(float)
        shape = [x.size, y.size, z.size]
        inds, vals = _interpmat3D(locs, x, y, z)

    I = np.repeat(range(npts), 2 ** len(shape))
    J = sub2ind(shape, inds)
    Q = sp.csr_matrix((vals, (I, J)), shape=(npts, np.prod(shape)))
    return Q


def volume_average(mesh_in, mesh_out, values=None, output=None):
    """Volume averaging interpolation between meshes.

    This volume averaging function looks for overlapping cells in each mesh,
    and weights the output values by the partial volume ratio of the overlapping
    input cells. The volume average operation should result in an output such that
    ``np.sum(mesh_in.cell_volumes*values)`` = ``np.sum(mesh_out.cell_volumes*output)``, when the input
    and output meshes have the exact same extent. When the output mesh extent goes
    beyond the input mesh, it is assumed to have constant values in that direction.
    When the output mesh extent is smaller than the input mesh, only the overlapping
    extent of the input mesh contributes to the output.

    This function operates in three different modes. If only *mesh_in* and
    *mesh_out* are given, the returned value is a ``scipy.sparse.csr_matrix``
    that represents this operation (so it could potentially be applied repeatedly).
    If *values* is given, the volume averaging is performed right away (without
    internally forming the matrix) and the returned value is the result of this.
    If *output* is given as well, it will be filled with the values of the
    operation and then returned (assuming it has the correct ``dtype``).

    Parameters
    ----------
    mesh_in : ~discretize.TensorMesh or ~discretize.TreeMesh
        Input mesh (the mesh you are interpolating from)
    mesh_out : ~discretize.TensorMesh or ~discretize.TreeMesh
        Output mesh (the mesh you are interpolating to)
    values : numpy.ndarray (optional)
        Array with values defined at the cells of ``mesh_in``
    output : numpy.ndarray (optional)
        Output array to be overwritten of length *mesh_out.nC* and type ``np.float64``

    Returns
    -------
    scipy.sparse.csr_matrix or numpy.ndarray
        If *values* = *None* , the returned value is a matrix representing this operation,
        otherwise it is a :class:`numpy.ndarray` of the result of the operation.

    Examples
    --------
    Create two meshes with the same extent, but different divisions (the meshes
    do not have to be the same extent).

    >>> import numpy as np
    >>> from discretize import TensorMesh
    >>> h1 = np.ones(32)
    >>> h2 = np.ones(16)*2
    >>> mesh_in = TensorMesh([h1, h1])
    >>> mesh_out = TensorMesh([h2, h2])

    Create a random model defined on the input mesh, and use volume averaging to
    interpolate it to the output mesh.

    >>> from discretize.utils import volume_average
    >>> model1 = np.random.rand(mesh_in.nC)
    >>> model2 = volume_average(mesh_in, mesh_out, model1)

    Because these two meshes' cells are perfectly aligned, but the output mesh
    has 1 cell for each 4 of the input cells, this operation should effectively
    look like averaging each of those cells values

    >>> import matplotlib.pyplot as plt
    >>> plt.figure(figsize=(6, 3))
    >>> ax1 = plt.subplot(121)
    >>> mesh_in.plot_image(model1, ax=ax1)
    >>> ax2 = plt.subplot(122)
    >>> mesh_out.plot_image(model2, ax=ax2)
    >>> plt.show()

    """
    try:
        in_type = mesh_in._meshType
        out_type = mesh_out._meshType
    except AttributeError:
        raise TypeError("Both input and output mesh must be valid discetize meshes")

    valid_meshs = ["TENSOR", "TREE"]
    if in_type not in valid_meshs or out_type not in valid_meshs:
        raise NotImplementedError(
            f"Volume averaging is only implemented for TensorMesh and TreeMesh, "
            f"not {type(mesh_in).__name__} and/or {type(mesh_out).__name__}"
        )

    if mesh_in.dim != mesh_out.dim:
        raise ValueError("Both meshes must have the same dimension")

    if values is not None and len(values) != mesh_in.nC:
        raise ValueError(
            "Input array does not have the same length as the number of cells in input mesh"
        )
    if output is not None and len(output) != mesh_out.nC:
        raise ValueError(
            "Output array does not have the same length as the number of cells in output mesh"
        )

    if values is not None:
        values = np.asarray(values, dtype=np.float64)
    if output is not None:
        output = np.asarray(output, dtype=np.float64)

    if in_type == "TENSOR":
        if out_type == "TENSOR":
            return _vol_interp(mesh_in, mesh_out, values, output)
        elif out_type == "TREE":
            return mesh_out._vol_avg_from_tens(mesh_in, values, output)
    elif in_type == "TREE":
        if out_type == "TENSOR":
            return mesh_in._vol_avg_to_tens(mesh_out, values, output)
        elif out_type == "TREE":
            return mesh_out._vol_avg_from_tree(mesh_in, values, output)
    else:
        raise TypeError("Unsupported mesh types")


interpmat = deprecate_function(
    interpolation_matrix, "interpmat", removal_version="1.0.0"
)
