import uuid
from collections import namedtuple
from itertools import chain

import numpy as np
from boututils import datafile as bdata

from zoidberg import __version__

from . import fieldtracer
from .grid import Grid
from .poloidal_grid import StructuredPoloidalGrid

try:
    from tqdm.auto import tqdm
except ModuleNotFoundError:
    tqdm = None
    from .progress import update_progress

# PyEVTK might be called pyevtk or evtk, depending on where it was
# installed from
have_evtk = True
try:
    from pyevtk.hl import gridToVTK
except ImportError:
    try:
        from evtk.hl import gridToVTK
    except ImportError:
        have_evtk = False


def parallel_slice_field_name(field, offset):
    """Form a unique, backwards-compatible name for field at a given offset

    Parameters
    ----------
    field : str
        Name of the field to convert
    offset : int
        Parallel slice offset

    """
    prefix = "forward" if offset > 0 else "backward"
    suffix = "_{}".format(abs(offset)) if abs(offset) > 1 else ""
    return "{}_{}{}".format(prefix, field, suffix)


def make_maps(grid, magnetic_field, nslice=1, quiet=False, field_tracer=None, **kwargs):
    """Make the forward and backward FCI maps

    Parameters
    ----------
    grid : :py:obj:`zoidberg.grid.Grid`
        Grid generated by Zoidberg
    magnetic_field : :py:obj:`zoidberg.field.MagneticField`
        Zoidberg magnetic field object
    nslice : int
        Number of parallel slices in each direction
    quiet : bool
        Don't display progress bar
    kwargs
        Optional arguments for field line tracing, etc.

    Returns
    -------
    dict
        Dictionary containing the forward/backward field line maps

    """

    # Get number of points
    # Note: Assumes that all poloidal grids have the same number of x and z(y) points
    ny = grid.numberOfPoloidalGrids()
    pol, _ = grid.getPoloidalGrid(0)
    nx = pol.nx
    nz = pol.nz

    # Get number of guard cells (default 2)
    mxg = kwargs.get("MXG", 2)

    shape = (nx, ny, nz)

    # Coordinates of each grid point
    R = np.zeros(shape)
    Z = np.zeros(shape)

    for j in range(ny):
        pol, _ = grid.getPoloidalGrid(j)
        R[:, j, :] = pol.R
        Z[:, j, :] = pol.Z

    field_tracer = field_tracer or fieldtracer.FieldTracer(magnetic_field)

    rtol = kwargs.get("rtol", None)

    # The field line maps and coordinates, etc.
    maps = {
        "R": R,
        "Z": Z,
        "MXG": mxg,
        "MYG": nslice,
    }

    # A helper data structure that groups the various field line maps along with the offset
    ParallelSlice = namedtuple(
        "ParallelSlice", ["offset", "R", "Z", "xt_prime", "zt_prime"]
    )
    # A list of the above data structures for each offset we want
    parallel_slices_list = []

    # Loop over offsets {1, ... nslice, -1, ... -nslice}
    for direction in 1, -1:
        parallel_slices_list.append([])
        for absoffset in range(1, nslice + 1):
            offset = absoffset * direction
            # Unique names of the field line maps for this offset
            field_names = [
                parallel_slice_field_name(field, offset)
                for field in ["R", "Z", "xt_prime", "zt_prime"]
            ]

            # Initialise the field arrays -- puts them straight into the result dict
            for field in field_names:
                maps[field] = np.zeros(shape)

            # Get the field arrays we just made and wrap them up in our helper tuple
            fields = map(lambda x: maps[x], field_names)
            parallel_slices_list[-1].append(ParallelSlice(offset, *fields))

    # Total size of the progress bar
    total_work = len(parallel_slices_list) * ny

    # TODO: if axisymmetric, don't loop, do one slice and copy
    # TODO: restart tracing for adjacent offsets
    if (not quiet) and (ny > 1):
        if tqdm:
            prog = tqdm(total=total_work, desc="Tracing")
        else:
            update_progress(0, **kwargs)
    for slice_index, parallel_slices in enumerate(parallel_slices_list):
        for j in range(ny):
            # Get this poloidal grid
            pol, ycoord = grid.getPoloidalGrid(j)

            # Get the next poloidal grid
            pol_slices = []
            y_slices = [ycoord]
            for parallel_slice in parallel_slices:
                pol_slice, y_slice = grid.getPoloidalGrid(j + parallel_slice.offset)
                pol_slices.append(pol_slice)
                y_slices.append(y_slice)

            # We only want the end point, as [0,...] is the initial position
            coords = field_tracer.follow_field_lines(pol.R, pol.Z, y_slices, rtol=rtol)[
                1:, ...
            ]

            for parallel_slice, coord, pol_slice in zip(
                parallel_slices, coords, pol_slices
            ):
                # Store the coordinates in real space
                parallel_slice.R[:, j, :] = coord[:, :, 0]
                parallel_slice.Z[:, j, :] = coord[:, :, 1]

                # Get the indices into the slice poloidal grid
                if pol_slice is None:
                    # No slice grid, so hit a boundary
                    xind = -1
                    zind = -1
                else:
                    # Find the indices for these new locations on the slice poloidal grid
                    xcoord = coord[:, :, 0]
                    zcoord = coord[:, :, 1]
                    xind, zind = pol_slice.findIndex(xcoord, zcoord)

                    # Check boundary defined by the field
                    outside = magnetic_field.boundary.outside(xcoord, y_slice, zcoord)
                    xind[outside] = -1
                    zind[outside] = -1

                parallel_slice.xt_prime[:, j, :] = xind
                parallel_slice.zt_prime[:, j, :] = zind

            if (not quiet) and (ny > 1):
                if tqdm:
                    prog.update()
                else:
                    update_progress((slice_index * ny + j + 1) / total_work, **kwargs)

    return maps


def update_metric_names(metric):
    # Translate between output variable names and metric names
    # Map from new to old names. Anything not in this dict
    # is unchanged
    name_changes = {
        "g_yy": "g_22",
        "gyy": "g22",
        "gxx": "g11",
        "gxz": "g13",
        "gzz": "g33",
        "g_xx": "g_11",
        "g_xz": "g_13",
        "g_zz": "g_33",
    }
    return {name_changes.get(key, key): value for key, value in metric.items()}


def get_metric(grid, magnetic_field):
    nx, ny, nz = grid.shape
    # Get metric tensor
    metric = grid.metric()

    # Check if the magnetic field is in cylindrical coordinates
    # If so, we need to change the gyy and g_yy metrics
    pol_grid, ypos = grid.getPoloidalGrid(0)
    Rmaj = magnetic_field.Rfunc(pol_grid.R, pol_grid.Z, ypos)
    if Rmaj is not None:
        assert np.all(np.isfinite(Rmaj))
        # In cylindrical coordinates
        Rmaj = np.zeros(grid.shape)
        for yindex in range(grid.numberOfPoloidalGrids()):
            pol_grid, ypos = grid.getPoloidalGrid(yindex)
            Rmaj[:, yindex, :] = magnetic_field.Rfunc(pol_grid.R, pol_grid.Z, ypos)
        metric["gyy"] = 1.0 / Rmaj**2
        metric["g_yy"] = Rmaj**2

    # Get magnetic field and pressure
    Bmag = np.zeros(grid.shape)
    pressure = np.zeros(grid.shape)

    for yindex in range(grid.numberOfPoloidalGrids()):
        pol_grid, ypos = grid.getPoloidalGrid(yindex)
        Bmag[:, yindex, :] = magnetic_field.Bmag(pol_grid.R, pol_grid.Z, ypos)
        pressure[:, yindex, :] = magnetic_field.pressure(pol_grid.R, pol_grid.Z, ypos)
        By = magnetic_field.Byfunc(pol_grid.R, pol_grid.Z, ypos)
        metric["g_yy"][:, yindex, :] *= (Bmag[:, yindex, :] / By) ** 2
        metric["gyy"][:, yindex, :] *= (By / Bmag[:, yindex, :]) ** 2

    metric["Bxy"] = Bmag
    metric["B"] = Bmag
    metric["pressure"] = pressure

    # B * J / sqrt(g22)
    BJg = Bmag * np.sqrt(metric["g_xx"] * metric["g_zz"] - metric["g_xz"] ** 2)

    return metric, BJg


class MapWriter:
    """
    with MapWriter(...) as mw:
        mw.add_grid_field(grid, field)
        ...
        # Add some additional things
        mw.add_dict(dict(a=f3d, b=f3d))
        ...
        mw.add_maps(maps)
        mw.add_dagp()

    Unlike map_write() this allows to delete objects when they are not needed
    any more. This can help avoid OOM issues for generating large grids.

    It also allows to update grids, if new fields are added to zoidberg, that
    old grids are missing.
    """

    def __init__(
        self,
        gridfile="fci.grid.nc",
        new_names=False,
        metric2d=False,
        quiet=False,
        create=True,
    ):
        self.fn = gridfile
        self.new_names = new_names
        self.metric2d = metric2d
        self.quiet = quiet
        self.is_open = False
        self.create = create

        self.BJg = None
        self.grid = None
        self.field = None
        self.metric_done = False
        self.maps_done = False
        self.final_done = False

    def __enter__(self):
        f = bdata.DataFile(self.fn, write=True, create=self.create)
        self.f = f.__enter__()
        self.is_open = True
        if not self.create:
            if "B" in self.f.list():
                B = self.f["B"]
                self.nxyz = B.shape
        return self

    def open(self):
        return self.__enter__()

    def __exit__(self, *exc_details):
        self.is_open = False
        self.f.__exit__(*exc_details)

    def add_grid(self, grid):
        """Add information from the grid to the grid file"""
        self.grid = grid
        self.nxyz = grid.shape

        shape = self.nxyz
        R = np.zeros(shape)
        Z = np.zeros(shape)
        for j in range(shape[1]):
            pol, _ = self.grid.getPoloidalGrid(j)
            R[:, j, :] = pol.R
            Z[:, j, :] = pol.Z
        self.write_dict(dict(R=R, Z=Z))
        del R
        del Z

        if self.field:
            self._write_metric()

    def add_dagp(self):
        """Add the coefficient for the finite-volume div-a-grad-perp implementation suitable for FCI."""
        from .stencil_dagp_fv import doit

        assert self.grid, "The grid is needed to compute the DAGP. Set the grid first."
        assert self.is_open, "The grid file needs to be open. Call open first."
        handles = {}

        poloidal_grids = self.grid.poloidal_grids
        if tqdm:
            poloidal_grids = tqdm(poloidal_grids, desc="compute DAGP terms")

        def getHandle(k, t=None, dims=("x", "y", "z"), init=None):
            try:
                # See if the variable already exists
                return self.f.impl.handle.variables[k]
            except KeyError:
                # t = v.dtype.str
                # dims = ("x", "y", "z")
                var = self.f.impl.handle.createVariable(k, t, dims)
                if init is not None:
                    var[:] = init
                return var

        i = np.int32(0)
        prog = getHandle("_dagp_generation_progress", i.dtype.str, (), init=i)

        for ind, pol in enumerate(poloidal_grids):
            if ind < prog[0]:
                continue
            dagp = doit([pol])
            for k, v in dagp.items():
                if k not in handles:
                    handles[k] = getHandle(k, v.dtype.str)

                handles[k][:, ind, :] = v
            prog[0] = ind + 1

    def add_field(self, field):
        """Add the information from the field to the grid"""
        self.field = field
        if self.grid:
            self._write_metric()

    def add_grid_field(self, grid, field):
        """Add the information from the grid and the field to the grid file.

        Short for
        mw.add_grid(grid)
        mw.add_field(field)
        """
        self.add_field(field)
        self.add_grid(grid)

    def _write_metric(self):
        if self.metric_done:
            return

        metric, self.BJg = get_metric(self.grid, self.field)

        if not self.new_names:
            metric = update_metric_names(metric)

        self.write_dict(metric, always3d=False)

        self.grid._metric_cache = None

        # Get attributes from magnetic field (e.g. psi)
        attributes = {}
        for name in self.field.attributes:
            attribute = np.zeros(self.grid.shape)
            for yindex in range(self.grid.numberOfPoloidalGrids()):
                pol_grid, ypos = self.grid.getPoloidalGrid(yindex)
                attribute[:, yindex, :] = self.field.attributes[name](
                    pol_grid.R, pol_grid.Z, ypos
                )
            attributes[name] = attribute
        self.write_dict(attributes)

    def write_dict(self, metric, always3d=True):
        """
        Add  data to the grid file"""

        assert self.is_open, "File needs to be open. Call open first."
        # Metric is now 3D
        if self.metric2d and not always3d:
            # Remove the Z dimension from metric components
            if not self.quiet:
                print("WARNING: Outputting 2D metrics, discarding metric information.")
            for key in metric:
                try:
                    metric[key] = metric[key][:, :, 0]
                except TypeError:
                    pass
            # Make dz a constant
            if "dz" in metric:
                metric["dz"] = metric["dz"][0, 0]

        for k, v in metric.items():
            if isinstance(v, np.ndarray):
                assert np.all(
                    np.isfinite(v)
                ), f"{k} is not finite in {v.size - np.sum(np.isfinite(v))} of {v.size} cells"
            self.f.write(k, v)

    def _write_par_metric(self, maps, nslice, ypar):
        # Loop over offsets {1, ... nslice, -1, ... -nslice}
        ny = len(ypar)
        meandy = np.mean(np.diff(ypar))
        Ly = self.grid.Ly if self.grid else meandy * ny
        if self.grid:
            assert np.isclose(
                Ly, meandy * ny
            ), f"Ly of grid (Ly={Ly}) does not seem to match the average dy (ny * dy = {ny} * {meandy} = {ny * meandy}"
        yperiodic = self.grid.yperiodic if self.grid else True
        for offset in chain(range(1, nslice + 1), range(-1, -(nslice + 1), -1)):
            pypar = np.roll(ypar, -offset)  # TODO: is the sign correct?
            if yperiodic:
                for i in range(1, ny):
                    if meandy * (pypar[i] - pypar[i - 1]) < 0:
                        pypar[i] += np.sign(meandy) * Ly
                assert np.isclose(
                    np.mean(np.diff(pypar)), meandy
                ), f"Mean of dy changes from {meandy} to {np.mean(np.diff(pypar))}. Values: {ypar} -> {pypar}"

            RZ = np.array([maps[parallel_slice_field_name(k, offset)] for k in "RZ"])
            par_pgrids = [StructuredPoloidalGrid(*RZ[:, :, i]) for i in range(ny)]
            par_grid = Grid(
                par_pgrids,
                pypar,
                Ly=Ly,
                yperiodic=yperiodic,
            )

            par_metric, par_BJg = get_metric(par_grid, self.field)

            # Check flux conservation
            if self.BJg is not None and self.BJg.shape[0] > 4:
                mymax = np.max(np.abs(par_BJg / self.BJg - 1)[2:-2], axis=(0, 2))
                if np.max(mymax) > 1e-6:
                    print("FluxError", mymax)

            if not self.new_names:
                par_metric = update_metric_names(par_metric)

            self.write_dict(
                {parallel_slice_field_name(k, offset): v for k, v in par_metric.items()}
            )

    def add_maps(self, maps):
        """Add mapping information to the grid field.

        Also adds additionally needed data like parallel metric components and
        grid id.
        """
        if "Rxy" not in maps:
            maps["Rxy"] = maps["R"]
        self.write_dict(maps)
        keep = {}
        nslice = maps["MYG"]
        for offset in chain(range(1, nslice + 1), range(-1, -(nslice + 1), -1)):
            for k in "RZ":
                n = parallel_slice_field_name(k, offset)
                keep[n] = maps[n]
        del maps
        ypar = self.grid.ycoords
        self._write_par_metric(keep, nslice, ypar)
        self._final()

    def _final(self, update=False):
        f = self.f

        grid_id = str(uuid.uuid1())
        attrs = {
            "title": "BOUT++ grid file",
            "software_name": "zoidberg",
            "software_version": __version__,
            "id": grid_id,  # conventional name
            "grid_id": grid_id,  # BOUT++ specific name
        }
        if update:
            oldattrs = f.list_file_attributes()
            for k in list(attrs.keys()):
                if k in oldattrs:
                    oldattr = f"old_{k}"
                    j = 0
                    while oldattr in oldattrs:
                        oldattr = f"old{j}_{k}"
                        j += 1
                    attrs[oldattr] = f.read_file_attribute(k)
        for kv in attrs.items():
            f.write_file_attribute(*kv)

        nx, ny, nz = self.nxyz
        f.write("nx", nx)
        f.write("ny", ny)
        f.write("nz", nz)

        f.write("ixseps1", nx + 1)
        f.write("ixseps2", nx + 1)


def write_maps(
    grid,
    magnetic_field,
    maps,
    gridfile="fci.grid.nc",
    new_names=False,
    metric2d=True,
    format=None,
    quiet=False,
):
    """Write FCI maps to BOUT++ grid file

    Parameters
    ----------
    grid : :py:obj:`zoidberg.grid.Grid`
        Grid generated by Zoidberg
    magnetic_field : :py:obj:`zoidberg.field.MagneticField`
        Zoidberg magnetic field object
    maps : dict
        Dictionary of FCI maps
    gridfile : str, optional
        Output filename
    new_names : bool, optional
        Write "g_yy" rather than "g_22"
    metric2d : bool, optional
        Output only 2D metrics
    format : optional
        deprecated and ignored
    quiet : bool, optional
        Don't warn about 2D metrics

    Returns
    -------

    Writes the following variables to the grid file


    """
    with MapWriter(gridfile, new_names=new_names, metric2d=metric2d, quiet=quiet) as mw:
        mw.add_grid_field(grid, magnetic_field)
        mw.add_maps(maps)


def write_Bfield_to_vtk(
    grid, magnetic_field, scale=5, vtkfile="fci_zoidberg", psi=True
):
    """Write the magnetic field to a VTK file

    Parameters
    ----------
    grid : :py:obj:`zoidberg.grid.Grid`
        Grid generated by Zoidberg
    magnetic_field : :py:obj:`zoidberg.field.MagneticField`
        Zoidberg magnetic field object
    scale : int, optional
        Factor to scale x, z dimensions by [5]
    vtkfile : str, optional
        Output filename without extension ["fci_zoidberg"]
    psi : bool, optional
        Write psi?

    Returns
    -------
    path           - Full path to vtkfile
    """

    point_data = {
        "B": (magnetic_field.bx * scale, magnetic_field.by, magnetic_field.bz * scale)
    }

    if psi:
        psi = make_surfaces(grid, magnetic_field)
        point_data["psi"] = psi

    path = gridToVTK(
        vtkfile,
        grid.xarray * scale,
        grid.yarray,
        grid.zarray * scale,
        pointData=point_data,
    )

    return path


def fci_to_vtk(infile, outfile, scale=5):
    if not have_evtk:
        return

    with bdata.DataFile(infile, write=False, create=False) as f:
        dx = f.read("dx")
        dy = f.read("dy")

        bx = f.read("bx")
        by = np.ones(bx.shape)
        bz = f.read("bz")
        if bx is None:
            xt_prime = f.read("forward_xt_prime")
            zt_prime = f.read("forward_zt_prime")
            array_indices = np.indices(xt_prime.shape)
            bx = xt_prime - array_indices[0, ...]
            by = by * dy
            bz = zt_prime - array_indices[2, ...]

        nx, ny, nz = bx.shape
        dz = nx * dx / nz

    x = np.linspace(0, nx * dx, nx)
    y = np.linspace(0, ny * dy, ny, endpoint=False)
    z = np.linspace(0, nz * dz, nz, endpoint=False)

    gridToVTK(
        outfile, x * scale, y, z * scale, pointData={"B": (bx * scale, by, bz * scale)}
    )


def make_surfaces(grid, magnetic_field, nsurfaces=10, revs=100):
    """Essentially interpolate a poincare plot onto the grid mesh

    Parameters
    ----------
    grid : :py:obj:`zoidberg.grid.Grid`
        Grid generated by Zoidberg
    magnetic_field : :py:obj:`zoidberg.field.MagneticField`
        Zoidberg magnetic field object
    nsurfaces : int, optional
        Number of surfaces to interpolate to [10]
    revs : int, optional
        Number of points on each surface [100]

    Returns
    -------
    surfaces
        Array of psuedo-psi on the grid mesh

    """

    from scipy.interpolate import griddata

    # initial x, z points in surface
    xpos = grid.xcentre + np.linspace(0, 0.5 * np.max(grid.xarray), nsurfaces)
    zpos = grid.zcentre

    phi_values = grid.yarray[:]
    # Extend the domain from [0,grid.Ly] to [0,revs*grid.Ly]
    for n in np.arange(1, revs):
        phi_values = np.append(phi_values, n * grid.Ly + phi_values[: grid.ny])

    # Get field line tracer and trace out surfaces
    tracer = fieldtracer.FieldTracer(magnetic_field)
    points = tracer.follow_field_lines(xpos, zpos, phi_values)

    # Reshape to be easier to work with
    points = points.reshape((revs, grid.ny, nsurfaces, 2))

    # Arbitarily number the surfaces from 0 to 1
    psi_points = np.zeros((revs, grid.ny, nsurfaces))
    for surf in range(nsurfaces):
        psi_points[:, :, surf] = float(surf) / float(nsurfaces - 1)

    x_2d, z_2d = np.meshgrid(grid.xarray, grid.zarray, indexing="ij")

    psi = np.zeros_like(grid.x_3d)
    for y_slice in range(grid.ny):
        points_2d = np.column_stack(
            (points[:, y_slice, :, 0].flatten(), points[:, y_slice, :, 1].flatten())
        )
        psi[:, y_slice, :] = griddata(
            points_2d,
            psi_points[:, y_slice, :].flatten(),
            (x_2d, z_2d),
            method="linear",
            fill_value=1,
        )

    return psi


def upscale(field, maps, upscale_factor=4, quiet=True):
    """Increase the resolution in y of field along the FCI maps.

    First, interpolate onto the (forward) field line end points, as in
    normal FCI technique. Then interpolate between start and end
    points. We also need to interpolate the xt_primes and
    zt_primes. This gives a cloud of points along the field lines,
    which we can finally interpolate back onto a regular grid.

    Parameters
    ----------
    field : array_like
        3D field to be upscaled
    maps : dict
        Zoidberg field line maps
    upscale_factor : int, optional
        Factor to increase resolution by [4]
    quiet : bool, optional
        Don't show progress bar [True]

    Returns
    -------
    Field with y-resolution increased *upscale_factor* times. Shape is
    (nx, upscale_factor*ny, nz).

    """

    from scipy.interpolate import griddata
    from scipy.ndimage.interpolation import map_coordinates

    xt_prime = maps["forward_xt_prime"]
    zt_prime = maps["forward_zt_prime"]

    # The field should be the same shape as the grid
    if field.shape != xt_prime.shape:
        try:
            field = field.reshape(xt_prime.T.shape).T
        except ValueError:
            raise ValueError(
                "Field, {}, must be same shape as grid, {}".format(
                    field.shape, xt_prime.shape
                )
            )

    # Get the shape of the grid
    nx, ny, nz = xt_prime.shape
    index_coords = np.mgrid[0:nx, 0:ny, 0:nz]

    # We use the forward maps, so get the y-index of the *next* y-slice
    yup_3d = index_coords[1, ...] + 1
    yup_3d[:, -1, :] = 0

    # Index space coordinates of the field line end points
    end_points = np.array([xt_prime, yup_3d, zt_prime])

    # Interpolation of the field at the end points
    field_prime = map_coordinates(field, end_points)

    # This is a 4D array where the first dimension is the start/end of
    # the field line
    field_aligned = np.array([field, field_prime])

    # x, z coords at start/end of field line
    x_start_end = np.array([index_coords[0, ...], xt_prime])
    z_start_end = np.array([index_coords[2, ...], zt_prime])

    # Parametric points along the field line
    midpoints = np.linspace(0, 1, upscale_factor, endpoint=False)
    # Need to make this 4D as well
    new_points = np.tile(midpoints[:, np.newaxis, np.newaxis, np.newaxis], [nx, ny, nz])

    # Index space coordinates of our upscaled field
    index_4d = np.mgrid[0:upscale_factor, 0:nx, 0:ny, 0:nz]
    hires_points = np.array(
        [new_points, index_4d[1, ...], index_4d[2, ...], index_4d[3, ...]]
    )

    # Upscale the field
    hires_field = map_coordinates(field_aligned, hires_points)

    # Linearly interpolate the x, z coordinates of the field lines
    hires_x = map_coordinates(x_start_end, hires_points)
    hires_z = map_coordinates(z_start_end, hires_points)

    def twizzle(array):
        """Transpose and reshape the output of map_coordinates to
        be 3D
        """
        return array.transpose((1, 2, 0, 3)).reshape((nx, upscale_factor * ny, nz))

    # Rearrange arrays to be 3D
    hires_field = twizzle(hires_field)
    hires_x = twizzle(hires_x)
    hires_z = twizzle(hires_z)

    # Interpolate from field line sections onto grid
    hires_grid_field = np.zeros((nx, upscale_factor * ny, nz))
    hires_index_coords = np.mgrid[0:nx, 0 : ny : 1.0 / upscale_factor, 0:nz]
    grid_points = (hires_index_coords[0, :, 0, :], hires_index_coords[2, :, 0, :])

    def y_first(array):
        """Put the middle index first"""
        return array.transpose((0, 2, 1))

    # The hires data is unstructed only in (x,z), interpolate onto
    # (x,z) grid for each y-slice individually
    for k, (x_points, z_points, f_slice) in enumerate(
        zip(y_first(hires_x).T, y_first(hires_z).T, y_first(hires_field).T)
    ):
        points = np.column_stack((x_points.flat, z_points.flat))
        hires_grid_field[:, k, :] = griddata(
            points, f_slice.flat, grid_points, method="linear", fill_value=0.0
        )
        if not quiet:
            update_progress(float(k) / float(ny - 1))

    return hires_grid_field
