"""
Two variables in one colour, decodably.

Blending two translucent layers also puts two variables in one colour, but by
accident: the result belongs to neither colour bar, so a reader who sees it
cannot get back to either number. A bivariate scheme does the same thing on
purpose. Both fields are cut into a few classes, every pair of classes is given
one colour, and the legend is the whole square of them - so every colour on the
map has a named cell, and "high CO2 with low H2O" is a colour one can point at.

The cost is honest and stated: the values are quantized into bins, and it works
for exactly two variables. Three has no square.

Pure numpy, so the scheme and its key can be tested without a figure - the same
doctrine `overlay.py` follows.
"""

import numpy as np

# How many classes each variable is cut into. Three by three is nine colours,
# which is about as many as a reader can hold; four by four is sixteen and is
# already a colour-matching exercise rather than a map.
CLASSES = 3

# The four corners of the square, in RGB. Low-low is near-white so an empty
# region reads as empty; the two single-variable corners are a blue and a red
# that stay apart under the common colour-vision deficiencies; high-high is
# their mixture, dark enough to read as "both".
CORNERS = {
    'low_low':   (0.92, 0.92, 0.92),
    'high_x':    (0.20, 0.45, 0.75),
    'high_y':    (0.80, 0.30, 0.20),
    'high_both': (0.25, 0.15, 0.35),
}


def classify(values, classes=CLASSES):
    """
    Cut a field into `classes` bins by quantile, as an integer array.

    Quantiles rather than even intervals: these fields are routinely spread over
    decades, and even intervals would put almost every cell in the bottom class
    and leave the rest of the square unused.

    Cells that are not finite come back as -1 and are drawn as nothing.
    """
    values = np.asarray(values, dtype=float)
    finite = values[np.isfinite(values)]
    out = np.full(values.shape, -1, dtype=int)
    if finite.size == 0:
        return out

    edges = np.quantile(finite, np.linspace(0, 1, classes + 1)[1:-1])
    # A field flat enough that its quantiles coincide has one class, not an
    # arbitrary split of identical values
    edges = np.unique(edges)
    binned = np.digitize(values, edges)
    out[np.isfinite(values)] = binned[np.isfinite(values)]
    return np.clip(out, -1, classes - 1)


def square(classes=CLASSES):
    """
    The colour of every (x class, y class) pair: an (classes, classes, 3) array.

    Indexed [y, x] so it can be handed straight to `imshow` as the legend key,
    with x increasing rightwards and y upwards.
    """
    grid = np.zeros((classes, classes, 3), dtype=float)
    span = max(classes - 1, 1)
    low_low = np.array(CORNERS['low_low'])
    high_x = np.array(CORNERS['high_x'])
    high_y = np.array(CORNERS['high_y'])
    high_both = np.array(CORNERS['high_both'])

    for yi in range(classes):
        for xi in range(classes):
            u, v = xi / span, yi / span
            # Bilinear between the four corners: each axis carries one variable,
            # and the diagonal is where they are both high.
            grid[yi, xi] = ((1 - u) * (1 - v) * low_low
                            + u * (1 - v) * high_x
                            + (1 - u) * v * high_y
                            + u * v * high_both)
    return np.clip(grid, 0.0, 1.0)


def colours(x_values, y_values, classes=CLASSES):
    """
    An (M, N, 3) RGB image of the two fields together.

    Where either field has nothing to say the cell comes back white, so a hole
    in one variable does not get a colour from the other alone.
    """
    xi = classify(x_values, classes)
    yi = classify(y_values, classes)
    key = square(classes)

    rgb = np.ones(xi.shape + (3,), dtype=float)
    known = (xi >= 0) & (yi >= 0)
    rgb[known] = key[yi[known], xi[known]]
    return rgb


def legend_axes(fig, ax, x_label, y_label, classes=CLASSES, size=0.16):
    """
    Draw the square key on the plot: the whole scheme, so every colour on the
    map has a cell a reader can point at.

    Placed in figure coordinates rather than as a fraction of the axes. A map
    axes is wide and short - cartopy locks its aspect - so the same fraction is
    a very different number of inches across than it is tall, and the square
    came out a cramped sliver with labels nothing could read. In figure
    coordinates it stays square and stays legible, in the empty canvas the map's
    own aspect leaves below it.

    This is what a bivariate scheme has instead of two colour bars, and the
    reason it is worth having at all: the same mixing, without the key, is just
    blending.
    """
    width, height = fig.get_size_inches()
    side = min(size * width, size * height) / max(width, height)
    key = fig.add_axes([0.80, 0.08, side * width / height * 0.9, side * 0.9])
    key.imshow(square(classes), origin='lower', interpolation='nearest',
               extent=(0, classes, 0, classes), aspect='auto')
    key.set_xticks([])
    key.set_yticks([])
    key.set_xlabel(_short(x_label), fontsize=7)
    key.set_ylabel(_short(y_label), fontsize=7)
    for spine in key.spines.values():
        spine.set_linewidth(0.4)
    return key


def _short(label, limit=18):
    """
    A key this small has room for a name, not for a name and its units.
    """
    label = str(label).split(' [')[0].split(' (')[0]
    return label if len(label) <= limit else label[:limit - 1] + '\u2026'