"""
Several variables in one picture: `--overlay`.

Some questions are about a variable and some are about how two of them sit
together. Where the CO2 frost gives way to water ice, whether the dust is under
the cloud or over it, which of them reaches the pole first - none of these can
be read from two figures side by side, because the eye cannot register a
boundary it has to remember from the other page.

Every layer keeps its own colour scale, because they measure different things,
and every layer is translucent, so an overlap reads as an overlap rather than as
one variable winning. That is a deliberate trade: a cell holding both CO2 and
H2O comes out a colour that is in neither bar. What it buys is that coexistence
is visible at all, which a priority order would hide by construction.

The hue carries which variable and the depth carries how much, which is why the
layer colormaps are single-hue ramps rather than perceptual ones - see
`colors.LAYER_COLORMAPS`.
"""

from dataclasses import dataclass, field

import numpy as np
from matplotlib.colors import Normalize

from .colors import LAYER_COLORMAPS


# Opacity of one layer at the top of its own scale. Below one, so that a layer
# underneath still shows through the densest part of the one above it.
LAYER_OPACITY = 0.85

# How far up its scale a layer has to be before it is worth drawing at all.
# Without it, a field that is faintly nonzero everywhere - which is most of them
# on a log scale - lays a flat wash over everything below.
LAYER_FLOOR = 0.02


@dataclass
class Layer:
    """
    One variable's contribution to a composite.
    """
    varname: str
    data: np.ndarray
    colormap: str
    norm: object = None
    units: str = None
    long_name: str = None
    frames: np.ndarray = None
    # The raw text after the colon in `--overlay VAR:SPEC`. A blended layer
    # reads it as a colormap; a contour, a hatch or a curve has to read the same
    # word as a single colour, so the unresolved form is kept alongside.
    requested: str = None
    # Explicit contour levels or hatch threshold, from --overlay-threshold.
    thresholds: tuple = None

    @property
    def label(self):
        base = self.long_name or self.varname
        return base + (f" [{self.units}]" if self.units else "")


@dataclass
class Composite:
    """
    The layers of one figure, bottom first.
    """
    layers: list = field(default_factory=list)

    def __bool__(self):
        return len(self.layers) > 1

    @property
    def names(self):
        return [layer.varname for layer in self.layers]


def parse_specs(values):
    """
    Turn `--overlay` arguments into (variable, colormap or None) pairs.

    `h2o_ice` takes the next colour in the rota; `h2o_ice:Blues` names one. The
    colon is the separator rather than a comma because a comma already means
    "two components of one vector" in `--vector`, and the two would be read
    wrongly by anyone who knows the other.
    """
    specs = []
    for value in values or ():
        for part in str(value).split(','):
            part = part.strip()
            if not part:
                continue
            name, _, colormap = part.partition(':')
            specs.append((name.strip(), colormap.strip() or None))
    return specs


def assign_colormaps(specs, taken=()):
    """
    Give every unnamed layer a colormap no other layer is using.

    Two layers in the same colours is the one outcome a composite cannot
    survive, so the rota skips whatever is already spoken for - including the
    base variable's own map, which was chosen before any of this ran.
    """
    used = [c for c in taken if c]
    out = []
    for name, colormap in specs:
        if colormap is None:
            spare = [c for c in LAYER_COLORMAPS if c not in used]
            if spare:
                colormap = spare[0]
            else:
                # The rota is out, so from here on two layers share a colour -
                # the one thing this function exists to prevent. Said out loud
                # rather than fixed, because there is no fixing it: a seventh
                # translucent field on one map cannot be read whatever palette
                # it is given.
                colormap = LAYER_COLORMAPS[len(used) % len(LAYER_COLORMAPS)]
                print(f"Warning: --overlay '{name}' reuses the '{colormap}' "
                      f"colours; only {len(LAYER_COLORMAPS)} layers can be told "
                      f"apart. Name one with '{name}:CMAP', or plot fewer.")
        used.append(colormap)
        out.append((name, colormap))
    return out


def layer_alpha(values, norm, opacity=LAYER_OPACITY, floor=LAYER_FLOOR):
    """
    Per-cell transparency for one layer of a composite.

    Absent is invisible and full is nearly opaque, with everything below `floor`
    of the layer's own scale treated as absent. Layers are drawn over one
    another, so a layer that renders its own floor as a faint wash does not
    merely look wrong - it hides whatever is beneath it across the whole map.

    The square root lifts the middle, the same way a single translucent shell
    does: without it a layer is either invisible or saturated, with very little
    in between where the interesting boundaries are.
    """
    values = np.asarray(values, dtype=float)
    scale = norm if norm is not None else _own_scale(values)

    scaled = np.ma.filled(np.ma.masked_invalid(scale(values)), np.nan)
    scaled = np.asarray(scaled, dtype=float)

    alpha = np.sqrt(np.clip(scaled, 0.0, 1.0)) * opacity
    alpha[~np.isfinite(scaled)] = 0.0
    alpha[scaled < floor] = 0.0
    return alpha


def _own_scale(values):
    finite = values[np.isfinite(values)]
    if finite.size == 0:
        return Normalize(vmin=0.0, vmax=1.0)
    low, high = float(finite.min()), float(finite.max())
    return Normalize(vmin=low, vmax=high if high > low else low + 1.0)


def resolve_thresholds(values, spec, default=None):
    """
    Turn a threshold spec into a sorted list of values in the field's own units.

    A bare number is used as-is; 'pNN' is that percentile of the finite values,
    which is what makes one setting work across fields whose magnitudes differ
    by decades. Several may be given, comma-separated.

    Resolved against each layer's own values rather than once for the figure:
    the p90 of a CO2 field and the p90 of a temperature are different numbers,
    and a single setting has to mean "the top tenth of each" or it means nothing.
    """
    finite = np.asarray(values, dtype=float)
    finite = finite[np.isfinite(finite)]
    if finite.size == 0:
        return []

    parts = [p.strip().lower() for p in str(spec).split(',')] if spec else list(default or ())
    out = []
    for part in parts:
        if not part:
            continue
        try:
            if str(part).startswith('p'):
                out.append(float(np.percentile(finite, float(str(part)[1:]))))
            else:
                out.append(float(part))
        except ValueError:
            print(f"Warning: '{part}' is not a value or a percentile like 'p90'; "
                  f"ignoring it.")
    # Sorted so that "inner" always means "later", whatever order they were given
    return sorted(set(out))


def describe(composite):
    """
    The one-line summary of what is stacked, for the title and the terminal.
    """
    return ' over '.join(reversed(composite.names))


def axis_groups(layers):
    """
    Which curves may share one axis, from their units.
    Returns (primary, secondary, refused).

    One unit is one axis. Two get a twin axis opposite it, which is as far as a
    reader can follow: past that there is no third side to the figure, so the
    odd ones out are refused by name rather than plotted against a scale that is
    not theirs. A curve in W/m2 read off a kelvin axis is not a smaller error
    than no curve at all - and unlike the 2-D styles, a 1-D plot has no
    per-layer colour bar to say which scale a line belongs to.

    Units are compared as written. Whether "K" and "kelvin" are the same unit is
    the file's business, not ours, and a variable with no units at all is its
    own group, because "unknown" cannot be claimed to match "K".
    """
    order = []
    grouped = {}
    for layer in layers:
        key = layer.units or ''
        if key not in grouped:
            grouped[key] = []
            order.append(key)
        grouped[key].append(layer)

    primary = grouped[order[0]] if order else []
    secondary = grouped[order[1]] if len(order) > 1 else []
    refused = [layer for key in order[2:] for layer in grouped[key]]
    return primary, secondary, refused


def units_of(layers):
    """
    The unit a group of curves shares, for the axis they share.
    """
    for layer in layers:
        if layer.units:
            return layer.units
    return ''


def name_list(composite):
    """
    The variables named side by side, in the order they were given.

    `describe` says "b over a", which is stacking language: right for layers
    that sit on top of one another, wrong for panels beside each other or for
    two curves on one axes, where nothing is over anything.
    """
    return ', '.join(composite.names)