#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Figure A4: extension and validation of the perimeter method.

Input: ../Data/FigA5_ellipses_by_latitude.npy (the ellipse test of panel b).
Output: ../Figures/FigA5_Perimeter_Method_v11.png and .pdf.

Run from this folder; ../Data holds the inputs and the figure is written to ../Figures.
"""
import math
import os

import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle

mpl.rcParams.update({
    "font.family":      "Helvetica",
    "font.size":         9,
    "axes.titlesize":   10,
    "mathtext.fontset": "cm",
    "axes.linewidth":  0.6,
    "axes.unicode_minus": True,
})

HERE = os.path.dirname(os.path.abspath(__file__))
DATA = os.path.join(HERE, "..", "Data", "FigA5_ellipses_by_latitude.npy")
OUT_DIR = os.path.join(HERE, "..", "Figures")
PIX, CELL = "#eb6834", "#2a78d6"        # traced on pixels / on cells (FlowMorph 1.1), as Figure R2
INK, MUTED, GRID = "#0b0b0b", "#52514e", "#d9d8d4"
FILL = "#e9e7e1"
DASH = (0, (4, 2))
A_WGS, E2 = 6378137.0, 0.00669437999014
W_AXIAL, W_DIAG, W_CORNER = 0.980, 1.406, 0.091
ROW = (-1, -1, 0, 1, 1, 1, 0, -1)
COL = (0, 1, 1, 1, 0, -1, -1, -1)


def moore(mask):
    """the cells visited by the Moore trace from the topmost-leftmost cell, Jacob's stopping rule; and the step codes
    and the index of the first return to the start"""
    nrow, ncol = mask.shape
    inside = lambda r, c: 0 <= r < nrow and 0 <= c < ncol and mask[r, c]   # noqa: E731
    r0 = int(np.nonzero(mask.any(axis=1))[0][0])
    c0 = int(np.nonzero(mask[r0])[0][0])
    back = next(d for d in range(8) if not inside(r0 + ROW[d], c0 + COL[d]))
    path, codes = [(r0, c0)], []
    r, c, first, first_return = r0, c0, None, None
    for _ in range(1000000):
        for turn in range(1, 9):
            code = (back + turn) % 8
            nr, nc = r + ROW[code], c + COL[code]
            if not inside(nr, nc):
                continue
            if codes and (r, c) == (r0, c0):
                if first_return is None:
                    first_return = len(path) - 1
                if (nr, nc) == first:
                    return path, codes, first_return
            if first is None:
                first = (nr, nc)
            codes.append(code)
            path.append((nr, nc))
            r, c, back = nr, nc, (code + 4) % 8
            break
    raise RuntimeError("the trace did not close")


def vs_length(codes, dx, dy):
    """Vossepoel & Smeulders (1982) along a piece of a trace on cells dx by dy"""
    total = 0.0
    for i, code in enumerate(codes):
        if code % 2 == 0:
            total += W_AXIAL * (dy if code in (0, 4) else dx)
        else:
            total += W_DIAG / math.sqrt(2) * math.hypot(dx, dy)
        if i > 0 and code % 2 != codes[i - 1] % 2:
            total -= W_CORNER * 0.5 * (dx + dy)
    return total


def pixel_edges(latitude):
    phi = math.radians(latitude)
    w = math.sqrt(1 - E2 * math.sin(phi) ** 2)
    return (A_WGS / w * math.cos(phi) * math.radians(1 / 1200),
            A_WGS * (1 - E2) / w ** 3 * math.radians(1 / 1200))


def draw_cells(ax, mask, dx, dy, fill, rows=None, cols=None, lw=0.6):
    rows = range(mask.shape[0]) if rows is None else rows
    cols = range(mask.shape[1]) if cols is None else cols
    for r in rows:
        for c in cols:
            ax.add_patch(Rectangle((c * dx, -(r + 1) * dy), dx, dy, facecolor=fill if mask[r, c] else "white",
                                   edgecolor=GRID, linewidth=lw))


def centres(path, dx, dy):
    return ([(c + 0.5) * dx for r, c in path], [-(r + 0.5) * dy for r, c in path])


def edge_ratio_and_window(dx, dy, k):
    """a 45-degree edge on the ground (the right side of a tall triangle), traced on cells of k columns"""
    nrow = 300
    ncol = (int(nrow * dy / dx) + 6 * k) // k * k
    r = np.arange(nrow)[:, None]
    c = np.arange(ncol)[None, :]
    pixels = (c + 0.5) * dx <= (r + 0.5) * dy + 0.37 * dx
    cells = pixels.reshape(nrow, ncol // k, k).any(axis=2) if k > 1 else pixels
    path, codes, _ = moore(cells)
    rows = [p[0] for p in path]
    i1, i2 = rows.index(50), rows.index(250)                # on the right side, going down
    (r1, c1), (r2, c2) = path[i1], path[i2]
    ratio = vs_length(codes[i1:i2], k * dx, dy) / math.hypot((c2 - c1) * k * dx, (r2 - r1) * dy)
    return ratio, cells, path[i1:i2 + 1]


def panel_edge(ax, latitude, k, color, ls, heading, what):
    dx, dy = pixel_edges(latitude)
    ratio, cells, path = edge_ratio_and_window(dx, dy, k)
    r_top, r_bottom = 100, 104
    x_left, x_right = (r_top - 1.6) * dy, (r_bottom + 1.6) * dy
    cols = range(int(x_left // (k * dx)), int(x_right // (k * dx)) + 1)
    draw_cells(ax, cells, k * dx, dy, FILL, rows=range(r_top, r_bottom + 1), cols=cols, lw=0.5)
    ys = np.array([-r_top * dy, -(r_bottom + 1) * dy])
    ax.plot(-ys + 0.37 * dx, ys, color=INK, lw=1.0, ls=":", zorder=3)
    x, y = centres(path, k * dx, dy)
    ax.plot(x, y, color=color, lw=2.0, ls=ls, zorder=4)
    ax.set_xlim(cols[0] * k * dx, (cols[-1] + 1) * k * dx)
    ax.set_ylim(-(r_bottom + 1) * dy, -r_top * dy)
    ax.set_aspect("equal")
    ax.axis("off")
    ax.set_title(heading, loc="left", fontsize=8.5, pad=3)
    err = f"{100 * (ratio - 1):+.1f} %".replace("-", "−")
    ax.text(0.5, -0.05, f"{what}: {err}", transform=ax.transAxes, ha="center", va="top", fontsize=8.5,
            color=color, fontweight="bold")
    return ratio


def panel_known(ax, rows):
    lat = rows[:, 0]
    ax.axhline(0.0, color=MUTED, lw=0.8)
    for lo, mid, hi, color, ls, label in ((2, 3, 4, PIX, DASH, "traced on the pixels"),
                                          (5, 6, 7, CELL, "-", "traced on cells of merged columns, "
                                                                "interpolated in $k$")):
        ax.fill_between(lat, 100 * (rows[:, lo] - 1), 100 * (rows[:, hi] - 1), color=color, alpha=0.14, lw=0)
        ax.plot(lat, 100 * (rows[:, mid] - 1), color=color, lw=2.0, ls=ls, label=label)
    for mid, color, dy in ((3, PIX, 0.0), (6, CELL, 0.0)):
        v = 100 * (rows[-1, mid] - 1)
        ax.text(lat[-1] + 1.2, v + dy, f"{v:+.1f} %".replace("-", "−"), color=color, fontsize=8.5,
                va="center", fontweight="bold")
    ax.set_xlim(0, 92)
    ax.set_ylim(-3, 16)
    ax.set_xticks([0, 20, 40, 60, 80])
    ax.set_xticklabels(["0°", "20°", "40°", "60°", "80°"])
    for side in ("top", "right"):
        ax.spines[side].set_visible(False)
    ax.grid(axis="y", color=GRID, lw=0.6)
    ax.set_axisbelow(True)
    ax.set_xlabel("Latitude (north or south)")
    ax.set_ylabel("Error of the perimeter (%)")
    shapes = [0.0, 45.0, 60.0, 70.5439, 80.0]
    top = ax.secondary_xaxis("top")
    top.set_xticks(shapes)
    top.set_xticklabels([f"{pixel_edges(v)[1] / pixel_edges(v)[0]:.1f} : 1" for v in shapes])
    top.tick_params(labelsize=8)
    top.set_xlabel("Pixel height : width", fontsize=8.5)
    ax.legend(loc="upper left", frameon=False, fontsize=8.5)
    ax.text(0.0, 1.20, "(b) Ellipses of known perimeter on the 3$''$ grid (median, 10th–90th percentiles)",
            transform=ax.transAxes, fontsize=10, fontweight="bold", va="bottom")


def main():
    rows = np.load(DATA)
    assert rows.shape[1] == 8
    fig = plt.figure(figsize=(7.4, 7.2))
    spec = fig.add_gridspec(3, 1, height_ratios=[0.62, 0.62, 0.95], hspace=0.0)
    mid = spec[0].subgridspec(1, 3, wspace=0.10)
    axes_a = [fig.add_subplot(mid[i]) for i in range(3)]
    dx30, dy30 = pixel_edges(30.0)
    dx70, dy70 = pixel_edges(70.5439)
    ratios = [panel_edge(axes_a[0], 30.0, 1, PIX, DASH, f"30° N: pixels {dx30:.0f} m × {dy30:.0f} m", "on the pixels"),
              panel_edge(axes_a[1], 70.5439, 1, PIX, DASH, f"70.5° N: pixels {dx70:.0f} m × {dy70:.0f} m",
                         "on the pixels"),
              panel_edge(axes_a[2], 70.5439, 3, CELL, "-", f"70.5° N: cells of 3 columns, {3 * dx70:.0f} m × {dy70:.0f} m",
                         "on the cells")]
    axes_a[0].text(0.0, 1.30, "(a) A boundary at 45° on the ground (dotted), traced on pixels and on cells of merged "
                   "columns", transform=axes_a[0].transAxes, fontsize=10, fontweight="bold", va="bottom")
    axes_a[1].text(0.5, -0.24, "Cells of $k$ columns: a cell is in the catchment when any of its $k$ pixels is; "
                   "cells start at global columns that are multiples of $k$." "\n"
                   r"Any latitude: $a = dy/dx$ at the catchment's middle row, $k_0 = \max(1, \lfloor a \rfloor)$, "
                   r"$w = \min(1, \max(0, a - k_0))$;" "\n"
                   r"$P = (1 - w)\,P(k_0) + w\,P(k_0 + 1)$, where $P(k)$ is the same trace and weights on cells of "
                   r"$k$ columns",
                   transform=axes_a[1].transAxes, ha="center", va="top", fontsize=8.3, linespacing=1.5)

    ax_b = fig.add_subplot(spec[2])
    panel_known(ax_b, rows)
    for ext in ("png", "pdf"):
        out = os.path.join(OUT_DIR, f"FigA5_Perimeter_Method_v11.{ext}")
        if os.path.exists(out):
            raise SystemExit(f"{out} is there already")
        fig.savefig(out, dpi=300, bbox_inches="tight", facecolor="white")
        print(f"Saved: {out}")
    print("  (a) errors " + ", ".join(f"{100 * (r - 1):+.2f} %" for r in ratios)
          + f"; (b) at {rows[-1, 0]:.0f}°: pixels {100 * (rows[-1, 3] - 1):+.1f} %, FlowMorph 1.1 "
          f"{100 * (rows[-1, 6] - 1):+.2f} %; FlowMorph 1.1 medians {100 * (rows[:, 6].min() - 1):+.2f} to "
          f"{100 * (rows[:, 6].max() - 1):+.2f} %")
    plt.close(fig)


if __name__ == "__main__":
    main()