"""Převod dvojice výškových polí na uzavřený (vodotěsný) trojúhelníkový mesh.

Těleso je definováno jako  { (x,y,z) : zb(x,y) <= z <= zt(x,y) }  na doméně
zadané maskou buněk.  Konstrukce je z principu 2-manifold a vodotěsná:
každá hranová (boční) stěna je generována právě jednou a spodní i horní
plocha používají tutéž triangulaci buněk.
"""
from __future__ import annotations

import numpy as np

from .raster import Grid


# --------------------------------------------------------------------------- #
#  Úpravy masky
# --------------------------------------------------------------------------- #
def fill_holes(mask: np.ndarray) -> np.ndarray:
    try:
        from scipy.ndimage import binary_fill_holes
        return binary_fill_holes(mask)
    except Exception:
        return mask


def largest_component(mask: np.ndarray, keep_all: bool = True) -> np.ndarray:
    if keep_all:
        return mask
    try:
        from scipy.ndimage import label
        lab, n = label(mask)
        if n <= 1:
            return mask
        sizes = np.bincount(lab.ravel())
        sizes[0] = 0
        return lab == int(sizes.argmax())
    except Exception:
        return mask


def fix_pinches(mask: np.ndarray, max_iter: int = 8) -> np.ndarray:
    """Odstraní diagonální „skřípnutí“ masky (nemanifoldní vrcholy).

    Vzor  [[1,0],[0,1]] nebo [[0,1],[1,0]] doplní na plný čtverec.
    """
    m = mask.copy()
    for _ in range(max_iter):
        a = m[:-1, :-1]
        b = m[:-1, 1:]
        c = m[1:, :-1]
        d = m[1:, 1:]
        p1 = a & d & ~b & ~c
        p2 = b & c & ~a & ~d
        if not (p1.any() or p2.any()):
            break
        m[:-1, 1:] |= p1
        m[1:, :-1] |= p1
        m[:-1, :-1] |= p2
        m[1:, 1:] |= p2
    return m


def _boundary_nodes(cm: np.ndarray) -> np.ndarray:
    """Uzly, které sousedí alespoň s jednou buňkou mimo masku."""
    ny, nx = cm.shape[0] + 1, cm.shape[1] + 1
    inside = np.zeros((ny, nx), dtype=bool)
    full = np.zeros((ny, nx), dtype=bool)
    full[:-1, :-1] |= cm
    full[:-1, 1:] |= cm
    full[1:, :-1] |= cm
    full[1:, 1:] |= cm
    inside[1:-1, 1:-1] = (cm[:-1, :-1] & cm[:-1, 1:] &
                          cm[1:, :-1] & cm[1:, 1:])
    return full & ~inside


def cells_from_nodes(node_mask: np.ndarray) -> np.ndarray:
    """Buňka je zahrnuta, pokud je zahrnut alespoň jeden její roh."""
    return (node_mask[:-1, :-1] | node_mask[:-1, 1:] |
            node_mask[1:, :-1] | node_mask[1:, 1:])


def cells_from_nodes_all(node_mask: np.ndarray) -> np.ndarray:
    return (node_mask[:-1, :-1] & node_mask[:-1, 1:] &
            node_mask[1:, :-1] & node_mask[1:, 1:])


# --------------------------------------------------------------------------- #
#  Meshování
# --------------------------------------------------------------------------- #
def heightfields_to_mesh(grid: Grid, z_bottom: np.ndarray, z_top: np.ndarray,
                         cell_mask: np.ndarray, tiny: float = 1e-9):
    """Vrátí (vertices, faces) uzavřeného tělesa mezi z_bottom a z_top."""
    ny, nx = grid.shape
    assert z_bottom.shape == (ny, nx) and z_top.shape == (ny, nx)
    assert cell_mask.shape == (ny - 1, nx - 1)

    thin_real = (z_top - z_bottom) <= tiny
    all_thin = (thin_real[:-1, :-1] & thin_real[:-1, 1:] &
                thin_real[1:, :-1] & thin_real[1:, 1:])

    cm = cell_mask.astype(bool)
    for _ in range(3):
        cm = fix_pinches(cm) & ~all_thin      # buňky s nulovou tloušťkou
        if not (fix_pinches(cm) & all_thin).any():
            break
    if not cm.any():
        return np.zeros((0, 3)), np.zeros((0, 3), dtype=np.int64)
    thin_node = thin_real & _boundary_nodes(cm)

    node_used = np.zeros((ny, nx), dtype=bool)
    node_used[:-1, :-1] |= cm
    node_used[:-1, 1:] |= cm
    node_used[1:, :-1] |= cm
    node_used[1:, 1:] |= cm

    xs = grid.xs()
    ys = grid.ys()
    X = np.broadcast_to(xs[None, :], (ny, nx))
    Y = np.broadcast_to(ys[:, None], (ny, nx))

    order = np.flatnonzero(node_used.ravel())
    n_used = order.size
    bidx = np.full(ny * nx, -1, dtype=np.int64)
    bidx[order] = np.arange(n_used)

    zb = z_bottom.ravel()[order]
    zt = z_top.ravel()[order]
    thin = thin_node.ravel()[order]

    tidx = np.full(ny * nx, -1, dtype=np.int64)
    tidx[order[thin]] = bidx[order[thin]]
    extra = order[~thin]
    tidx[extra] = n_used + np.arange(extra.size)

    Vb = np.stack([X.ravel()[order], Y.ravel()[order], zb], axis=1)
    Vt = np.stack([X.ravel()[extra], Y.ravel()[extra],
                   z_top.ravel()[extra]], axis=1)
    vertices = np.concatenate([Vb, Vt], axis=0)

    bidx = bidx.reshape(ny, nx)
    tidx = tidx.reshape(ny, nx)

    jc, ic = np.nonzero(cm)

    b00 = bidx[jc, ic]
    b10 = bidx[jc, ic + 1]
    b11 = bidx[jc + 1, ic + 1]
    b01 = bidx[jc + 1, ic]
    t00 = tidx[jc, ic]
    t10 = tidx[jc, ic + 1]
    t11 = tidx[jc + 1, ic + 1]
    t01 = tidx[jc + 1, ic]

    # Volba úhlopříčky buňky: preferujeme tu, která spojuje méně „splynulých“
    # uzlů (kde se dno a vrch dotýkají).  Jinak by úhlopříčná hrana patřila
    # čtyřem stěnám -> nemanifoldní mesh.
    tn = thin_node
    cA = tn[jc, ic].astype(np.int8) + tn[jc + 1, ic + 1].astype(np.int8)
    cB = tn[jc, ic + 1].astype(np.int8) + tn[jc + 1, ic].astype(np.int8)
    A = cA <= cB
    B = ~A

    faces = [
        # úhlopříčka 00-11
        np.stack([b00[A], b11[A], b10[A]], axis=1),   # dno, normála -Z
        np.stack([b00[A], b01[A], b11[A]], axis=1),
        np.stack([t00[A], t10[A], t11[A]], axis=1),   # vrch, normála +Z
        np.stack([t00[A], t11[A], t01[A]], axis=1),
        # úhlopříčka 10-01
        np.stack([b00[B], b01[B], b10[B]], axis=1),
        np.stack([b10[B], b01[B], b11[B]], axis=1),
        np.stack([t00[B], t10[B], t01[B]], axis=1),
        np.stack([t10[B], t11[B], t01[B]], axis=1),
    ]

    def _wall(sel_j, sel_i, nb_node_a, nb_node_b):
        """Boční stěna mezi uzly A a B (A->B tak, aby normála mířila ven)."""
        if sel_j.size == 0:
            return
        ja, ia = nb_node_a
        jb, ib = nb_node_b
        ba = bidx[sel_j + ja, sel_i + ia]
        bb = bidx[sel_j + jb, sel_i + ib]
        ta = tidx[sel_j + ja, sel_i + ia]
        tb = tidx[sel_j + jb, sel_i + ib]
        faces.append(np.stack([ba, bb, tb], axis=1))
        faces.append(np.stack([ba, tb, ta], axis=1))

    nbR = np.zeros_like(cm)
    nbR[:, :-1] = cm[:, 1:]
    nbL = np.zeros_like(cm)
    nbL[:, 1:] = cm[:, :-1]
    nbU = np.zeros_like(cm)
    nbU[:-1, :] = cm[1:, :]
    nbD = np.zeros_like(cm)
    nbD[1:, :] = cm[:-1, :]

    j, i = np.nonzero(cm & ~nbR)
    _wall(j, i, (0, 1), (1, 1))            # pravá hrana, normála +X
    j, i = np.nonzero(cm & ~nbL)
    _wall(j, i, (1, 0), (0, 0))            # levá hrana, normála -X
    j, i = np.nonzero(cm & ~nbU)
    _wall(j, i, (1, 1), (1, 0))            # horní hrana, normála +Y
    j, i = np.nonzero(cm & ~nbD)
    _wall(j, i, (0, 0), (0, 1))            # dolní hrana, normála -Y

    F = np.concatenate(faces, axis=0)
    ok = (F[:, 0] != F[:, 1]) & (F[:, 1] != F[:, 2]) & (F[:, 0] != F[:, 2])
    F = F[ok]

    used = np.unique(F)
    remap = np.full(vertices.shape[0], -1, dtype=np.int64)
    remap[used] = np.arange(used.size)
    return vertices[used], remap[F]


# --------------------------------------------------------------------------- #
#  Adaptivní (kvadrantové) meshování — výrazně menší STL při stejné přesnosti
# --------------------------------------------------------------------------- #
def _block_all(arr: np.ndarray, B: int) -> np.ndarray:
    """Logické AND přes bloky B×B (zbytek mřížky se zahodí)."""
    nj, ni = arr.shape
    nbj, nbi = nj // B, ni // B
    if nbj == 0 or nbi == 0:
        return np.zeros((0, 0), dtype=bool)
    return arr[:nbj * B, :nbi * B].reshape(nbj, B, nbi, B).all(axis=(1, 3))


def _planar_dev(Z: np.ndarray, B: int, nbj: int, nbi: int) -> np.ndarray:
    """Max. odchylka uzlů bloku od bilineární plochy proložené jeho rohy."""
    jj = np.arange(nbj)[:, None] * B + np.arange(B + 1)[None, :]
    ii = np.arange(nbi)[:, None] * B + np.arange(B + 1)[None, :]
    Zb = Z[jj[:, :, None, None], ii[None, None, :, :]]     # (nbj,B+1,nbi,B+1)
    z00 = Zb[:, :1, :, :1]
    z01 = Zb[:, :1, :, -1:]
    z10 = Zb[:, -1:, :, :1]
    z11 = Zb[:, -1:, :, -1:]
    a = (np.arange(B + 1) / B)[None, :, None, None]        # podél j
    b = (np.arange(B + 1) / B)[None, None, None, :]        # podél i
    interp = (z00 * (1 - a) * (1 - b) + z01 * (1 - a) * b +
              z10 * a * (1 - b) + z11 * a * b)
    return np.abs(Zb - interp).max(axis=(1, 3))


def _level_map(cm: np.ndarray, zb: np.ndarray, zt: np.ndarray,
               thin_node: np.ndarray, tol: float, max_level: int) -> np.ndarray:
    """Úroveň zjednodušení pro každou buňku (0 = plné rozlišení)."""
    nj, ni = cm.shape
    L = np.zeros((nj, ni), dtype=np.int16)

    # bloky úrovně >= 1 musí ležet i s jednobuňkovým lemem uvnitř masky
    # (hranice masky se tak vždy řeší v plném rozlišení) a nesmí obsahovat
    # uzly s nulovou tloušťkou
    inner = cm.copy()
    inner[1:, :] &= cm[:-1, :]
    inner[:-1, :] &= cm[1:, :]
    inner[:, 1:] &= cm[:, :-1]
    inner[:, :-1] &= cm[:, 1:]
    inner[1:, 1:] &= cm[:-1, :-1]
    inner[:-1, :-1] &= cm[1:, 1:]
    inner[1:, :-1] &= cm[:-1, 1:]
    inner[:-1, 1:] &= cm[1:, :-1]
    # buňky na okraji pole nemají souseda, kterým by se dala ověřit vnitřnost
    inner[0, :] = inner[-1, :] = False
    inner[:, 0] = inner[:, -1] = False
    fat_cell = ~(thin_node[:-1, :-1] | thin_node[:-1, 1:] |
                 thin_node[1:, :-1] | thin_node[1:, 1:])
    base = inner & fat_cell

    prev_elig = None
    for lev in range(1, max_level + 1):
        B = 1 << lev
        nbj, nbi = nj // B, ni // B
        if nbj == 0 or nbi == 0:
            break
        okm = _block_all(base, B)
        if not okm.any():
            break
        dev = np.maximum(_planar_dev(zb, B, nbj, nbi),
                         _planar_dev(zt, B, nbj, nbi))
        elig = okm & (dev <= tol)
        if prev_elig is not None:                 # vnořenost kvadrantů
            child = _block_all(prev_elig, 2)
            elig = elig[:child.shape[0], :child.shape[1]] & child
            tmp = np.zeros((nbj, nbi), dtype=bool)
            tmp[:elig.shape[0], :elig.shape[1]] = elig
            elig = tmp
        if not elig.any():
            break
        big = np.kron(elig, np.ones((B, B), dtype=bool))
        L[:big.shape[0], :big.shape[1]] = np.where(
            big, lev, L[:big.shape[0], :big.shape[1]])
        prev_elig = elig

    # 2:1 vyvážení: blok, jehož soused je o více než jednu úroveň jemnější,
    # se rozdělí (sníží o jednu úroveň).  Opakuje se do ustálení.
    for _ in range(4 * max_level + 4):
        changed = False
        for lev in range(max_level, 0, -1):
            B = 1 << lev
            nbj, nbi = nj // B, ni // B
            if nbj == 0 or nbi == 0:
                continue
            sub = L[:nbj * B, :nbi * B].reshape(nbj, B, nbi, B)
            blk = (sub == lev).all(axis=(1, 3))
            if not blk.any():
                continue
            rows = np.arange(nbj) * B
            cols = np.arange(nbi) * B
            span = np.arange(B)
            cspan = cols[:, None] + span[None, :]              # (nbi,B)
            rspan = rows[:, None] + span[None, :]              # (nbj,B)
            jm = np.clip(rows - 1, 0, nj - 1)
            jp = np.clip(rows + B, 0, nj - 1)
            im = np.clip(cols - 1, 0, ni - 1)
            ip = np.clip(cols + B, 0, ni - 1)
            s0 = L[jm[:, None, None], cspan[None, :, :]].min(axis=2)
            s2 = L[jp[:, None, None], cspan[None, :, :]].min(axis=2)
            s1 = L[rspan[:, :, None], ip[None, None, :]].min(axis=1)
            s3 = L[rspan[:, :, None], im[None, None, :]].min(axis=1)
            bad = blk & (np.minimum(np.minimum(s0, s1),
                                    np.minimum(s2, s3)) < lev - 1)
            if bad.any():
                changed = True
                rep = np.kron(bad, np.ones((B, B), dtype=bool))
                reg = L[:nbj * B, :nbi * B]
                L[:nbj * B, :nbi * B] = np.where(rep, lev - 1, reg)
        if not changed:
            break
    return L


def heightfields_to_mesh_adaptive(grid: Grid, z_bottom: np.ndarray,
                                  z_top: np.ndarray, cell_mask: np.ndarray,
                                  tol: float = 0.02, max_level: int = 6,
                                  tiny: float = 1e-9):
    """Vodotěsný mesh s adaptivním zjednodušením rovinných oblastí.

    Používá omezený kvadrantový strom (2:1 vyvážení) s vějířovou triangulací
    a vkládáním středů hran sousedících s jemnějším blokem — mesh je proto
    bez prasklin a T-uzlů.  Všechny vrcholy leží v uzlech mřížky.
    """
    if tol <= 0:
        return heightfields_to_mesh(grid, z_bottom, z_top, cell_mask, tiny)

    ny, nx = grid.shape
    thin_real = (z_top - z_bottom) <= tiny
    all_thin = (thin_real[:-1, :-1] & thin_real[:-1, 1:] &
                thin_real[1:, :-1] & thin_real[1:, 1:])
    cm = cell_mask.astype(bool)
    for _ in range(3):
        cm = fix_pinches(cm) & ~all_thin
        if not (fix_pinches(cm) & all_thin).any():
            break
    if not cm.any():
        return np.zeros((0, 3)), np.zeros((0, 3), dtype=np.int64)

    # Dno a vrch se smí ztotožnit jen v uzlech na hranici masky. Uvnitř
    # (např. drážka zaříznutá až k dělící rovině) by společný vrchol udělal
    # z hrany mezi dvěma nulovými uzly hranu se čtyřmi stěnami = nemanifold.
    thin_node = thin_real & _boundary_nodes(cm)

    L = _level_map(cm, z_bottom, z_top, thin_real, tol, max_level)
    L[~cm] = -1

    faces = []          # seznam polí (F,3) s ID = node_flat*2 + vrstva

    def nid(j, i, layer):
        flat = j * nx + i
        return flat * 2 + np.where(thin_node.ravel()[flat], 0, layer)

    # ---- bloky úrovně >= 1: vějířová triangulace -------------------------
    for lev in range(1, max_level + 1):
        B = 1 << lev
        nbj, nbi = (ny - 1) // B, (nx - 1) // B
        if nbj == 0 or nbi == 0:
            continue
        sub = L[:nbj * B, :nbi * B].reshape(nbj, B, nbi, B)
        blk = (sub == lev).all(axis=(1, 3))
        if not blk.any():
            continue
        bj, bi = np.nonzero(blk)
        j0 = bj * B
        i0 = bi * B
        H = B // 2

        def side_finer(axis, sign):
            """Je soused podél dané strany jemnější?"""
            out = np.zeros(bj.size, dtype=bool)
            for k in range(B):
                if axis == 0:                      # strany -y / +y
                    jj = j0 - 1 if sign < 0 else j0 + B
                    ii = i0 + k
                else:                              # strany -x / +x
                    jj = j0 + k
                    ii = i0 - 1 if sign < 0 else i0 + B
                jj = np.clip(jj, 0, L.shape[0] - 1)
                ii = np.clip(ii, 0, L.shape[1] - 1)
                out |= (L[jj, ii] >= 0) & (L[jj, ii] < lev)
            return out

        f = [side_finer(0, -1), side_finer(1, +1),
             side_finer(0, +1), side_finer(1, -1)]
        code = (f[0].astype(np.int8) | (f[1].astype(np.int8) << 1) |
                (f[2].astype(np.int8) << 2) | (f[3].astype(np.int8) << 3))

        for cval in np.unique(code):
            sel = code == cval
            jj0, ii0 = j0[sel], i0[sel]
            ring = [(0, 0)]
            if cval & 1:
                ring.append((0, H))
            ring.append((0, B))
            if cval & 2:
                ring.append((H, B))
            ring.append((B, B))
            if cval & 4:
                ring.append((B, H))
            ring.append((B, 0))
            if cval & 8:
                ring.append((H, 0))
            cj, ci = jj0 + H, ii0 + H
            for k in range(len(ring)):
                aj, ai = ring[k]
                bjj, bii = ring[(k + 1) % len(ring)]
                A = (jj0 + aj, ii0 + ai)
                Bv = (jj0 + bjj, ii0 + bii)
                faces.append(np.stack([nid(cj, ci, 1), nid(*A, 1),
                                       nid(*Bv, 1)], axis=1))
                faces.append(np.stack([nid(cj, ci, 0), nid(*Bv, 0),
                                       nid(*A, 0)], axis=1))

    # ---- buňky úrovně 0: klasické dva trojúhelníky + boční stěny ---------
    zero = (L == 0)
    if zero.any():
        jc, ic = np.nonzero(zero)
        tn = thin_node
        cA = tn[jc, ic].astype(np.int8) + tn[jc + 1, ic + 1].astype(np.int8)
        cB = tn[jc, ic + 1].astype(np.int8) + tn[jc + 1, ic].astype(np.int8)
        A = cA <= cB
        for lay, flip in ((1, False), (0, True)):
            for msk, quads in (
                (A, [((0, 0), (0, 1), (1, 1)), ((0, 0), (1, 1), (1, 0))]),
                (~A, [((0, 0), (0, 1), (1, 0)), ((0, 1), (1, 1), (1, 0))]),
            ):
                if not msk.any():
                    continue
                jm, im = jc[msk], ic[msk]
                for tri in quads:
                    idx = [nid(jm + dj, im + di, lay) for dj, di in tri]
                    if flip:
                        idx = idx[::-1]
                    faces.append(np.stack(idx, axis=1))

    # ---- boční stěny na hranici masky (vždy v plném rozlišení) -----------
    nbR = np.zeros_like(cm)
    nbR[:, :-1] = cm[:, 1:]
    nbL = np.zeros_like(cm)
    nbL[:, 1:] = cm[:, :-1]
    nbU = np.zeros_like(cm)
    nbU[:-1, :] = cm[1:, :]
    nbD = np.zeros_like(cm)
    nbD[1:, :] = cm[:-1, :]

    def wall(sel, a_off, b_off):
        j, i = np.nonzero(sel)
        if j.size == 0:
            return
        aj, ai = j + a_off[0], i + a_off[1]
        bj_, bi_ = j + b_off[0], i + b_off[1]
        ba, bb = nid(aj, ai, 0), nid(bj_, bi_, 0)
        ta, tb = nid(aj, ai, 1), nid(bj_, bi_, 1)
        faces.append(np.stack([ba, bb, tb], axis=1))
        faces.append(np.stack([ba, tb, ta], axis=1))

    wall(cm & ~nbR, (0, 1), (1, 1))
    wall(cm & ~nbL, (1, 0), (0, 0))
    wall(cm & ~nbU, (1, 1), (1, 0))
    wall(cm & ~nbD, (0, 0), (0, 1))

    if not faces:
        return np.zeros((0, 3)), np.zeros((0, 3), dtype=np.int64)
    Fid = np.concatenate(faces, axis=0)
    ok = (Fid[:, 0] != Fid[:, 1]) & (Fid[:, 1] != Fid[:, 2]) & \
         (Fid[:, 0] != Fid[:, 2])
    Fid = Fid[ok]

    used, inv = np.unique(Fid, return_inverse=True)
    node = used // 2
    layer = used % 2
    xs = grid.xs()
    ys = grid.ys()
    vj, vi = node // nx, node % nx
    zsel = np.where(layer == 1, z_top.ravel()[node], z_bottom.ravel()[node])
    vertices = np.stack([xs[vi], ys[vj], zsel], axis=1)
    return vertices, inv.reshape(Fid.shape).astype(np.int64)
