"""Field gain model: gain mapping, weighted reduction and unloaded baseline."""

from __future__ import annotations

import hashlib
from copy import deepcopy

import numpy as np
import pytest

from orisys.calibration import CalibrationError, parse, validate_for_device
from orisys.calibration_context import build_context
from orisys.calibration_models import parse_model
from orisys.field_calibration import (
    BASELINE_POLICY,
    FIELD_IDENTITY,
    MAPPING_POLICY,
    build_field_gain_map,
    coerce_field,
    reduce_field,
)

FULL_SQUARE = [[0.0, 0.0], [3.0, 0.0], [0.0, 3.0], [3.0, 3.0]]
SMALL_SQUARE = [[0.0, 0.0], [2.0, 0.0], [2.0, 2.0], [0.0, 2.0]]


def _artifact(tmp_path, *, shape=(4, 4), uv=None, valid=None):
    height, width = shape
    if uv is None:
        rows, cols = np.mgrid[0:height, 0:width]
        uv = np.stack((cols, rows), axis=-1).astype(np.float64)
    if valid is None:
        valid = np.ones((height, width), dtype=bool)
    path = tmp_path / "map.npz"
    np.savez(path, uv_by_output_px=uv, valid_mask=valid, origin_xy=np.array([0, 0]))
    return path, hashlib.sha256(path.read_bytes()).hexdigest()


def _pipeline(*, digest):
    pipeline = {
        "stitch_resolution": [4, 4],
        "dis_mode": 1,
        "hhd_downsample": 1,
        "filters": {"measurement": {"flow_ema_alpha": 1.0}},
        "source_scale": 0.5,
        "geometry_sha256": "b" * 64,
        "coordinates": {"frame": "autocal_flat", "unit": "px",
                        "origin_xy": [0, 0], "map_sha256": digest},
        "feature_signature_sha256": "a" * 64,
        "field": deepcopy(FIELD_IDENTITY),
    }
    return pipeline


def _document(tmp_path, *, gains=None, points=None, triangles=None, uv=None,
              coord_valid=None, artifact=True):
    height, width = (4, 4) if uv is None else uv.shape[:2]
    path, digest = _artifact(tmp_path, shape=(height, width), uv=uv, valid=coord_valid)
    pipeline = _pipeline(digest=digest)
    params = {
        "points_uv": points if points is not None else FULL_SQUARE,
        "gains_N_per_raw": gains if gains is not None else [1.0, 2.0, 3.0, 4.0],
        "triangles": triangles if triangles is not None else [[0, 1, 2], [1, 3, 2]],
        "valid_force_range_N": [0.0, 100.0],
        "mapping": deepcopy(MAPPING_POLICY),
        "baseline": deepcopy(BASELINE_POLICY),
    }
    document = {
        "schema_version": 3,
        "unit": "N",
        "identity": {"sensor_type": "finger", "sn": "DD02-GL"},
        "pipeline": pipeline,
        "model": {
            "type": "field_gain/1",
            "input_profile": "hhd_normal_l1_field/1",
            "output_frame": "sdk_force/1",
            "outputs": ["fn"],
            "params": params,
        },
    }
    if artifact:
        document["model"]["artifacts"] = [
            {"role": "coordinate_map", "path": path.name, "sha256": digest}
        ]
    return document


def _field(dx=0.5, dy=0.5, *, shape=(2, 2), area_scale=4.0, mask=None, dtype=np.float32):
    d = np.zeros((*shape, 2), dtype=dtype)
    d[:, :, 0] = dx
    d[:, :, 1] = dy
    if mask is None:
        mask = np.ones(shape, dtype=bool)
    mask = np.asarray(mask, dtype=bool)
    d[~mask] = 0.0
    return {"d": d, "valid_mask": mask, "area_scale": area_scale, "grid_shape": tuple(shape)}


def _features(field, *, mask_shape=(4, 4), output_mask=None):
    if output_mask is None:
        output_mask = np.ones(mask_shape, dtype=bool)
    return {"field": field, "output_valid_mask": output_mask}


def _model_with_gain(tmp_path, gain=1.0):
    document = _document(tmp_path, gains=[gain] * 4)
    return parse_model(document, tmp_path / "calibration.json")


def _feed_ready(model, frames=30, field=None):
    model.begin_reference("unloaded")
    field = field if field is not None else _field()
    for _ in range(frames):
        model.predict(1, 0.0, _features(field))
    assert model.field_baseline_N is not None


def test_parse_accepts_field_model_and_declares_profile(tmp_path):
    model = parse_model(_document(tmp_path), tmp_path / "calibration.json")
    assert model.model_type == "field_gain/1"
    assert model.input_profile == "hhd_normal_l1_field/1"
    assert model.channels == {"fn": "N"}
    assert model.to_dict()["model"]["params"]["mapping"] == MAPPING_POLICY
    assert len(model.fingerprint) == 64
    gain_map = model.field_gain_map((2, 2), np.ones((4, 4), dtype=bool))
    assert gain_map.gain.shape == (2, 2)
    assert np.allclose(gain_map.coverage, 1.0)


def test_constant_gain_matches_hand_computed_weighted_sum(tmp_path):
    model = _model_with_gain(tmp_path, gain=1.0)
    _feed_ready(model)
    result = model.predict(31, 0.0, _features(_field()))
    # q=1 on each of 4 cells, coverage 1, gain 1, A=4: P=16, baseline=16, FN=0.
    assert result.values_N["fn"] == pytest.approx(0.0)
    assert result.status["fn"] == "ok"
    assert result.diagnostics["baseline_N"] == pytest.approx(16.0)
    result = model.predict(32, 0.0, _features(_field(1.5, 1.5)))
    # q=3 per cell -> P=48, FN=48-16=32.
    assert result.values_N["fn"] == pytest.approx(32.0)


def test_two_regions_contribute_and_vectors_combine_before_l1(tmp_path):
    model = _model_with_gain(tmp_path, gain=1.0)
    _feed_ready(model)
    d = np.zeros((2, 2, 2), dtype=np.float32)
    d[0, 0] = (1.0, 0.0)
    d[1, 1] = (0.0, 0.5)
    mask = np.ones((2, 2), dtype=bool)
    field = {"d": d, "valid_mask": mask, "area_scale": 4.0, "grid_shape": (2, 2)}
    result = model.predict(31, 0.0, _features(field))
    # q sums 1.5, two cells -> P=6; baseline 16 -> -10.
    assert result.values_N["fn"] == pytest.approx(-10.0)

    combined = np.zeros((2, 2, 2), dtype=np.float32)
    combined[0, 0] = (0.75, 0.0)
    # A negative and a positive vector at the same cell combine to 0.5, not 1.0.
    first = combined.copy()
    second = np.zeros_like(combined)
    second[0, 0] = (-0.25, 0.0)
    first_map = {"d": first, "valid_mask": mask, "area_scale": 4.0, "grid_shape": (2, 2)}
    second_map = {"d": second, "valid_mask": mask, "area_scale": 4.0, "grid_shape": (2, 2)}
    combined_map = {"d": first + second, "valid_mask": mask, "area_scale": 4.0, "grid_shape": (2, 2)}
    gain = model.field_gain_map((2, 2), np.ones((4, 4), dtype=bool))
    reduction = reduce_field(coerce_field(combined_map), gain)
    assert reduction.weighted_N == pytest.approx(4.0 * 0.5)
    assert reduction.weighted_N != pytest.approx(
        reduce_field(coerce_field(first_map), gain).weighted_N
        + reduce_field(coerce_field(second_map), gain).weighted_N
    )


def test_boundary_projection_and_mask_coverage(tmp_path):
    model = parse_model(_document(tmp_path, points=SMALL_SQUARE,
                                  triangles=[[0, 1, 2], [0, 2, 3]]),
                        tmp_path / "calibration.json")
    gains = model.field_gain_map((4, 4), np.ones((4, 4), dtype=bool))
    # (1,1) is inside the small hull; (3,3) projects to the nearest edge.
    assert gains.gain[1, 1] == pytest.approx(2.0)
    assert gains.coverage[1, 1] == pytest.approx(1.0)
    assert gains.inside[1, 1] == pytest.approx(1.0)
    assert gains.gain[3, 3] == pytest.approx(3.0)
    assert gains.inside[3, 3] == pytest.approx(0.0)
    assert gains.coverage[3, 3] == pytest.approx(1.0)

    # A physical mask that keeps only the left half removes the right cells entirely.
    constant = parse_model(_document(tmp_path, points=SMALL_SQUARE, gains=[1.0] * 4,
                                     triangles=[[0, 1, 2], [0, 2, 3]]),
                           tmp_path / "calibration.json")
    left_half = np.zeros((4, 4), dtype=bool)
    left_half[:, :2] = True
    gains = constant.field_gain_map((2, 2), left_half)
    assert np.allclose(gains.coverage, [[1.0, 0.0], [1.0, 0.0]])
    field = _field(mask=np.array([[True, False], [True, False]]))
    reduction = reduce_field(coerce_field(field), gains)
    # Two active cells with q=1, A=4: P=8.  Inactive cells never contribute.
    assert reduction.weighted_N == pytest.approx(8.0)


def test_partial_coverage_constant_gain_still_matches_raw_sum(tmp_path):
    """R1: coverage is diagnostic, not a second mask weight.

    HHD raw counts a partially covered mask-edge cell in full; folding coverage
    into ``P`` would discount it again and break ``P == gain * raw``.
    """
    model = _model_with_gain(tmp_path, gain=0.5)
    mask = np.ones((4, 4), dtype=bool)
    mask[0, 0] = False
    gain_map = model.field_gain_map((2, 2), mask)
    assert gain_map.coverage[0, 0] == pytest.approx(0.75)
    field = _field(0.5, 0.5, shape=(2, 2), area_scale=4.0)
    raw = 4.0 * float(np.sum(np.abs(field["d"])))
    reduction = reduce_field(coerce_field(field), gain_map)
    assert reduction.weighted_N == pytest.approx(0.5 * raw)
    assert reduction.coverage_fraction == pytest.approx((0.75 + 1.0 + 1.0 + 1.0) / 4.0)


def test_active_hhd_cell_without_coverage_is_a_mapping_error(tmp_path):
    """R1: an active cell with no mapped pixel must not be silently counted as zero."""
    model = _model_with_gain(tmp_path, gain=1.0)
    left_half = np.zeros((4, 4), dtype=bool)
    left_half[:, :2] = True
    gain_map = model.field_gain_map((2, 2), left_half)
    assert np.allclose(gain_map.coverage, [[1.0, 0.0], [1.0, 0.0]])
    with pytest.raises(ValueError, match="mapping error"):
        reduce_field(coerce_field(_field(mask=np.ones((2, 2), dtype=bool))), gain_map)
    result = model.predict(1, 0.0, _features(
        _field(mask=np.ones((2, 2), dtype=bool)), output_mask=left_half))
    assert result.status == {"fn": "coordinate_unmapped"}
    assert result.values_N == {"fn": None}


def test_mapping_and_input_errors_are_explicit(tmp_path):
    uv = np.stack(np.mgrid[0:4, 0:4][::-1], axis=-1).astype(np.float64)
    valid = np.ones((4, 4), dtype=bool)
    valid[0, 0] = False
    model = parse_model(_document(tmp_path, uv=uv, coord_valid=valid),
                        tmp_path / "calibration.json")
    with pytest.raises(ValueError, match="mapping error"):
        model.field_gain_map((2, 2), np.ones((4, 4), dtype=bool))
    result = model.predict(1, 0.0, _features(_field()))
    assert result.status == {"fn": "coordinate_unmapped"}
    assert result.values_N == {"fn": None}

    broken = _field()
    broken["d"][0, 0, 0] = np.nan
    result = model.predict(2, 0.0, _features(broken))
    assert result.status == {"fn": "invalid_input"}
    assert result.values_N == {"fn": None}

    wrong_dtype = _field(dtype=np.float64)
    with pytest.raises(ValueError, match="dtype"):
        coerce_field(wrong_dtype)
    wrong_shape = _field()
    wrong_shape["grid_shape"] = (3, 2)
    with pytest.raises(ValueError, match="grid_shape"):
        coerce_field(wrong_shape)


def test_empty_field_reports_zero_and_negative_baseline_difference(tmp_path):
    model = _model_with_gain(tmp_path, gain=1.0)
    _feed_ready(model)
    empty = _field(0.0, 0.0)
    result = model.predict(31, 0.0, _features(empty))
    assert result.status["fn"] == "out_of_range"
    assert result.values_N["fn"] == pytest.approx(-16.0)
    assert result.diagnostics["extrapolated_fraction"] == pytest.approx(0.0)


def test_weighted_baseline_needs_declaration_and_thirty_consecutive_frames(tmp_path):
    model = _model_with_gain(tmp_path, gain=1.0)
    result = model.predict(1, 0.0, _features(_field()))
    assert result.status == {"fn": "zero_unavailable"}
    assert result.diagnostics["pending_reference_frames"] == 0
    model.begin_reference("unloaded")
    for _ in range(29):
        result = model.predict(1, 0.0, _features(_field()))
        assert result.status == {"fn": "zero_unavailable"}
    result = model.predict(1, 0.0, _features(_field()))
    assert result.status == {"fn": "zero_unavailable"}
    assert result.diagnostics["baseline_N"] == pytest.approx(16.0)
    assert result.diagnostics["reference_id"] == "unloaded"
    ready = model.predict(2, 0.0, _features(_field()))
    assert ready.status["fn"] == "ok"

    model.begin_reference("second")
    invalid = _field()
    invalid["d"][0, 0, 0] = np.nan
    for _ in range(29):
        model.predict(3, 0.0, _features(_field()))
    assert model.predict(3, 0.0, _features(invalid)).status == {"fn": "invalid_input"}
    result = model.predict(4, 0.0, _features(_field()))
    assert result.diagnostics["pending_reference_frames"] == 1
    assert result.status == {"fn": "zero_unavailable"}


def test_reset_and_model_identity_void_the_weighted_baseline(tmp_path):
    document = _document(tmp_path, gains=[1.0] * 4)
    model = parse_model(document, tmp_path / "calibration.json")
    _feed_ready(model)
    assert model.field_baseline_N is not None
    model.reset(None)
    assert model.field_baseline_N is None
    result = model.predict(31, 0.0, _features(_field()))
    assert result.status == {"fn": "zero_unavailable"}
    fresh = parse_model(document, tmp_path / "calibration.json")
    assert fresh.field_baseline_N is None
    assert fresh.predict(31, 0.0, _features(_field())).status == {"fn": "zero_unavailable"}


def test_field_document_rejections(tmp_path):
    mutations = [
        lambda data: data["model"].update(input_profile="fn_delta_autocal_uv/1"),
        lambda data: data["model"].update(outputs=["fn", "fx"]),
        lambda data: data["model"].pop("artifacts"),
        lambda data: data["model"]["params"].update(mapping={"inside": "nearest"}),
        lambda data: data["model"]["params"].update(baseline={"frames": 29, "statistic": "median"}),
        lambda data: data["pipeline"].pop("field"),
        lambda data: data["pipeline"].update(field={"grid": "other"}),
        lambda data: data["model"]["params"].update(extra=True),
        lambda data: data["model"]["params"].update(gains_N_per_raw=[1.0, 2.0, 3.0, 0.0]),
        lambda data: data["pipeline"].pop("geometry_sha256"),
    ]
    for mutate in mutations:
        data = _document(tmp_path)
        mutate(data)
        with pytest.raises(ValueError):
            parse_model(data, tmp_path / "calibration.json")


def test_field_model_is_loadable_through_calibration_entry_points(tmp_path):
    document = _document(tmp_path)
    model = parse(document, source_path=tmp_path / "calibration.json")
    assert model.schema_version == 3
    assert parse_model(document, tmp_path / "calibration.json").model_type == "field_gain/1"
    with pytest.raises(CalibrationError, match="identity mismatch"):
        validate_for_device(model, "finger", "OTHER-SN")


def test_context_includes_field_identity_and_rejects_scalar_profile(tmp_path):
    path, digest = _artifact(tmp_path)
    coordinates = {"frame": "autocal_flat", "unit": "px",
                   "origin_xy": [0, 0], "map_sha256": digest}
    config = {
        "type": "compound_eye",
        "camera": {"resolution_raw": [4, 4], "resolution_stitch": [4, 4]},
        "tracker": {"dis_mode": 1, "hhd_downsample": 1},
        "filters": {"measurement": {"flow_ema_alpha": 1.0}},
        "contact": {},
    }
    context = build_context(config, config["filters"], geometry_sha256="b" * 64,
                            source_scale=0.5, coordinates=coordinates,
                            processing_backend="uv_unroll",
                            preprocessing={"flip_source_vertical": True, "mirror": False},
                            field=FIELD_IDENTITY)
    document = _document(tmp_path)
    document["pipeline"] = dict(context)
    model = parse_model(document, tmp_path / "calibration.json")
    model.validate_for_context(context)
    scalar_context = build_context(config, config["filters"], geometry_sha256="b" * 64,
                                   source_scale=0.5, coordinates=coordinates,
                                   processing_backend="uv_unroll")
    with pytest.raises(ValueError, match="field"):
        model.validate_for_context(scalar_context)


def test_gain_map_is_deterministic_across_builds(tmp_path):
    model = _model_with_gain(tmp_path, gain=2.0)
    first = model.field_gain_map((2, 2), np.ones((4, 4), dtype=bool))
    model._field_maps.clear()
    second = model.field_gain_map((2, 2), np.ones((4, 4), dtype=bool))
    assert np.array_equal(first.gain, second.gain)
    assert np.array_equal(first.coverage, second.coverage)
    assert first.physical_mask_sha256 == second.physical_mask_sha256
    with pytest.raises(ValueError, match="mask"):
        build_field_gain_map(model._points, model._gains, model._triangles,
                             model.coordinate_map, model.coordinate_valid_mask,
                             (2, 2), np.ones((3, 3), dtype=bool))


def test_field_model_ignores_contact_state(tmp_path):
    model = _model_with_gain(tmp_path, gain=1.0)
    model.begin_reference("unloaded")
    for _ in range(30):
        model.predict(1, 0.0, _features(_field()))
    features = _features(_field())
    features.update(contact_found=False, zero_valid=False)
    result = model.predict(31, 0.0, features)
    assert result.status["fn"] == "ok"
    assert result.values_N["fn"] == pytest.approx(0.0)


def _hhd(grid_shape, reference_shape, downsample):
    from orisys.hhd import nHHD

    return nHHD(cuda=False, grid=grid_shape, spacings=(18 * downsample, 18 * downsample),
                poisson_backend="cpu", force_reference_grid=reference_shape)


def _identity_uv(shape):
    rows, cols = np.mgrid[0:shape[0], 0:shape[1]]
    return np.stack((cols, rows), axis=-1).astype(np.float64)


def test_hhd_field_payload_uses_actual_reference_area_scale():
    reference = (9, 8)
    for downsample in (1, 2, 3, 4):
        grid = (reference[0] // downsample, reference[1] // downsample)
        solver = _hhd(grid, reference, downsample)
        solver.get_force(np.zeros((*grid, 2), dtype=np.float32))
        payload = solver.get_field()
        assert payload["area_scale"] == pytest.approx(
            reference[0] * reference[1] / (grid[0] * grid[1]))
        assert payload["grid_shape"] == grid
        assert payload["reference_shape"] == reference
        assert payload["d"].dtype == np.float32
        assert payload["d"].shape == (*grid, 2)
        payload["d"][0, 0, 0] = 123.0
        assert solver.d[0, 0, 0] != 123.0


def test_constant_gain_matches_raw_hhd_sum(tmp_path):
    shape = (6, 5)
    rng = np.random.default_rng(7)
    solver = _hhd(shape, shape, 1)
    _, forces = solver.get_force(rng.normal(size=(*shape, 2)).astype(np.float32))
    raw_sn = float(forces[0])
    payload = solver.get_field()
    uv = _identity_uv(shape)
    document = _document(tmp_path, gains=[0.5] * 4, uv=uv,
                         points=[[0.0, 0.0], [shape[1] - 1.0, 0.0],
                                 [0.0, shape[0] - 1.0], [shape[1] - 1.0, shape[0] - 1.0]],
                         triangles=[[0, 1, 2], [1, 3, 2]])
    model = parse_model(document, tmp_path / "calibration.json")
    gain_map = model.field_gain_map(shape, np.ones(shape, dtype=bool))
    assert np.allclose(gain_map.coverage, 1.0)
    reduction = reduce_field(coerce_field(payload), gain_map)
    # Constant gain: the field-weighted sum equals the raw HHD sum times the gain.
    assert reduction.weighted_N == pytest.approx(0.5 * raw_sn, rel=1e-5, abs=1e-6)


def test_constant_field_integral_across_downsample_and_nondivisible_grids(tmp_path):
    reference = (9, 8)
    expected = 0.5 * 2.0 * reference[0] * reference[1]
    for downsample in (1, 2, 3, 4):
        grid = (reference[0] // downsample, reference[1] // downsample)
        solver = _hhd(grid, reference, downsample)
        solver.get_force(np.zeros((*grid, 2), dtype=np.float32))
        payload = solver.get_field()
        payload["d"] = np.ones((*grid, 2), dtype=np.float32)
        payload["valid_mask"] = np.ones(grid, dtype=bool)
        uv = _identity_uv(reference)
        document = _document(
            tmp_path, gains=[0.5] * 4, uv=uv,
            points=[[0.0, 0.0], [reference[1] - 1.0, 0.0],
                    [0.0, reference[0] - 1.0], [reference[1] - 1.0, reference[0] - 1.0]],
            triangles=[[0, 1, 2], [1, 3, 2]],
        )
        model = parse_model(document, tmp_path / "calibration.json")
        gain_map = model.field_gain_map(grid, np.ones(reference, dtype=bool))
        reduction = reduce_field(coerce_field(payload), gain_map)
        # q=2 per cell; A*cell_count == reference pixel count for every integer
        # downsample, including grids that do not divide the reference exactly.
        assert reduction.weighted_N == pytest.approx(expected, rel=1e-5, abs=1e-6)