"""HHD downsample compatibility diagnostics and non-divisible grid integration.

The pipeline contract allows a model calibrated at one ``hhd_downsample`` to run at
another, but the difference is an accuracy risk that must be visible instead of a
silent pass. These tests pin the report, the unchanged refusal of every other
identity mismatch, and the analytic block integration on grids that do not divide
the output image evenly.
"""

from __future__ import annotations

import numpy as np
import pytest

from orisys.calibration import CalibrationError, parse
from orisys.calibration_context import build_context
from orisys.calibration_models import ModelCalibration
from orisys.field_calibration import FieldFrame, _block_starts, build_field_gain_map, reduce_field
from orisys.force_readout import OutputSettings
from orisys.force_session import ForceSession, FrameStamp
from orisys.sensor import FrameMeasurement

FRAME = np.zeros((2, 2, 3), dtype=np.uint8)


def _config(downsample=1, *, dis_mode=4):
    return {
        "camera": {"resolution_raw": [1280, 720], "resolution_stitch": [700, 750]},
        "tracker": {"dis_mode": dis_mode, "hhd_downsample": downsample},
    }


def _document(config) -> dict:
    return {
        "schema_version": 3,
        "unit": "N",
        "identity": {"sensor_type": "box", "sn": "A1"},
        "pipeline": dict(build_context(config)),
        "model": {
            "type": "global_gain/1",
            "input_profile": "force_delta/1",
            "output_frame": "sdk_force/1",
            "outputs": ["fn"],
            "params": {
                "gains_N_per_raw": {"fn": 0.001},
                "valid_force_range_N": [0, 3],
            },
        },
    }


def _model(config) -> ModelCalibration:
    loaded = parse(_document(config))
    assert isinstance(loaded, ModelCalibration)
    return loaded


# ------------------------------------------------------------------ the report
def test_downsample_only_difference_runs_and_is_reported():
    model = _model(_config(2))
    runtime = build_context(_config(3))

    model.validate_for_context(runtime)  # allowed: no exception

    report = model.pipeline_compatibility(runtime)
    assert report["ok"] is True
    downsample = report["downsample"]
    assert downsample["declared"] == 2
    assert downsample["actual"] == 3
    assert downsample["match"] is False
    assert downsample["allowed"] is True
    assert downsample["accuracy_risk"]
    assert model.downsample_report(runtime).known is True


def test_a_matching_downsample_has_no_accuracy_risk():
    model = _model(_config(3))
    report = model.pipeline_compatibility(build_context(_config(3)))
    assert report["downsample"]["match"] is True
    assert report["downsample"]["accuracy_risk"] == ""


def test_an_unknown_runtime_downsample_is_unknown_not_a_match():
    model = _model(_config(2))
    report = model.downsample_report({})
    assert report.actual is None
    assert report.match is False
    assert report.allowed is True
    assert "unknown" in report.accuracy_risk


def test_every_other_identity_mismatch_is_still_refused():
    model = _model(_config(2))
    runtime = build_context(_config(2, dis_mode=5))
    report = model.pipeline_compatibility(runtime)
    assert report["ok"] is False
    assert "calibration pipeline mismatch: dis_mode" in report["mismatches"].values()
    with pytest.raises(CalibrationError, match="calibration pipeline mismatch: dis_mode"):
        model.validate_for_context(runtime)


@pytest.mark.parametrize("raw, expected", [(None, 1), ("bad", 1), (0, 1), (-3, 1), ("3", 3), (3.7, 3)])
def test_context_reports_its_actual_downsample(raw, expected):
    config = _config(1)
    config["tracker"]["hhd_downsample"] = raw
    assert build_context(config).actual_downsample == expected


# ------------------------------------------------- non-divisible grid integration
def _identity_coordinates(rows: int, cols: int) -> np.ndarray:
    uv = np.empty((rows, cols, 2), dtype=np.float64)
    uv[..., 0] = np.arange(cols)[None, :]
    uv[..., 1] = np.arange(rows)[:, None]
    return uv


@pytest.mark.parametrize("grid_shape", [(1, 1), (2, 3), (3, 4), (4, 4)])
def test_a_constant_field_integrates_exactly_on_a_non_divisible_grid(grid_shape):
    rows, cols = 37, 53
    uv = _identity_coordinates(rows, cols)
    mask = np.ones((rows, cols), dtype=bool)
    points = np.array(
        [
            [0.0, 0.0],
            [cols - 1.0, 0.0],
            [cols - 1.0, rows - 1.0],
            [0.0, rows - 1.0],
        ]
    )
    gains = np.ones(4, dtype=np.float64)
    triangles = np.array([[0, 1, 2], [0, 2, 3]], dtype=int)

    # The block layout is an exact partition of the output image, for every grid
    # size that does not divide it: no lost pixels and no double counting.
    for divisions, size in ((grid_shape[0], rows), (grid_shape[1], cols)):
        starts, ends = _block_starts(divisions, size)
        assert starts[0] == 0
        assert ends[-1] == size
        assert np.array_equal(starts[1:], ends[:-1])
        assert np.all(ends > starts)

    gain_map = build_field_gain_map(points, gains, triangles, uv, mask, grid_shape, mask)
    row_starts, row_ends = _block_starts(grid_shape[0], rows)
    col_starts, col_ends = _block_starts(grid_shape[1], cols)
    block_area = np.outer(row_ends - row_starts, col_ends - col_starts)
    # A constant unit gain and a fully mapped mask: the area-weighted coverage is
    # the whole image, exactly (analytic area, not a sampled sum).
    assert float(np.sum(gain_map.coverage * block_area)) == pytest.approx(rows * cols)
    assert np.allclose(gain_map.gain, 1.0)
    assert np.allclose(gain_map.inside, 1.0)

    field = FieldFrame(
        d=np.ones((grid_shape[0], grid_shape[1], 2), dtype=np.float32),
        valid_mask=np.ones(grid_shape, dtype=bool),
        area_scale=1.0,
        grid_shape=grid_shape,
        reference_shape=None,
    )
    reduction = reduce_field(field, gain_map)
    # q = |d_x| + |d_y| = 2 per active HHD cell; the gain map's block sums are exact
    # whatever the block boundaries are, so the integral is exactly 2 * grid cells.
    cells = grid_shape[0] * grid_shape[1]
    assert reduction.weighted_N == pytest.approx(2.0 * cells)
    assert reduction.extrapolated_fraction == 0.0
    assert reduction.coverage_fraction == pytest.approx(1.0)


# ------------------------------------------------------- session output diagnostics
class _Source:
    """The measurement-source surface ForceSession needs; no camera involved."""

    def __init__(self) -> None:
        self.session_enabled = False
        self.reset_count = 0

    def enable_measurement_session(self) -> None:
        self.session_enabled = True

    def disable_measurement_session(self) -> None:
        self.session_enabled = False

    def reset_tracker(self) -> None:
        self.reset_count += 1

    def read_measurement_frame(
        self, source_frame, *, frame_id, timestamp_s=None, want_field=False, want_contact=False
    ) -> FrameMeasurement:
        del source_frame, timestamp_s, want_field, want_contact
        return FrameMeasurement(
            raw=(10.0, 20.0, 30.0),
            raw_valid=(True, True, True),
        )


def test_the_session_output_states_the_downsample_risk():
    model = _model(_config(2))
    session = ForceSession(
        source=_Source(),
        sensor_family="box",
        context=build_context(_config(3)),
        calibration=model,
        settings=OutputSettings(sample_rate_hz=30.0),
    )
    output = session.process_frame(
        FRAME,
        FrameStamp(
            session_id="s1",
            frame_id=0,
            source_frame_id=0,
            signal_time_s=0.0,
            time_basis="recorded_signal",
        ),
    )
    diagnostics = dict(output.measurement.diagnostics)
    assert diagnostics["hhd_downsample_calibrated"] == 2.0
    assert diagnostics["hhd_downsample_actual"] == 3.0
    assert diagnostics["hhd_downsample_mismatch"] == 1.0
    assert diagnostics["downsample_accuracy_risk"]
    assert session.calibration_compatibility["downsample"]["actual"] == 3