"""
Deciding how to plot a variable, from what its dimensions mean.

The old code chose a plot type from `data.ndim` and then guessed the axis
ordering from dimension names. Here the decision is made from resolved axis
roles and expressed as a PlotPlan, so adding a plot type is adding a row to the
table below and an entry in render.RENDERERS, not another branch in a ladder.
"""

from dataclasses import dataclass

from .coords import Axis

# Every plot kind the renderers know about
KINDS = ('scalar', 'line', 'profile', 'timeseries', 'section', 'geomap', 'globe')


@dataclass(frozen=True)
class PlotPlan:
    """
    What to draw and how to orient it.

    For `profile` the data goes on X and the coordinate on Y, which is why `y`
    can be set while `x` is None. `globe` is the only kind with three axes: it
    keeps the vertical coordinate instead of asking for a slice through it.
    """
    kind: str
    x: Axis = None
    y: Axis = None
    z: Axis = None
    invert_y: bool = False
    cyclic: bool = False
    reason: str = ''

    def describe(self):
        bits = [f"kind={self.kind}"]
        if self.x is not None:
            bits.append(f"x={self.x.dim}[{self.x.role}]")
        if self.y is not None:
            bits.append(f"y={self.y.dim}[{self.y.role}]")
        if self.z is not None:
            bits.append(f"z={self.z.dim}[{self.z.role}]")
        if self.invert_y:
            bits.append("y inverted (positive down)")
        if self.cyclic:
            bits.append("cyclic longitude")
        return ', '.join(bits) + (f" - {self.reason}" if self.reason else '')


class SubsurfaceShell(Exception):
    """
    Raised when a globe was asked for over a vertical axis that measures depth.

    A soil column cannot be an atmospheric shell, and silently drawing one
    inside out would be worse than refusing, so the caller reports it and
    suggests the cross-section instead.
    """

    def __init__(self, axis):
        self.axis = axis
        super().__init__(axis.dim)


class TooManyDimensions(Exception):
    """
    Raised when more than two dimensions remain, so the caller can report the
    remaining dimensions and exit rather than guess.
    """

    def __init__(self, axes):
        self.axes = axes
        super().__init__(', '.join(a.dim for a in axes))


def _order_two(a, b):
    """
    Which of two axes goes on X, and which on Y.

    Vertical coordinates always take Y, time takes X, and longitude precedes
    latitude. This is the same intent as the old choose_x_dim, but expressed
    over roles rather than over name tables, so it works for `rlonu` too.
    """
    roles = (a.role, b.role)
    if 'Z' in roles:
        z, other = (a, b) if a.role == 'Z' else (b, a)
        return other, z, "vertical coordinate stays on Y"
    if 'T' in roles:
        t, other = (a, b) if a.role == 'T' else (b, a)
        return t, other, "time on X"
    if a.role == 'X' and b.role == 'Y':
        return a, b, "longitude on X, latitude on Y"
    if a.role == 'Y' and b.role == 'X':
        return b, a, "longitude on X, latitude on Y"
    return a, b, "dimension order"


def infer(axes, x_dim=None, plot_kind=None):
    """
    Build a PlotPlan for the axes remaining after selection.

    `x_dim` forces which dimension goes on X; `plot_kind` forces the renderer.
    Raises TooManyDimensions when more than two dimensions are left.
    """
    axes = list(axes)
    n = len(axes)

    if n == 0:
        return PlotPlan(kind='scalar', reason="no dimensions left")

    if n == 3 and plot_kind == 'globe':
        return _globe_plan(axes)

    if n > 2:
        raise TooManyDimensions(axes)

    if n == 1:
        axis = axes[0]
        if axis.role == 'T':
            plan = PlotPlan(kind='timeseries', x=axis, reason="single time axis")
        elif axis.role == 'Z':
            plan = PlotPlan(kind='profile', y=axis,
                            invert_y=(axis.positive == 'down'),
                            reason="vertical profile: value on X, depth on Y")
        else:
            plan = PlotPlan(kind='line', x=axis, reason=f"single {axis.role} axis")
        return _apply_overrides(plan, axes, x_dim, plot_kind)

    a, b = axes
    x, y, reason = _order_two(a, b)
    if {a.role, b.role} == {'X', 'Y'}:
        plan = PlotPlan(kind='geomap', x=x, y=y, cyclic=bool(x.is_cyclic),
                        reason="latitude/longitude map")
    else:
        plan = PlotPlan(kind='section', x=x, y=y,
                        invert_y=(y.positive == 'down'), reason=reason)
    return _apply_overrides(plan, axes, x_dim, plot_kind)


def _globe_plan(axes):
    """
    The plan for a 3D atmospheric shell: longitude, latitude and a vertical
    coordinate that points up.

    This is the one kind that keeps three axes. Anything else with three
    dimensions still has to be reduced by the caller, and a vertical axis
    measuring depth is refused outright rather than drawn as an atmosphere.
    """
    by_role = {a.role: a for a in axes}
    if set(by_role) != {'X', 'Y', 'Z'}:
        raise TooManyDimensions(axes)

    z = by_role['Z']
    if z.positive == 'down':
        raise SubsurfaceShell(z)

    return PlotPlan(kind='globe', x=by_role['X'], y=by_role['Y'], z=z,
                    cyclic=bool(by_role['X'].is_cyclic),
                    reason="longitude, latitude and a vertical axis: 3D shell")


def _apply_overrides(plan, axes, x_dim, plot_kind):
    """
    Apply --x-dim and --plot-kind on top of the inferred plan.
    """
    if x_dim:
        names = [a.dim for a in axes]
        if x_dim not in names:
            print(f"Warning: --x-dim '{x_dim}' is not among {names}; ignoring it.")
        elif plan.kind == 'geomap' and x_dim != plan.x.dim:
            # A geographic map has a fixed orientation; honouring --x-dim here
            # used to produce a transposed plot with no topography.
            print(f"Note: --x-dim '{x_dim}' is ignored on geographic maps; "
                  f"pass --plot-kind section for a transposed heatmap.")
        elif len(axes) == 2:
            chosen = next(a for a in axes if a.dim == x_dim)
            other = next(a for a in axes if a.dim != x_dim)
            plan = PlotPlan(kind=plan.kind, x=chosen, y=other,
                            invert_y=(other.positive == 'down'),
                            cyclic=bool(chosen.is_cyclic),
                            reason=f"--x-dim {x_dim}")

    if plot_kind and plot_kind != 'auto' and plot_kind != plan.kind:
        if plot_kind not in KINDS:
            print(f"Warning: unknown --plot-kind '{plot_kind}'; ignoring it.")
        else:
            plan = PlotPlan(kind=plot_kind, x=plan.x, y=plan.y,
                            invert_y=plan.invert_y, cyclic=plan.cyclic,
                            reason=f"--plot-kind {plot_kind}")
    return plan