"""
The sphere the globe views are drawn on: relief, coordinate conversion, and the
surface mesh.

Everything here is plain numpy. The vedo-facing code lives in `scene.py` and
`shell.py`, so this module can be tested without a render window.
"""

import numpy as np
from matplotlib.colors import LogNorm, Normalize

from ...paths import planet_radius
from ...topography import load_topography

try:
    from scipy.interpolate import RegularGridInterpolator
    scipy_available = True
except ImportError:
    scipy_available = False


# Vertical exaggeration of the relief. On Mars, MOLA's -8201..18168 m over a
# 3389.5 km radius then spans 0.976 R .. 1.054 R, enough to read the shield
# volcanoes and Hellas without visibly distorting the sphere.
RELIEF_EXAGGERATION = 10

# Vertical exaggeration of the atmosphere above it, which is a larger number for
# a reason. The weather is in the bottom scale height: a water ice cloud in the
# sample file reaches 10 km, which at the relief's own factor is 100 km against
# a 3390 km sphere - a film on the ground, and one drawn *inside* terrain
# exaggerated to 182 km. At x30 the same cloud stands 300 km off the surface,
# a tenth of the radius, which is where it can be seen for what it is.
#
# The two factors differing is the price: a cloud over Olympus Mons can be drawn
# below a summit it is physically above. The figure says both numbers, and
# --shell-exaggeration sets this one - pass 10 to get the old shared scale back.
AIR_EXAGGERATION = 30

# Colour for cells the data does not cover. VTK's default is a dark red that
# reads as a value rather than as a gap, which matters here: a regional field
# leaves most of the globe uncovered.
NO_DATA_COLOR = (0.35, 0.35, 0.35, 1.0)

# What an uncovered cell looks like on a shell, which has nothing to be neutral
# against: a hole rather than a grey. The globe behind it is the answer.
NO_DATA_TRANSPARENT = (0.0, 0.0, 0.0, 0.0)

# Peak height of a relief built from data rather than from topography, as a
# fraction of the planet radius. Such a field carries no metric height, so it is
# normalized into a range that reads like terrain instead of like a spike.
DATA_RELIEF_FRACTION = 0.05


def finite_range(values):
    """
    (min, max) over the finite entries, or (0, 1) when there are none.
    """
    finite = np.asarray(values)[np.isfinite(values)]
    if finite.size == 0:
        return 0.0, 1.0
    return float(finite.min()), float(finite.max())


def _cell_fraction(samples, axis, periodic):
    """
    (index of the cell each sample is in, how far across it), for a regular axis.

    The axis holds cell centres and may run either way; a periodic one wraps a
    turn later, so a sample past the last longitude belongs to the cell that
    closes onto the first. A sample outside a non-periodic axis is held at the
    nearest edge, which is what `elevation_at` does with a line drawn to the
    pole - the grid stops at 89.5 and the cap above it is built from that row.
    """
    step = float(axis[1] - axis[0]) if axis.size > 1 else 1.0
    position = (samples - axis[0]) / step
    if periodic:
        position = np.mod(position, axis.size)
        index = np.floor(position).astype(int) % axis.size
    else:
        position = np.clip(position, 0.0, axis.size - 1)
        index = np.clip(np.floor(position).astype(int), 0, max(axis.size - 2, 0))
    return index, position - index


def densify(lats, lons, max_step):
    """
    A polyline resampled so no segment spans more than `max_step` degrees.

    A line is straight between its vertices while the ground under it is not, so
    a segment that crosses a whole cell cuts the corner off whatever is inside -
    including the diagonal where the two triangles of that cell meet. Splitting
    it leaves only the twist of half a cell, tens of metres, and costs one
    interpolation over a few thousand points.

    The graticule needs none of this: its vertices are grid nodes and its
    segments run along the cell edges, which are the mesh's own edges.
    """
    lats = np.asarray(lats, dtype=float)
    lons = np.asarray(lons, dtype=float)
    if lats.size < 2:
        return lats, lons

    # Longitudes are unrolled first, so a segment across the date line is one
    # short step rather than a sweep back round the planet
    steps = np.diff(lons)
    unrolled = np.concatenate([lons[:1], lons[0] + np.cumsum(
        steps - 360.0 * np.round(steps / 360.0))])

    spans = np.maximum(np.abs(np.diff(lats)), np.abs(np.diff(unrolled)))
    cuts = np.maximum(1, np.ceil(spans / float(max_step)).astype(int))
    if not np.any(cuts > 1):
        return lats, lons

    out_lat, out_lon = [], []
    for start in range(lats.size - 1):
        fraction = np.linspace(0.0, 1.0, cuts[start] + 1)[:-1]
        out_lat.append(lats[start] + fraction * (lats[start + 1] - lats[start]))
        out_lon.append(unrolled[start] + fraction * (unrolled[start + 1] - unrolled[start]))
    out_lat.append(lats[-1:])
    out_lon.append(unrolled[-1:])
    return np.concatenate(out_lat), np.concatenate(out_lon)


def normalize_relief(values, norm=None, radius=planet_radius):
    """
    Turn an arbitrary field into elevations in metres.

    `--globe-relief <var>` shapes the globe by a field whose units are not
    lengths at all - an ice mass, a pressure - so the values are mapped onto a
    fixed fraction of the radius. The result is readable as terrain; it is not
    an altitude, which is why the caller says so in the figure text.

    `norm` is the colour scale of the field, used when the relief and the colour
    show the same variable. It matters: `h2oice` spans ten decades, so a linear
    stretch puts all but one cell on the floor and the globe comes out smooth,
    while the log scale that colours it produces relief that matches what the
    eye already reads from the colour.
    """
    values = np.asarray(values, dtype=float)
    if norm is not None:
        scaled = np.ma.filled(np.ma.masked_invalid(norm(values)), 0.0)
        scaled = np.clip(np.asarray(scaled, dtype=float), 0.0, 1.0)
    else:
        low, high = finite_range(values)
        if high <= low:
            return np.zeros_like(values)
        scaled = (values - low) / (high - low)
    return np.nan_to_num(scaled, nan=0.0) * radius * DATA_RELIEF_FRACTION / RELIEF_EXAGGERATION


class Globe:
    """
    A sphere of a given radius carrying relief, and the conversions between
    geographic and cartesian coordinates on it.

    `elevation` is always in metres above the datum and is multiplied by
    `exaggeration` on the way into cartesian space, so a caller that knows its
    own elevations never has to think about the exaggeration itself. An
    `altitude` above the ground is multiplied by `air_exaggeration` instead: the
    atmosphere is stretched harder than the terrain, for the reason written up
    at AIR_EXAGGERATION.
    """

    def __init__(self, relief=None, lats=None, lons=None,
                 radius=planet_radius, exaggeration=RELIEF_EXAGGERATION,
                 air_exaggeration=None):
        self.radius = float(radius)
        self.exaggeration = float(exaggeration)
        self.air_exaggeration = float(air_exaggeration if air_exaggeration is not None
                                      else AIR_EXAGGERATION)
        self.relief = relief
        self.lats = lats
        self.lons = lons
        self._interp = None
        if relief is not None and lats is not None and lons is not None and scipy_available:
            # The first column repeated a turn later, so the seam is a cell like
            # any other. Without it the half degree between the last longitude
            # and the first is off the end of the axis, and a line crossing the
            # date line - the 180 degree meridian, for one - is placed at the
            # fill value, sea level, whatever the ground there is doing.
            self._interp = RegularGridInterpolator(
                (lats, np.append(lons, lons[0] + 360.0)),
                np.concatenate([relief, relief[:, :1]], axis=1),
                bounds_error=False, fill_value=0.0)

    @classmethod
    def from_topography(cls, **kwargs):
        """
        The MOLA globe, or a bare sphere when the topography file is missing.
        """
        topo = load_topography()
        if topo is None:
            return cls(**kwargs)
        return cls(relief=topo.values, lats=topo.lats, lons=topo.lons, **kwargs)

    @classmethod
    def bare(cls, **kwargs):
        """
        A globe with no relief at all, for `--globe-relief none`.
        """
        return cls(**kwargs)

    def elevation_at(self, lat, lon):
        """
        Relief in metres at the given samples, zero when there is no relief.

        Latitudes are held inside the grid before being asked for. The grid
        holds cell centres, so it stops at +/-89.5, and anything past that -
        a line drawn to the pole, a terminator crossing it - fell out of bounds
        and came back as the interpolator's fill value, sea level, in the middle
        of terrain kilometres high. Clamping reads the nearest real cell
        instead, which is what the pole cap on the surface mesh is built from
        too.
        """
        if self._interp is None:
            return np.zeros(np.shape(lat), dtype=float)
        edges = (min(self.lats[0], self.lats[-1]), max(self.lats[0], self.lats[-1]))
        start = float(self.lons[0])
        wrapped = start + np.mod(np.asarray(lon, dtype=float) - start, 360.0)
        return self._interp((np.clip(lat, *edges), wrapped))

    def mesh_elevation_at(self, lat, lon):
        """
        The elevation of the surface as it is *drawn*, in metres.

        `elevation_at` gives the smooth bilinear surface through the grid nodes;
        the globe is drawn as flat triangles between those same nodes, and the
        two differ inside a cell by its twist, `|z00 + z11 - z01 - z10| / 4`,
        which over MOLA reaches 1977 m - twenty kilometres once exaggerated.
        That gap is why anything laid on the surface used to be lifted so far
        clear of it that it read as floating rather than as lying on the ground.

        So this reads the triangles instead. `build_surface` splits each cell
        along the p1-p2 anti-diagonal, which in cell fractions `s` (down the
        latitudes) and `t` (along the longitudes) is the line s + t = 1: below
        it the triangle spans the two edges out of the top-left corner, above it
        the two out of the bottom-right one. A point interpolated inside its own
        triangle is *on* the surface, not near it.
        """
        if self.relief is None or self.lats is None or self.lons is None:
            return np.zeros(np.shape(lat), dtype=float)

        relief = np.asarray(self.relief, dtype=float)
        cols = relief.shape[1]
        lats = np.asarray(self.lats, dtype=float)
        lons = np.asarray(self.lons, dtype=float)

        # Cell fractions along each axis. The latitudes usually descend, so the
        # row index runs with the array rather than with the latitude.
        i, s = _cell_fraction(np.asarray(lat, dtype=float), lats, periodic=False)
        j, t = _cell_fraction(np.asarray(lon, dtype=float), lons, periodic=True)

        z00 = relief[i, j]
        z01 = relief[i, (j + 1) % cols]
        z10 = relief[i + 1, j]
        z11 = relief[i + 1, (j + 1) % cols]

        lower = z00 + t * (z01 - z00) + s * (z10 - z00)
        upper = z11 + (1.0 - t) * (z10 - z11) + (1.0 - s) * (z01 - z11)
        return np.where(s + t <= 1.0, lower, upper)

    def on_surface(self, lat, lon, lift):
        """
        Cartesian points on the ground, `lift` metres above the mesh below them.

        Anything drawn *on* the globe - a meridian, a contour, the terminator -
        has to clear a mesh of flat triangles, and `mesh_elevation_at` reads
        those triangles, so the lift is only what keeps one line from fighting
        the surface, or another line, for the depth buffer. It used to have to
        cover the difference between the mesh and the smooth surface as well,
        which is why it was ten times larger and why the lines read as hovering.
        """
        elevation = self.mesh_elevation_at(lat, lon) + lift
        return self.to_cartesian(lat, lon, elevation=elevation)

    def radius_at(self, lat, lon, scale=1.0, elevation=None):
        """
        Distance from the centre, with the relief exaggerated.
        """
        if elevation is None:
            elevation = self.elevation_at(lat, lon)
        return (self.radius + np.asarray(elevation, dtype=float) * self.exaggeration) * scale

    def to_cartesian(self, lat, lon, scale=1.0, elevation=None, altitude=None):
        """
        Map lat/lon samples onto the exaggerated sphere, returning (N, 3).

        `elevation` overrides the relief lookup, which is what the surface mesh
        does with its own grid: exact, and it skips interpolating the
        topography back onto the grid it came from. `altitude` adds metres above
        the relief, which is how the atmospheric shell is lifted off the
        surface; it carries the air exaggeration rather than the relief's own,
        and may be an array of any shape the lat/lon samples broadcast against -
        one number per level, or one per column of a terrain-following shell.
        """
        lat = np.asarray(lat, dtype=float)
        lon = np.asarray(lon, dtype=float)
        phi = np.deg2rad(90.0 - lat)
        theta = np.deg2rad(lon)
        rr = self.radius_at(lat, lon, scale=scale, elevation=elevation)
        if altitude is not None:
            rr = rr + np.asarray(altitude, dtype=float) * self.air_exaggeration * scale
        return np.column_stack([
            (rr * np.sin(phi) * np.cos(theta)).ravel(),
            (rr * np.sin(phi) * np.sin(theta)).ravel(),
            (rr * np.cos(phi)).ravel()
        ])

    @staticmethod
    def spherical(xyz):
        """
        (latitude, longitude, distance from the centre) of cartesian points.

        The half of the inverse that is pure trigonometry, shared by the two
        readings of the other half: a height above the datum and a height above
        the ground are the same radius divided by different things.
        """
        xyz = np.atleast_2d(np.asarray(xyz, dtype=float))
        r = np.linalg.norm(xyz, axis=1)
        safe = np.where(r > 0, r, 1.0)
        lat = 90.0 - np.rad2deg(np.arccos(np.clip(xyz[:, 2] / safe, -1.0, 1.0)))
        lon = np.rad2deg(np.arctan2(xyz[:, 1], xyz[:, 0]))
        return lat, lon, r

    def to_geographic(self, xyz):
        """
        The inverse of to_cartesian: (lat, lon, height above the datum in metres).

        Used by the click readout, which gets a point back from vedo and has to
        say where on the planet it is.
        """
        lat, lon, r = self.spherical(xyz)
        height = (r - self.radius) / self.exaggeration
        if r.size == 1:
            return float(lat[0]), float(lon[0]), float(height[0])
        return lat, lon, height

    def altitude_at(self, lat, lon, radius):
        """
        Metres of air under a point at `radius`: the inverse of
        to_cartesian(altitude=), in geographic coordinates.

        Not the same quantity as to_geographic's third value, which is a height
        above the datum on the relief's scale. A point of the shell carries both
        exaggerations - the ground it stands on at one, its own altitude at the
        other - so the relief has to come off before the air factor is divided
        out. Reading it the other way, as `--shell-color height` used to, called
        an isosurface over Olympus Mons three times as high as the same surface
        over the plain beside it.

        Takes lat/lon rather than points because the cloud sampler has them
        already, in slabs, and converting those to points and back is the whole
        cost of a 256^3 box.
        """
        ground = self.elevation_at(lat, lon) * self.exaggeration
        return (np.asarray(radius, dtype=float) - self.radius - ground) / self.air_exaggeration

    def height_above_ground(self, xyz):
        """
        `altitude_at` for cartesian points, which is what a picked mesh gives.
        """
        lat, lon, r = self.spherical(xyz)
        return self.altitude_at(lat, lon, r)

    def highest_ground(self):
        """
        Metres of the highest ground drawn above the datum, zero for a globe
        with no relief at all.
        """
        if self.relief is None:
            return 0.0
        relief = np.asarray(self.relief, dtype=float)
        if not np.any(np.isfinite(relief)):
            return 0.0
        return float(max(np.nanmax(relief), 0.0))

    def clearance_altitude(self, margin=1.0):
        """
        The altitude, measured from the datum, at which a shell clears every
        summit on the planet.

        The relief and the air are stretched by different amounts, so the answer
        is not the height of the ground: at x10 against x30, Olympus Mons is
        cleared by a third of its own height in air.
        """
        return margin * self.highest_ground() * self.exaggeration / self.air_exaggeration


def build_surface(globe, lat2d, lon2d, elevation, scalars):
    """
    Vertices, triangles and per-vertex scalars for the surface globe.

    The longitude seam is closed by wrapping the last column onto the first, and
    each pole is capped with a triangle fan: the topography grid holds cell
    centres, so it stops at +/-89.5 and would otherwise leave a hole about 30 km
    across at each pole. The two fans wind in opposite directions so both end up
    facing outwards, like the body faces.
    """
    nlat, nlon = elevation.shape
    pts = globe.to_cartesian(lat2d, lon2d, elevation=elevation)

    i = np.arange(nlat - 1)[:, None]
    j = np.arange(nlon)[None, :]
    p0 = i * nlon + j
    p1 = i * nlon + (j + 1) % nlon
    p2 = p0 + nlon
    p3 = p1 + nlon
    faces = np.concatenate([
        np.stack(np.broadcast_arrays(p0, p2, p1), axis=-1).reshape(-1, 3),
        np.stack(np.broadcast_arrays(p1, p2, p3), axis=-1).reshape(-1, 3),
    ])

    north, south = nlat * nlon, nlat * nlon + 1
    pts = np.vstack([
        pts,
        globe.to_cartesian(90.0, 0.0, elevation=elevation[0].mean()),
        globe.to_cartesian(-90.0, 0.0, elevation=elevation[-1].mean()),
    ])

    jj = np.arange(nlon)
    jn = (jj + 1) % nlon
    last = (nlat - 1) * nlon
    faces = np.concatenate([
        faces,
        np.stack([np.full(nlon, north), jj, jn], axis=-1),
        np.stack([np.full(nlon, south), last + jn, last + jj], axis=-1),
    ])

    return pts, faces, surface_scalars(scalars)


def surface_scalars(values):
    """
    Per-vertex scalars for the surface mesh, with a value for each pole apex.

    Split out from build_surface because an animation re-colours the same
    geometry every frame and only needs this part.
    """
    values = np.asarray(values, dtype=float)
    return np.concatenate([
        values.ravel(),
        [_ring_mean(values[0]), _ring_mean(values[-1])],
    ])


def _ring_mean(row):
    """
    Mean of a pole-adjacent row, NaN when the whole row is missing.
    """
    row = np.asarray(row, dtype=float)
    finite = row[np.isfinite(row)]
    return float(finite.mean()) if finite.size else np.nan


def apply_colormap(mesh, colormap, values, norm, title, alpha=1.0,
                   no_data=NO_DATA_COLOR):
    """
    Colour a vedo mesh the way the 2D map that preceded it was coloured, and
    return the scalar-bar title.

    vedo drives a VTK lookup table, which offers a linear and a log scale and
    nothing else. The two norms it can express are handed over as vmin/vmax
    (plus `logscale`); any other norm is applied to the data here instead, and
    the title then carries the real range because the bar itself runs over the
    normalized 0..1.

    `alpha` may be a per-vertex array. `Mesh.alpha()` only takes a scalar, so
    varying transparency has to ride along with the colours.

    `no_data` is what an uncovered cell comes out as. It belongs here rather
    than in the callers because VTK's own answer - an opaque dark red - is a
    value-looking colour that has to be overridden every single time, and doing
    it by convention meant a mode that forgot drew a red wedge over the pole.
    A per-vertex alpha cannot help: a NaN never reaches the alpha array, it
    reaches the lookup table's NaN colour.
    """
    values = np.asarray(values, dtype=float)
    bar_title = _colour(mesh, colormap, values, norm, title, alpha)
    if no_data is not None:
        set_no_data_color(mesh, no_data)
    return bar_title


def _colour(mesh, colormap, values, norm, title, alpha):
    """
    The colouring itself, split out so `apply_colormap` has one exit.
    """
    if norm is None:
        mesh.cmap(colormap, values, alpha=alpha)
        return title

    data_low, data_high = finite_range(values)
    low = data_low if norm.vmin is None else float(norm.vmin)
    high = data_high if norm.vmax is None else float(norm.vmax)

    if isinstance(norm, LogNorm):
        mesh.cmap(colormap, values, vmin=low, vmax=high, logscale=True, alpha=alpha)
        return title

    # `type` rather than `isinstance`: every Normalize subclass other than the
    # plain one bends the scale in a way VTK cannot follow, so they belong on
    # the pre-normalized path below
    if type(norm) is Normalize:
        mesh.cmap(colormap, values, vmin=low, vmax=high, alpha=alpha)
        return title

    scaled = np.ma.filled(np.ma.masked_invalid(norm(values)), np.nan)
    mesh.cmap(colormap, np.asarray(scaled, dtype=float), vmin=0.0, vmax=1.0, alpha=alpha)
    return f"{title} ({low:.3g} to {high:.3g})"


def add_bar(actor, title, pos=None, **kwargs):
    """
    Attach a scalar bar, letting vedo place it when no position is given.

    vedo takes the length of `pos`, so a None has to be left out of the call
    rather than passed through as "wherever you like".

    The 'nan' swatch vedo adds whenever the lookup table carries a no-data
    colour is turned back off. On the surface it repeats what the figure text
    already says; on a shell, whose gaps are transparent, it would draw a solid
    black patch to stand for something invisible.
    """
    if pos is not None:
        kwargs['pos'] = pos
    actor.add_scalarbar(title=title, c='white', **kwargs)
    bar = getattr(actor, 'scalarbar', None)
    if bar is not None and hasattr(bar, 'DrawNanAnnotationOff'):
        bar.DrawNanAnnotationOff()
    return actor


def set_no_data_color(mesh, color=NO_DATA_COLOR):
    """
    Paint the cells the data does not reach in a neutral grey.
    """
    lut = mesh.mapper.GetLookupTable()
    if lut is not None:
        lut.SetNanColor(*color)