"""Routines for generating structured meshes on poloidal domains

Classes
-------
`RectangularPoloidalGrid`
    Simple rectangles in R-Z
`StructuredPoloidalGrid`
    Curvilinear structured grids in R-Z

Functions
---------
`grid_annulus`
    Create a `StructuredPoloidalGrid` from inner and outer RZLine
    objects using a simple algorithm
`grid_elliptic`
    Create a `StructuredPoloidalGrid` from inner and outer RZLine
    objects using elliptic meshing method

"""

import warnings

import numpy as np
from numpy import linspace, pi, zeros
from scipy.interpolate import RectBivariateSpline
from scipy.spatial import cKDTree as KDTree

try:
    import matplotlib.pyplot as plt

    plotting_available = True
except ImportError:
    warnings.warn("Couldn't import matplotlib, plotting not available.")
    plotting_available = False

from . import rzline


class PoloidalGrid(object):
    """Represents a poloidal grid

    Note: Here the 2D plane (R,Z) is labelled by (x,z) indices

    Attributes
    ----------
    nx, nz : int
        Number of points in x and z
    R : ndarray
        2D Numpy array of R coordinates
    Z : ndarray
        2D Numpy array of Z coordinates

    """

    def plot(self, axis=None, show=True):
        """Plot grid using matplotlib

        Parameters
        ----------
        axis : matplotlib axis, optional
            A matplotlib axis to plot on. By default a new figure
            is created
        show : bool, optional
            Calls plt.show() at the end

        Returns
        -------
        axis
            The matplotlib axis that was used

        """

        if not plotting_available:
            warnings.warn("matplotlib not available, unable to plot")
            return None

        if axis is None:
            fig = plt.figure()
            axis = fig.add_subplot(1, 1, 1)

        axis.plot(self.R, self.Z, "k-")
        axis.plot(self.R.T, self.Z.T, "k-")

        if show:
            plt.show()

        return axis


class RectangularPoloidalGrid(PoloidalGrid):
    """Represents a poloidal grid consisting of a rectangular domain

    Note: Here the 2D plane (R,Z) is labelled by (x,z) indices

    Attributes
    ----------
    nx, nz : int
        Number of points in x and z
    R : ndarray
        2D Numpy array of R coordinates
    Z : ndarray
        2D Numpy array of Z coordinates

    Parameters
    ----------
    nx : int
        Number of points in major radius (including boundaries)
    nz : int
        Number of points in height (including boundaries)
    Lx : float
        Radial domain size  [m]
    Lz : float
        Vertical domain size [m]
    Rcentre : float, optional
        Coordinate at the middle of the domain
    Zcentre : float, optional
        Coordinate at the middle of the domain
    MXG : int, optional
        Number of guard cells in X. The boundary is put half-way
        between the guard cell and the domain

    """

    def __init__(self, nx, nz, Lx, Lz, Rcentre=0.0, Zcentre=0.0, MXG=2):
        self.nx = nx
        self.nz = nz

        self.Lx = Lx
        self.Lz = Lz

        self.Rcentre = Rcentre
        self.Zcentre = Zcentre

        # Some useful derived quantities
        # Note: index at the middle of the domain is (nx - 1)/2
        #       e.g. nx=5 :  0 1 | 2 | 3 4
        self.dR = self.Lx / (self.nx - 2 * MXG)
        self.dZ = self.Lz / self.nz
        self.Rmin = self.Rcentre - self.dR * (self.nx - 1.0) / 2.0
        self.Zmin = self.Zcentre - self.dZ * (self.nz - 1.0) / 2.0

        # Generate 2D arrays
        # Using getCoordinate to ensure consistency
        xind, zind = np.meshgrid(np.arange(nx), np.arange(nz), indexing="ij")
        self.R, self.Z = self.getCoordinate(xind, zind)

    def __repr__(self):
        return (
            "RectangularPoloidalGrid({0},{1},{2},{3},Rcentre={4},Zcentre={5})".format(
                self.nx, self.nz, self.Lx, self.Lz, self.Rcentre, self.Zcentre
            )
        )

    def getCoordinate(self, xind, zind, dx=0, dz=0):
        """Get coordinates (R,Z) at given (xind,zind) index

        Parameters
        ----------
        xind, zind : array_like
            Indices in X and Z. These should be the same shape
        dx : int, optional
            Order of x derivative
        dz : int, optional
            Order of z derivative

        Returns
        -------
        R, Z : (ndarray, ndarray)
            Locations of point or derivatives of R,Z with respect to
            indices if dx,dz != 0

        """

        # Convert to NumPy arrays if not already
        xind = np.asarray(xind)
        zind = np.asarray(zind)
        # Make sure dx and dz are integers
        dx = int(dx)
        dz = int(dz)

        assert xind.shape == zind.shape
        assert dx >= 0
        assert dz >= 0

        shape = xind.shape

        if dx + dz > 2:
            # Second derivatives and mixed derivatives all zero
            return np.zeros(shape), np.zeros(shape)

        if dx == 1:
            # dR/dx, dZ/dx
            return np.full(shape, self.dR), np.zeros(shape)
        elif dz == 1:
            # dR/dz, dZ/dz
            return np.zeros(shape), np.full(shape, self.dZ)
        # Return (R,Z) location
        return self.Rmin + xind * self.dR, self.Zmin + zind * self.dZ

    def findIndex(self, R, Z):
        """Finds the (x,z) index corresponding to the given (R,Z) coordinate

        Parameters
        ----------
        R, Z : array_like
            Locations to find indices for

        Returns
        -------
        x, z : (ndarray, ndarray)
            Index as a float, same shape as R,Z

        """

        # Make sure inputs are NumPy arrays
        R = np.asarray(R)
        Z = np.asarray(Z)

        # Check that they have the same shape
        assert R.shape == Z.shape

        xind = (R - self.Rmin) / self.dR
        zind = (Z - self.Zmin) / self.dZ

        # Note: These indices may be outside the domain,
        # but this is handled in BOUT++, and useful for periodic
        # domains.

        return xind, zind

    def metric(self):
        """Return the metric tensor, dx and dz

        For this rectangular grid the metric is the identity

        Returns
        -------
        dict
            Dictionary containing:
            - **dx, dz**: Grid spacing
            - **gxx, gxz, gzz**: Covariant components
            - **g_xx, g_xz, g_zz**: Contravariant components
        """
        return {
            "dx": self.dR,
            "dz": self.dZ,  # Grid spacing
            "gxx": 1.0,
            "g_xx": 1.0,
            "gxz": 0.0,
            "g_xz": 0.0,
            "gzz": 1.0,
            "g_zz": 1.0,
        }


class StructuredPoloidalGrid(PoloidalGrid):
    """Represents a structured poloidal grid in R-Z

    Attributes
    ----------
    nx, nz : int
        Number of points in x and z
    R : ndarray
        2D Numpy array of R coordinates
    Z : ndarray
        2D Numpy array of Z coordinates

    Parameters
    ----------
    R, Z : ndarray
        2D Numpy arrays of R,Z points

        .. note:: R,Z are not copied, so these arrays should not be
                  modified afterwards

    """

    def __init__(self, R, Z):
        assert R.shape == Z.shape

        self.R = R
        self.Z = Z

        # Create a KDTree for quick lookup of nearest points
        n = R.size
        data = np.concatenate((R.reshape((n, 1)), Z.reshape((n, 1))), axis=1)
        self.tree = KDTree(data)

        # Create splines for quick interpolation of coordinates
        nx, nz = R.shape

        self.nx = nx
        self.nz = nz

        xinds = np.arange(nx)
        zinds = np.arange(nz * 3)
        # Repeat the data in z, to approximate periodicity
        R_ext = np.concatenate((R, R, R), axis=1)
        Z_ext = np.concatenate((Z, Z, Z), axis=1)

        self._spl_r = RectBivariateSpline(xinds, zinds, R_ext)
        self._spl_z = RectBivariateSpline(xinds, zinds, Z_ext)

    def __repr__(self):
        return "StructuredPoloidalGrid()"

    def getCoordinate(self, xind, zind, dx=0, dz=0):
        """Get coordinates (R, Z) at given (xind, zind) index

        Parameters
        ----------
        xind, zind : array_like
            Indices in X and Z. These should be the same shape
        dx : int, optional
            Order of x derivative
        dz : int, optional
            Order of z derivative

        Returns
        -------
        R, Z : (ndarray, ndarray)
            Locations of point or derivatives of R,Z with respect to
            indices if dx,dz != 0

        """
        nx, nz = self.R.shape
        if (np.amin(xind) < 0) or (np.amax(xind) > nx - 1):
            raise ValueError("x index out of range")

        # Periodic in y
        zind = np.remainder(zind, nz)

        R = self._spl_r(xind, zind + self.nz, dx=dx, dy=dz, grid=False)
        Z = self._spl_z(xind, zind + self.nz, dx=dx, dy=dz, grid=False)

        return R, Z

    def findIndex(self, R, Z, tol=1e-10, show=False):
        """Finds the (x, z) index corresponding to the given (R, Z) coordinate

        Parameters
        ----------
        R, Z : array_like
            Locations. Can be scalar or array, must be the same shape
        tol : float, optional
            Maximum tolerance on the square distance

        Returns
        -------
        x, z : (ndarray, ndarray)
            Index as a float, same shape as R, Z

        """

        # Make sure inputs are NumPy arrays
        R = np.asarray(R)
        Z = np.asarray(Z)

        # Check that they have the same shape
        assert R.shape == Z.shape

        input_shape = R.shape  # So output has same shape as input

        # Get distance and index into flattened data
        # Note ind can be an integer, or an array of ints
        # with the same number of elements as the input (R,Z) arrays
        n = R.size
        position = np.concatenate((R.reshape((n, 1)), Z.reshape((n, 1))), axis=1)

        R = R.reshape((n,))
        Z = Z.reshape((n,))

        dists, ind = self.tree.query(position)

        # Calculate (x,y) index
        nx, nz = self.R.shape
        xind = np.floor_divide(ind, nz)
        zind = ind - xind * nz

        # Convert indices to float
        xind = np.asarray(xind, dtype=float)
        zind = np.asarray(zind, dtype=float)

        # Create a mask for the positions
        mask = np.ones(xind.shape)
        mask[np.logical_or((xind < 0.5), (xind > (nx - 1.5)))] = (
            0.0  # Set to zero if near the boundary
        )

        if show and plotting_available:
            plt.plot(self.R, self.Z, ".")
            plt.plot(R, Z, "x")

        cnt = 0
        underrelax = 1

        while True:
            # Use Newton iteration to find the index
            # dR, dZ are the distance away from the desired point
            Rpos, Zpos = self.getCoordinate(xind, zind)
            if show and plotting_available:
                plt.plot(Rpos, Zpos, "o")
            dR = Rpos - R
            dZ = Zpos - Z

            # Check if close enough
            # Note: only check the points which are not in the boundary
            val = np.amax(mask * (dR**2 + dZ**2))
            if val < tol:
                break
            cnt += 1
            if cnt == 10:
                underrelax = 1.5
            if cnt == 100:
                underrelax = 2
            if cnt == 300:
                underrelax = 2.5
            if cnt == 700:
                underrelax = 3
            if cnt == 1000:
                raise RuntimeError("Failed to converge")

            # Calculate derivatives
            dRdx, dZdx = self.getCoordinate(xind, zind, dx=1)
            dRdz, dZdz = self.getCoordinate(xind, zind, dz=1)

            # Invert 2x2 matrix to get change in coordinates
            #
            # (x) -=  ( dR/dx   dR/dz )^-1  (dR)
            # (y)     ( dZ/dx   dZ/dz )     (dz)
            #
            #
            # (x) -=  ( dZ/dz  -dR/dz ) (dR)
            # (y)     (-dZ/dx   dR/dx ) (dZ) / (dR/dx*dZ/dy - dR/dy*dZ/dx)
            determinant = dRdx * dZdz - dRdz * dZdx

            xind -= mask * ((dZdz * dR - dRdz * dZ) / determinant / underrelax)
            zind -= mask * ((dRdx * dZ - dZdx * dR) / determinant / underrelax)

            # Re-check for boundary
            in_boundary = xind < 0.5
            mask[in_boundary] = 0.0  # Set to zero if near the boundary
            xind[in_boundary] = 0.0
            out_boundary = xind > (nx - 1.5)
            mask[out_boundary] = 0.0  # Set to zero if near the boundary
            xind[out_boundary] = nx - 1

        if show and plotting_available:
            plt.show()

        # Set xind to -1 if in the inner boundary, nx if in outer boundary
        in_boundary = xind < 0.5
        xind[in_boundary] = -1
        out_boundary = xind > (nx - 1.5)
        xind[out_boundary] = nx

        return xind.reshape(input_shape), zind.reshape(input_shape)

    def metric(self):
        """Return the metric tensor, dx and dz

        Returns
        -------
        dict
            Dictionary containing:
            - **dx, dz**: Grid spacing
            - **gxx, gxz, gzz**: Covariant components
            - **g_xx, g_xz, g_zz**: Contravariant components

        """

        # Get arrays of indices
        xind, zind = np.meshgrid(np.arange(self.nx), np.arange(self.nz), indexing="ij")

        # Calculate the gradient along each coordinate
        dolddnew = np.array(
            [self.getCoordinate(xind, zind, dx=a, dz=b) for a, b in ((1, 0), (0, 1))]
        )
        # dims: 0 : dx or dz?
        #       1 : R or z?
        #       2 : spatial: r
        #       3 : spatial: \theta
        ddist = np.sqrt(np.sum(dolddnew**2, axis=1))  # sum R + z
        nx, nz = ddist.shape[1:]
        ddist[0] = 1 / nx
        ddist[1] = 1 / nz
        dolddnew /= ddist[:, None, ...]

        # g_ij = J_ki J_kj
        # (2.5.27) from D'Haeseleer 1991
        # Note: our J is transposed
        J = dolddnew
        g = np.sum(
            np.array(
                [
                    [[J[j, i] * J[k, i] for i in range(2)] for j in range(2)]
                    for k in range(2)
                ]
            ),
            axis=2,
        )

        assert np.all(
            g[0, 0] > 0
        ), f"g[0, 0] is expected to be positive, but some values are not (minimum {np.min(g[0, 0])})"
        assert np.all(
            g[1, 1] > 0
        ), f"g[1, 1] is expected to be positive, but some values are not (minimum {np.min(g[1, 1])})"
        g = g.transpose(2, 3, 0, 1)
        assert np.all(
            np.linalg.det(g) > 0
        ), f"All determinants of g should be positive, but some are not (minimum {np.min(np.linalg.det(g))})"
        ginv = np.linalg.inv(g)
        # Jacobian from BOUT++
        JB = self.R * (J[0, 0] * J[1, 1] - J[0, 1] * J[1, 0])
        return {
            "dx": ddist[0],
            "dz": ddist[1],  # Grid spacing
            "gxx": ginv[..., 0, 0],
            "g_xx": g[..., 0, 0],
            "gxz": ginv[..., 0, 1],
            "g_xz": g[..., 0, 1],
            "gzz": ginv[..., 1, 1],
            "g_zz": g[..., 1, 1],
            # "J": JB,
        }


def grid_annulus(inner, outer, nx, nz, show=True, return_coords=False):
    """Grid an annular region, given inner and outer boundaries both of
    which are RZline objects

    This is a very simple algorithm which just draws straight lines
    between inner and outer boundaries.

    Parameters
    ----------
    inner, outer : `RZline`
        Inner and outer boundaries of the domain
    nx : int
        The required radial resolution, including boundaries
    nz : int
        The required poloidal resolution
    show : bool, optional
        If True, plot the resulting grid
    return_coords : bool, optional
        If True, return the R, Z coordinates of the grid points,
        instead of a `StructuredPoloidalGrid`

    Returns
    -------
    StructuredPoloidalGrid
        A grid of the region

    """

    assert nx >= 2
    assert nz > 1

    R = zeros((nx, nz))
    Z = zeros((nx, nz))

    # Generate angle values, which should now be equally spaced
    # in distance along inner and outer boundaries
    thetavals = linspace(0, 2 * pi, nz, endpoint=False)

    # Radial coordinate
    xvals = linspace(0, 1.0, nx, endpoint=True)

    innerR = inner.Rvalue(thetavals)
    innerZ = inner.Zvalue(thetavals)

    outerR = outer.Rvalue(thetavals)
    outerZ = outer.Zvalue(thetavals)
    for i, x in enumerate(xvals):
        # Get the R and Z coordinates of this line
        R[i, :] = x * outerR + (1.0 - x) * innerR
        Z[i, :] = x * outerZ + (1.0 - x) * innerZ

    if show and plotting_available:
        plt.plot(inner.R, inner.Z, "-o")
        plt.plot(outer.R, outer.Z, "-o")

        plt.plot(R, Z, "x")

        plt.show()

    if return_coords:
        return R, Z
    return StructuredPoloidalGrid(R, Z)


def grid_elliptic(
    inner,
    outer,
    nx,
    nz,
    show=False,
    tol=1e-10,
    align=True,
    restrict_size=20,
    restrict_factor=2,
    return_coords=False,
    nx_outer=0,
    nx_inner=0,
    legacy_align=False,
    inner_ort=True,
    maxfac_inner=None,
    dz_relax=None,
):
    """Create a structured grid between inner and outer boundaries using
    elliptic method

    Coordinates x = x(R, Z) and z = z(R,Z) obey an elliptic equation:

    .. math::

        d^2x/dR^2 + d^2x/dZ^2 = 0

        d^2z/dR^2 + d^2z/dZ^2 = 0

    where here x is in in the domain (0, 1) and z in (0, 2pi)

    The above equations are inverted, giving:

    .. math::

        a*R_xx - 2*b*R_xz + c*R_zz = 0

        a*Z_xx - 2*b*Z_xz + c*Z_zz = 0

    where

    .. math::

        a &= R_z^2 + Z_z^2

        b &= R_z*R_x + Z_x*Z_z

        c &= R_x^2 + Z_x^2

    This is a nonlinear system of equations which is solved
    iteratively.

    Parameters
    ----------
    inner, outer : `RZline`
        Inner and outer boundaries of the domain
    nx : int
        The required radial resolution, including boundaries
    nz : int
        The required poloidal resolution
    show : bool, optional
        Display plots of intermediate results
    tol : float, optional
        Controls when iteration stops
    align : bool, optional
        Attempt to align the inner and outer boundaries
    restrict_size : int, optional
        The size (nx or nz) above which the grid is coarsened
    restrict_factor : int, optional
        The factor by which the grid is divided if coarsened
    nx_outer: int, optional
        The number of additional points outside of the outer boundary
    nx_inner: int, optional
        The number of additional points outside of the inner boundary
    inner_ort: bool, optional
        Whether to place the inner points as close as possible to the
        corresponding outer ones. That increases orthogonality of the grid.
    maxfac_inner : int, optional
        If given, ensure the spacing between smalles and largest distance on
        the inner boundary. Only used with inner_ort. Must be larger than 1.
    return_coords : bool, optional
        If True, return the R, Z coordinates of the grid points,
        instead of a `StructuredPoloidalGrid`

    Returns
    -------
    If return_coords is true, returns R,Z as arrays.
    If return_coords is false, returns a `StructuredPoloidalGrid` object

    References
    ----------
    https://web.archive.org/web/20040530002835/https://www.nada.kth.se/kurser/kth/2D1263/l2.pdf
    https://en.wikipedia.org/wiki/Principles_of_grid_generation

    """
    if nx_outer:
        assert nx_outer > 0
        nx -= nx_outer
    if nx_inner:
        assert nx_inner > 0
        nx -= nx_inner

    assert nx > 1
    assert nz > 1

    # Generate angle values (y coordinate),
    # which should now be equally spaced
    # in distance along inner and outer boundaries
    thetavals = linspace(0, 2 * pi, nz, endpoint=False)

    # Radial coordinate
    xvals = linspace(0, 1.0, nx, endpoint=True)

    if align and not inner_ort:
        # Align inner and outer boundaries
        # Easiest way is to roll both boundaries
        # so that index 0 is on the outboard midplane.
        # This is useless if inner_ort is used, where the inner points are
        # based on a project of the outer points.
        if len(inner.R) < len(outer.R):
            shorter = inner
            longer = outer
        else:
            shorter = outer
            longer = inner
        ind = np.argmax(shorter.R)
        shorter = rzline.RZline(np.roll(shorter.R, -ind), np.roll(shorter.Z, -ind))
        if legacy_align:
            ind = np.argmax(longer.R)
            longer = rzline.RZline(np.roll(longer.R, -ind), np.roll(longer.Z, -ind))
        else:
            dr = shorter.R - longer.R[:, None]
            dz = shorter.Z - longer.Z[:, None]
            delta = dr**2 + dz**2
            fac = len(longer.R) / len(shorter.R)
            j = np.arange(len(shorter.R), dtype=int)
            sums = [
                np.sum(delta[np.round(j * fac).astype(int) - i, j])
                for i in range(len(longer.R))
            ]
            ind = -np.argmin(sums)
            longer = rzline.RZline(np.roll(longer.R, -ind), np.roll(longer.Z, -ind))
        if len(inner.R) < len(outer.R):
            inner = shorter
            outer = longer
        else:
            outer = shorter
            inner = longer

    def laplace(x):
        def dx(x, sign):
            dxv = np.roll(x, sign) - x
            dxv = np.remainder(dxv + np.pi, 2 * np.pi) - np.pi
            return dxv

        fac = 0.1 * nz / 192
        steps = 10
        if fac > 0.1:
            steps = int(steps * fac / 0.1)
            fac = 0.1
        x0 = x.copy()
        for i in range(steps):
            x += fac * (dx(x, 1) + dx(x, -1))

        extra = 0
        if maxfac_inner:

            def getfac(x):
                d = dx(x, -1)
                return np.max(d) / np.min(d)

            while not np.all(dx(x, -1) > 0):
                x += fac * (dx(x, 1) + dx(x, -1))
                extra += 1

            if maxfac_inner == 1:
                x1 = x.copy()
                for i in range(1, len(x0)):
                    if x1[i - 1] > x1[i]:
                        x1[i] += 2 * np.pi
                x2 = (
                    np.mean(x1)
                    + np.linspace(-np.pi, np.pi, len(x0), endpoint=False)
                    + np.pi / len(x0)
                )
                gf = getfac(x0)
                ratio = (maxfac_inner - 1) / (gf - 1)
                x = x2 * (1 - ratio) + np.array(x1) * (ratio)
                x %= 2 * np.pi
            else:
                while getfac(x) > maxfac_inner:
                    extra += 1
                    x += fac * (dx(x, 1) + dx(x, -1))
        else:
            while not np.all(dx(x, -1) > 0):
                x += fac * (dx(x, 1) + dx(x, -1))
                extra += 1
            if extra:
                print(f"Required {extra} extra steps in smoothing!")

        if not np.all(dx(x, -1) > 0):
            plt.plot(x, label="result")
            plt.plot(x0, label="init")
            plt.plot(dx(x0, -1))
            plt.plot(dx(x, -1))
            plt.legend()
            plt.title("Not monotonic!")
            plt.show()
            assert False
        return x

    if (nx > restrict_size) or (nz > restrict_size):
        # Create a coarse grid first to get a starting guess
        # Only restrict the dimensions which exceed restrict_size
        # Note that this might result in multiple levels of resolution

        nx_r = nx
        if nx > restrict_size:
            nx_r = int(nx / restrict_factor)

        nz_r = nz
        if nz > restrict_size:
            nz_r = int(nz / restrict_factor)

        # Create the coarse mesh
        R_r, Z_r = grid_elliptic(
            inner,
            outer,
            nx_r,
            nz_r,
            align=False,
            tol=tol,
            restrict_size=restrict_size,
            restrict_factor=restrict_factor,
            return_coords=True,
            inner_ort=inner_ort,
            maxfac_inner=maxfac_inner,
            dz_relax=dz_relax,
        )

        # Note: Lower case x,z are indices
        z_r = linspace(
            0, 2 * pi, nz_r + 1, endpoint=True
        )  # Add on the final point duplicating the first
        x_r = linspace(0, 1.0, nx_r, endpoint=True)

        # Note: Upper case R,Z are real-space locations
        R_r = np.concatenate((R_r, np.reshape(R_r[:, 0], (nx_r, 1))), axis=1)
        Z_r = np.concatenate((Z_r, np.reshape(Z_r[:, 0], (nx_r, 1))), axis=1)

        # Now interpolate
        spl = RectBivariateSpline(x_r, z_r, R_r)
        R = spl(xvals, thetavals, grid=True)
        spl = RectBivariateSpline(x_r, z_r, Z_r)
        Z = spl(xvals, thetavals, grid=True)

        # Make sure that the inner and outer boundaries are on the
        # inner and outer RZline, not interpolated
        R[-1, :] = outer.Rvalue(thetavals)
        Z[-1, :] = outer.Zvalue(thetavals)

        if inner_ort:
            thetavals_inner = [inner.closestPoint(*x) for x in zip(R[-1], Z[-1])]
            thetavals_inner = laplace(thetavals_inner)

        else:
            thetavals_inner = thetavals

        R[0, :] = inner.Rvalue(thetavals_inner)
        Z[0, :] = inner.Zvalue(thetavals_inner)

    else:
        # Interpolate coordinates of inner and outer boundary
        Router = outer.Rvalue(thetavals)
        Zouter = outer.Zvalue(thetavals)

        if inner_ort:
            thetavals_inner = [inner.closestPoint(*x) for x in zip(Router, Zouter)]
            thetavals_inner = laplace(thetavals_inner)
        else:
            thetavals_inner = thetavals

        Rinner = inner.Rvalue(thetavals_inner)
        Zinner = inner.Zvalue(thetavals_inner)

        # Interpolate in x between inner and outer
        # to get starting guess for a grid
        R = xvals[:, None] * Router[None, :] + (1 - xvals[:, None]) * Rinner[None, :]
        Z = xvals[:, None] * Zouter[None, :] + (1 - xvals[:, None]) * Zinner[None, :]

    dx = xvals[1] - xvals[0]
    dz = thetavals[1] - thetavals[0]

    if show and plotting_available:
        # Markers on original points on inner and outer boundaries
        plt.plot(inner.R, inner.Z, "-o")
        plt.plot(outer.R, outer.Z, "-o")

        # Black lines through inner and outer boundaries
        r, z = inner.position(np.linspace(0, 2 * np.pi, 10 * nz))
        plt.plot(r, z, "k")
        r, z = outer.position(np.linspace(0, 2 * np.pi, 10 * nz))
        plt.plot(r, z, "k")

        # Red dots to mark the inner and outer boundaries
        plt.plot(R[0, :], Z[0, :], "ro")
        plt.plot(R[-1, :], Z[-1, :], "ro")

    # Start solver loop
    for _ in range(R.size * 2):
        # Calculate coefficients, which exclude boundary points
        # Note that the domain is periodic in y so roll arrays

        R_xm = R[:-2, :]  # R(x-1,z)
        R_xp = R[2:, :]  # R(x+1,z)
        R_zm = np.roll(R, 1, axis=1)  # R(x, z-1)
        R_zp = np.roll(R, -1, axis=1)  # R(x, z+1)
        R_xmzm = R_zm[:-2, :]  # R(x-1, z-1)
        R_xpzm = R_zm[2:, :]  # R(x+1, z-1)
        R_xmzp = R_zp[:-2, :]  # R(x-1, z+1)
        R_xpzp = R_zp[2:, :]  # R(x+1, z+1)
        R_zm = R_zm[1:-1, :]  # Now chop off x boundaries
        R_zp = R_zp[1:-1, :]  # This is done to minimise number of rolls

        Z_xm = Z[:-2, :]
        Z_xp = Z[2:, :]
        Z_zm = np.roll(Z, 1, axis=1)
        Z_zp = np.roll(Z, -1, axis=1)
        Z_xmzm = Z_zm[:-2, :]
        Z_xpzm = Z_zm[2:, :]
        Z_xmzp = Z_zp[:-2, :]
        Z_xpzp = Z_zp[2:, :]
        Z_zm = Z_zm[1:-1, :]
        Z_zp = Z_zp[1:-1, :]

        dz_relax = dz_relax or 10 / 3
        dRdz = dz_relax * (R_zp - R_zm) / (2.0 * dz)
        dRdx = (R_xp - R_xm) / (2.0 * dx)

        dZdz = dz_relax * (Z_zp - Z_zm) / (2.0 * dz)
        dZdx = (Z_xp - Z_xm) / (2.0 * dx)

        a = dRdz**2 + dZdz**2
        b = dRdz * dRdx + dZdx * dZdz
        c = dRdx**2 + dZdx**2

        # Now solve a*R_xx - 2*b*R_xz + c*R_zz = 0
        # For now using Jacobi update

        a_dx2 = a / dx**2
        b_dxdz = b / (2.0 * dx * dz)
        c_dz2 = c / dz**2
        inv_diag = 1.0 / (2 * a / dx**2 + 2 * c / dz**2)

        Rold = R.copy()
        Zold = Z.copy()

        R[1:-1, :] = (
            a_dx2 * (R_xm + R_xp)
            - b_dxdz * (R_xpzp - R_xmzp - R_xpzm + R_xmzm)
            + c_dz2 * (R_zm + R_zp)
        ) * inv_diag

        Z[1:-1, :] = (
            a_dx2 * (Z_xm + Z_xp)
            - b_dxdz * (Z_xpzp - Z_xmzp - Z_xpzm + Z_xmzm)
            + c_dz2 * (Z_zm + Z_zp)
        ) * inv_diag

        maxchange_sq = np.amax((R - Rold) ** 2 + (Z - Zold) ** 2)

        if maxchange_sq < tol:
            break

    else:
        print("Convergence failure, trying to plot ...")
        if plotting_available:
            if not show:
                # Markers on original points on inner and outer boundaries
                plt.plot(inner.R, inner.Z, "-o", label="inner")
                plt.plot(outer.R, outer.Z, "-o", label="outer")

                # Black lines through inner and outer boundaries
                r, z = inner.position(np.linspace(0, 2 * np.pi, 10 * nz))
                plt.plot(r, z, "k")
                r, z = outer.position(np.linspace(0, 2 * np.pi, 10 * nz))
                plt.plot(r, z, "k")

                # Red dots to mark the inner and outer boundaries
                plt.plot(R[0, :], Z[0, :], "ro")
                plt.plot(R[-1, :], Z[-1, :], "ro")
            plt.show()

        raise RuntimeError(
            f"Failed to converge - grid shape {R.shape}: Error was {maxchange_sq} > {tol}"
        )

    if nx_outer:
        dn = np.arange(nx_outer) + 1
        dR = R[-1] - R[-2]
        dZ = Z[-1] - Z[-2]
        R = np.vstack([R, R[-1] + dR[None, :] * dn[:, None]])
        Z = np.vstack([Z, Z[-1] + dZ[None, :] * dn[:, None]])
    if nx_inner:
        dn = np.arange(nx_inner)[::-1] + 1
        dR = R[0] - R[1]
        dZ = Z[0] - Z[1]
        R = np.vstack([R[0] + dR[None, :] * dn[:, None], R])
        Z = np.vstack([Z[0] + dZ[None, :] * dn[:, None], Z])

    if show and plotting_available:
        plt.plot(R, Z)
        plt.plot(np.transpose(R), np.transpose(Z))
        plt.show()

    if return_coords:
        return R, Z
    return StructuredPoloidalGrid(R, Z)


if __name__ == "__main__":
    inner = rzline.shaped_line(R0=3.0, a=0.5, elong=1.0, triang=0.0, indent=1.0, n=50)
    outer = rzline.shaped_line(R0=2.8, a=1.5, elong=1.0, triang=0.0, indent=0.2, n=50)

    grid = grid_elliptic(inner, outer, 100, 100, show=True)

    x, z = grid.findIndex([2.0, 1.9], [1.5, 2.0])
    print(x, z)
