# module for inperpolating / smoothing data
import numpy as np
from scipy.interpolate import CloughTocher2DInterpolator, Rbf
from scipy.spatial import cKDTree
try:
    import matplotlib.pyplot as plt
except ImportError:
    print('could not import matplotlib, no plotting functions enabled. Set plotonly = False in function call')


def MeshSplineToMesh2d(x, y, data, xi, yi, tol=1e-6, fill_nans=False, **kwargs):  #{{{
    '''
    Piecewise cubic, C1 smooth, curvature-minimizing interpolant in 2D.
    The interpolant is guaranteed to be continuously differentiable,
    and the gradients are chosen such that the curvature of the interpolant
    is approximately minimized.

    Uses scipy.interpolate.CloughTocher2DInterpolator

    x, y:            data point coordinates
    data:            data to be interpolated (same length as x, y)
    xi, yi:        coordintes to interpolate data onto
    tol:            tolerance for gradient estimation (default 1e-6)
    fill_nans:    fill nan's (holes) in data using the spline fit?
    **kwargs:    optional keywork arguments:
                    maxiter: maximum iterations in gradient estimation

    Returns interpolated data at given x, y coordinates.

    Usage:
        interpdata = CloughToucher2d(x, y, data)

    Examples:
        interpdata = CloughToucher2d(md.mesh.x, md.mesh.y, data)
        interpdata = CloughToucher2d(md.mesh.x, md.mesh.y, data, tol = 1e-3, maxiter = 100)
    '''

    # unpack kwargs
    maxiter = kwargs.pop('maxiter', None)
    if 'maxiter' in kwargs:
        del kwargs['maxiter']
    if maxiter:
        assert type(maxiter) == int, 'error, maxiter should be an integer'
    assert len(kwargs) == 0, 'error, unexpected or misspelled kwargs'

    # create sub - vectors that just cover the limits of xi and yi
    # TODO x, y not necessarily a grid, so need a better definition of dx, dy (e.g. average element size)
    dx = 500
    dy = 500
    #dx = x[1] - x[0]
    #dy = y[1] - y[0]
    xlim = [min(xi) - dx, max(xi) + dx]
    ylim = [min(yi) - dy, max(yi) + dy]
    xflag = np.logical_and(x > xlim[0], x < xlim[1])
    yflag = np.logical_and(y > ylim[0], y < ylim[1])
    bothind = np.nonzero(np.logical_and(xflag, yflag))
    subdata = data[bothind]
    subx = x[bothind]
    suby = y[bothind]
    points = np.array([subx, suby]).T

    # mask out any nan's in the data and corresponding coordinate points
    mask = np.isnan(subdata)
    ind = np.nonzero(mask)[0]
    if len(ind) and fill_nans:
        print("        WARNING: filling nans using spline fit through good data points, which may or may not be appropriate. Check results carefully.")
    subdata = np.delete(subdata, ind)
    points = np.delete(points, ind, axis=0)

    if maxiter:
        spline = CloughTocher2DInterpolator(points, subdata, tol, maxiter=maxiter)
    else:
        spline = CloughTocher2DInterpolator(points, subdata, tol)

    interpdata = spline(xi, yi)

    if not fill_nans:
        # identify nan's in xi, yi using nearest neighbors
        xyinterp = np.dstack([xi, yi])[0]
        xg, yg = np.meshgrid(subx, suby)
        xydata = np.dstack([subx, suby])[0]
        tree = cKDTree(xydata)
        nearest = tree.query(xyinterp)[1]
        pos = np.nonzero(np.isnan(subdata[nearest]))
        interpdata[pos] = subdata[nearest][pos]

    return interpdata
    #}}}


def GridSplineToMesh2d(x, y, data, xi, yi, default_value=np.nan, plotonly=False, fill_nans=False):  #{{{
    '''
    python analog to InterpFromGridToMesh.  This routine uses
    scipy.interpolate.CloughTocher2dInterpolator to create a bivariate spline
    interpolation of the input data and then return values of the spline
    on the x, y coordinates of the model mesh.  The interpolant is piece-wise
    cubic, C1 smooth (continuously differentiable) and has approximately
    minimized curvature.  See "help(scipy.interpolate.CloughTocher2dInterpolator)"
    for more information on the routine.

    NOTE: this routine will not be appropriate if there are large holes (nan's) in
    the input data.  A non - spline interpolation scheme should be used in that case.

    x, y:                vectors defining the coordinates of the input data
    data:                2D array of input data
    xi, yi:            x and y coordinates to be interpolated onto
    default_value:    default value if points lie outside the convex hull of input
                        points (defaults to nan if not specified)
    plotonly:        plot the data to be interpolated using imshow (useful for
    fill_nans:        fill nan's (holes) in data using the spline fit?

    Usage:
        interpdata = GridToMesh(x, y, data, xi, yi, default_value = np.nan, plotonly = False, fill_nans = False)

    Examples:
        interpdata = GridToMesh(x_m, y_m, data, md.mesh.x, md.mesh.y, 0)
    '''

    if np.ndim(x) == 2:
        x = x.reshape(-1, )
    if np.ndim(y) == 2:
        y = y.reshape(-1, )
    if len(x) != data.shape[1] + 1 and len(x) != data.shape[1]:
        raise ValueError('x should have same length as ncols(data) or ncols(data) + 1')
    if len(y) != data.shape[0] + 1 and len(y) != data.shape[0]:
        raise ValueError('y should have same length as nrows(data) or nrows(data) + 1')

    # create sub - grid that just covers the limits of xi and yi
    dx = x[1] - x[0]
    dy = y[1] - y[0]
    xlim = [min(xi) - dx, max(xi) + dx]
    ylim = [min(yi) - dy, max(yi) + dy]

    # TODO create grid differently depending on whether data is defined at x, y
    # or at the center of a grid cell with corner coordinates defined by xi, yi
    # create points array and flattened data array
    if len(x) == data.shape[1] and len(y) == data.shape[0]:
        print('        x, y taken to define the center of data grid cells')
        xind = np.nonzero(np.logical_and(x > xlim[0], x < xlim[1]))[0]
        yind = np.nonzero(np.logical_and(y > ylim[0], y < ylim[1]))[0]
        xg, yg = np.meshgrid(x[xind], y[yind])
        subdata = data[yind[0]:yind[-1] + 1, xind[0]:xind[-1] + 1]
    elif len(x) == data.shape[1] + 1 and len(y) == data.shape[0] + 1:
        print('        x, y taken to define the corners of data grid cells')
        xcenter = np.fromiter(((x[i] + x[i + 1]) / 2 for i in range(len(x) - 1)), np.float)
        ycenter = np.fromiter(((y[i] + y[i + 1]) / 2 for i in range(len(y) - 1)), np.float)
        xind = np.nonzero(np.logical_and(xcenter > xlim[0], xcenter < xlim[1]))[0]
        yind = np.nonzero(np.logical_and(ycenter > ylim[0], ycenter < ylim[1]))[0]
        xg, yg = np.meshgrid(xcenter[xind], ycenter[yind])
        subdata = data[yind[0]:yind[-1] + 1, xind[0]:xind[-1] + 1]
    else:
        raise ValueError('x and y have inconsistent sizes: both should have length ncols(data) / nrows(data) or ncols(data) + 1 / nrows(data) + 1')

    points = np.array([xg.ravel(), yg.ravel()]).T
    flatsubdata = subdata.ravel()

    if plotonly:
        plt.imshow(np.flipud(subdata), origin='upper')
        plt.show()
        return

    # mask out any nan's in the data and corresponding coordinate points
    mask = np.isnan(flatsubdata)
    ind = np.nonzero(mask)[0]
    if len(ind) and fill_nans:
        print("        WARNING: filling nans using spline fit through good data points, which may or may not be appropriate. Check results carefully.")
    goodsubdata = np.delete(flatsubdata, ind)
    goodpoints = np.delete(points, ind, axis=0)

    # create spline and index spline at mesh points
    spline = CloughTocher2DInterpolator(goodpoints, goodsubdata)
    interpdata = spline(xi, yi)

    if not fill_nans:
        # identify nan's in xi, yi using nearest neighbors
        xyinterp = np.dstack([xi, yi])[0]
        xydata = np.dstack([xg.ravel(), yg.ravel()])[0]
        tree = cKDTree(xydata)
        nearest = tree.query(xyinterp)[1]
        pos = np.nonzero(np.isnan(flatsubdata[nearest]))
        interpdata[pos] = flatsubdata[nearest][pos]

    return interpdata
    #}}}


def RadialInterp(x, y, data, xi, yi, **kwargs):  #{{{
    '''
    Interpolation using a radial basis function in 2 or 3 dimensions.
    Useful for smoothing input data after interpolation.

    Uses scipy.interpolate.Rbf

    x, y:            data point coordinates
    data:            data to be interpolated (same length as x, y)
    xi, yi:        coordinates to interpolate onto
    function:    form of radial basis function for interpolation:
                    'multiquadric': sqrt((r / self.epsilon)**2 + 1) (default)
                    'inverse': 1.0 / sqrt((r / self.epsilon)**2 + 1)
                    'gaussian': exp(-(r / self.epsilon)**2)
                    'linear': r
                    'cubic': r**3
                    'quintic': r**5
                    'thin_plate': r**2 * log(r)
    epsilon:        adjustable constant for scaling radial distance.  Defaults to
                    approximate average distance between nodes.
    smooth:        float > 0, adjusts the amount of smoothing applied.  Defaults to 0,
                    such that the function always passes through nodal points.
    z:                coordinate array if interpolating in 3 dimensions
    zi:            coordinate array if interpolating in 3 dimensions

    Usage:
        interpdata = RadialInterp(x, y, data,**kwargs)

    Examples:
        interpdata = RadialInterp(md.mesh.x, md.mesh.y, data)
        interpdata = RadialInterp(md.mesh.x, md.mesh.y, data, function = 'gaussian', epsilon = 100, smooth = 1)
    '''

    # unpack kwargs
    function = kwargs.pop('function', 'gaussian')
    if 'function' in kwargs:
        del kwargs['function']
    epsilon = kwargs.pop('epsilon', None)
    if 'epsilon' in kwargs:
        del kwargs['epsilon']
    smooth = kwargs.pop('smooth', 0)
    if 'smooth' in kwargs:
        del kwargs['smooth']
    z = kwargs.pop('z', None)
    if 'z' in kwargs:
        del kwargs['z']
    assert len(kwargs) == 0, 'error, unexpected or misspelled kwargs'

    if z:
        if epsilon:
            rbfi = Rbf(x, y, z, data, function=function, smooth=smooth, epsilon=epsilon)
        else:
            rbfi = Rbf(x, y, z, data, function=function, smooth=smooth)
        interpdata = rbfi(xi, yi, zi)
    else:
        if epsilon:
            rbfi = Rbf(x, y, data, function=function, smooth=smooth, epsilon=epsilon)
        else:
            rbfi = Rbf(x, y, data, function=function, smooth=smooth)
        interpdata = rbfi(xi, yi)

    return interpdata
    #}}}
