"""Rasterizace meshe do výškové mapy (horní obálka) pomocí z-bufferu.

Klíčová myšlenka celého Babofkátoru
-----------------------------------
Těleso je odformovatelné jednodílnou formou ve směru +Z **právě tehdy**,
když je popsatelné výškovým polem h(x, y):

    Solid = { (x, y, z) : 0 <= z <= h(x, y) }

Důkaz (náčrt): řezy tělesa rovinami z = konst tvoří vnořenou klesající
posloupnost { h >= z }.  Doplněk (dutina formy) tedy směrem nahoru pouze
roste, takže odlitek lze vysunout podél +Z bez kolize.  Naopak, není-li
těleso výškovým polem, existuje svislá přímka protínající těleso ve dvou
oddělených intervalech — mezi nimi je materiál formy, který vysunutí brání.

Horní obálka (max z přes všechny trojúhelníky nad daným bodem) je proto
minimální možná úprava, která odstraní **všechny** podkosy — lokální i
globální (zastínění jinou částí modelu).
"""
from __future__ import annotations

from dataclasses import dataclass

import numpy as np


@dataclass
class Grid:
    x0: float
    y0: float
    res: float
    nx: int
    ny: int

    @property
    def shape(self):
        return (self.ny, self.nx)

    def xs(self):
        return self.x0 + np.arange(self.nx) * self.res

    def ys(self):
        return self.y0 + np.arange(self.ny) * self.res

    @property
    def cell_area(self):
        return self.res * self.res


def make_grid(vmin, vmax, res: float, pad_mm: float) -> Grid:
    x0 = float(vmin[0]) - pad_mm
    y0 = float(vmin[1]) - pad_mm
    nx = int(np.ceil((float(vmax[0]) + pad_mm - x0) / res)) + 1
    ny = int(np.ceil((float(vmax[1]) + pad_mm - y0) / res)) + 1
    return Grid(x0, y0, res, nx, ny)


def _scatter_max(dst_flat: np.ndarray, idx: np.ndarray, val: np.ndarray):
    """dst[idx] = max(dst[idx], val) i pro opakující se indexy."""
    if idx.size == 0:
        return
    if idx.size < 200_000:
        np.maximum.at(dst_flat, idx, val)
        return
    order = np.argsort(idx, kind="stable")
    si = idx[order]
    sv = val[order]
    starts = np.flatnonzero(np.r_[True, si[1:] != si[:-1]])
    grouped = np.maximum.reduceat(sv, starts)
    keys = si[starts]
    np.maximum(dst_flat[keys], grouped, out=grouped)
    dst_flat[keys] = grouped


def _scatter_add(dst_flat: np.ndarray, idx: np.ndarray):
    if idx.size == 0:
        return
    dst_flat += np.bincount(idx, minlength=dst_flat.size).astype(dst_flat.dtype)


def rasterize(vertices: np.ndarray, faces: np.ndarray, grid: Grid,
              eps_bary: float = 1e-9, max_batch_cells: int = 6_000_000):
    """Z-buffer rasterizace.

    Vrací (hmax, hmin, hits):
        hmax  (ny,nx) horní obálka (-inf mimo model)
        hmin  (ny,nx) dolní obálka (+inf mimo model)
        hits  (ny,nx) počet zásahů trojúhelníkem (heuristika vícevrstevnatosti)
    """
    ny, nx = grid.shape
    hmax = np.full(ny * nx, -np.inf, dtype=np.float64)
    hmin = np.full(ny * nx, np.inf, dtype=np.float64)
    hits = np.zeros(ny * nx, dtype=np.int32)

    a = vertices[faces[:, 0]]
    b = vertices[faces[:, 1]]
    c = vertices[faces[:, 2]]

    # znaménková plocha průmětu do XY; svislé trojúhelníky (d ~ 0) vynecháme,
    # jejich průmět má nulovou plochu
    d = (b[:, 0] - a[:, 0]) * (c[:, 1] - a[:, 1]) - \
        (b[:, 1] - a[:, 1]) * (c[:, 0] - a[:, 0])
    scale = np.maximum(np.abs(d), 0.0)
    keep = scale > (1e-12 * max(1.0, float(np.abs(vertices).max()) ** 2))
    if not np.any(keep):
        return (hmax.reshape(ny, nx), hmin.reshape(ny, nx), hits.reshape(ny, nx))

    a, b, c, d = a[keep], b[keep], c[keep], d[keep]

    xmin = np.minimum.reduce([a[:, 0], b[:, 0], c[:, 0]])
    xmax = np.maximum.reduce([a[:, 0], b[:, 0], c[:, 0]])
    ymin = np.minimum.reduce([a[:, 1], b[:, 1], c[:, 1]])
    ymax = np.maximum.reduce([a[:, 1], b[:, 1], c[:, 1]])

    i0 = np.ceil((xmin - grid.x0) / grid.res - 1e-9).astype(np.int64)
    i1 = np.floor((xmax - grid.x0) / grid.res + 1e-9).astype(np.int64)
    j0 = np.ceil((ymin - grid.y0) / grid.res - 1e-9).astype(np.int64)
    j1 = np.floor((ymax - grid.y0) / grid.res + 1e-9).astype(np.int64)
    np.clip(i0, 0, nx - 1, out=i0)
    np.clip(i1, 0, nx - 1, out=i1)
    np.clip(j0, 0, ny - 1, out=j0)
    np.clip(j1, 0, ny - 1, out=j1)

    w = i1 - i0 + 1
    hgt = j1 - j0 + 1
    ok = (w > 0) & (hgt > 0)
    if not np.any(ok):
        return (hmax.reshape(ny, nx), hmin.reshape(ny, nx), hits.reshape(ny, nx))

    idx_all = np.flatnonzero(ok)
    # oddělené zaokrouhlení šířky a výšky na mocninu dvou — dlouhé úzké
    # trojúhelníky by při společném čtvercovém okně plýtvaly pamětí i časem
    def _p2(v):
        return np.maximum(1, 1 << np.ceil(np.log2(np.maximum(v, 1))).astype(int))

    bw = _p2(w[idx_all])
    bh = _p2(hgt[idx_all])
    key = bw.astype(np.int64) * 100000 + bh

    for kv in np.unique(key):
        sel = idx_all[key == kv]
        Sx = int(kv // 100000)
        Sy = int(kv % 100000)
        per = max(1, int(max_batch_cells // (Sx * Sy)))
        for s in range(0, sel.size, per):
            t = sel[s:s + per]
            nT = t.size
            ax, ay, az = a[t, 0], a[t, 1], a[t, 2]
            bx, by, bz = b[t, 0], b[t, 1], b[t, 2]
            cx, cy, cz = c[t, 0], c[t, 1], c[t, 2]
            dd = d[t]

            ii = i0[t][:, None, None] + np.arange(Sx)[None, None, :]
            jj = j0[t][:, None, None] + np.arange(Sy)[None, :, None]
            # nepatrné iracionální posunutí vzorkovacích bodů zabrání tomu,
            # aby uzel padl přesně na hranu nebo vrchol meshe (dvojí započtení)
            px = grid.x0 + ii * grid.res + 2.7e-4 * grid.res
            py = grid.y0 + jj * grid.res + 4.3e-4 * grid.res

            rx = px - ax[:, None, None]
            ry = py - ay[:, None, None]
            e1x = (bx - ax)[:, None, None]
            e1y = (by - ay)[:, None, None]
            e2x = (cx - ax)[:, None, None]
            e2y = (cy - ay)[:, None, None]
            inv = (1.0 / dd)[:, None, None]
            u = (rx * e2y - ry * e2x) * inv
            v = (e1x * ry - e1y * rx) * inv

            m = (u >= -eps_bary) & (v >= -eps_bary) & (u + v <= 1.0 + eps_bary)
            m &= (ii <= i1[t][:, None, None]) & (jj <= j1[t][:, None, None])
            if not m.any():
                continue
            z = az[:, None, None] + u * (bz - az)[:, None, None] + \
                v * (cz - az)[:, None, None]

            flat = (np.broadcast_to(jj, m.shape) * nx +
                    np.broadcast_to(ii, m.shape))[m]
            zz = np.broadcast_to(z, m.shape)[m]
            _scatter_max(hmax, flat, zz)
            _scatter_max(hmin, flat, -zz)   # trik: min přes max(-z)
            _scatter_add(hits, flat)

    hmin = -hmin
    return (hmax.reshape(ny, nx), hmin.reshape(ny, nx), hits.reshape(ny, nx))


def height_field(vertices, faces, grid: Grid, z_plane: float = 0.0):
    """Horní obálka jako výškové pole h >= 0 + maska půdorysu.

    Vrací dict:
        h        (ny,nx) výškové pole, 0 mimo model
        covered  (ny,nx) bool, půdorys modelu
        hits     (ny,nx) počet zásahů (>2 ~ vícevrstevnatost)
        hmin     dolní obálka (nan mimo model)
    """
    hmax, hmin, hits = rasterize(vertices, faces, grid)
    covered = np.isfinite(hmax)
    h = np.where(covered, np.maximum(hmax - z_plane, 0.0), 0.0)
    hmin_out = np.where(covered, hmin - z_plane, np.nan)
    return {"h": h, "covered": covered, "hits": hits, "hmin": hmin_out}


def mesh_volume(vertices: np.ndarray, faces: np.ndarray) -> float:
    """Objem uzavřeného tělesa (divergenční věta).  Kladný pro normály ven."""
    v0 = vertices[faces[:, 0]]
    v1 = vertices[faces[:, 1]]
    v2 = vertices[faces[:, 2]]
    return float(np.einsum("ij,ij->i", v0, np.cross(v1, v2)).sum() / 6.0)
