"""Generátor testovacích STL modelů se záludnými podkosy.

Používá scikit-image (marching cubes) — pouze pro testy, běhová část
Babofkátoru scikit-image nepotřebuje.
"""
from __future__ import annotations

import os
import sys

import numpy as np
from skimage.measure import marching_cubes

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
from babofkator.stlio import write_stl_binary          # noqa: E402
from babofkator.raster import mesh_volume              # noqa: E402


def implicit_to_mesh(fn, lo, hi, n=160):
    xs = np.linspace(lo[0], hi[0], n)
    ys = np.linspace(lo[1], hi[1], n)
    zs = np.linspace(lo[2], hi[2], n)
    X, Y, Z = np.meshgrid(xs, ys, zs, indexing="ij")
    Fv = fn(X, Y, Z)
    Fv = np.pad(Fv, 1, mode="constant", constant_values=1e3)
    sp = (xs[1] - xs[0], ys[1] - ys[0], zs[1] - zs[0])
    verts, faces, _, _ = marching_cubes(Fv, level=0.0, spacing=sp)
    verts -= np.array(sp)                       # kompenzace paddingu
    verts += np.array([lo[0], lo[1], lo[2]])
    if mesh_volume(verts, faces) < 0:
        faces = faces[:, ::-1]
    # posadit na z = 0
    verts[:, 2] -= verts[:, 2].min()
    return verts, faces.astype(np.int64)


def _cut_base(f, Z):
    return np.maximum(f, -Z)


def hrib(X, Y, Z):
    """Hřib: tenká noha + široký kulový klobouk = brutální podkos."""
    r = np.sqrt(X ** 2 + Y ** 2)
    stem = np.maximum(r - 4.0, np.abs(Z - 5.0) - 5.0)
    cap = np.sqrt(X ** 2 + Y ** 2 + (Z - 11.0) ** 2) - 9.0
    return _cut_base(np.minimum(stem, cap), Z)


def kostka(X, Y, Z):
    """Kvádr — svislé stěny, nulový úkos."""
    b = np.maximum.reduce([np.abs(X) - 10.0, np.abs(Y) - 8.0,
                           np.abs(Z - 6.0) - 6.0])
    return _cut_base(b, Z)


def polokoule(X, Y, Z):
    """Polokoule — u paty svislá tečna."""
    s = np.sqrt(X ** 2 + Y ** 2 + Z ** 2) - 12.0
    return _cut_base(s, Z)


def rybina(X, Y, Z):
    """Rybinová drážka / lišta rozšiřující se vzhůru = podkos."""
    half = 2.0 + 0.45 * Z
    ridge = np.maximum.reduce([np.abs(X) - half, np.abs(Y) - 18.0,
                               np.abs(Z - 5.0) - 5.0])
    slab = np.maximum.reduce([np.abs(X) - 14.0, np.abs(Y) - 20.0,
                              np.abs(Z - 1.0) - 1.0])
    return _cut_base(np.minimum(ridge, slab), Z)


def relief(X, Y, Z):
    """Reliéf: vlnitá deska + převislý pysk + dvě jehly."""
    surf = 3.0 + 1.6 * np.sin(X * 0.55) * np.cos(Y * 0.5)
    plate = np.maximum.reduce([np.abs(X) - 18.0, np.abs(Y) - 14.0, Z - surf])
    lip = np.maximum.reduce([np.abs(X - 10.0) - 4.0, np.abs(Y) - 10.0,
                             np.abs(Z - 7.0) - 1.2])
    neck = np.maximum.reduce([np.abs(X - 10.0) - 1.0, np.abs(Y) - 10.0,
                              np.abs(Z - 5.0) - 2.0])
    p1 = np.maximum(np.sqrt((X + 8) ** 2 + (Y - 6) ** 2) - 1.2,
                    np.abs(Z - 6.0) - 6.0)
    p2 = np.maximum(np.sqrt((X + 4) ** 2 + (Y + 7) ** 2) - 2.5,
                    np.abs(Z - 4.0) - 4.0)
    return _cut_base(np.minimum.reduce([plate, lip, neck, p1, p2]), Z)


MODELS = {
    "hrib":      (hrib,      (-12, -12, -2), (12, 12, 22)),
    "kostka":    (kostka,    (-13, -11, -2), (13, 11, 15)),
    "polokoule": (polokoule, (-15, -15, -2), (15, 15, 15)),
    "rybina":    (rybina,    (-17, -22, -2), (17, 22, 13)),
    "relief":    (relief,    (-21, -17, -2), (21, 17, 12)),
}


def main(out_dir):
    os.makedirs(out_dir, exist_ok=True)
    for name, (fn, lo, hi) in MODELS.items():
        v, f = implicit_to_mesh(fn, lo, hi, n=180)
        path = os.path.join(out_dir, f"{name}.stl")
        write_stl_binary(path, v, f, header=f"test {name}")
        print(f"{name:10s} {f.shape[0]:7d} trojúhelníků  ->  {path}")


if __name__ == "__main__":
    main(sys.argv[1] if len(sys.argv) > 1 else
         os.path.join(os.path.dirname(__file__), "..", "examples"))
