"""Stage 3 acceptance: the field path, composition, and the frame schedule.

The three rows that matter most are the ones that are easy to get subtly wrong and
impossible to notice on screen: the *colour* path must never resample the field onto
the vertices (that is what painted 30.15 % of the island grey), the display gain must
not touch the measurement, and a producer must never be able to make the display queue
grow.
"""

from __future__ import annotations

import numpy as np
import pytest
from orisys_app_shared.scene3d import synthetic
from orisys_app_shared.scene3d.composer import DISPLAY_ILLUSTRATIVE, SceneComposer
from orisys_app_shared.scene3d.contracts import (
    DisplaySettings,
    FieldDescriptor,
    FrameIdentity,
    SurfaceFieldFrame,
    ThreeDError,
    freeze,
)
from orisys_app_shared.scene3d.sampling import (
    remap_to_domain,
    sample_vertex_displacement,
)


# --------------------------------------------------------------------- helpers
def _frame(scene, scalar, valid=None, *, displacement=None, quality=None,
           identity=None, domain=None, shape=None):
    scalar = np.asarray(scalar, dtype=np.float32)
    if valid is None:
        valid = np.ones(scalar.shape, dtype=bool)
    frame = SurfaceFieldFrame(
        identity=identity or FrameIdentity("live", 1, 1, 0.1, generation=1),
        geometry_id=scene.geometry_id,
        field_domain_id=domain or scene.field_domain_id,
        scalar=freeze(scalar, dtype=np.float32),
        scalar_descriptor=FieldDescriptor("flow_magnitude", "px", "approximate", False),
        valid=freeze(np.asarray(valid, dtype=bool), dtype=np.bool_),
        displacement_sensor_m=(
            freeze(np.asarray(displacement, dtype=np.float32), dtype=np.float32)
            if displacement is not None
            else None
        ),
        displacement_quality=quality,
    )
    if shape is not None:
        object.__setattr__(frame, "scalar", freeze(np.asarray(shape, dtype=np.float32), dtype=np.float32))
    return frame


def _binding(shape=(3, 3), indices=None, weights=None):
    """The identity binding of the PoC scene, optionally overridden."""
    scene = synthetic.poc_scene()
    binding = scene.binding
    if indices is None and weights is None and shape == (3, 3):
        return binding
    from orisys_app_shared.scene3d.contracts import SurfaceBinding

    count = int(indices.shape[0]) if indices is not None else shape[0] * shape[1]
    uvs = np.zeros((count, 2), dtype=np.float32)
    if count == shape[0] * shape[1]:
        grid = np.stack(
            np.meshgrid(np.arange(shape[1]), np.arange(shape[0])), axis=-1
        ).reshape(-1, 2)
        uvs = grid.astype(np.float32)
    return SurfaceBinding(
        part_id="sensing_surface",
        geometry_id=scene.geometry_id,
        field_domain_id=scene.field_domain_id,
        shape_hw=shape,
        uvs=freeze(uvs, dtype=np.float32),
        sample_indices=freeze(np.asarray(indices, dtype=np.int32), dtype=np.int32),
        sample_weights=freeze(np.asarray(weights, dtype=np.float32), dtype=np.float32),
    )


# -------------------------------------------------------- the colour field path
def test_the_colour_path_never_gathers_the_field_onto_the_vertices():
    """The property the whole texture-first change exists for.

    The render frame must carry the acquisition frame's *own* ``[H,W]`` grid — not a
    per-vertex resampling of it. This is asserted by identity, which is the strongest
    form available: a gather cannot produce the same object.
    """
    scene = synthetic.poc_scene()
    frame = _frame(scene, np.arange(9, dtype=np.float32).reshape(3, 3))
    composed = SceneComposer().compose(scene, frame, DisplaySettings())
    assert composed.field.scalar is frame.scalar
    assert composed.field.valid is frame.valid
    assert composed.field.shape_hw == (3, 3)
    assert composed.field.domain_id == scene.field_domain_id
    assert composed.field.descriptor is frame.scalar_descriptor
    # ...and there is no per-vertex array left to read by mistake.
    assert not hasattr(composed, "scalar")
    assert not hasattr(composed, "valid")


def test_a_hole_in_the_field_stays_a_hole_and_is_never_a_vertex_nan():
    """A gap is not a measurement of zero, and it is not a grey *triangle* either.

    Under the old sampler one invalid tap made a whole vertex ``NaN``, and the shader
    painted every triangle touching it grey. Now the hole travels as a hole: exactly
    the pixels the device called invalid are invalid, and a value of 0.0 inside them
    is not read at all.
    """
    scene = synthetic.poc_scene()
    valid = np.ones((3, 3), dtype=bool)
    valid[1, 1] = False
    field = np.zeros((3, 3), dtype=np.float32)
    field[1, 1] = 99.0
    composed = SceneComposer().compose(scene, _frame(scene, field, valid), DisplaySettings())
    assert composed.field.valid.tolist() == valid.tolist()
    assert int(composed.field.valid.sum()) == 8
    # The scalar is carried through untouched, `NaN` or not: nothing rewrites it into
    # a per-vertex absence, so the fragment reads the *device's* answer.
    assert np.array_equal(np.asarray(composed.field.scalar), field)


def test_a_zero_measurement_is_valid_and_stays_zero():
    """Zero is a value. Only `valid` decides whether it means anything."""
    scene = synthetic.poc_scene()
    composed = SceneComposer().compose(
        scene, _frame(scene, np.zeros((3, 3), dtype=np.float32)), DisplaySettings()
    )
    assert composed.field.valid.all()
    assert np.all(np.asarray(composed.field.scalar) == 0.0)
    assert not np.isnan(np.asarray(composed.field.scalar)).any()


def test_the_composer_refuses_a_frame_from_another_domain():
    scene = synthetic.poc_scene()
    frame = _frame(scene, np.zeros((3, 3), dtype=np.float32), domain="some-other-domain")
    with pytest.raises(ThreeDError) as excinfo:
        SceneComposer().compose(scene, frame, DisplaySettings())
    assert excinfo.value.diagnostic.code == "field_domain_mismatch"


def test_the_composer_refuses_a_frame_of_the_wrong_shape():
    scene = synthetic.poc_scene()
    frame = _frame(scene, np.zeros((4, 4), dtype=np.float32), shape=np.zeros((4, 4), dtype=np.float32))
    with pytest.raises(ThreeDError) as excinfo:
        SceneComposer().compose(scene, frame, DisplaySettings())
    assert excinfo.value.diagnostic.code == "field_domain_mismatch"


def test_composing_a_frame_costs_no_per_vertex_work():
    """The measured CPU cost of the colour path, over a realistic grid.

    The old path gathered 1032 vertices x 4 taps every frame *and* the panel's own
    tests measured it as the bulk of a compose. The field path is a shape check and a
    dataclass: the number below is the evidence.
    """
    scene = synthetic.poc_scene()
    composer = SceneComposer()
    frame = _frame(scene, np.zeros((3, 3), dtype=np.float32))
    composer.compose(scene, frame, DisplaySettings())
    samples = []
    for _ in range(20):
        composer.compose(scene, frame, DisplaySettings())
        samples.append(composer.last_stats.sample_ms)
    assert max(samples) < 2.0, f"the colour path still does per-vertex work: {samples}"


# -------------------------------------------- the vertex displacement (geometry)
def test_a_constant_displacement_gathers_to_the_same_constant():
    scene = synthetic.poc_scene()
    displacement = np.full((3, 3, 3), 0.25, dtype=np.float32)
    frame = _frame(
        scene, np.zeros((3, 3), dtype=np.float32), displacement=displacement
    )
    values = sample_vertex_displacement(frame, scene.binding)
    assert values is not None
    assert np.allclose(np.asarray(values), 0.25)


def test_an_invalid_tap_leaves_that_vertex_with_no_displacement():
    """Absent is absent: the vertex does not move at all, and the composer reads NaN."""
    scene = synthetic.poc_scene()
    displacement = np.ones((3, 3, 3), dtype=np.float32) * 0.5
    valid = np.ones((3, 3), dtype=bool)
    valid[1, 1] = False
    frame = _frame(
        scene, np.zeros((3, 3), dtype=np.float32), valid, displacement=displacement
    )
    values = np.asarray(sample_vertex_displacement(frame, scene.binding))
    assert np.isnan(values[4]).all()
    assert np.allclose(values[0], 0.5)


def test_a_frame_without_a_displacement_gathers_nothing():
    scene = synthetic.poc_scene()
    frame = _frame(scene, np.zeros((3, 3), dtype=np.float32))
    assert sample_vertex_displacement(frame, scene.binding) is None


def test_the_displacement_path_never_reads_the_scalar_field():
    """It is geometry, not colour: the gathered value may not depend on the scalar."""
    scene = synthetic.poc_scene()
    displacement = np.ones((3, 3, 3), dtype=np.float32) * 0.125
    first = sample_vertex_displacement(
        _frame(scene, np.full((3, 3), 1.0, dtype=np.float32), displacement=displacement),
        scene.binding,
    )
    second = sample_vertex_displacement(
        _frame(scene, np.full((3, 3), 900.0, dtype=np.float32), displacement=displacement),
        scene.binding,
    )
    assert np.array_equal(np.asarray(first), np.asarray(second))


# ----------------------------------------------------------------------- remap
def test_a_remap_gathers_the_source_pixels_the_map_names():
    field = np.arange(12, dtype=np.float32).reshape(3, 4)
    map_x = np.array([[0, 1], [2, 3]])
    map_y = np.array([[0, 0], [2, 2]])
    values, valid = remap_to_domain(field, map_x=map_x, map_y=map_y, shape_hw=(2, 2))
    assert np.array_equal(np.asarray(values), [[0.0, 1.0], [10.0, 11.0]])
    assert np.asarray(valid).all()


def test_a_remap_that_has_no_source_pixel_leaves_a_hole_not_a_zero():
    field = np.full((3, 3), 5.0, dtype=np.float32)
    map_x = np.array([[0, -1], [1, 1]])
    map_y = np.array([[0, -1], [1, 1]])
    values, valid = remap_to_domain(field, map_x=map_x, map_y=map_y, shape_hw=(2, 2))
    assert np.asarray(valid).tolist() == [[True, False], [True, True]]
    assert np.isnan(np.asarray(values)[0, 1]), "a hole must not read as zero"
    assert np.asarray(values)[0, 0] == 5.0


def test_a_remap_keeps_a_trailing_axis_for_vector_fields():
    flow = np.zeros((3, 3, 2), dtype=np.float32)
    flow[..., 0] = 1.0
    flow[..., 1] = 2.0
    map_x = np.array([[0, 1]])
    map_y = np.array([[1, 2]])
    values, valid = remap_to_domain(flow, map_x=map_x, map_y=map_y, shape_hw=(1, 2))
    assert np.asarray(values).shape == (1, 2, 2)
    assert np.asarray(values)[0, 0].tolist() == [1.0, 2.0]
    assert np.asarray(valid).all()


def test_a_remap_of_the_wrong_shape_is_refused():
    field = np.zeros((3, 3), dtype=np.float32)
    with pytest.raises(ThreeDError) as excinfo:
        remap_to_domain(
            field, map_x=np.zeros((2, 2)), map_y=np.zeros((2, 2)), shape_hw=(3, 3)
        )
    assert excinfo.value.diagnostic.code == "field_domain_mismatch"


def test_a_remap_leaves_the_result_read_only():
    field = np.zeros((2, 2), dtype=np.float32)
    values, valid = remap_to_domain(
        field, map_x=np.zeros((2, 2)), map_y=np.zeros((2, 2)), shape_hw=(2, 2)
    )
    assert not np.asarray(values).flags["WRITEABLE"]
    assert not np.asarray(valid).flags["WRITEABLE"]


# ------------------------------------------------------------------- composition
def test_the_gain_moves_positions_and_never_the_measurement():
    """The row the whole colour scale depends on."""
    scene = synthetic.poc_scene()
    count = 9
    field = np.linspace(0.0, 1.0, count, dtype=np.float32).reshape(3, 3)
    displacement = np.zeros((3, 3, 3), dtype=np.float32)
    displacement[:, :, 2] = 0.001  # a millimetre of illustrative relief
    frame = _frame(scene, field, displacement=displacement, quality="approximate")
    composer = SceneComposer()
    quiet = composer.compose(scene, frame, DisplaySettings(deform_enabled=False))
    small = composer.compose(
        scene, frame, DisplaySettings(deform_enabled=True, display_gain=1.0)
    )
    large = composer.compose(
        scene, frame, DisplaySettings(deform_enabled=True, display_gain=50.0)
    )
    # The measurement is identical in all three.
    assert np.array_equal(np.asarray(quiet.field.scalar), np.asarray(small.field.scalar))
    assert np.array_equal(np.asarray(small.field.scalar), np.asarray(large.field.scalar))
    assert np.array_equal(np.asarray(quiet.field.valid), np.asarray(large.field.valid))
    # The positions are not.
    rest = np.asarray(small.positions_sensor_m)
    moved = np.asarray(large.positions_sensor_m)
    rest_scene = np.asarray(scene.sensing_part().positions_sensor_m)
    assert np.allclose((rest - rest_scene)[:, :2], 0.0, atol=1e-7), "the gain moved sideways"
    assert np.allclose((rest - rest_scene)[:, 2], 0.001, atol=1e-7)
    assert np.allclose((moved - rest_scene)[:, 2], 0.05, atol=1e-7)


def test_an_illustrative_deformation_says_so():
    scene = synthetic.poc_scene()
    displacement = np.zeros((3, 3, 3), dtype=np.float32)
    displacement[:, :, 2] = 0.001
    frame = _frame(scene, np.zeros((3, 3), dtype=np.float32), displacement=displacement,
                   quality="approximate")
    composed = SceneComposer().compose(
        scene, frame, DisplaySettings(deform_enabled=True, display_gain=3.0)
    )
    codes = [item.code for item in composed.diagnostics]
    assert DISPLAY_ILLUSTRATIVE in codes, codes
    note = next(item for item in composed.diagnostics if item.code == DISPLAY_ILLUSTRATIVE)
    assert note.severity == "info"
    assert "illustration" in note.message and "×3" in note.message


def test_a_calibrated_displacement_does_not_need_the_caveat():
    scene = synthetic.poc_scene()
    displacement = np.zeros((3, 3, 3), dtype=np.float32)
    displacement[:, :, 2] = 0.0005
    frame = _frame(scene, np.zeros((3, 3), dtype=np.float32), displacement=displacement,
                   quality="calibrated")
    composed = SceneComposer().compose(
        scene, frame, DisplaySettings(deform_enabled=True, display_gain=2.0)
    )
    assert DISPLAY_ILLUSTRATIVE not in [item.code for item in composed.diagnostics]


def test_deformation_without_a_displacement_colours_in_place_and_says_so():
    scene = synthetic.poc_scene()
    frame = _frame(scene, np.full((3, 3), 2.0, dtype=np.float32))
    composed = SceneComposer().compose(
        scene, frame, DisplaySettings(deform_enabled=True, display_gain=10.0)
    )
    assert np.allclose(
        np.asarray(composed.positions_sensor_m),
        np.asarray(scene.sensing_part().positions_sensor_m),
    )
    assert any(item.code == "field_unavailable" for item in composed.diagnostics)


def test_normals_are_recomputed_from_the_deformed_positions():
    """A bump must change the shading on its slopes."""
    scene = synthetic.poc_scene()
    displacement = np.zeros((3, 3, 3), dtype=np.float32)
    # Only the centre vertex is lifted: the slopes around it have to tilt. (The
    # apex itself stays vertical, which is the correct normal for a symmetric
    # peak — asserting otherwise would be asserting a bug.)
    displacement[1, 1, 2] = 0.004
    frame = _frame(scene, np.zeros((3, 3), dtype=np.float32), displacement=displacement,
                   quality="approximate")
    composed = SceneComposer().compose(
        scene, frame, DisplaySettings(deform_enabled=True, display_gain=1.0)
    )
    normals = np.asarray(composed.normals_sensor)
    assert np.allclose(np.linalg.norm(normals, axis=1), 1.0, atol=1e-5)
    rest_normals = np.asarray(scene.sensing_part().normals_sensor)
    assert not np.allclose(normals, rest_normals), "the normals ignored the deformation"
    # Only the vertices whose triangles touch the lifted one change; the two
    # diagonal corners (2 and 6) share no triangle with it and correctly stay flat.
    neighbours = [0, 1, 3, 5, 7, 8]
    untouched = [2, 6]
    tilted = np.array([float(normals[index, 2]) for index in neighbours])
    assert (tilted < 0.999).all(), f"the slopes did not tilt: {tilted}"
    flat = np.array([float(normals[index, 2]) for index in untouched])
    assert np.allclose(flat, 1.0, atol=1e-6), flat
    assert normals[4, 2] == pytest.approx(1.0, abs=1e-3), (
        "the apex of a symmetric bump is vertical"
    )
    # ...and the reverse bump tilts the other way.
    upside_down = -displacement
    flipped = SceneComposer().compose(
        scene,
        _frame(scene, np.zeros((3, 3), dtype=np.float32), displacement=upside_down,
               quality="approximate"),
        DisplaySettings(deform_enabled=True, display_gain=1.0),
    )
    flipped_normals = np.asarray(flipped.normals_sensor)
    assert np.allclose(
        flipped_normals[neighbours, 0], -normals[neighbours, 0], atol=1e-6
    ), "the reverse bump must tilt the slopes the other way"


def test_the_composer_refuses_a_frame_from_another_geometry():
    scene = synthetic.poc_scene()
    frame = _frame(scene, np.zeros((3, 3), dtype=np.float32))
    object.__setattr__(frame, "geometry_id", "someone-elses")
    with pytest.raises(ThreeDError) as excinfo:
        SceneComposer().compose(scene, frame, DisplaySettings())
    assert excinfo.value.diagnostic.code == "geometry_mismatch"


def test_a_reference_frame_measures_nothing():
    scene = synthetic.poc_scene()
    rest = SceneComposer().rest_frame(scene, FrameIdentity("live", 1, 0, 0.0, 1))
    assert not np.asarray(rest.field.valid).any()
    assert np.isnan(np.asarray(rest.field.scalar)).all()
    assert rest.field.descriptor.quantity == "unavailable"
    assert rest.field.descriptor.quality == "display_only"
    assert np.allclose(
        np.asarray(rest.positions_sensor_m),
        np.asarray(scene.sensing_part().positions_sensor_m),
    )


def test_the_composed_frame_carries_read_only_arrays():
    scene = synthetic.poc_scene()
    composed = SceneComposer().compose(
        scene, _frame(scene, np.zeros((3, 3), dtype=np.float32)), DisplaySettings()
    )
    for name in ("positions_sensor_m", "normals_sensor"):
        value = getattr(composed, name)
        assert not value.flags["WRITEABLE"], f"{name} is writable"
        assert value.flags["C_CONTIGUOUS"]
    for name in ("scalar", "valid"):
        value = getattr(composed.field, name)
        assert not value.flags["WRITEABLE"], f"field.{name} is writable"
        assert value.flags["C_CONTIGUOUS"]


def test_the_composer_does_not_rebuild_the_binding_or_the_indices():
    """Composition must be cheap enough to run per frame."""
    scene = synthetic.poc_scene()
    composer = SceneComposer()
    frame = _frame(scene, np.zeros((3, 3), dtype=np.float32))
    for _ in range(3):
        composer.compose(scene, frame, DisplaySettings())
    assert composer.last_stats.total_ms < 50.0, composer.last_stats


# --------------------------------------------- stage 9: reuse and real numbers
def test_colouring_in_place_reuses_the_assets_own_geometry():
    """No displacement to draw means the picture *is* the asset: nothing to recompute.

    Normals were recomputed from the same untouched positions on every frame, and the
    measured cost was the bulk of a 20k-vertex compose. The reused arrays are the
    asset's own (frozen) ones, so nothing is copied either.
    """
    scene = synthetic.poc_scene()
    part = scene.sensing_part()
    composer = SceneComposer()
    frame = _frame(scene, np.zeros((3, 3), dtype=np.float32))
    composed = composer.compose(scene, frame, DisplaySettings())
    assert composed.positions_sensor_m is part.positions_sensor_m
    assert composed.normals_sensor is part.normals_sensor
    # A deform request with nothing to deform is the same case, and says so.
    composed = composer.compose(
        scene, frame, DisplaySettings(deform_enabled=True, display_gain=5.0)
    )
    assert composed.positions_sensor_m is part.positions_sensor_m
    assert any(item.code == "field_unavailable" for item in composed.diagnostics)


def test_a_calibrated_displacement_still_recomputes_positions_and_normals():
    """The cache is for the colour-in-place case only, never for a deformed surface."""
    scene = synthetic.poc_scene()
    part = scene.sensing_part()
    displacement = np.zeros((3, 3, 3), dtype=np.float32)
    displacement[1, 1, 2] = 0.5
    frame = _frame(
        scene,
        np.ones((3, 3), dtype=np.float32),
        displacement=displacement,
        quality="calibrated",
    )
    composed = SceneComposer().compose(
        scene, frame, DisplaySettings(deform_enabled=True, display_gain=2.0)
    )
    assert composed.deformed
    assert composed.positions_sensor_m is not part.positions_sensor_m
    assert composed.normals_sensor is not part.normals_sensor
    assert not np.allclose(
        np.asarray(composed.positions_sensor_m), np.asarray(part.positions_sensor_m)
    )


def test_the_composed_frame_carries_its_own_cost():
    """The panel's "compose ms" reads this, not the composer's shared counter."""
    scene = synthetic.poc_scene()
    composer = SceneComposer()
    frame = _frame(scene, np.zeros((3, 3), dtype=np.float32))
    composed = composer.compose(scene, frame, DisplaySettings())
    assert composed.compose_ms > 0.0
    assert composed.compose_ms == pytest.approx(composer.last_stats.total_ms)
    # A second compose rewrites ``last_stats``; the first frame's number must not move.
    first = composed.compose_ms
    composer.compose(scene, frame, DisplaySettings())
    assert composed.compose_ms == first