"""
One panel per variable, instead of one figure carrying all of them.

Not one plot, so not what `--overlay` was originally for - but frequently the
honest answer. Two fields in different units have no shared colour scale and no
shared axis, and stacking them anyway produces a picture whose colours belong to
neither. Side by side on the same axes they can at least be compared by eye,
which is what the question usually was.

It is also the only style that works for *every* plot kind. A scalar, a curve, a
profile, a cross-section and a map are all things a renderer can draw into an
axes it was handed, so nothing here has to know which. That is what lets the
`--overlay` gate offer an alternative rather than merely refusing.
"""

import math
from dataclasses import dataclass, field

import matplotlib.pyplot as plt

from ..colors import choose_style


@dataclass
class Grid:
    """
    A figure of several axes, handed out one at a time.

    `figure.axes_for` asks this for its axes when a context carries one, so a
    renderer draws into the next cell without ever knowing it is in a grid.
    """
    fig: object
    cells: list = field(default_factory=list)
    used: int = 0

    def next_axes(self, **subplot_kw):
        """
        The next free cell, replaced by one with the right projection when the
        renderer asks for one - a map needs a GeoAxes, and the placeholder made
        at layout time is a plain one.

        The figure is marked unfinished until the last cell has been handed out.
        Renderers save or show as their last act, and without this each panel
        would write the half-drawn figure over the one before it, or pop up its
        own window - the renderer cannot know it is one of several.
        """
        position = self.cells[min(self.used, len(self.cells) - 1)]
        self.used += 1
        self.fig._dispnc_incomplete = self.used < len(self.cells)
        if subplot_kw:
            # A cartopy GeoAxes cannot be converted after the fact, so the
            # placeholder goes and a properly projected one takes its geometry.
            spec = position.get_subplotspec()
            position.remove()
            position = self.fig.add_subplot(spec, **subplot_kw)
            self.cells[min(self.used - 1, len(self.cells) - 1)] = position
        return self.fig, position


def layout(count, figsize=None, per_panel=(6.2, 4.2)):
    """
    A near-square grid of `count` panels.

    Near-square rather than one row: eight variables in a row is a figure nobody
    can read, and the eye compares a 3x3 block far better than a 1x8 strip.
    """
    columns = max(1, math.ceil(math.sqrt(count)))
    rows = max(1, math.ceil(count / columns))
    size = figsize or (per_panel[0] * columns, per_panel[1] * rows)

    # Constrained layout rather than a fixed grid: the panels are filled in one
    # at a time by renderers that each add their own colour bar and labels, and
    # only matplotlib knows how much room those took once they are there.
    fig, axes = plt.subplots(rows, columns, figsize=size, squeeze=False,
                             layout='constrained')
    # Extra padding because a cartopy gridliner draws its degree labels outside
    # the axes without telling the layout engine, so a neighbouring colour bar's
    # label lands on them unless the gap is opened by hand.
    # The `rect` margin is the other half of the same problem: the leftmost
    # panel's latitude labels hang outside its axes, off the canvas, and no
    # amount of inter-panel padding reaches them.
    fig.get_layout_engine().set(w_pad=0.10, h_pad=0.06, wspace=0.06, hspace=0.05,
                                rect=(0.035, 0.0, 0.94, 0.95))
    cells = [ax for row in axes for ax in row]
    # Any cell the variables do not reach is removed rather than left as an
    # empty frame, which reads as a panel whose data failed to draw.
    for spare in cells[count:]:
        spare.remove()
    return Grid(fig=fig, cells=cells[:count])


def shared_style(layers, varname):
    """
    One colour scale over every panel, or None to let each keep its own.

    Shared is right for one variable at several times, where a colour has to
    mean the same thing across the row. It is wrong for different variables,
    where the shared scale is set by whichever happens to be largest and the
    rest come out flat - which is why `own` is the default.
    """
    stacked = [layer.data for layer in layers if layer.data is not None]
    if not stacked:
        return None, None
    import numpy as np
    _, norm, _ = choose_style(np.concatenate([s.ravel() for s in stacked]), varname)
    return norm, 'one scale over every panel'