"""
Two-dimensional cross-sections and Hovmoller diagrams: any 2D field that is not
a geographic map.
"""

import numpy as np
import matplotlib.pyplot as plt

from .. import overlay
from ..figure import axes_for, finish_figure, attach_format_coord
from . import layers as layer_draw


def render_section(ctx):
    """
    A pcolormesh on plain axes, with the vertical coordinate inverted when it
    measures depth.

    A section is a pcolormesh like a map is, so `--overlay` costs it nothing
    beyond passing the layers to the same drawing code, on plain axes instead of
    a projection.
    """
    plan = ctx.plan
    x = plan.x.values if plan.x.values is not None else np.arange(ctx.data.shape[1])
    y = plan.y.values if plan.y.values is not None else np.arange(ctx.data.shape[0])

    prepared = [(layer, np.asarray(layer.data, dtype=float), layer.frames)
                for layer in (ctx.composite.layers[1:] if ctx.composite else ())]

    # Wider per colour bar that will actually be drawn. A contour or a hatch
    # layer needs none, so those styles give the width back to the plot.
    extra = layer_draw.bar_count(ctx.composite, ctx.overlay_style) - 1
    fig, ax = axes_for(ctx, (8 + 1.1 * extra, 6))
    mesh, wants_bar = layer_draw.draw_base(
        fig, ax, ctx.data, prepared, x, y, ctx.colormap, ctx.norm, ctx.label,
        style=ctx.overlay_style, interpolate=ctx.interpolate)
    try:
        attach_format_coord(ax, ctx.data, x, y, plan.x.dim, plan.y.dim, ctx.varname)
    except ValueError:
        # A single-point or non-monotonic axis cannot be mapped back to cells;
        # the plot is still fine, only the hover readout is unavailable.
        pass

    if prepared:
        # Inset bars only once there is something to stack beside: a plain
        # section keeps the colorbar that takes its room out of the axes, so a
        # single-variable figure comes out exactly as it always has.
        if extra:
            fig.subplots_adjust(right=0.86 - 0.07 * extra)
        if wants_bar:
            fig.colorbar(mesh, cax=layer_draw.bar_axes(ax, 0)).set_label(ctx.label)
        layer_draw.draw(fig, ax, prepared, x, y, style=ctx.overlay_style,
                        interpolate=ctx.interpolate)
    else:
        fig.colorbar(mesh, ax=ax, pad=0.02).set_label(ctx.label)

    ax.set_xlabel(plan.x.label)
    ax.set_ylabel(plan.y.label)
    if plan.invert_y:
        ax.invert_yaxis()
    title = ctx.titled(f'({plan.y.dim} vs {plan.x.dim})')
    if prepared:
        title = f"{title}\n{overlay.describe(ctx.composite)}"
    ax.set_title(title, fontweight='bold')
    ax.grid(True)
    return finish_figure(fig, ctx.output_path, dpi=ctx.dpi)