"""Vizualizace: mapy úkosů, podkosů, změn a odlehčené meshe pro 3D náhled."""
from __future__ import annotations

import base64
import io

import numpy as np

from .draft import draft_from_slope, slope_from_draft
from .raster import Grid


# --------------------------------------------------------------------------- #
def surface_draft_map(h: np.ndarray, res: float) -> np.ndarray:
    """Mapa úkosu (°) po částech lineární plochy nad výškovým polem."""
    ny, nx = h.shape
    z00 = h[:-1, :-1]
    z10 = h[:-1, 1:]
    z01 = h[1:, :-1]
    z11 = h[1:, 1:]
    p = z10 - z00
    q = z11 - z10
    u = z11 - z01
    v = z01 - z00
    g = np.sqrt(np.maximum(p * p, u * u) + np.maximum(q * q, v * v)) / res
    out = np.full((ny, nx), np.nan)
    ang = np.degrees(np.arctan2(1.0, np.maximum(g, 1e-12)))
    out[:-1, :-1] = ang
    out[-1, :] = out[-2, :]
    out[:, -1] = out[:, -2]
    return out


def _png(fig) -> str:
    buf = io.BytesIO()
    fig.savefig(buf, format="png", dpi=110, bbox_inches="tight",
                facecolor=fig.get_facecolor())
    buf.seek(0)
    import matplotlib.pyplot as plt
    plt.close(fig)
    return "data:image/png;base64," + base64.b64encode(buf.read()).decode()


def make_maps(result, required_draft: float | None = None) -> dict:
    """Vrátí slovník {název: data-URL PNG}.  Vyžaduje matplotlib."""
    try:
        import matplotlib
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
        from matplotlib.colors import ListedColormap, BoundaryNorm
    except Exception:
        return {}

    g: Grid = result.grid
    h0 = result.h0
    h = result.h
    req = required_draft or result.params.get("draft_deg", 3.0)
    ext = [g.x0, g.x0 + (g.nx - 1) * g.res, g.y0, g.y0 + (g.ny - 1) * g.res]
    out = {}

    def new_fig():
        fig, ax = plt.subplots(figsize=(4.6, 4.1), facecolor="#101318")
        ax.set_facecolor("#101318")
        ax.tick_params(colors="#8b95a5", labelsize=7)
        for s in ax.spines.values():
            s.set_color("#2a3240")
        return fig, ax

    # 1) výška původního reliéfu
    fig, ax = new_fig()
    im = ax.imshow(np.where(h0 > 0, h0, np.nan), origin="lower", extent=ext,
                   cmap="viridis")
    ax.set_title("Výška reliéfu h₀ (mm)", color="#dfe6f0", fontsize=9)
    cb = fig.colorbar(im, ax=ax, fraction=0.046)
    cb.ax.tick_params(colors="#8b95a5", labelsize=7)
    out["vyska"] = _png(fig)

    # 2) úkos původního povrchu
    dm = surface_draft_map(h0, g.res)
    dm = np.where(h0 > 1e-9, dm, np.nan)
    fig, ax = new_fig()
    cmap = ListedColormap(["#c0392b", "#e67e22", "#f1c40f",
                           "#2ecc71", "#16a085"])
    norm = BoundaryNorm([0, req * 0.5, req, req * 2, req * 5, 90], cmap.N)
    im = ax.imshow(dm, origin="lower", extent=ext, cmap=cmap, norm=norm)
    ax.set_title(f"Úkos původního povrchu (°), práh {req}°",
                 color="#dfe6f0", fontsize=9)
    cb = fig.colorbar(im, ax=ax, fraction=0.046)
    cb.ax.tick_params(colors="#8b95a5", labelsize=7)
    out["ukos"] = _png(fig)

    # 3) mapa podkosů (vícevrstevnatost) a nedostatečných úkosů
    prob = np.zeros(h0.shape)
    prob[(result.hits > 2)] = 2.0
    bad = (dm < req) & np.isfinite(dm)
    prob[bad & (prob == 0)] = 1.0
    fig, ax = new_fig()
    cmap = ListedColormap(["#151a21", "#e67e22", "#c0392b"])
    im = ax.imshow(prob, origin="lower", extent=ext, cmap=cmap, vmin=0, vmax=2)
    ax.set_title("Oranžově malý úkos, červeně podkos",
                 color="#dfe6f0", fontsize=9)
    out["podkosy"] = _png(fig)

    # 4) změna po korekci
    d = h - h0
    fig, ax = new_fig()
    lim = float(np.abs(d).max()) or 1.0
    im = ax.imshow(np.where(np.abs(d) > 1e-6, d, np.nan), origin="lower",
                   extent=ext, cmap="coolwarm", vmin=-lim, vmax=lim)
    ax.set_title("Změna výšky h − h₀ (mm)", color="#dfe6f0", fontsize=9)
    cb = fig.colorbar(im, ax=ax, fraction=0.046)
    cb.ax.tick_params(colors="#8b95a5", labelsize=7)
    out["zmena"] = _png(fig)

    # 5) histogram úkosů původní sítě
    an = result.report.get("draft_analysis", {})
    edges = an.get("hist_edges")
    vals = an.get("hist_area_mm2")
    if edges and vals:
        fig, ax = plt.subplots(figsize=(5.4, 2.8), facecolor="#101318")
        ax.set_facecolor("#101318")
        centers = np.arange(len(vals))
        colors = ["#c0392b" if edges[i + 1] <= 0 else
                  ("#e67e22" if edges[i + 1] <= req else "#2ecc71")
                  for i in range(len(vals))]
        ax.bar(centers, vals, color=colors)
        ax.set_xticks(centers)
        ax.set_xticklabels([f"{edges[i]:g}…{edges[i+1]:g}"
                            for i in range(len(vals))],
                           rotation=60, fontsize=6, color="#8b95a5")
        ax.set_ylabel("plocha (mm²)", color="#8b95a5", fontsize=8)
        ax.set_title("Rozdělení úkosů původního STL (°)",
                     color="#dfe6f0", fontsize=9)
        ax.tick_params(colors="#8b95a5", labelsize=7)
        for s in ax.spines.values():
            s.set_color("#2a3240")
        out["histogram"] = _png(fig)

    return out


# --------------------------------------------------------------------------- #
def decimate(grid: Grid, arrays, max_cells: int = 260):
    """Podvzorkuje mřížku i pole pro rychlý náhled v prohlížeči."""
    step = max(1, int(np.ceil(max(grid.nx, grid.ny) / max_cells)))
    if step == 1:
        return grid, arrays
    sub = Grid(grid.x0, grid.y0, grid.res * step,
               len(range(0, grid.nx, step)), len(range(0, grid.ny, step)))
    return sub, [a[::step, ::step] for a in arrays]
