#  **************************************************************************  #
#   BornAgain: simulate and fit reflection and scattering
#
#   @file      Wrap/Python/src/bornagain/ba_view3d.py
#   @brief     Interactive 3D sample visualization using VTK.
#
#   @homepage  http://www.bornagainproject.org
#   @license   GNU General Public License v3 or higher (see COPYING)
#   @copyright Forschungszentrum Juelich GmbH 2026
#   @authors   Scientific Computing Group at MLZ (see CITATION, AUTHORS)
#  **************************************************************************  #

import hashlib
import math
import os
import pickle
import struct
import subprocess
import sys
import warnings
from types import SimpleNamespace
import numpy as np

_MESH_RESOLUTION = 16  # angular resolution for smooth shapes (phi/theta/circle segments)
_DISPLAY_R = 15.0  # display half-size of the visible sample square
_MAX_ROUGHNESS_RESOLUTION = 2048
_OPEN_WINDOWS = []
_LAYER_FIELDS = (
    "material", "z_top", "z_bottom", "semi_infinite",
    "roughness_nx", "roughness_ny", "top_roughness",
)
_STRUCT_FIELDS = ("type", "placement", "ref_z", "pos_x", "pos_y", "pos_z", "mixture_weights")
_COLOR_FIELDS = ("r", "g", "b", "a")
_PARTICLE_FIELDS = (
    "shape", "param_names", "param_values",
    "pos_x", "pos_y", "pos_z",
    "euler_alpha", "euler_beta", "euler_gamma",
    "material", "alignment",
    "vert_x", "vert_y", "vert_z",
    "face_vertex_counts", "face_vertex_indices",
    "mixture_groups", "mixture_variants",
)

# Rotation matrix: RotateX(+90°), used to align VTK's Y-polar-axis to Z.
# vtkSphereSource and vtkCylinderSource use Y as their principal axis.
_RX90 = np.array([[1, 0, 0], [0, 0, -1], [0, 1, 0]], dtype=float)

# Rotation matrix: RotateZ(-90°), used to align Y-axis cylinder to X-axis.
_RZ_NEG90 = np.array([[0, 1, 0], [-1, 0, 0], [0, 0, 1]], dtype=float)


def _require_vtk():
    try:
        import vtkmodules.all as vtk
        return vtk
    except ImportError:
        raise ImportError(
            "vtk is required for 3D visualization.\n"
            "Install it with: pip install vtk"
        )


def _rotation_matrix(alpha, beta, gamma):
    """Builds a 3×3 rotation matrix from ZXZ Euler angles (radians)."""
    ca, sa = math.cos(alpha), math.sin(alpha)
    cb, sb = math.cos(beta), math.sin(beta)
    cg, sg = math.cos(gamma), math.sin(gamma)
    return np.array([
        [ca * cg - sa * cb * sg, -ca * sg - sa * cb * cg, sa * sb],
        [sa * cg + ca * cb * sg, -sa * sg + ca * cb * cg, -ca * sb],
        [sb * sg, sb * cg, cb],
    ])


def _np_to_vtk_transform(vtk, mat4):
    """Convert a 4×4 numpy matrix to a vtkTransform."""
    vtk_mat = vtk.vtkMatrix4x4()
    for i in range(4):
        for j in range(4):
            vtk_mat.SetElement(i, j, float(mat4[i, j]))
    t = vtk.vtkTransform()
    t.SetMatrix(vtk_mat)
    return t


def _apply_transform(vtk, polydata, mat4):
    """Apply a 4×4 numpy transform to a vtkPolyData; return new vtkPolyData."""
    tf = vtk.vtkTransformPolyDataFilter()
    tf.SetInputData(polydata)
    tf.SetTransform(_np_to_vtk_transform(vtk, mat4))
    tf.Update()
    return tf.GetOutput()


def _mat4_rot(rot3):
    """Embed a 3×3 rotation into a 4×4 identity matrix."""
    m = np.eye(4)
    m[:3, :3] = rot3
    return m


def _mat4_scale(sx, sy, sz):
    """4×4 scale matrix."""
    return np.diag([sx, sy, sz, 1.0])


def _mat4_translate(tx, ty, tz):
    """4×4 translation matrix."""
    m = np.eye(4)
    m[0, 3] = tx
    m[1, 3] = ty
    m[2, 3] = tz
    return m


# ---- Mesh creation from ParticleData ----

def _spherical_segment_polydata(vtk, r, rm_top, rm_bottom, sx=1.0, sy=1.0):
    """Creates a truncated sphere polydata with the bottom face at z = 0."""
    z_high = r - rm_top
    z_low = -(r - rm_bottom)
    if z_high - z_low <= 0:
        return None
    eps = 1e-9
    start_phi = math.degrees(math.acos(min(1, max(-1, z_high / (r + eps)))))
    end_phi = math.degrees(math.acos(min(1, max(-1, z_low / (r + eps)))))

    src = vtk.vtkSphereSource()
    src.SetRadius(r)
    src.SetCenter(0, 0, 0)
    src.SetStartPhi(start_phi)
    src.SetEndPhi(end_phi)
    src.SetPhiResolution(_MESH_RESOLUTION)
    src.SetThetaResolution(_MESH_RESOLUTION)
    src.Update()

    # vtkSphereSource: Y = polar axis.  RotateX(90) → Z = polar axis.
    # Then scale x/y for ellipsoidal variant, then lift bottom to z = 0.
    mat = _mat4_translate(0, 0, -z_low) @ _mat4_scale(sx, sy, 1.0) @ _mat4_rot(_RX90)
    return _apply_transform(vtk, src.GetOutput(), mat)


def _make_smooth_polydata(vtk, pd):
    """Creates a VTK polydata for smooth (non-polyhedral) particle shapes."""
    params = dict(zip(pd.param_names, pd.param_values))
    name = pd.shape

    if name == "Sphere":
        r = params.get("Radius", 1.0)
        src = vtk.vtkSphereSource()
        src.SetRadius(r)
        src.SetCenter(0, 0, 0)
        src.SetPhiResolution(_MESH_RESOLUTION)
        src.SetThetaResolution(_MESH_RESOLUTION)
        src.Update()
        mat = _mat4_translate(0, 0, r) @ _mat4_rot(_RX90)
        return _apply_transform(vtk, src.GetOutput(), mat)

    if name in ("Cylinder", "EllipsoidalCylinder"):
        rx = params.get("RadiusX", params.get("Radius", 1.0))
        ry = params.get("RadiusY", params.get("Radius", 1.0))
        h = params.get("Height", 1.0)
        src = vtk.vtkCylinderSource()
        src.SetRadius(1.0)
        src.SetHeight(h)
        src.SetResolution(_MESH_RESOLUTION)
        src.SetCenter(0, 0, 0)
        src.Update()
        # RotateX(90): Y-axis cylinder → Z-axis; scale radii; lift bottom to z=0.
        mat = (_mat4_translate(0, 0, h / 2)
               @ _mat4_scale(rx, ry, 1.0)
               @ _mat4_rot(_RX90))
        return _apply_transform(vtk, src.GetOutput(), mat)

    if name == "Cone":
        r = params.get("Radius", 1.0)
        h = params.get("Height", 1.0)
        src = vtk.vtkConeSource()
        src.SetRadius(r)
        src.SetHeight(h)
        src.SetResolution(_MESH_RESOLUTION)
        src.SetDirection(0, 0, 1)
        src.SetCenter(0, 0, 0)
        src.Update()
        # vtkConeSource with direction=(0,0,1), center=(0,0,0):
        #   base at z = -h/2, tip at z = +h/2 → translate by +h/2.
        mat = _mat4_translate(0, 0, h / 2)
        return _apply_transform(vtk, src.GetOutput(), mat)

    if name in ("Ellipsoid", "Spheroid"):
        rx = params.get("Radius X", params.get("Radius XY", 1.0))
        ry = params.get("Radius Y", params.get("Radius XY", 1.0))
        rz = params.get("Radius Z", 1.0)
        src = vtk.vtkSphereSource()
        src.SetRadius(1.0)
        src.SetCenter(0, 0, 0)
        src.SetPhiResolution(_MESH_RESOLUTION)
        src.SetThetaResolution(_MESH_RESOLUTION)
        src.Update()
        mat = (_mat4_translate(0, 0, rz)
               @ _mat4_scale(rx, ry, rz)
               @ _mat4_rot(_RX90))
        return _apply_transform(vtk, src.GetOutput(), mat)

    if name in ("SphericalSegment", "EllipsoidalSegment", "SpheroidalSegment"):
        r = params.get("Radius", params.get("Radius Z", 1.0))
        rm_t = params.get("Top cut", 0.0)
        rm_b = params.get("Bottom cut", 0.0)
        rx = params.get("Radius X", params.get("Radius XY", r)) / r
        ry = params.get("Radius Y", params.get("Radius XY", r)) / r
        return _spherical_segment_polydata(vtk, r, rm_t, rm_b, sx=rx, sy=ry)

    if name == "HorizontalCylinder":
        r = params.get("Radius", 1.0)
        L = params.get("Length", 1.0)
        src = vtk.vtkCylinderSource()
        src.SetRadius(r)
        src.SetHeight(L)
        src.SetResolution(_MESH_RESOLUTION)
        src.SetCenter(0, 0, 0)
        src.Update()
        # RotateZ(-90): Y-axis cylinder → X-axis; translate center to z = r.
        mat = _mat4_translate(0, 0, r) @ _mat4_rot(_RZ_NEG90)
        return _apply_transform(vtk, src.GetOutput(), mat)

    if name == "FuzzySphere":
        r = params.get("Radius", 1.0) + params.get("Fuzziness", 0.0)
        src = vtk.vtkSphereSource()
        src.SetRadius(r)
        src.SetCenter(0, 0, 0)
        src.SetPhiResolution(_MESH_RESOLUTION)
        src.SetThetaResolution(_MESH_RESOLUTION)
        src.Update()
        mat = _mat4_translate(0, 0, r) @ _mat4_rot(_RX90)
        return _apply_transform(vtk, src.GetOutput(), mat)

    if name == "GaussSphere":
        r = params.get("MeanRadius", 1.0)
        src = vtk.vtkSphereSource()
        src.SetRadius(r)
        src.SetCenter(0, 0, 0)
        src.SetPhiResolution(_MESH_RESOLUTION)
        src.SetThetaResolution(_MESH_RESOLUTION)
        src.Update()
        mat = _mat4_translate(0, 0, r) @ _mat4_rot(_RX90)
        return _apply_transform(vtk, src.GetOutput(), mat)

    # Bar/LongBox: rendered as a plain cuboid; the Gaussian/Lorentzian profile
    # affects scattering but not the 3D shape.
    if name in ("BarGauss", "BarLorentz", "LongBoxGauss", "LongBoxLorentz"):
        L = params.get("Length", 1.0)
        W = params.get("Width", 1.0)
        H = params.get("Height", 1.0)
        src = vtk.vtkCubeSource()
        src.SetXLength(L)
        src.SetYLength(W)
        src.SetZLength(H)
        src.SetCenter(0, 0, H / 2)
        src.Update()
        return src.GetOutput()

    if name in ("CosineRippleBox", "CosineRippleGauss", "CosineRippleLorentz"):
        L = params.get("Length", 1.0)
        W = params.get("Width", 1.0)
        H = params.get("Height", 1.0)
        return _ripple_polydata(vtk, L, W, H, profile="cosine")

    if name in ("SawtoothRippleBox", "SawtoothRippleGauss", "SawtoothRippleLorentz"):
        L = params.get("Length", 1.0)
        W = params.get("Width", 1.0)
        H = params.get("Height", 1.0)
        a = params.get("Asymmetry", 0.0)
        return _ripple_polydata(vtk, L, W, H, profile="sawtooth", asymmetry=a)

    return None


def _ripple_polydata(vtk, L, W, H, profile, asymmetry=0.0):
    """Build a ripple polydata extruded along X.

    Cross-section in the YZ plane:
      - cosine:   z(y) = (H/2) * (1 + cos(2π y / W)),  y ∈ [-W/2, +W/2]
      - sawtooth: triangle with base from (-W/2,0) to (+W/2,0), apex at (asymmetry, H)

    Bottom edge is z = 0; the ridge runs along X from -L/2 to +L/2.
    """
    if profile == "sawtooth":
        cross = [(-W / 2, 0.0),
                 (asymmetry, H),
                 (+W / 2, 0.0)]
        cy, cz = asymmetry, H / 2
    else:  # cosine
        n = _MESH_RESOLUTION * 4
        cross = []
        for s in range(n):
            th = math.pi * s / (n + 1)
            y = -(W / 2) * math.cos(th)
            z = (H / 2) * (1 + math.cos(2 * math.pi * y / W))
            cross.append((y, z))
        cy, cz = 0.0, H / 2

    points = vtk.vtkPoints()
    front_ids = []
    back_ids = []
    for y, z in cross:
        front_ids.append(points.InsertNextPoint(-L / 2, y, z))
        back_ids.append(points.InsertNextPoint(+L / 2, y, z))
    front_center = points.InsertNextPoint(-L / 2, cy, cz)
    back_center = points.InsertNextPoint(+L / 2, cy, cz)

    cells = vtk.vtkCellArray()

    def add_tri(a, b, c):
        cell = vtk.vtkIdList()
        cell.InsertNextId(a)
        cell.InsertNextId(b)
        cell.InsertNextId(c)
        cells.InsertNextCell(cell)

    n_cross = len(cross)
    for s in range(n_cross):
        s1, s2 = s, (s + 1) % n_cross
        # Front face: outward normal points to -X. Triangle fan from front_center.
        add_tri(front_ids[s1], front_ids[s2], front_center)
        # Back face: outward normal points to +X. Reversed winding.
        add_tri(back_ids[s1], back_center, back_ids[s2])
        # Side surface: quad (vfront[s2], vfront[s1], vback[s1], vback[s2]) split.
        add_tri(front_ids[s2], front_ids[s1], back_ids[s1])
        add_tri(front_ids[s2], back_ids[s1], back_ids[s2])

    polydata = vtk.vtkPolyData()
    polydata.SetPoints(points)
    polydata.SetPolys(cells)

    normals = vtk.vtkPolyDataNormals()
    normals.SetInputData(polydata)
    normals.AutoOrientNormalsOn()
    normals.ConsistencyOn()
    normals.Update()
    return normals.GetOutput()


def _make_polyhedron_polydata(vtk, pd):
    """Creates a VTK polydata from polyhedron vertex/face data in a ParticleData."""
    verts = np.column_stack([pd.vert_x, pd.vert_y, pd.vert_z]).astype(float)
    verts[:, 2] -= verts[:, 2].min()

    points = vtk.vtkPoints()
    for v in verts:
        points.InsertNextPoint(float(v[0]), float(v[1]), float(v[2]))

    counts = list(pd.face_vertex_counts)
    indices = list(pd.face_vertex_indices)
    cells = vtk.vtkCellArray()
    offset = 0
    for n in counts:
        cell = vtk.vtkIdList()
        for idx in indices[offset:offset + n]:
            cell.InsertNextId(int(idx))
        cells.InsertNextCell(cell)
        offset += n

    polydata = vtk.vtkPolyData()
    polydata.SetPoints(points)
    polydata.SetPolys(cells)

    normals = vtk.vtkPolyDataNormals()
    normals.SetInputData(polydata)
    normals.ComputePointNormalsOn()
    normals.Update()
    return normals.GetOutput()


def _make_mesh(vtk, pd):
    """Creates a VTK polydata for a ParticleData."""
    if pd.vert_x:
        return _make_polyhedron_polydata(vtk, pd)
    return _make_smooth_polydata(vtk, pd)


ALIGN_BOTTOM, ALIGN_CENTER, ALIGN_TOP = 0, 1, 2


def _prepare_template(vtk, pd):
    """Creates a template polydata with alignment translation and Euler rotation applied."""
    mesh = _make_mesh(vtk, pd)
    if mesh is None:
        return None

    bounds = mesh.GetBounds()  # (xmin, xmax, ymin, ymax, zmin, zmax)
    height = bounds[5] - bounds[4]

    has_alignment = pd.alignment != ALIGN_BOTTOM
    has_rotation = any(abs(a) > 1e-12
                       for a in (pd.euler_alpha, pd.euler_beta, pd.euler_gamma))

    if not has_alignment and not has_rotation:
        return mesh

    # Build combined transform: first apply alignment translation, then Euler rotation.
    align = np.eye(4)
    if pd.alignment == ALIGN_TOP:
        align[2, 3] = -height
    elif pd.alignment == ALIGN_CENTER:
        align[2, 3] = -height / 2

    rot4 = np.eye(4)
    if has_rotation:
        rot4[:3, :3] = _rotation_matrix(pd.euler_alpha, pd.euler_beta, pd.euler_gamma)

    return _apply_transform(vtk, mesh, rot4 @ align)


def _particle_key(pd):
    """Returns a hashable key grouping geometrically identical particles."""
    return (
        pd.shape,
        pd.material,
        tuple(pd.param_values),
        pd.color.r, pd.color.g, pd.color.b,
        pd.color.a,
        pd.alignment,
        pd.euler_alpha, pd.euler_beta, pd.euler_gamma,
    )


# ---- Scene rendering ----


def _layer_top_roughness(ld):
    """Returns the layer top roughness map as a 2D array, or None for flat layers."""
    nx = int(getattr(ld, "roughness_nx", 0) or 0)
    ny = int(getattr(ld, "roughness_ny", 0) or 0)
    values = getattr(ld, "top_roughness", [])
    if nx < 2 or ny < 2:
        return None
    values = list(values)
    if len(values) != nx * ny:
        return None
    return np.asarray(values, dtype=float).reshape((ny, nx))


def _vtk_points_from_numpy(vtk, coordinates):
    """Creates vtkPoints from an Nx3 numpy coordinate array."""
    from vtkmodules.util.numpy_support import numpy_to_vtk

    points = vtk.vtkPoints()
    points.SetData(numpy_to_vtk(np.ascontiguousarray(coordinates, dtype=np.float64), deep=True))
    return points


def _vtk_cell_array_from_quads(vtk, quads):
    """Creates vtkCellArray from an Nx4 numpy quad-connectivity array."""
    from vtkmodules.util.numpy_support import numpy_to_vtkIdTypeArray

    quads = np.ascontiguousarray(quads, dtype=np.int64)
    cells = vtk.vtkCellArray()
    if quads.size == 0:
        return cells

    offsets = np.arange(0, quads.size + 1, 4, dtype=np.int64)
    connectivity = quads.reshape(-1)
    cells.SetData(
        numpy_to_vtkIdTypeArray(offsets, deep=True),
        numpy_to_vtkIdTypeArray(connectivity, deep=True),
    )
    return cells


def _rough_layer_polydata(vtk, ld, scale, half, top_map, bottom_map, draw_bottom):
    """Creates a layer mesh from optional top and bottom roughness maps."""
    base = top_map if top_map is not None else bottom_map
    if base is None:
        return None
    ny, nx = base.shape
    if nx < 2 or ny < 2:
        return None
    if top_map is None or top_map.shape != base.shape:
        top_map = np.zeros_like(base)
    if bottom_map is None or bottom_map.shape != base.shape:
        bottom_map = np.zeros_like(base)

    n_grid = nx * ny
    x = np.linspace(-half, half, nx)
    y = np.linspace(-half, half, ny)
    xx, yy = np.meshgrid(x, y)

    top_points = np.column_stack((
        xx.reshape(-1),
        yy.reshape(-1),
        ((ld.z_top + top_map) * scale).reshape(-1),
    ))
    bottom_points = np.column_stack((
        xx.reshape(-1),
        yy.reshape(-1),
        ((ld.z_bottom + bottom_map) * scale).reshape(-1),
    ))
    points = np.vstack((top_points, bottom_points))

    top = np.arange(n_grid, dtype=np.int64).reshape(ny, nx)
    bottom = top + n_grid

    top_quads = np.stack((
        top[:-1, :-1],
        top[:-1, 1:],
        top[1:, 1:],
        top[1:, :-1],
    ), axis=-1).reshape(-1, 4)

    left_wall = np.column_stack((bottom[:-1, 0], top[:-1, 0], top[1:, 0], bottom[1:, 0]))
    right_wall = np.column_stack((bottom[1:, -1], top[1:, -1], top[:-1, -1],
                                  bottom[:-1, -1]))
    front_wall = np.column_stack((bottom[0, 1:], top[0, 1:], top[0, :-1], bottom[0, :-1]))
    back_wall = np.column_stack((bottom[-1, :-1], top[-1, :-1], top[-1, 1:],
                                 bottom[-1, 1:]))

    quads = [top_quads, left_wall, right_wall, front_wall, back_wall]
    if draw_bottom:
        bottom_quads = np.stack((
            bottom[:-1, 1:],
            bottom[:-1, :-1],
            bottom[1:, :-1],
            bottom[1:, 1:],
        ), axis=-1).reshape(-1, 4)
        quads.append(bottom_quads)

    polydata = vtk.vtkPolyData()
    polydata.SetPoints(_vtk_points_from_numpy(vtk, points))
    polydata.SetPolys(_vtk_cell_array_from_quads(vtk, np.vstack(quads)))

    normals = vtk.vtkPolyDataNormals()
    normals.SetInputData(polydata)
    normals.ConsistencyOn()
    normals.SplittingOff()
    normals.Update()
    return normals.GetOutput()


def _set_phong_material(prop, color, opacity, ambient, diffuse, specular, power):
    prop.SetColor(*color)
    prop.SetOpacity(opacity)
    prop.SetAmbient(ambient)
    prop.SetDiffuse(diffuse)
    prop.SetSpecular(specular)
    prop.SetSpecularPower(power)
    if hasattr(prop, "SetSpecularColor"):
        prop.SetSpecularColor(1.0, 1.0, 1.0)
    prop.SetInterpolationToPhong()


def _add_layer(renderer, vtk, ld, scale, layer_half, bottom_map=None, draw_bottom=False):
    """Adds a semi-transparent layer slab to the renderer."""
    half = layer_half
    top_map = _layer_top_roughness(ld)
    polydata = _rough_layer_polydata(vtk, ld, scale, half, top_map, bottom_map, draw_bottom)
    rough_layer = polydata is not None

    mapper = vtk.vtkPolyDataMapper()
    if rough_layer:
        mapper.SetInputData(polydata)
    else:
        src = vtk.vtkCubeSource()
        src.SetBounds(-half, half, -half, half, ld.z_bottom * scale, ld.z_top * scale)
        src.Update()
        mapper.SetInputConnection(src.GetOutputPort())

    actor = vtk.vtkActor()
    actor.SetMapper(mapper)
    prop = actor.GetProperty()
    layer_color = (ld.color.r, ld.color.g, ld.color.b)
    if rough_layer:
        _set_phong_material(prop, layer_color, 0.30, 0.18, 0.58, 0.72, 24)
    else:
        _set_phong_material(prop, layer_color, 0.27, 0.25, 0.66, 0.42, 30)
    renderer.AddActor(actor)


def _add_lateral_clipping_planes(vtk, mapper, half):
    """Clips rendered particle geometry to the lateral display patch."""
    if half is None or half <= 0:
        return
    for normal, origin in (
            ((1, 0, 0), (-half, 0, 0)),
            ((-1, 0, 0), (half, 0, 0)),
            ((0, 1, 0), (0, -half, 0)),
            ((0, -1, 0), (0, half, 0)),
    ):
        plane = vtk.vtkPlane()
        plane.SetNormal(*normal)
        plane.SetOrigin(*origin)
        mapper.AddClippingPlane(plane)


def _site_pick_u01(scene_seed, struct_idx, group_id, lx, ly, lz):
    """Returns a stable pseudo-random float in [0, 1) for one site/group.

    The key includes the scene seed, struct, mixture group, and site position so
    Mixture realizations do not change when visible sites are clipped or reordered.
    """
    key = struct.pack(
        "<QII3d",
        int(scene_seed) & 0xffffffffffffffff,
        int(struct_idx) & 0xffffffff,
        int(group_id) & 0xffffffff,
        float(lx), float(ly), float(lz),
    )
    digest = hashlib.blake2b(key, digest_size=8).digest()
    return int.from_bytes(digest, "little") / float(1 << 64)


def _weighted_variant(weights, u01):
    """Maps one uniform draw to a variant index using normalized Mixture weights."""
    total = float(sum(weights))
    if total <= 0.0:
        raise ValueError("Mixture variant weights must have positive total")

    threshold = u01 * total
    cumulative = 0.0
    for variant, weight in enumerate(weights):
        cumulative += float(weight)
        if threshold < cumulative:
            return variant
    return len(weights) - 1


def _pick_active_variants(scene_seed, struct_idx, lx, ly, lz, mixture_weights):
    """Picks one variant per mixture group for a layout site.

    All groups are sampled independently, including nested Mixture groups.  A
    leaf particle is rendered only when its full activation path matches these
    choices, so groups whose parent variant was not selected are simply ignored.
    This is equivalent to recursive conditional sampling but stable under view
    pruning and layout reordering.
    """
    return {
        group_id: _weighted_variant(
            weights, _site_pick_u01(scene_seed, struct_idx, group_id, lx, ly, lz)
        )
        for group_id, weights in enumerate(mixture_weights)
    }


def _particle_matches_mixture_choices(pd, choices):
    """Checks whether a particle leaf is active under the sampled Mixture choices.

    Each exported leaf carries the full path of group/variant decisions that must
    match for this site; leaves outside the chosen path are skipped.
    """
    groups = getattr(pd, "mixture_groups", ())
    variants = getattr(pd, "mixture_variants", ())
    if len(groups) != len(variants):
        raise ValueError("ParticleData mixture path has mismatched group and variant lengths")
    for group_id, variant_id in zip(groups, variants):
        if choices.get(int(group_id)) != int(variant_id):
            return False
    return True


def _add_struct(renderer, vtk, sd, scale=1.0, view_half=None, clip_half=None,
                scene_seed=0, struct_idx=0):
    """Renders all particles of a struct using vtkGlyph3D instancing.

    A layout instance is kept if its whole bbox intersects ±view_half (display
    units).  Composite particles such as mesocrystals must not be clipped
    leaf-by-leaf after expansion into basis particles.  Particle geometry itself
    is clipped laterally at the mapper level for display only.
    """
    if not sd.particles:
        return

    view_limit = float("inf") if view_half is None else view_half
    rel_x_min = rel_y_min = float("inf")
    rel_x_max = rel_y_max = float("-inf")
    glyph_groups = {}
    prepared = []
    for pd in sd.particles:
        key = _particle_key(pd)
        if key not in glyph_groups:
            tmpl = _prepare_template(vtk, pd)
            if tmpl is None:
                print(f"Warning: unsupported shape '{pd.shape}', skipping")
                glyph_groups[key] = None
                continue
            scaled_tmpl = _apply_transform(vtk, tmpl, _mat4_scale(scale, scale, scale))
            glyph_groups[key] = {
                "template": scaled_tmpl,
                "positions": [],
                "color": (pd.color.r, pd.color.g, pd.color.b),
                "alpha": pd.color.a,
                "label": pd.shape,
            }

        if glyph_groups[key] is None:
            continue

        if view_half is not None:
            bounds = glyph_groups[key]["template"].GetBounds()
            rel_x_min = min(rel_x_min, pd.pos_x * scale + bounds[0])
            rel_x_max = max(rel_x_max, pd.pos_x * scale + bounds[1])
            rel_y_min = min(rel_y_min, pd.pos_y * scale + bounds[2])
            rel_y_max = max(rel_y_max, pd.pos_y * scale + bounds[3])
        prepared.append((pd, glyph_groups[key]))

    visible_layout = []
    for lx, ly, lz in zip(sd.pos_x, sd.pos_y, sd.pos_z):
        instance_x = lx * scale
        instance_y = ly * scale
        if (instance_x + rel_x_max < -view_limit or instance_x + rel_x_min > view_limit
                or instance_y + rel_y_max < -view_limit or instance_y + rel_y_min > view_limit):
            continue
        visible_layout.append((lx, ly, lz))

    mixture_weights = getattr(sd, "mixture_weights", [])
    for lx, ly, lz in visible_layout:
        mixture_choices = _pick_active_variants(
            scene_seed, struct_idx, lx, ly, lz, mixture_weights
        ) if mixture_weights else {}
        for pd, grp in prepared:
            if not _particle_matches_mixture_choices(pd, mixture_choices):
                continue
            x = (pd.pos_x + lx) * scale
            y = (pd.pos_y + ly) * scale
            z = (pd.pos_z + lz) * scale
            grp["positions"].append([x, y, z])

    for grp in glyph_groups.values():
        if grp is None or not grp["positions"]:
            continue

        pts = vtk.vtkPoints()
        for pos in grp["positions"]:
            pts.InsertNextPoint(float(pos[0]), float(pos[1]), float(pos[2]))
        cloud = vtk.vtkPolyData()
        cloud.SetPoints(pts)

        glyph = vtk.vtkGlyph3D()
        glyph.SetSourceData(grp["template"])
        glyph.SetInputData(cloud)
        glyph.OrientOff()
        glyph.ScalingOff()
        glyph.Update()

        mapper = vtk.vtkPolyDataMapper()
        mapper.SetInputConnection(glyph.GetOutputPort())
        _add_lateral_clipping_planes(vtk, mapper, clip_half)

        actor = vtk.vtkActor()
        actor.SetMapper(mapper)
        prop = actor.GetProperty()
        _set_phong_material(prop, grp["color"], grp["alpha"], 0.20, 0.64, 0.60, 36)
        renderer.AddActor(actor)


def _scene_to_dict(scene):
    """Converts the SWIG scene description to a plain dict for IPC transport."""
    def plain(value):
        """Recursively converts SWIG vector-like values to JSON-pickleable lists."""
        if isinstance(value, (str, int, float, bool)):
            return value
        return [plain(v) for v in value]

    def fields(obj, names):
        """Copies selected attributes from a SWIG object into plain Python data."""
        return {name: plain(getattr(obj, name)) for name in names}

    def particle(pd):
        """Serializes one particle template, including its material color."""
        data = fields(pd, _PARTICLE_FIELDS)
        data["color"] = fields(pd.color, _COLOR_FIELDS)
        return data

    def layer(ld):
        """Serializes one layer record, including its material color."""
        data = fields(ld, _LAYER_FIELDS)
        data["color"] = fields(ld.color, _COLOR_FIELDS)
        return data

    return {
        "lateral_half_extent": float(scene.lateral_half_extent),
        "layers": [layer(ld) for ld in scene.layers],
        "structs": [
            {**fields(sd, _STRUCT_FIELDS), "particles": [particle(pd) for pd in sd.particles]}
            for sd in scene.structs
        ],
    }


def _dict_to_scene(data):
    """Converts plain scene data back to attribute-based objects."""
    def namespace(value):
        if isinstance(value, dict):
            return SimpleNamespace(**{k: namespace(v) for k, v in value.items()})
        if isinstance(value, list):
            return [namespace(v) for v in value]
        return value

    return namespace(data)


def _viewer_kwargs(kwargs):
    """Keeps only subprocess-safe viewer options for non-blocking rendering."""
    result = {}
    if "window_size" in kwargs:
        result["window_size"] = list(kwargs["window_size"])
    if "background" in kwargs:
        result["background"] = kwargs["background"]
    if "zoom" in kwargs:
        result["zoom"] = float(kwargs["zoom"])
    result["seed"] = int(kwargs.get("seed", 0))
    result["blocking"] = True
    return result


def _reap_open_viewers():
    _OPEN_WINDOWS[:] = [proc for proc in _OPEN_WINDOWS if proc.poll() is None]


def _child_env():
    env = os.environ.copy()
    parent = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    env["PYTHONPATH"] = parent + os.pathsep + env.get("PYTHONPATH", "")
    return env


def _launch_scene_process(scene, kwargs):
    """Starts a separate process whose main thread owns the VTK event loop."""
    _reap_open_viewers()
    payload = {"scene": _scene_to_dict(scene), "kwargs": _viewer_kwargs(kwargs)}
    proc = subprocess.Popen(
        [sys.executable, "-m", "bornagain.ba_view3d"],
        stdin=subprocess.PIPE,
        env=_child_env(),
    )
    pickle.dump(payload, proc.stdin)
    proc.stdin.close()
    _OPEN_WINDOWS.append(proc)
    return proc


def _show_scene(scene, **kwargs):
    vtk = _require_vtk()
    blocking = bool(kwargs.get("blocking", kwargs.get("block", True)))
    lateral_half_extent = scene.lateral_half_extent
    if lateral_half_extent <= 0:
        raise ValueError(
            "Scene lateral_half_extent must be positive; pass sample_size > 0 to showSample3D"
        )

    scale = _DISPLAY_R / lateral_half_extent
    view_half = _DISPLAY_R

    layers_to_show = [(i, ld) for i, ld in enumerate(scene.layers)
                      if not (i == 0 and ld.semi_infinite)]

    layer_half = view_half

    renderer = vtk.vtkRenderer()
    named_colors = vtk.vtkNamedColors()
    renderer.SetBackground(named_colors.GetColor3d(kwargs.get("background", "white")))

    for pos, (i, ld) in enumerate(layers_to_show):
        # C++ generates rough interfaces as one correlated stack. The top map of the layer below is
        # therefore also the rough bottom boundary of the current layer.
        bottom_map = (_layer_top_roughness(scene.layers[i + 1])
                      if i + 1 < len(scene.layers) else None)
        _add_layer(renderer, vtk, ld, scale, layer_half, bottom_map,
                   draw_bottom=pos + 1 == len(layers_to_show))
    seed = int(kwargs.get("seed", 0))
    for struct_idx, sd in enumerate(scene.structs):
        _add_struct(renderer, vtk, sd, scale, view_half, view_half, seed, struct_idx)

    window_size = kwargs.get("window_size", [1330, 1075])
    render_window = vtk.vtkRenderWindow()
    render_window.AddRenderer(renderer)
    render_window.SetSize(*window_size)
    render_window.SetWindowName("BornAgain 3D View")
    if hasattr(render_window, "SetMultiSamples"):
        render_window.SetMultiSamples(8)

    interactor = vtk.vtkRenderWindowInteractor()
    interactor.SetRenderWindow(render_window)
    interactor.SetInteractorStyle(vtk.vtkInteractorStyleTrackballCamera())

    # Axes orientation marker in the bottom-left corner
    axes_actor = vtk.vtkAxesActor()
    axes_widget = vtk.vtkOrientationMarkerWidget()
    axes_widget.SetOrientationMarker(axes_actor)
    axes_widget.SetInteractor(interactor)
    axes_widget.SetViewport(0.0, 0.0, 0.2, 0.2)
    axes_widget.EnabledOn()
    axes_widget.InteractiveOff()

    # Camera: frame the displayed sample patch and all visible z extents.
    z_values = []
    for i, ld in layers_to_show:
        z_values.extend((ld.z_bottom * scale, ld.z_top * scale))
        top_map = _layer_top_roughness(ld)
        if top_map is not None:
            z_values.extend(((ld.z_top + float(np.min(top_map))) * scale,
                             (ld.z_top + float(np.max(top_map))) * scale))
        bottom_map = (_layer_top_roughness(scene.layers[i + 1])
                      if i + 1 < len(scene.layers) else None)
        if bottom_map is not None:
            z_values.extend(((ld.z_bottom + float(np.min(bottom_map))) * scale,
                             (ld.z_bottom + float(np.max(bottom_map))) * scale))

    for sd in scene.structs:
        if not sd.pos_z:
            continue
        visible_lz = [
            lz for lx, ly, lz in zip(sd.pos_x, sd.pos_y, sd.pos_z)
            if abs(lx * scale) <= view_half and abs(ly * scale) <= view_half
        ]
        if not visible_lz:
            continue
        for pd in sd.particles:
            tmpl = _prepare_template(vtk, pd)
            if tmpl is None:
                continue
            bounds = tmpl.GetBounds()
            z_low = bounds[4] * scale
            z_high = bounds[5] * scale
            for lz in visible_lz:
                base_z = (pd.pos_z + lz) * scale
                z_values.extend((base_z + z_low, base_z + z_high))

    z_min = min(z_values) if z_values else 0.0
    z_max = max(z_values) if z_values else 0.0
    z_mid = (z_min + z_max) / 2.0

    az = math.radians(45)
    el = math.radians(30)
    view_angle = math.radians(30)
    aspect = window_size[0] / max(window_size[1], 1)
    half_angle = min(view_angle / 2, math.atan(aspect * math.tan(view_angle / 2)))
    radius = math.sqrt((2 * view_half)**2 + (2 * view_half)**2 + (z_max - z_min)**2) / 2
    dist = radius / max(math.sin(half_angle), 1e-6) * 1.05

    cam = renderer.GetActiveCamera()
    cam.SetViewAngle(math.degrees(view_angle))
    cam.SetFocalPoint(0.0, 0.0, z_mid)
    cam.SetPosition(dist * math.cos(el) * math.cos(az),
                    dist * math.cos(el) * math.sin(az),
                    z_mid + dist * math.sin(el))
    cam.SetViewUp(0, 0, 1)
    cam.OrthogonalizeViewUp()
    zoom = float(kwargs.get("zoom", 1.2))
    if zoom <= 0:
        raise ValueError("zoom must be positive")
    cam.Zoom(zoom)
    renderer.ResetCameraClippingRange()

    # Lighting: warm key light, cool fill, and grazing/camera glints for glossy surfaces.
    renderer.RemoveAllLights()

    sun = vtk.vtkLight()
    sun.SetLightTypeToSceneLight()
    sun.SetPosition(dist * 0.6, -dist * 0.3, dist * 1.2)
    sun.SetFocalPoint(0.0, 0.0, z_mid)
    sun.SetIntensity(1.08)
    sun.SetColor(1.0, 0.97, 0.93)
    renderer.AddLight(sun)

    fill = vtk.vtkLight()
    fill.SetLightTypeToSceneLight()
    fill.SetPosition(-dist * 0.5, dist * 0.6, dist * 0.3)
    fill.SetFocalPoint(0.0, 0.0, z_mid)
    fill.SetIntensity(0.36)
    fill.SetColor(0.82, 0.88, 1.0)
    renderer.AddLight(fill)

    glint = vtk.vtkLight()
    glint.SetLightTypeToSceneLight()
    glint.SetPosition(-dist * 1.0, -dist * 0.8, z_mid + dist * 0.18)
    glint.SetFocalPoint(0.0, 0.0, z_mid)
    glint.SetIntensity(0.55)
    glint.SetColor(1.0, 1.0, 1.0)
    renderer.AddLight(glint)

    headlight = vtk.vtkLight()
    headlight.SetLightTypeToHeadlight()
    headlight.SetIntensity(0.14)
    headlight.SetColor(1.0, 1.0, 1.0)
    renderer.AddLight(headlight)

    render_window.Render()
    interactor.Initialize()

    if blocking:
        interactor.Start()
        return interactor

    return interactor


def plot_sample(sample, sample_size, seed=0, **kwargs):
    """Shows an interactive 3D visualization of a BornAgain Sample.

    The visible sample area is an explicit square patch of side sample_size,
    in nm.  Particles extending beyond this patch are clipped for display.

    Args:
        sample:       a bornagain.Sample instance.
        sample_size   (float): side length of the visible square sample patch,
                               in nm.
        seed          (int):   base seed for random particle positions in
                               disordered layouts (default 0). Change this
                               to regenerate a different position configuration
                               and Mixture realization.
        **kwargs:
            window_size (list): [width, height] of the render window.
            background  (str):  background color name.
            zoom        (float): initial camera zoom. Default is 1.2.
            roughness   (bool): render rough interfaces. Default is True.
            roughness_resolution (int): points along each roughness-map axis.
                                        Default is 256. Use 0 to disable.
            blocking    (bool): run the VTK event loop in the current process
                                before returning. Default is False. Alias: block.
    """
    import bornagain as ba

    seed = int(seed)
    roughness = bool(kwargs.pop("roughness", True))
    roughness_resolution = int(kwargs.pop("roughness_resolution", 256))
    if roughness_resolution < 0:
        raise ValueError("roughness_resolution must be non-negative")
    if roughness_resolution > _MAX_ROUGHNESS_RESOLUTION:
        warnings.warn(
            f"roughness_resolution={roughness_resolution} is too high for the 3D view; "
            f"limiting to {_MAX_ROUGHNESS_RESOLUTION}.",
            RuntimeWarning,
            stacklevel=2,
        )
        roughness_resolution = _MAX_ROUGHNESS_RESOLUTION
    roughness_points = roughness_resolution if roughness else 0
    kwargs = {**kwargs, "seed": seed}
    sample_size = float(sample_size)
    if sample_size <= 0:
        raise ValueError(
            "sample_size must be positive; it is the side length of the visible square patch in nm"
        )
    scene = ba.buildScene3D(sample, sample_size / 2.0, seed, roughness_points)

    blocking = bool(kwargs.get("blocking", kwargs.get("block", False)))
    if blocking:
        return _show_scene(scene, **kwargs)
    return _launch_scene_process(scene, kwargs)


if __name__ == "__main__":
    payload = pickle.load(sys.stdin.buffer)
    _show_scene(_dict_to_scene(payload["scene"]), **payload["kwargs"])
