"""Korekce úkosů — Lipschitzovy obálky výškového pole.

Matematické jádro
-----------------
Výškové pole h má v každém bodě úkos

    delta(p) = arctan( 1 / |grad h(p)| )        [odklon stěny od osy Z]

Podmínka „všude alespoň delta_req“ je proto ekvivalentní

    |grad h| <= s,      s = cot(delta_req)

tedy tomu, že h je s-Lipschitzovská funkce.  Minimální úpravy jsou dvě:

  * **horní obálka** (přidání materiálu)
        h_up(p) = max_q ( h(q) - s * |p - q| )
    = nejmenší s-Lipschitzovská funkce >= h.

  * **dolní obálka** (odebrání materiálu)
        h_lo(p) = min_q ( h(q) + s * |p - q| )
    = největší s-Lipschitzovská funkce <= h.

  * jejich aritmetický průměr je opět s-Lipschitzovský (vyvážený režim).

Obě obálky měníme jen tam, kde je to nutné; boční dosah úpravy je
(delta h)/s, tedy pro 3° a výškový skok 10 mm pouhých 0.52 mm — úprava je
proto z hlediska celkového tvaru nenápadná.

Výpočet: chamfer „fast sweeping“ s 16-okolím doplněný o **exaktní**
neomezené 1D scany ve směru X (trik s kumulativním maximem).  Reziduální
anizotropie metriky je kompenzována bezpečnostním koeficientem, takže
výsledek podmínku |grad h| <= s vždy splňuje.
"""
from __future__ import annotations

import numpy as np

# stencil (di, dj>0); pro dj se sweepuje po řádcích
_STENCIL = [(0, 1), (-1, 1), (1, 1), (-2, 1), (2, 1), (-3, 1), (3, 1),
            (-1, 2), (1, 2), (-3, 2), (3, 2), (-1, 3), (1, 3),
            (-2, 3), (2, 3)]


def slope_from_draft(draft_deg: float) -> float:
    """s = cot(delta).  Pro delta -> 0 vrací velmi velké číslo (svislá stěna)."""
    d = float(np.clip(draft_deg, 1e-4, 89.99))
    return float(1.0 / np.tan(np.radians(d)))


def draft_from_slope(slope: float) -> float:
    if slope <= 0:
        return 90.0
    return float(np.degrees(np.arctan(1.0 / slope)))


def _decay_max_x(row: np.ndarray, step: float) -> np.ndarray:
    """Exaktní 1D horní obálka podél osy X: max_k(row[k] - step*|i-k|)."""
    n = row.size
    k = np.arange(n, dtype=np.float64) * step
    b = row + k
    np.maximum.accumulate(b, out=b)
    left = b - k
    c = (row - k)[::-1].copy()
    np.maximum.accumulate(c, out=c)
    right = c[::-1] + k
    return np.maximum(left, right)


def cone_dilate(h: np.ndarray, res: float, slope: float,
                rounds: int = 2) -> np.ndarray:
    """Horní s-Lipschitzovská obálka pole h (kuželová dilatace)."""
    g = np.ascontiguousarray(h, dtype=np.float64).copy()
    ny, nx = g.shape
    step_x = slope * res
    costs = [(di, dj, slope * res * float(np.hypot(di, dj)))
             for di, dj in _STENCIL]

    for _ in range(int(rounds)):
        for direction in (1, -1):
            rows = range(ny) if direction == 1 else range(ny - 1, -1, -1)
            for j in rows:
                row = g[j]
                for di, dj, cost in costs:
                    jp = j - direction * dj
                    if jp < 0 or jp >= ny:
                        continue
                    src = g[jp]
                    if di == 0:
                        np.maximum(row, src - cost, out=row)
                    elif di > 0:
                        np.maximum(row[di:], src[:nx - di] - cost,
                                   out=row[di:])
                    else:
                        k = -di
                        np.maximum(row[:nx - k], src[k:] - cost,
                                   out=row[:nx - k])
                g[j] = _decay_max_x(row, step_x)
    return g


def cone_erode(h: np.ndarray, res: float, slope: float,
               rounds: int = 2) -> np.ndarray:
    """Dolní s-Lipschitzovská obálka pole h (kuželová eroze)."""
    return -cone_dilate(-h, res, slope, rounds)


def max_slope(h: np.ndarray, res: float) -> float:
    """Maximální diskrétní |grad h| (4- i 8-okolí)."""
    gx = np.abs(np.diff(h, axis=1)) / res
    gy = np.abs(np.diff(h, axis=0)) / res
    gd1 = np.abs(h[1:, 1:] - h[:-1, :-1]) / (res * np.sqrt(2.0))
    gd2 = np.abs(h[1:, :-1] - h[:-1, 1:]) / (res * np.sqrt(2.0))
    vals = [a.max() for a in (gx, gy, gd1, gd2) if a.size]
    return float(max(vals)) if vals else 0.0


def max_slope_pl(h: np.ndarray, res: float) -> float:
    """Maximální sklon **po částech lineární** plochy nad mřížkou.

    Rozhodující není rozdíl hodnot mezi uzly, ale gradient jednotlivých
    trojúhelníků.  Pro buňku s rohy 00,10,11,01 a obě možné úhlopříčky
    vycházejí gradienty (p,q), (u,v), (p,v), (u,q), kde

        p = z10-z00,  q = z11-z10,  u = z11-z01,  v = z01-z00

    Maximum přes všechny čtyři je sqrt(max(p²,u²) + max(q²,v²)) / res.
    V nejhorším případě (sedlo) je až sqrt(2)krát větší než rozdíl mezi
    sousedními uzly — proto se musí kontrolovat takto, ne po hranách.
    """
    if h.shape[0] < 2 or h.shape[1] < 2:
        return 0.0
    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
    m = np.sqrt(np.maximum(p * p, u * u) + np.maximum(q * q, v * v))
    return float(m.max()) / res


def achieved_draft(h: np.ndarray, res: float) -> float:
    """Skutečně dosažený minimální úkos vygenerované plochy ve stupních."""
    return draft_from_slope(max_slope_pl(h, res))


def correct_draft(h: np.ndarray, res: float, draft_deg: float,
                  mode: str = "add", safety: float = 1.03,
                  smooth_px: float = 0.0, rounds: int = 2,
                  max_iter: int = 6):
    """Upraví výškové pole tak, aby všude platilo delta >= draft_deg.

    mode: 'add' (přidat materiál) | 'cut' (ubrat) | 'balanced' (půl na půl)
    safety: počáteční rezerva (>1); dále se dolaďuje automaticky
    smooth_px: volitelné zjemnění přechodů; po vyhlazení se obálka aplikuje
               znovu, takže podmínka zůstává splněna.

    Po každé iteraci se měří **skutečný** sklon trojúhelníkové plochy
    (max_slope_pl).  Pokud požadavek není splněn, zpřísní se cílový sklon a
    obálka se přepočítá — výsledek proto dodržuje zadaný úkos i po
    triangulaci, a to s minimální nutnou rezervou.
    """
    s_target = slope_from_draft(draft_deg)
    s_eff = s_target / max(1.0, float(safety))

    def _apply(field, s):
        up = cone_dilate(field, res, s, rounds)
        if mode == "add":
            return np.maximum(up, 0.0)
        lo = cone_erode(field, res, s, rounds)
        if mode == "cut":
            return np.maximum(lo, 0.0)
        return np.maximum(0.5 * (up + lo), 0.0)

    def _post(field, s):
        if smooth_px and smooth_px > 0:
            try:
                from scipy.ndimage import gaussian_filter
                field = gaussian_filter(field, float(smooth_px),
                                        mode="nearest")
                field = _apply(np.maximum(field, 0.0), s)
            except Exception:
                pass
        return field

    out = _post(_apply(h, s_eff), s_eff)
    for _ in range(int(max_iter)):
        g = max_slope_pl(out, res)
        if g <= s_target * (1.0 + 1e-9):
            break
        s_eff *= max(0.5, s_target / g) * 0.999
        out = _post(_apply(h, s_eff), s_eff)
    return out
