from __future__ import annotations

from typing import Optional, Tuple

import math

import numpy as np
try:
    import numba as nb
    _HAVE_NUMBA = True
except Exception:
    nb = None
    _HAVE_NUMBA = False

try:
    from scipy.spatial import cKDTree, ConvexHull
except Exception:
    cKDTree = None
    ConvexHull = None


# ----------------------------------------------------------------------
# Utility helpers
# ----------------------------------------------------------------------


def _extract_ca_coordinates_from_atoms(atoms: np.ndarray) -> Optional[np.ndarray]:
    """Return CA atom coordinates if available, else ``None``.

    The ``atoms`` array is expected to expose ``'atom_name'`` and ``'xyz'``
    fields compatible with :mod:`chisurf.structure`.
    """

    coords, _, _ = _extract_ca_trace(atoms)
    return coords


def _compute_center_radius(xyz: np.ndarray) -> tuple[np.ndarray, float]:
    """Return (center, radius) bounding sphere approximation."""

    arr = np.asarray(xyz, dtype=float)
    center = arr.mean(axis=0)
    # Use max distance from center as radius
    diffs = arr - center
    dist2 = np.sum(diffs * diffs, axis=1)
    radius = float(math.sqrt(float(dist2.max()))) if dist2.size else 1.0
    if not np.isfinite(radius) or radius <= 0:
        radius = 1.0
    return center, radius


def _build_sphere_mesh(radius: float) -> Optional[dict[str, np.ndarray]]:
    """Return a procedural sphere approximation with normals."""

    r = float(radius)
    if not np.isfinite(r) or r <= 0.0:
        return None

    rows = 10
    cols = 20
    phi = np.linspace(0.0, np.pi, rows)
    theta = np.linspace(0.0, 2.0 * np.pi, cols, endpoint=False)
    phi, theta = np.meshgrid(phi, theta, indexing="ij")
    sin_phi = np.sin(phi)
    cos_phi = np.cos(phi)
    cos_theta = np.cos(theta)
    sin_theta = np.sin(theta)

    x = r * sin_phi * cos_theta
    y = r * sin_phi * sin_theta
    z = r * cos_phi
    vertices = np.stack([x, y, z], axis=-1).reshape(-1, 3)

    normals = np.stack([sin_phi * cos_theta, sin_phi * sin_theta, cos_phi], axis=-1)
    normals = normals.reshape(-1, 3)
    norms = np.linalg.norm(normals, axis=1, keepdims=True)
    normals = normals / np.clip(norms, 1e-8, None)

    faces = []
    for i in range(rows - 1):
        for j in range(cols):
            k0 = i * cols + j
            k1 = i * cols + (j + 1) % cols
            k2 = (i + 1) * cols + j
            k3 = (i + 1) * cols + (j + 1) % cols
            faces.append([k0, k2, k1])
            faces.append([k1, k2, k3])
    faces_arr = np.asarray(faces, dtype=np.int32)
    return {"vertices": vertices, "normals": normals, "faces": faces_arr}


_CYLINDER_TEMPLATE_CACHE: dict[int, dict[str, np.ndarray]] = {}


def _get_cylinder_template(segments_circle: int = 12) -> Optional[dict[str, np.ndarray]]:
    """Return cached unit-length cylinder template aligned with the +Z axis."""

    seg = max(3, int(segments_circle))
    template = _CYLINDER_TEMPLATE_CACHE.get(seg)
    if template is not None:
        return template

    angles = np.linspace(0.0, 2.0 * np.pi, seg, endpoint=False, dtype=float)
    cos = np.cos(angles)
    sin = np.sin(angles)

    bottom = np.column_stack((cos, sin, np.zeros(seg, dtype=float)))
    top = np.column_stack((cos, sin, np.ones(seg, dtype=float)))

    vertices = np.vstack((bottom, top)).astype(np.float32, copy=False)
    normals = np.vstack(
        (
            np.column_stack((cos, sin, np.zeros(seg, dtype=float))),
            np.column_stack((cos, sin, np.zeros(seg, dtype=float))),
        )
    ).astype(np.float32, copy=False)

    faces: list[list[int]] = []
    for i in range(seg):
        j = (i + 1) % seg
        faces.append([i, j, seg + i])
        faces.append([seg + i, j, seg + j])

    faces_arr = np.asarray(faces, dtype=np.int32)
    template = {
        "vertices": vertices,
        "normals": normals,
        "faces": faces_arr,
        "z": vertices[:, 2].astype(np.float32, copy=False),
    }
    _CYLINDER_TEMPLATE_CACHE[seg] = template
    return template


def _rotation_from_z(direction: np.ndarray) -> np.ndarray:
    """Return rotation matrix that aligns the +Z axis with ``direction``."""

    dir_vec = np.asarray(direction, dtype=float)
    norm = float(np.linalg.norm(dir_vec))
    if norm <= 1e-8 or not np.isfinite(norm):
        return np.eye(3, dtype=float)
    dir_unit = dir_vec / norm
    z_axis = np.array([0.0, 0.0, 1.0], dtype=float)
    c = float(np.dot(z_axis, dir_unit))
    if c >= 0.9999:
        return np.eye(3, dtype=float)
    if c <= -0.9999:
        return np.array(
            [[1.0, 0.0, 0.0], [0.0, -1.0, 0.0], [0.0, 0.0, -1.0]],
            dtype=float,
        )

    axis = np.cross(z_axis, dir_unit)
    sin_theta = float(np.linalg.norm(axis))
    if sin_theta <= 1e-8 or not np.isfinite(sin_theta):
        return np.eye(3, dtype=float)
    axis_unit = axis / sin_theta

    kx, ky, kz = axis_unit
    K = np.array(
        [
            [0.0, -kz, ky],
            [kz, 0.0, -kx],
            [-ky, kx, 0.0],
        ],
        dtype=float,
    )

    rot = np.eye(3, dtype=float) + sin_theta * K + (1.0 - c) * (K @ K)
    return rot


def _build_stick_mesh(
    bonds: np.ndarray,
    atom_positions: np.ndarray,
    atom_colors: Optional[np.ndarray],
    radius: float,
    segments_circle: int = 12,
) -> Optional[Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]]:
    """Build a mesh for cylindrical sticks between bonded atom pairs."""

    template = _get_cylinder_template(segments_circle)
    if template is None:
        return None

    pts = np.asarray(atom_positions, dtype=float)
    if pts.ndim != 2 or pts.shape[1] != 3:
        return None

    bonds_arr = np.asarray(bonds, dtype=int)
    if bonds_arr.ndim != 2 or bonds_arr.shape[1] != 2:
        return None

    colors_arr: Optional[np.ndarray]
    if atom_colors is not None:
        try:
            colors_arr = np.asarray(atom_colors, dtype=float)
            if colors_arr.shape[0] != pts.shape[0]:
                colors_arr = None
        except Exception:
            colors_arr = None
    else:
        colors_arr = None

    base_color = np.array([0.8, 0.8, 0.8, 1.0], dtype=float)
    base_vertices = np.asarray(template["vertices"], dtype=float)
    base_normals = np.asarray(template["normals"], dtype=float)
    base_faces = np.asarray(template["faces"], dtype=np.int32)
    base_z = np.asarray(template.get("z", base_vertices[:, 2]), dtype=float)

    verts_list: list[np.ndarray] = []
    norms_list: list[np.ndarray] = []
    cols_list: list[np.ndarray] = []
    faces_list: list[np.ndarray] = []
    idx_offset = 0

    radius_val = float(radius)
    if not np.isfinite(radius_val) or radius_val <= 0.0:
        return None

    verts_per_cyl = base_vertices.shape[0]

    for pair in bonds_arr:
        i0 = int(pair[0])
        i1 = int(pair[1])
        if (
            i0 < 0
            or i1 < 0
            or i0 >= pts.shape[0]
            or i1 >= pts.shape[0]
            or i0 == i1
        ):
            continue

        start = pts[i0]
        end = pts[i1]
        vec = end - start
        length = float(np.linalg.norm(vec))
        if not np.isfinite(length) or length <= 1e-5:
            continue

        rot = _rotation_from_z(vec)

        verts_local = base_vertices.copy()
        verts_local[:, :2] *= radius_val
        verts_local[:, 2] *= length
        verts_world = verts_local @ rot.T + start
        verts_list.append(verts_world)

        normals_world = base_normals @ rot.T
        norms_list.append(normals_world)

        if colors_arr is not None:
            c0 = colors_arr[i0]
            c1 = colors_arr[i1]
        else:
            c0 = base_color
            c1 = base_color

        z = base_z.reshape(-1, 1)
        col = c0 * (1.0 - z) + c1 * z
        col[:, 3] = 1.0
        cols_list.append(col)

        faces_list.append(base_faces + idx_offset)
        idx_offset += verts_per_cyl

    if not verts_list:
        return None

    positions = np.vstack(verts_list).astype(np.float32, copy=False)
    normals = np.vstack(norms_list).astype(np.float32, copy=False)
    colors = np.vstack(cols_list).astype(np.float32, copy=False)
    faces = np.vstack(faces_list).astype(np.int32, copy=False)

    return positions, normals, faces, colors


if _HAVE_NUMBA and nb is not None:
    @nb.jit(nopython=True, nogil=True, cache=True)
    def _estimate_ambient_occlusion_nb(
        pts: np.ndarray,
        radius: float,
        max_neighbors: int,
    ) -> np.ndarray:
        n = pts.shape[0]
        occ = np.zeros(n, dtype=np.float64)
        r2 = radius * radius
        if n <= 1 or r2 <= 0.0:
            return occ

        for i in range(n):
            count = 0
            x0 = pts[i, 0]
            y0 = pts[i, 1]
            z0 = pts[i, 2]
            for j in range(n):
                if i == j:
                    continue
                dx = pts[j, 0] - x0
                dy = pts[j, 1] - y0
                dz = pts[j, 2] - z0
                if dx * dx + dy * dy + dz * dz < r2:
                    count += 1
                    if max_neighbors > 0 and count >= max_neighbors:
                        break

            if max_neighbors > 0:
                if count > max_neighbors:
                    count = max_neighbors
                occ[i] = float(count) / float(max_neighbors)

        return occ


def _estimate_ambient_occlusion(
    points: np.ndarray,
    radius: float = 4.0,
    max_neighbors: int = 32,
) -> Optional[np.ndarray]:
    pts = np.asarray(points, dtype=float)
    if pts.ndim != 2 or pts.shape[0] == 0:
        return None

    n = pts.shape[0]
    if n == 1:
        return np.zeros(1, dtype=float)

    r = float(radius)
    if not np.isfinite(r) or r <= 0.0:
        return None

    if _HAVE_NUMBA and nb is not None:
        try:
            occ_nb = _estimate_ambient_occlusion_nb(pts, r, int(max_neighbors))
            return np.clip(occ_nb, 0.0, 1.0)
        except Exception:
            pass

    cell = r
    inv_cell = 1.0 / cell

    centered = pts - pts.mean(axis=0)
    ijk = np.floor(centered * inv_cell).astype(np.int32)

    grid: dict[tuple[int, int, int], list[int]] = {}
    for idx, key in enumerate(map(tuple, ijk)):
        grid.setdefault(key, []).append(idx)

    neighbor_offsets = [
        (dx, dy, dz)
        for dx in (-1, 0, 1)
        for dy in (-1, 0, 1)
        for dz in (-1, 0, 1)
    ]

    occ = np.zeros(n, dtype=float)
    r2 = r * r

    for idx, key in enumerate(map(tuple, ijk)):
        ix, iy, iz = key
        cand_idx: list[int] = []
        for dx, dy, dz in neighbor_offsets:
            cand_idx.extend(grid.get((ix + dx, iy + dy, iz + dz), []))

        if not cand_idx:
            continue

        if len(cand_idx) > max_neighbors * 4:
            cand_idx = cand_idx[: max_neighbors * 4]

        cand = np.asarray(cand_idx, dtype=int)
        diffs = pts[cand] - pts[idx]
        dist2 = np.einsum("ij,ij->i", diffs, diffs)
        within = dist2 < r2
        count = int(np.count_nonzero(within))
        if count <= 0:
            continue

        if count > max_neighbors:
            count = max_neighbors
        occ[idx] = float(count)

    if max_neighbors > 0:
        occ /= float(max_neighbors)

    return np.clip(occ, 0.0, 1.0)


if _HAVE_NUMBA and nb is not None:

    @nb.jit(nopython=True, nogil=True)
    def _accumulate_gaussians_nb(
        pts: np.ndarray,
        sigmas: np.ndarray,
        grid: np.ndarray,
        origin: np.ndarray,
        spacing: float,
        cutoff_factor: float,
    ) -> None:
        nx, ny, nz = grid.shape
        for i in range(pts.shape[0]):
            sigma = sigmas[i]
            if sigma <= 0.0 or not np.isfinite(sigma):
                continue
            cutoff = cutoff_factor * sigma
            px = pts[i, 0]
            py = pts[i, 1]
            pz = pts[i, 2]
            inv_two_sigma2 = 1.0 / (2.0 * sigma * sigma)

            min_ix = max(int(math.floor((px - origin[0] - cutoff) / spacing)), 0)
            min_iy = max(int(math.floor((py - origin[1] - cutoff) / spacing)), 0)
            min_iz = max(int(math.floor((pz - origin[2] - cutoff) / spacing)), 0)
            max_ix = min(int(math.ceil((px - origin[0] + cutoff) / spacing)), nx - 1)
            max_iy = min(int(math.ceil((py - origin[1] + cutoff) / spacing)), ny - 1)
            max_iz = min(int(math.ceil((pz - origin[2] + cutoff) / spacing)), nz - 1)

            for ix in range(min_ix, max_ix + 1):
                dx = origin[0] + ix * spacing - px
                dx2 = dx * dx
                for iy in range(min_iy, max_iy + 1):
                    dy = origin[1] + iy * spacing - py
                    dy2 = dy * dy
                    for iz in range(min_iz, max_iz + 1):
                        dz = origin[2] + iz * spacing - pz
                        dist2 = dx2 + dy2 + dz * dz
                        grid[ix, iy, iz] += math.exp(-dist2 * inv_two_sigma2)
else:

    def _accumulate_gaussians_nb(
        pts: np.ndarray,
        sigmas: np.ndarray,
        grid: np.ndarray,
        origin: np.ndarray,
        spacing: float,
        cutoff_factor: float,
    ) -> None:
        nx, ny, nz = grid.shape
        for i in range(pts.shape[0]):
            sigma = sigmas[i]
            if sigma <= 0.0 or not np.isfinite(sigma):
                continue
            cutoff = cutoff_factor * sigma
            px, py, pz = pts[i]
            inv_two_sigma2 = 1.0 / (2.0 * sigma * sigma)

            min_ix = max(int(math.floor((px - origin[0] - cutoff) / spacing)), 0)
            min_iy = max(int(math.floor((py - origin[1] - cutoff) / spacing)), 0)
            min_iz = max(int(math.floor((pz - origin[2] - cutoff) / spacing)), 0)
            max_ix = min(int(math.ceil((px - origin[0] + cutoff) / spacing)), nx - 1)
            max_iy = min(int(math.ceil((py - origin[1] + cutoff) / spacing)), ny - 1)
            max_iz = min(int(math.ceil((pz - origin[2] + cutoff) / spacing)), nz - 1)

            for ix in range(min_ix, max_ix + 1):
                dx = origin[0] + ix * spacing - px
                dx2 = dx * dx
                for iy in range(min_iy, max_iy + 1):
                    dy = origin[1] + iy * spacing - py
                    dy2 = dy * dy
                    for iz in range(min_iz, max_iz + 1):
                        dz = origin[2] + iz * spacing - pz
                        dist2 = dx2 + dy2 + dz * dz
                        grid[ix, iy, iz] += math.exp(-dist2 * inv_two_sigma2)


def _build_gaussian_density_grid(
    pts: np.ndarray,
    sigmas: np.ndarray,
    grid_spacing: float = 1.0,
    padding: float = 2.5,
    cutoff_factor: float = 2.5,
    max_dim: int = 96,
) -> Optional[tuple[np.ndarray, np.ndarray, float]]:
    pts_arr = np.asarray(pts, dtype=float)
    if pts_arr.ndim != 2 or pts_arr.shape[0] < 4:
        return None
    sig_arr = np.asarray(sigmas, dtype=float)
    if sig_arr.shape[0] != pts_arr.shape[0]:
        return None

    spacing = max(float(grid_spacing), 0.2)
    pad = max(float(padding), spacing * 1.5)
    max_dim = max(int(max_dim), 16)

    mins = pts_arr.min(axis=0) - pad
    maxs = pts_arr.max(axis=0) + pad
    extent = maxs - mins
    dims_float = np.ceil(extent / spacing).astype(int) + 1
    max_axis = int(np.max(dims_float))
    if max_axis > max_dim:
        scale = max_axis / float(max_dim)
        spacing *= scale
        dims_float = np.ceil(extent / spacing).astype(int) + 1

    dims = np.maximum(dims_float, 3)
    grid = np.zeros(tuple(int(x) for x in dims), dtype=np.float32)

    sig_arr = np.clip(sig_arr, spacing * 0.25, spacing * 5.0)
    origin = mins.astype(np.float32)

    _accumulate_gaussians_nb(
        pts_arr.astype(np.float32),
        sig_arr.astype(np.float32),
        grid,
        origin,
        float(spacing),
        float(cutoff_factor),
    )

    if not np.isfinite(grid.max()) or grid.max() <= 0.0:
        return None

    return grid, origin, float(spacing)


def _generate_surface_mesh_from_gaussians(
    pts: np.ndarray,
    sigmas: np.ndarray,
    *,
    grid_spacing: float = 1.0,
    padding: float = 2.5,
    cutoff_factor: float = 2.5,
    iso_value: float = 0.2,
    max_dim: int = 96,
) -> Optional[tuple[np.ndarray, np.ndarray, np.ndarray]]:
    if not _HAVE_SKIMAGE:
        return None

    density_data = _build_gaussian_density_grid(
        pts,
        sigmas,
        grid_spacing=grid_spacing,
        padding=padding,
        cutoff_factor=cutoff_factor,
        max_dim=max_dim,
    )
    if density_data is None:
        return None

    grid, origin, spacing = density_data
    level = float(iso_value)
    if not np.isfinite(level) or level <= 0.0:
        level = 0.2 * float(grid.max())
    level = min(level, float(grid.max()) * 0.9)
    if level <= 0.0:
        return None

    try:
        verts, faces, norms, _ = _sk_measure.marching_cubes(
            grid,
            level=level,
            spacing=(spacing, spacing, spacing),
        )
    except Exception:
        return None

    verts = np.asarray(verts, dtype=np.float32)
    verts += origin
    faces = np.asarray(faces, dtype=np.int32)
    norms = np.asarray(norms, dtype=np.float32)
    return verts, faces, norms


def _build_trace_ups(
    atoms: np.ndarray,
    res_ids: Optional[np.ndarray],
    ca_coords: Optional[np.ndarray],
) -> Optional[np.ndarray]:
    """Estimate backbone up-vectors per CA from full atom data.

    Uses C-O vectors within each residue as an approximate "up" direction
    similar to pyball. Returns an array of shape ``(N, 3)`` aligned to
    ``res_ids`` / ``ca_coords`` or ``None`` if not available.
    """

    if res_ids is None or ca_coords is None:
        return None
    if not isinstance(atoms, np.ndarray):
        return None

    fields = set(atoms.dtype.fields or {})
    if not {"res_id", "atom_name", "xyz"}.issubset(fields):
        return None

    try:
        atom_res_id = np.asarray(atoms["res_id"])
        atom_names = np.char.strip(atoms["atom_name"].astype(str))
        atom_xyz = np.asarray(atoms["xyz"], dtype=float)
    except Exception:
        return None

    n = len(res_ids)
    ups = np.zeros((n, 3), dtype=float)

    for i, rid in enumerate(res_ids):
        mask = atom_res_id == rid
        if not np.any(mask):
            ups[i] = np.array([0.0, 0.0, 1.0], dtype=float)
            continue

        names = atom_names[mask]
        coords = atom_xyz[mask]
        idx_c = np.where(names == "C")[0]
        idx_o = np.where(names == "O")[0]

        if idx_c.size and idx_o.size:
            c = coords[idx_c[0]]
            o = coords[idx_o[0]]
            up_vec = c - o
        else:
            up_vec = np.array([0.0, 0.0, 1.0], dtype=float)

        norm = float(np.linalg.norm(up_vec))
        if norm > 0.0:
            up_vec = up_vec / norm
        ups[i] = up_vec

    # Make neighbouring up-vectors point in a consistent direction
    for i in range(1, n):
        if float(np.dot(ups[i - 1], ups[i])) < 0.0:
            ups[i] = -ups[i]

    return ups


def _generate_cartoon_tube_arrays(
    coords: np.ndarray,
    colors: Optional[np.ndarray],
    trace_ups: Optional[np.ndarray] = None,
    base_radius: float = 0.5,
    segments_circle: int = 14,
    subdivisions: int = 6,
    *,
    style: str = "tube",
    ss_codes: Optional[np.ndarray] = None,
    config: Optional[dict] = None,
) -> Optional[tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray]]]:
    """Return (vertices, normals, faces, colors) for the cartoon mesh."""

    style_l = str(style or "tube").lower()

    arr = np.asarray(coords, dtype=float)
    if arr.ndim != 2 or arr.shape[0] < 2:
        return None

    n = arr.shape[0]
    col_arr: Optional[np.ndarray]
    if colors is not None:
        col_arr = np.asarray(colors, dtype=float)
        if col_arr.shape[0] != n:
            col_arr = None
    else:
        col_arr = None

    try:
        path, path_colors = _smooth_backbone(arr, col_arr, subdivisions=subdivisions)
    except Exception:
        path = arr
        path_colors = col_arr

    m = path.shape[0]
    if m < 2:
        return None

    # Densify/up-sample up-vectors if provided
    ups_path: Optional[np.ndarray] = None
    if trace_ups is not None:
        try:
            ups_arr = np.asarray(trace_ups, dtype=float)
            if ups_arr.shape[0] == n:
                ups_path, _ = _smooth_backbone(ups_arr, None, subdivisions=subdivisions)
        except Exception:
            ups_path = None

    if style_l == "ribbon":
        return _generate_cartoon_ribbon_arrays(
            path,
            path_colors,
            ups_path,
            base_radius,
            subdivisions,
            ss_codes,
            config,
            n,
        )

    two_pi = 2.0 * math.pi
    segments_circle = max(int(segments_circle), 6)
    angles = np.linspace(0.0, two_pi, segments_circle, endpoint=False)
    cos_a = np.cos(angles)
    sin_a = np.sin(angles)

    verts = np.zeros((m * segments_circle, 3), dtype=float)
    norms = np.zeros_like(verts)
    cols: Optional[np.ndarray]
    if path_colors is not None:
        cols = np.zeros((m * segments_circle, 4), dtype=float)
    else:
        cols = None

    for i in range(m):
        p = path[i]
        if i == 0:
            t = path[i + 1] - path[i]
        elif i == m - 1:
            t = path[i] - path[i - 1]
        else:
            t = path[i + 1] - path[i - 1]

        tn = float(np.linalg.norm(t))
        if tn <= 0.0:
            t = np.array([0.0, 0.0, 1.0], dtype=float)
        else:
            t = t / tn

        if ups_path is not None and ups_path.shape[0] == m:
            up_vec = ups_path[i]
        else:
            up_vec = np.array([0.0, 0.0, 1.0], dtype=float)

        # Remove component parallel to tangent
        up_vec = up_vec - np.dot(up_vec, t) * t
        un = float(np.linalg.norm(up_vec))
        if un <= 0.0:
            if abs(t[2]) < 0.9:
                up_vec = np.array([0.0, 0.0, 1.0], dtype=float)
            else:
                up_vec = np.array([1.0, 0.0, 0.0], dtype=float)
            up_vec = up_vec - np.dot(up_vec, t) * t
            un = float(np.linalg.norm(up_vec))
        if un > 0.0:
            up_vec = up_vec / un

        side = np.cross(t, up_vec)
        sn = float(np.linalg.norm(side))
        if sn > 0.0:
            side = side / sn
        else:
            side = np.array([1.0, 0.0, 0.0], dtype=float)

        radius = float(base_radius)
        base_index = i * segments_circle
        for j in range(segments_circle):
            nx = cos_a[j] * up_vec[0] + sin_a[j] * side[0]
            ny = cos_a[j] * up_vec[1] + sin_a[j] * side[1]
            nz = cos_a[j] * up_vec[2] + sin_a[j] * side[2]
            normal = np.array([nx, ny, nz], dtype=float)
            normal_norm = float(np.linalg.norm(normal))
            if normal_norm > 0.0:
                normal /= normal_norm
            norms[base_index + j] = normal
            verts[base_index + j] = p + radius * normal
            if cols is not None and path_colors is not None:
                cols[base_index + j] = path_colors[i]

    faces = []
    for i in range(m - 1):
        i0 = i * segments_circle
        i1 = (i + 1) * segments_circle
        for j in range(segments_circle):
            k0 = i0 + j
            k1 = i0 + (j + 1) % segments_circle
            k2 = i1 + j
            k3 = i1 + (j + 1) % segments_circle
            faces.append((k0, k2, k1))
            faces.append((k1, k2, k3))

    if not faces:
        return None

    faces_arr = np.asarray(faces, dtype=np.int32)
    return verts, norms, faces_arr, cols


def _generate_cartoon_ribbon_arrays(
    path: np.ndarray,
    path_colors: Optional[np.ndarray],
    ups_path: Optional[np.ndarray],
    base_radius: float,
    subdivisions: int,
    ss_codes: Optional[np.ndarray],
    config: Optional[dict],
    n_samples: int,
) -> Optional[tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray]]]:
    m = path.shape[0]
    if m < 2:
        return None

    cfg = config or {}
    profile_segments = max(int(cfg.get("profile_segments", 20)), 4)
    ribbon_scale = float(cfg.get("ribbon_thickness_scale", 0.45))

    widths_res, thickness_res, profile_power_res = _compute_ss_profile_arrays(
        ss_codes,
        cfg,
        max(int(n_samples), 2),
        base_radius,
        ribbon_scale,
    )

    if (
        widths_res is None
        or thickness_res is None
        or profile_power_res is None
    ):
        widths_res = np.full(path.shape[0], base_radius, dtype=float)
        thickness_res = np.full(path.shape[0], base_radius * ribbon_scale, dtype=float)
        profile_power_res = np.ones(path.shape[0], dtype=float) * 2.0

    width_path = _subdivide_attribute_array(widths_res, subdivisions)
    thickness_path = _subdivide_attribute_array(thickness_res, subdivisions)
    profile_power_path = _subdivide_attribute_array(profile_power_res, subdivisions)

    if width_path.shape[0] != m:
        width_path = _resample_attribute_array(width_path, m)
    if thickness_path.shape[0] != m:
        thickness_path = _resample_attribute_array(thickness_path, m)
    if profile_power_path.shape[0] != m:
        profile_power_path = _resample_attribute_array(profile_power_path, m)

    two_pi = 2.0 * math.pi
    angles = np.linspace(0.0, two_pi, profile_segments, endpoint=False)
    cos_a = np.cos(angles)
    sin_a = np.sin(angles)

    verts = np.zeros((m * profile_segments, 3), dtype=float)
    norms = np.zeros_like(verts)
    cols: Optional[np.ndarray]
    if path_colors is not None:
        cols = np.zeros((m * profile_segments, 4), dtype=float)
    else:
        cols = None

    tangents, up_vectors = _build_parallel_transport_frames(path, ups_path)
    prev_side = None

    for i in range(m):
        p = path[i]
        t = tangents[i]
        up_vec = up_vectors[i]

        side = np.cross(t, up_vec)
        sn = float(np.linalg.norm(side))
        if sn <= 0.0:
            side = _default_side_from_up(up_vec)
        else:
            side = side / sn

        if prev_side is not None and float(np.dot(prev_side, side)) < 0.0:
            side = -side
        prev_side = side

        width = max(float(width_path[i]), 1e-4)
        thickness = max(float(thickness_path[i]), 1e-4)
        power = max(float(profile_power_path[i]), 1e-3)
        exponent = 2.0 / max(power, 1e-3)
        base_index = i * profile_segments

        for j in range(profile_segments):
            cx = math.copysign(abs(cos_a[j]) ** exponent, cos_a[j])
            cy = math.copysign(abs(sin_a[j]) ** exponent, sin_a[j])
            offset = (width * 0.5 * cx) * side + (thickness * 0.5 * cy) * up_vec
            verts[base_index + j] = p + offset
            norm_vec = offset
            norm_len = float(np.linalg.norm(norm_vec))
            if norm_len > 0.0:
                norms[base_index + j] = norm_vec / norm_len
            else:
                norms[base_index + j] = up_vec
            if cols is not None and path_colors is not None:
                cols[base_index + j] = path_colors[i]

    faces = []
    for i in range(m - 1):
        i0 = i * profile_segments
        i1 = (i + 1) * profile_segments
        for j in range(profile_segments):
            k0 = i0 + j
            k1 = i0 + (j + 1) % profile_segments
            k2 = i1 + j
            k3 = i1 + (j + 1) % profile_segments
            faces.append((k0, k2, k1))
            faces.append((k1, k2, k3))

    if not faces:
        return None

    faces_arr = np.asarray(faces, dtype=np.int32)
    return verts, norms, faces_arr, cols


def _subdivide_attribute_array(values: np.ndarray, subdivisions: int) -> np.ndarray:
    vals = np.asarray(values, dtype=float)
    n = vals.shape[0]
    if n == 0:
        return np.zeros(0, dtype=float)
    subdivisions = max(int(subdivisions), 1)
    if subdivisions <= 1 or n == 1:
        return vals.copy()

    out: list[np.ndarray] = []
    for i in range(n - 1):
        v0 = vals[i]
        v1 = vals[i + 1]
        for j in range(subdivisions):
            t = float(j) / float(subdivisions)
            out.append((1.0 - t) * v0 + t * v1)
    out.append(vals[-1])
    return np.asarray(out, dtype=float)


def _resample_attribute_array(values: np.ndarray, target_len: int) -> np.ndarray:
    if target_len <= 0:
        return np.zeros(0, dtype=float)
    src = np.asarray(values, dtype=float)
    if src.size == 0:
        return np.zeros(target_len, dtype=float)
    x_src = np.linspace(0.0, 1.0, src.size)
    x_dst = np.linspace(0.0, 1.0, target_len)
    return np.interp(x_dst, x_src, src)


def _compute_ss_profile_arrays(
    ss_codes: Optional[np.ndarray],
    config: dict,
    n_points: int,
    base_radius: float,
    base_thickness_scale: float,
) -> tuple[Optional[np.ndarray], Optional[np.ndarray], Optional[np.ndarray]]:
    if n_points <= 0:
        return None, None, None

    codes = None
    if ss_codes is not None:
        try:
            codes = np.asarray(ss_codes)
        except Exception:
            codes = None

    shapes_cfg = (config.get("ss_shapes") or {}) if config is not None else {}

    def _shape_for(code: Optional[str]) -> dict:
        key = str(code).strip().upper()[:1] if code is not None else "C"
        if key == "H":
            lookup = "helix"
        elif key == "E":
            lookup = "strand"
        else:
            lookup = "coil"
        return shapes_cfg.get(lookup, shapes_cfg.get("coil", {}))

    widths = np.zeros(n_points, dtype=float)
    thicknesses = np.zeros(n_points, dtype=float)
    strand_mask = np.zeros(n_points, dtype=bool)
    strand_arrow_scale = np.ones(n_points, dtype=float)
    profile_power = np.ones(n_points, dtype=float)

    for i in range(n_points):
        code = codes[i] if (codes is not None and i < codes.shape[0]) else None
        shape = _shape_for(code)
        width_scale = float(shape.get("width", 1.0))
        thick_scale = float(shape.get("thickness", 1.0))
        widths[i] = base_radius * width_scale
        thicknesses[i] = base_radius * base_thickness_scale * thick_scale
        profile_power[i] = float(shape.get("profile_power", 2.0))
        if str(code).strip().upper()[:1] == "E":
            strand_mask[i] = True
            strand_arrow_scale[i] = float(shape.get("arrow_scale", 1.0))

    if np.any(strand_mask):
        arrow_len = max(int(config.get("arrow_tip_residues", 0)), 0) if config else 0
        if arrow_len > 0:
            start = 0
            n = n_points
            while start < n:
                if not strand_mask[start]:
                    start += 1
                    continue
                end = start
                while end + 1 < n and strand_mask[end + 1]:
                    end += 1
                block_len = end - start + 1
                tip = min(block_len, arrow_len)
                if tip > 0:
                    arrow_scale = float(np.max(strand_arrow_scale[start : end + 1]))
                    if not math.isclose(arrow_scale, 1.0):
                        for offset in range(tip):
                            idx = end - tip + 1 + offset
                            frac = float(offset + 1) / float(tip)
                            widths[idx] *= 1.0 + (arrow_scale - 1.0) * frac
                            thicknesses[idx] *= max(0.2, 1.0 - 0.4 * frac)
                start = end + 1

    return widths, thicknesses, profile_power


def _build_parallel_transport_frames(
    path: np.ndarray,
    ups_hint: Optional[np.ndarray],
) -> tuple[np.ndarray, np.ndarray]:
    pts = np.asarray(path, dtype=float)
    m = pts.shape[0]
    tangents = np.zeros((m, 3), dtype=float)
    ups = np.zeros_like(tangents)
    hint = None
    if ups_hint is not None:
        try:
            hint = np.asarray(ups_hint, dtype=float)
            if hint.shape[0] != m:
                hint = None
        except Exception:
            hint = None

    prev_up = None
    for i in range(m):
        if i == 0:
            t = pts[i + 1] - pts[i]
        elif i == m - 1:
            t = pts[i] - pts[i - 1]
        else:
            t = pts[i + 1] - pts[i - 1]

        tn = float(np.linalg.norm(t))
        if tn <= 0.0:
            t = np.array([0.0, 0.0, 1.0], dtype=float)
        else:
            t = t / tn
        tangents[i] = t

        candidate = hint[i] if hint is not None else None
        up_vec = _project_perpendicular(candidate, t)
        if up_vec is None and prev_up is not None:
            up_vec = _project_perpendicular(prev_up, t)
        if up_vec is None:
            fallback = _default_up_from_tangent(t)
            up_vec = _project_perpendicular(fallback, t)
        if up_vec is None:
            up_vec = np.array([0.0, 1.0, 0.0], dtype=float)

        if prev_up is not None and float(np.dot(prev_up, up_vec)) < 0.0:
            up_vec = -up_vec

        ups[i] = up_vec
        prev_up = up_vec

    return tangents, ups


def _project_perpendicular(vec: Optional[np.ndarray], direction: np.ndarray) -> Optional[np.ndarray]:
    if vec is None:
        return None
    arr = np.asarray(vec, dtype=float)
    if arr.shape != (3,):
        return None
    proj = arr - np.dot(arr, direction) * direction
    norm = float(np.linalg.norm(proj))
    if norm <= 1e-8:
        return None
    return proj / norm


def _default_up_from_tangent(tangent: np.ndarray) -> np.ndarray:
    axis_candidates = (
        np.array([0.0, 0.0, 1.0], dtype=float),
        np.array([0.0, 1.0, 0.0], dtype=float),
        np.array([1.0, 0.0, 0.0], dtype=float),
    )
    for axis in axis_candidates:
        proj = axis - np.dot(axis, tangent) * tangent
        if float(np.linalg.norm(proj)) > 1e-6:
            return proj
    return np.array([0.0, 1.0, 0.0], dtype=float)


def _default_side_from_up(up_vec: np.ndarray) -> np.ndarray:
    axis = np.array([1.0, 0.0, 0.0], dtype=float)
    if abs(float(np.dot(axis, up_vec))) > 0.9:
        axis = np.array([0.0, 1.0, 0.0], dtype=float)
    side = np.cross(up_vec, axis)
    sn = float(np.linalg.norm(side))
    if sn <= 0.0:
        return np.array([0.0, 0.0, 1.0], dtype=float)
    return side / sn


def _build_cartoon_tube_item(
    coords: np.ndarray,
    colors: Optional[np.ndarray],
    trace_ups: Optional[np.ndarray] = None,
    base_radius: float = 0.5,
    segments_circle: int = 14,
    subdivisions: int = 6,
) -> Optional[gl.GLMeshItem]:
    """Legacy GL helper that wraps the backend-agnostic tube arrays."""

    if gl is None:
        return None

    arrays = _generate_cartoon_tube_arrays(
        coords,
        colors,
        trace_ups=trace_ups,
        base_radius=base_radius,
        segments_circle=segments_circle,
        subdivisions=subdivisions,
    )
    if arrays is None:
        return None

    verts, faces_arr, cols = arrays
    try:
        md = gl.MeshData(vertexes=verts, faces=faces_arr)
        if cols is not None:
            md.setVertexColors(cols)
        mi = gl.GLMeshItem(
            meshdata=md,
            smooth=True,
            shader="shaded",
            glOptions="opaque",
        )
    except Exception:
        return None

    return mi


def _smooth_backbone(
    coords: np.ndarray,
    colors: Optional[np.ndarray],
    subdivisions: int = 5,
) -> tuple[np.ndarray, Optional[np.ndarray]]:
    """Return a slightly densified backbone polyline and matching colors.

    A small number of linear subdivisions between consecutive CA atoms
    produces a smoother-looking cartoon without heavy spline machinery.
    """

    arr = np.asarray(coords, dtype=float)
    if arr.ndim != 2 or arr.shape[0] < 2:
        return arr, colors

    n = arr.shape[0]
    subdivs = max(int(subdivisions), 1)
    if subdivs <= 1:
        return arr, colors

    out_pos: list[np.ndarray] = []
    out_col: Optional[list[np.ndarray]] = None
    col_arr = None
    if colors is not None:
        col_arr = np.asarray(colors, dtype=float)
        if col_arr.shape[0] == n:
            out_col = []
        else:
            col_arr = None

    for i in range(n - 1):
        p0 = arr[i - 1] if i > 0 else arr[i]
        p1 = arr[i]
        p2 = arr[i + 1]
        p3 = arr[i + 2] if (i + 2) < n else arr[i + 1]

        if col_arr is not None:
            c0 = col_arr[i - 1] if i > 0 else col_arr[i]
            c1 = col_arr[i]
            c2 = col_arr[i + 1]
            c3 = col_arr[i + 2] if (i + 2) < n else col_arr[i + 1]

        for j in range(subdivs):
            if i > 0 and j == 0:
                continue
            t = float(j) / float(subdivs)
            pos = _catmull_rom(p0, p1, p2, p3, t)
            out_pos.append(pos)
            if out_col is not None and col_arr is not None:
                col = _catmull_rom(c0, c1, c2, c3, t)
                out_col.append(col)

    # Append final endpoint explicitly
    out_pos.append(arr[-1])
    if out_col is not None and col_arr is not None:
        out_col.append(col_arr[-1])

    pos_arr = np.asarray(out_pos, dtype=float)
    col_out_arr = np.asarray(out_col, dtype=float) if out_col is not None else None
    return pos_arr, col_out_arr


def _catmull_rom(p0: np.ndarray, p1: np.ndarray, p2: np.ndarray, p3: np.ndarray, t: float) -> np.ndarray:
    t = float(np.clip(t, 0.0, 1.0))
    t2 = t * t
    t3 = t2 * t
    return 0.5 * (
        (2 * p1)
        + (-p0 + p2) * t
        + (2 * p0 - 5 * p1 + 4 * p2 - p3) * t2
        + (-p0 + 3 * p1 - 3 * p2 + p3) * t3
    )


def _generate_trace_arrays(
    coords: np.ndarray,
    colors: Optional[np.ndarray],
    subdivisions: int = 5,
) -> tuple[np.ndarray, Optional[np.ndarray]]:
    """Return smoothed trace coordinates and colors for the CA backbone."""

    return _smooth_backbone(coords, colors, subdivisions=subdivisions)


def _extract_ca_trace(
    atoms: np.ndarray,
) -> tuple[Optional[np.ndarray], Optional[np.ndarray], Optional[np.ndarray]]:
    """Return (coords, res_ids, res_names) for CA atoms or (None, None, None).

    The ``atoms`` array is expected to expose at least ``'atom_name'`` and
    ``'xyz'`` fields; residue ids/names are optional.
    """

    if not isinstance(atoms, np.ndarray):
        return None, None, None
    fields = set(atoms.dtype.fields or {})
    if "atom_name" not in fields or "xyz" not in fields:
        return None, None, None

    names = atoms["atom_name"]
    try:
        text = np.char.strip(names.astype(str))
    except Exception:
        text = np.array([str(n).strip() for n in names])

    mask = text == "CA"
    if not mask.any():
        return None, None, None

    coords = np.asarray(atoms["xyz"][mask], dtype=float)

    res_ids = None
    res_names = None
    if "res_id" in fields:
        res_ids = np.asarray(atoms["res_id"][mask])
    if "res_name" in fields:
        try:
            res_names = np.char.strip(atoms["res_name"][mask].astype(str))
        except Exception:
            res_names = np.array(
                [str(x).strip() for x in atoms["res_name"][mask]]
            )
    return coords, res_ids, res_names


if _HAVE_NUMBA and nb is not None:
    @nb.jit(nopython=True, nogil=True, cache=True)
    def _build_bond_pairs_nb(pts: np.ndarray, max_length: float) -> np.ndarray:
        n = pts.shape[0]
        r2 = max_length * max_length
        if n < 2 or r2 <= 0.0:
            return np.zeros((0, 2), dtype=np.int64)

        count = 0
        for i in range(n - 1):
            x0 = pts[i, 0]
            y0 = pts[i, 1]
            z0 = pts[i, 2]
            for j in range(i + 1, n):
                dx = pts[j, 0] - x0
                dy = pts[j, 1] - y0
                dz = pts[j, 2] - z0
                if dx * dx + dy * dy + dz * dz <= r2:
                    count += 1

        if count == 0:
            return np.zeros((0, 2), dtype=np.int64)

        out = np.empty((count, 2), dtype=np.int64)
        k = 0
        for i in range(n - 1):
            x0 = pts[i, 0]
            y0 = pts[i, 1]
            z0 = pts[i, 2]
            for j in range(i + 1, n):
                dx = pts[j, 0] - x0
                dy = pts[j, 1] - y0
                dz = pts[j, 2] - z0
                if dx * dx + dy * dy + dz * dz <= r2:
                    out[k, 0] = i
                    out[k, 1] = j
                    k += 1

        return out


def _build_bond_pairs(coords: np.ndarray, max_length: float) -> np.ndarray:
    """Return an array of (i, j) index pairs for simple covalent bonds.

    Bonds are inferred purely from distance using a cutoff ``max_length`` in
    the *raw* coordinate frame. A simple grid-based neighbor search is used
    so the cost grows roughly linearly with the number of atoms. This is a
    lightweight approximation similar in spirit to pyball's stick geometry.
    """

    pts = np.asarray(coords, dtype=float)
    if pts.ndim != 2 or pts.shape[0] < 2:
        return np.zeros((0, 2), dtype=int)

    r = float(max_length)
    if not np.isfinite(r) or r <= 0.0:
        return np.zeros((0, 2), dtype=int)

    n = pts.shape[0]
    if _HAVE_NUMBA and nb is not None and n > 1:
        try:
            return _build_bond_pairs_nb(pts, r)
        except Exception:
            pass

    cell = r
    inv_cell = 1.0 / cell

    centered = pts - pts.mean(axis=0)
    ijk = np.floor(centered * inv_cell).astype(np.int32)

    grid: dict[tuple[int, int, int], list[int]] = {}
    for idx, key in enumerate(map(tuple, ijk)):
        grid.setdefault(key, []).append(idx)

    neighbor_offsets = [
        (dx, dy, dz)
        for dx in (-1, 0, 1)
        for dy in (-1, 0, 1)
        for dz in (-1, 0, 1)
    ]

    r2 = r * r
    bonds: list[tuple[int, int]] = []

    for i, key in enumerate(map(tuple, ijk)):
        ix, iy, iz = key
        cand_idx: list[int] = []
        for dx, dy, dz in neighbor_offsets:
            cand_idx.extend(grid.get((ix + dx, iy + dy, iz + dz), []))

        if not cand_idx:
            continue

        pi = pts[i]
        for j in cand_idx:
            if j <= i:
                continue
            d = pts[j] - pi
            if float(d[0] * d[0] + d[1] * d[1] + d[2] * d[2]) <= r2:
                bonds.append((i, j))

    if not bonds:
        return np.zeros((0, 2), dtype=int)

    return np.asarray(bonds, dtype=int)
