"""
Floating keys — IPF triangles and colour wheels
===============================================

Some plots are coloured by a *direction*, not by a magnitude, and a colorbar
cannot say what the colours mean.  An orientation map is coloured by which
crystal axis points at you; a polarization map by which way the moment lies in
the plane.  Both need a small picture as their legend: an inverse pole figure
triangle, or a hue wheel.

:meth:`~anyplotlib.Plot2D.add_key` pins that picture over the panel in *screen*
space — it does not pan or zoom with the data, exactly like the scale bar it is
modelled on.

This is deliberately not :meth:`~anyplotlib.Figure.add_inset`, which is a
draggable window with a title bar and its own canvas stack.  That is the right
tool when the overlay is a live plot; a key is a static picture that should
read as part of the figure.
"""
import numpy as np

import anyplotlib as apl

# %%
# The IPF colour key
# ------------------
# The standard cubic stereographic triangle.  Each pixel's colour is its
# barycentric distance to the three corners, which is the classic IPF key: a
# grain pointing [1 0 0] at the detector reads red, [1 1 0] green, [1 1 1] blue.
#
# The triangle is built as an ``(H, W, 4)`` RGBA array, and **alpha 0 outside
# the triangle** is what lets it sit on the map without a rectangular card
# around it.

KEY_N = 220


def ipf_triangle(n=KEY_N):
    """RGBA image of the 001–011–111 stereographic triangle."""
    yy, xx = np.mgrid[0:n, 0:n]
    u = xx / (n - 1)
    v = 1.0 - yy / (n - 1)
    inside = v <= u + 1e-9                       # lower-right half

    # Barycentric-ish weights: distance to each corner, normalised.
    d100 = np.hypot(u, v)                        # corner (0, 0)
    d110 = np.hypot(u - 1.0, v)                  # corner (1, 0)
    d111 = np.hypot(u - 1.0, v - 1.0)            # corner (1, 1)
    far = np.maximum.reduce([d100, d110, d111])
    rgb = np.stack([1 - d100 / far, 1 - d110 / far, 1 - d111 / far], -1)
    rgb /= rgb.max(-1, keepdims=True) + 1e-9     # full saturation at the corners

    img = np.zeros((n, n, 4), np.uint8)
    img[..., :3] = np.clip(rgb, 0, 1) * 255
    img[..., 3] = np.where(inside, 255, 0)
    return img


# %%
# A synthetic orientation map
# ---------------------------
# Voronoi grains, each with a random orientation, coloured through the same key
# so the map and its legend agree by construction.

rng = np.random.default_rng(11)
H, W, NGRAIN = 210, 280, 40

cy, cx = rng.uniform(0, H, NGRAIN), rng.uniform(0, W, NGRAIN)
yy, xx = np.mgrid[0:H, 0:W]
grain = np.hypot(yy[..., None] - cy, xx[..., None] - cx).argmin(-1)

# Each grain gets a point in the triangle, then reads its colour off the key.
gu = rng.uniform(0, 1, NGRAIN)
gv = rng.uniform(0, 1, NGRAIN) * gu              # keep it inside v <= u
key_img = ipf_triangle()
kx = np.clip((gu * (KEY_N - 1)).astype(int), 0, KEY_N - 1)
ky = np.clip(((1 - gv) * (KEY_N - 1)).astype(int), 0, KEY_N - 1)
grain_rgb = key_img[ky, kx, :3]
ipf_map = grain_rgb[grain]                       # (H, W, 3) true colour

# %%
# Pinning the key
# ---------------
# ``labels`` draws text *inside* the picture, positioned as fractions of the
# key image, so the corner indices stay on the corners at any ``size``.

fig, ax = apl.subplots(1, 1, figsize=(520, 420))
vmap = ax.imshow(ipf_map)
vmap.set_title("orientation map")

vmap.add_key(
    key_img,
    corner="bottom-right",
    size=0.34,
    # `align` keeps a label inside the key: centring text on a corner would
    # hang half of it off the edge, where the panel clips it.
    labels=[
        {"x": 0.02, "y": 0.93, "text": "[1 0 0]", "align": "left"},
        {"x": 0.98, "y": 0.93, "text": "[1 1 0]", "align": "right"},
        {"x": 0.98, "y": 0.08, "text": "[1 1 1]", "align": "right"},
    ],
    name="ipf",
)

fig

# %%
# A colour wheel over a polarization map
# --------------------------------------
# Same mechanism, different legend.  Here the key gets a translucent card
# (``bgcolor``) because the field underneath is saturated everywhere and a bare
# wheel would fight with it.

def hue_wheel(n=KEY_N):
    """RGBA colour wheel: hue = in-plane angle, value = magnitude."""
    yy, xx = np.mgrid[0:n, 0:n]
    ang = (np.arctan2(-(yy - n / 2), xx - n / 2) + np.pi) / (2 * np.pi)
    rad = np.hypot(yy - n / 2, xx - n / 2) / (n / 2)
    h6 = ang * 6.0
    chan = np.clip(
        np.abs(((h6 + np.array([0, 4, 2])[:, None, None]) % 6) - 3) - 1, 0, 1)
    img = np.zeros((n, n, 4), np.uint8)
    img[..., :3] = chan.transpose(1, 2, 0) * 255 * np.clip(rad, 0, 1)[..., None]
    img[..., 3] = np.where(rad <= 1.0, 255, 0)
    return img


# A vortex: the moment angle winds once around the centre.
ang = np.arctan2(yy - H / 2, xx - W / 2)
mag = np.clip(np.hypot(yy - H / 2, xx - W / 2) / (0.5 * min(H, W)), 0, 1)
a6 = ((ang + np.pi) / (2 * np.pi)) * 6.0
chan = np.clip(np.abs(((a6 + np.array([0, 4, 2])[:, None, None]) % 6) - 3) - 1, 0, 1)
polar_map = (chan.transpose(1, 2, 0) * 255 * (0.3 + 0.7 * mag)[..., None]).astype(np.uint8)

fig2, ax2 = apl.subplots(1, 1, figsize=(520, 420))
vpol = ax2.imshow(polar_map)
vpol.set_title("in-plane magnetic polarization")

vpol.add_key(
    hue_wheel(),
    corner="top-right",
    size=0.26,
    bgcolor="rgba(0,0,0,0.45)",     # legible over a busy field
    border="#ffffff",
    label="moment direction",
    labels=[
        (0.5, 0.06, "N"), (0.94, 0.5, "E"),
        (0.5, 0.94, "S"), (0.06, 0.5, "W"),
    ],
    name="wheel",
)

fig2

# %%
# Keeping it out of the way
# -------------------------
# ``hover_only=True`` shows the key only while the pointer is over the panel —
# a reading aid that does not sit on the data while you study it.  PNG export
# renders the panel as though the pointer were there, so an exported figure
# still carries the key.
#
# Everything is live: :meth:`~anyplotlib.KeyOverlay.set` restyles a key without
# re-sending the picture, and :meth:`~anyplotlib.KeyOverlay.set_image` swaps
# the picture without disturbing the placement::
#
#     key = vmap.get_key("ipf")
#     key.set(size=0.4, corner="top-left")
#     key.visible = False