"""
Vector fields drawn over a map.

`--vector U,V` names the two components explicitly rather than guessing them,
because the naming is not consistent across the models this tool reads: GCM
winds are `u`/`v`, XIOS output uses long descriptive names, and PEM restarts
carry slope-resolved fluxes.

Both components must share a grid. `start.nc` holds a real pair,
`ucov(Time, altitude, latitude, rlonu)` and `vcov(Time, altitude, rlatv, longitude)`,
on a staggered Arakawa C-grid: the two live on different dimensions and cannot
be drawn together without interpolating one onto the other's grid.
De-staggering is deliberately out of scope, so such a pair is refused with a
message naming both sets of dimensions rather than silently plotted wrong.
"""

import numpy as np
import matplotlib.pyplot as plt
import cartopy.crs as ccrs

from ..figure import axes_for, finish_figure
from ..topography import overlay_topography
from . import extras
from .geomap import _equally_spaced, _normalize_longitudes
from .layers import mesh_shading, paint_blank_cells, share_wrapped_norm


class VectorError(Exception):
    """
    A vector pair that cannot be drawn, reported rather than approximated.
    """


def check_pair(u_da, v_da):
    """
    Refuse component pairs that do not share a grid.
    """
    if u_da.dims != v_da.dims:
        raise VectorError(
            f"vector components live on different grids: "
            f"'{u_da.name}' has dimensions {tuple(u_da.dims)} while "
            f"'{v_da.name}' has {tuple(v_da.dims)}. "
            f"Staggered (Arakawa C) pairs such as ucov/vcov would need "
            f"de-staggering, which this tool does not do.")
    if u_da.shape != v_da.shape:
        raise VectorError(
            f"vector components have different shapes: {u_da.shape} vs {v_da.shape}")


def _subsample(n, target):
    """
    Stride that keeps roughly `target` arrows along an axis, so a 669-point grid
    does not become a black rectangle.
    """
    return max(1, int(np.ceil(n / float(target))))


def render_vectors(ctx, u, v, style='quiver', density=None, background='magnitude'):
    """
    Draw a vector field over a map. `ctx.data` carries the background scalar,
    or None when only the vectors are wanted.
    """
    plan = ctx.plan
    lons = plan.x.values if plan.x.values is not None else np.arange(u.shape[1], dtype=float)
    lats = plan.y.values if plan.y.values is not None else np.arange(u.shape[0], dtype=float)

    magnitude = np.hypot(u, v)
    field = magnitude if (background == 'magnitude' or ctx.data is None) else ctx.data

    lons_sorted, stacked = _normalize_longitudes(
        lons, np.stack([u, v, magnitude] + ([field] if field is not None else [])))
    u, v, magnitude = stacked[0], stacked[1], stacked[2]
    field = stacked[3] if field is not None else None

    proj = ccrs.PlateCarree()
    fig, ax = axes_for(ctx, (9, 6), projection=proj)

    if field is not None and background != 'none':
        paint_blank_cells(ax, lons_sorted, lats, field, ctx.colormap, transform=proj)
        mesh = share_wrapped_norm(
            ax.pcolormesh(lons_sorted, lats, field,
                          shading=mesh_shading(ctx.interpolate, lons_sorted,
                                               lats, field),
                          cmap=ctx.colormap, norm=ctx.norm, transform=proj))
        cbar = fig.colorbar(mesh, ax=ax, pad=0.02)
        cbar.set_label(ctx.label if background != 'magnitude'
                       else f"|({ctx.varname})|" + (f" ({ctx.units})" if ctx.units else ''))

    if style == 'stream':
        # streamplot needs a strictly ascending, evenly spaced grid
        if not _equally_spaced(lons_sorted):
            raise VectorError("--vector-style stream needs an evenly spaced "
                              "longitude axis; use quiver instead")
        order = np.argsort(lats)
        ax.streamplot(lons_sorted, np.asarray(lats)[order],
                      u[order], v[order],
                      color='k', linewidth=0.7, density=density or 1.5,
                      transform=proj)
    else:
        step_x = _subsample(lons_sorted.size, density or 30)
        step_y = _subsample(len(lats), density or 20)
        ax.quiver(lons_sorted[::step_x], np.asarray(lats)[::step_y],
                  u[::step_y, ::step_x], v[::step_y, ::step_x],
                  transform=proj, scale_units='xy', angles='xy',
                  width=0.0025, color='k')

    if ctx.show_topo:
        overlay_topography(ax, transform=proj, levels=10)

    gl = ax.gridlines(draw_labels=True, linewidth=0.4, color='gray',
                      alpha=0.6, linestyle='--')
    gl.top_labels = False
    gl.right_labels = False
    ax.set_title(ctx.title, fontweight='bold')

    # As on a map, a failure to write is carried past the secondary views
    # rather than cutting them off: they write their own files.
    status = finish_figure(fig, ctx.output_path, dpi=ctx.dpi)

    # The globe draws the same wind pair as arrows tangent to the sphere, which
    # is the one view where a flow over a pole is not torn apart by the map
    # projection
    if ctx.globe is not None:
        ctx.globe.wind_u, ctx.globe.wind_v = u, v
    lon2d, lat2d = np.meshgrid(lons_sorted, lats)
    extras.show_extra_views(lon2d, lat2d, magnitude if field is None else field,
                            ctx.colormap, ctx.varname, ctx.units,
                            ctx.interactive, ctx.show_polar, ctx.show_3d,
                            ctx.show_topo, ctx.output_path, norm=ctx.norm,
                            globe_options=ctx.globe)
    return status