"""n3xd_ocp.sample — the bulk UV-grid sampler.

The reference is the nested Python loop in cad/operations/surface_pattern.py's
``_sample_grid``, reproduced below.  It has to match exactly rather than
closely: the samples feed a B-spline fit, so a last-ulp difference in a sample
location moves the patterned surface and would make the corpus disagree with
itself for no reason.
"""

from __future__ import annotations

import pytest

np = pytest.importorskip("numpy")

import n3xd_ocp
from OCP.BRepAdaptor import BRepAdaptor_Surface
from OCP.BRepPrimAPI import BRepPrimAPI_MakeBox, BRepPrimAPI_MakeCylinder
from OCP.BRepTools import BRepTools
from OCP.gp import gp_Ax2, gp_Dir, gp_Pnt, gp_Vec
from OCP.TopAbs import TopAbs_FACE, TopAbs_REVERSED
from OCP.TopoDS import TopoDS

from .test_inc1_modeling import sub_shapes

_MIN_NORMAL = 1e-12


def _faces(shape):
    return [TopoDS.Face_s(f) for f in sub_shapes(shape, TopAbs_FACE)]


def _python_grid(face, n):
    """surface_pattern.py::_sample_grid, reduced to points and normals."""
    adaptor = BRepAdaptor_Surface(face)
    umin, umax, vmin, vmax = BRepTools.UVBounds_s(face)
    sign = -1.0 if face.Orientation() == TopAbs_REVERSED else 1.0
    us = np.linspace(umin, umax, n)
    vs = np.linspace(vmin, vmax, n)
    point = np.zeros((n, n, 3))
    normal = np.zeros((n, n, 3))
    pnt, du, dv = gp_Pnt(), gp_Vec(), gp_Vec()
    for i in range(n):
        up = float(us[i])
        for j in range(n):
            adaptor.D1(up, float(vs[j]), pnt, du, dv)
            point[i, j] = (pnt.X(), pnt.Y(), pnt.Z())
            nx = du.Y() * dv.Z() - du.Z() * dv.Y()
            ny = du.Z() * dv.X() - du.X() * dv.Z()
            nz = du.X() * dv.Y() - du.Y() * dv.X()
            mag = (nx * nx + ny * ny + nz * nz) ** 0.5
            if mag > _MIN_NORMAL:
                normal[i, j] = (sign * nx / mag, sign * ny / mag, sign * nz / mag)
    return point, normal, (umin, umax, vmin, vmax)


@pytest.fixture(scope="module")
def shapes():
    return {
        "box": BRepPrimAPI_MakeBox(10.0, 20.0, 30.0).Shape(),
        "cylinder": BRepPrimAPI_MakeCylinder(
            gp_Ax2(gp_Pnt(0, 0, 0), gp_Dir(0, 0, 1)), 5.0, 12.0
        ).Shape(),
    }


@pytest.mark.parametrize("name", ["box", "cylinder"])
@pytest.mark.parametrize("n", [5, 33, 65])
def test_matches_the_python_loop_exactly(shapes, name, n):
    for face in _faces(shapes[name]):
        points, normals, bounds = n3xd_ocp.sample.face_grid(face, n)
        ref_points, ref_normals, ref_bounds = _python_grid(face, n)
        assert bounds == ref_bounds
        assert np.array_equal(points, ref_points)
        assert np.array_equal(normals, ref_normals)


def test_reversed_face_flips_the_normal(shapes):
    """A box has both orientations; the sign has to follow each face."""
    seen = set()
    for face in _faces(shapes["box"]):
        _points, normals, _bounds = n3xd_ocp.sample.face_grid(face, 5)
        _rp, ref_normals, _rb = _python_grid(face, 5)
        seen.add(face.Orientation() == TopAbs_REVERSED)
        assert np.array_equal(normals, ref_normals)
    assert seen == {True, False}, "fixture must cover both orientations"


def test_normals_are_unit_length(shapes):
    for face in _faces(shapes["cylinder"]):
        _points, normals, _bounds = n3xd_ocp.sample.face_grid(face, 17)
        lengths = np.linalg.norm(normals.reshape(-1, 3), axis=1)
        assert np.allclose(lengths, 1.0, atol=1e-12)


def test_endpoints_hit_the_bounds_exactly(shapes):
    """linspace forces the last sample onto the bound; start + (n-1)*step can
    miss it by an ulp, which is enough to move a fitted surface."""
    face = _faces(shapes["cylinder"])[0]
    n = 33
    points, _normals, (umin, umax, vmin, vmax) = n3xd_ocp.sample.face_grid(face, n)
    adaptor = BRepAdaptor_Surface(face)
    for i, u in ((0, umin), (n - 1, umax)):
        for j, v in ((0, vmin), (n - 1, vmax)):
            expected = adaptor.Value(u, v)
            assert points[i, j, 0] == expected.X()
            assert points[i, j, 1] == expected.Y()
            assert points[i, j, 2] == expected.Z()


def test_single_sample_is_the_lower_bound(shapes):
    face = _faces(shapes["box"])[0]
    points, _normals, (umin, _umax, vmin, _vmax) = n3xd_ocp.sample.face_grid(face, 1)
    assert points.shape == (1, 1, 3)
    expected = BRepAdaptor_Surface(face).Value(umin, vmin)
    assert points[0, 0, 0] == expected.X()


def test_rejects_a_non_positive_grid(shapes):
    with pytest.raises(ValueError):
        n3xd_ocp.sample.face_grid(_faces(shapes["box"])[0], 0)


def test_arrays_own_their_memory(shapes):
    import gc

    points, normals, _bounds = n3xd_ocp.sample.face_grid(
        _faces(BRepPrimAPI_MakeBox(1.0, 2.0, 3.0).Shape())[0], 9
    )
    expected = float(points.sum())
    gc.collect()
    assert float(points.sum()) == expected
    assert normals.base is not None  # the capsule, not the face