import numpy as np
import pytest

from SR_analysis.utils.data_contracts import (
    CAUSAL_ALIGNMENT,
    HISTORICAL_ALIGNMENT,
    actions_to_alpha,
    build_batch_dataset,
    contiguous_block_slices,
)
from SR_analysis.utils.feature_builder import (
    AUG_LEVEL_4_KEYS,
    ILLUSION_AUG_LEVEL_3_KEYS,
    build_feature_matrix,
    compute_dimensionless,
    compute_features,
)
from SR_analysis.utils.feature_state import StatefulFeatureBuilder


def synthetic_trajectory(T=8):
    t = np.arange(T, dtype=np.float64)
    sensors = np.column_stack([
        0.01 + 0.001 * t, 0.002 * t,
        0.012 - 0.0002 * t, -0.001 * t,
        0.009 + 0.0004 * t, 0.0005 * t,
    ])
    forces = np.column_stack([
        0.001 + 1e-5 * t, 2e-5 * t,
        0.0007 - 2e-5 * t, -1e-5 * t,
        0.0008 + 3e-5 * t, 1.5e-5 * t,
    ])
    actions = np.column_stack([0.2 * t, -0.1 * t, 0.05 * t])
    target = np.column_stack([0.002 + 2e-5 * t, -3e-5 * t])
    return sensors, forces, actions, target


def test_batch_and_stateful_features_match_elementwise():
    sensors, forces, alpha, target = synthetic_trajectory()
    T = len(sensors)
    a1 = np.zeros_like(alpha)
    a2 = np.zeros_like(alpha)
    a1[1:] = alpha[:-1]
    a2[2:] = alpha[:-2]
    dim = compute_dimensionless(sensors, forces)
    batch = compute_features(
        dim, a1, a2, mu=0.02, alpha_mode=True, include_mu=True,
        include_cos_sin=True, dt_c=0.3, target_forces=target,
        sensors_raw=sensors, forces_raw=forces,
    )
    builder = StatefulFeatureBuilder(
        mu=0.02, dt_c=0.3, action_mode="alpha", require_target=True,
        include_mu=True, include_cos_sin=True,
    )
    rows = []
    keys = list(dict.fromkeys(AUG_LEVEL_4_KEYS + ILLUSION_AUG_LEVEL_3_KEYS + [
        "sin_ua", "cos_ua", "target_Cd", "target_Cl", "daF", "daB", "daT",
    ]))
    for i in range(T):
        previous = alpha[i - 1] if i >= 1 else np.zeros(3)
        previous2 = alpha[i - 2] if i >= 2 else np.zeros(3)
        _, vector = builder.step(
            sensors[i], forces[i], previous_action=previous,
            previous_action2=previous2, target_force=target[i], feat_keys=keys,
        )
        rows.append(vector)
    expected = build_feature_matrix(batch, keys, add_bias=False)
    np.testing.assert_allclose(np.vstack(rows), expected, rtol=1e-12, atol=1e-12)


@pytest.mark.parametrize("sample_interval", [400, 600, 800])
def test_linear_signal_derivatives_respect_dt_c(sample_interval):
    dt_c = sample_interval / 2000.0
    T = 6
    t_c = np.arange(T) * dt_c
    sensors = np.zeros((T, 6))
    sensors[:, 0] = 0.01 + 0.01 * 3.0 * t_c
    sensors[:, 2] = 0.01 + 0.01 * 1.0 * t_c
    sensors[:, 4] = 0.01 - 0.01 * 1.0 * t_c
    forces = np.zeros((T, 6))
    actions = np.zeros((T, 3))
    dim = compute_dimensionless(sensors, forces)
    features = compute_features(
        dim, actions, actions, 0.02, alpha_mode=True, include_mu=False,
        include_cos_sin=False, dt_c=dt_c, sensors_raw=sensors, forces_raw=forces,
    )
    np.testing.assert_allclose(features["du_m_dt"][1:], 1.0)
    np.testing.assert_allclose(features["du_a_dt"][1:], 2.0)
    assert features["du_m_dt"][0] == pytest.approx(0.0)
    assert features["du_a_dt"][0] == pytest.approx(0.0)


def test_alignment_fixture_has_explicit_state_and_target_indices():
    sensors, forces, alpha, _ = synthetic_trajectory(T=6)
    keys = ["aF_lag1", "daF"]
    historical = build_batch_dataset(
        sensors, forces, alpha, keys, mu=0.02, dt_c=0.4,
        alignment=HISTORICAL_ALIGNMENT, action_mode="alpha", warmup=2,
    )
    causal = build_batch_dataset(
        sensors, forces, alpha, keys, mu=0.02, dt_c=0.4,
        alignment=CAUSAL_ALIGNMENT, action_mode="alpha", warmup=1,
    )
    np.testing.assert_array_equal(historical.metadata["state_indices"], [2, 3, 4, 5])
    np.testing.assert_array_equal(historical.metadata["target_indices"], [2, 3, 4, 5])
    np.testing.assert_allclose(historical.X[:, 0], alpha[1:-1, 0])
    np.testing.assert_allclose(historical.Y, alpha[2:])
    np.testing.assert_array_equal(causal.metadata["state_indices"], [1, 2, 3, 4])
    np.testing.assert_array_equal(causal.metadata["target_indices"], [2, 3, 4, 5])
    np.testing.assert_allclose(causal.X[:, 0], alpha[1:-1, 0])
    np.testing.assert_allclose(causal.Y, alpha[2:])


def test_missing_target_and_feature_fail_strictly():
    sensors, forces, alpha, _ = synthetic_trajectory(T=4)
    with pytest.raises(ValueError, match="require target_forces"):
        build_batch_dataset(
            sensors, forces, alpha, ["Cd_err"], mu=0.02, dt_c=0.4,
            action_mode="alpha", illusion=True,
        )
    with pytest.raises(KeyError, match="missing required feature"):
        build_feature_matrix({"u_m": np.ones(3)}, ["u_m", "not_computed"])
    compatible = build_feature_matrix(
        {"u_m": np.ones(3)}, ["u_m", "not_computed"], strict=False
    )
    np.testing.assert_array_equal(compatible[:, -1], np.zeros(3))


def test_transient_and_steady_segments_are_disjoint():
    rows = 10
    sensors = np.zeros((rows, 6))
    forces = np.zeros((rows, 6))
    actions = np.arange(rows * 3, dtype=float).reshape(rows, 3)
    transient = build_batch_dataset(
        sensors, forces, actions, ["u_m"], mu=0.02, dt_c=0.4,
        action_mode="alpha", warmup=2, sample_segment="transient", transient_count=5,
    )
    steady = build_batch_dataset(
        sensors, forces, actions, ["u_m"], mu=0.02, dt_c=0.4,
        action_mode="alpha", warmup=2, sample_segment="steady", transient_count=5,
    )
    np.testing.assert_array_equal(transient.kept_indices, [2, 3, 4])
    np.testing.assert_array_equal(steady.kept_indices, [5, 6, 7, 8, 9])


def test_contiguous_offline_blocks_are_ordered_disjoint_and_complete():
    blocks = contiguous_block_slices(11)
    indices = {name: np.arange(11)[block] for name, block in blocks.items()}
    np.testing.assert_array_equal(indices["train"], np.arange(6))
    np.testing.assert_array_equal(indices["validation"], np.arange(6, 8))
    np.testing.assert_array_equal(indices["blind"], np.arange(8, 11))
    np.testing.assert_array_equal(np.concatenate(list(indices.values())), np.arange(11))


def test_normalized_action_alpha_conversion():
    actions = np.array([[1.0, 0.0, -1.0], [0.5, -0.5, 0.25]])
    alpha = actions_to_alpha(
        actions, action_mode="normalized", u0=0.01,
        action_scale=12.0, action_bias=(1.0, -2.0, 3.0),
    )
    np.testing.assert_allclose(alpha, actions * 12.0 + [1.0, -2.0, 3.0])
    physical = alpha * 0.01
    np.testing.assert_allclose(
        actions_to_alpha(physical, action_mode="physical", u0=0.01), alpha
    )
