from pathlib import Path

import numpy as np
import pytest

import src.drl_pinball.legacy_test.acquire as acquire
import src.drl_pinball.legacy_test.core.legacy_env_builder as builder


class RoleFlow:
    DATA_TYPE = np.float32

    def __init__(self, n_objects, obs_size, absolute_base=2_000_000):
        self.action = np.ones(n_objects, dtype=np.float32)
        self.obs = np.zeros(obs_size, dtype=np.float32)
        self.intervals = 0
        self.absolute_base = absolute_base
        self.commands = []
        self.capture_intervals = []

    def restore_ddf(self):
        pass

    def apply_ddf(self):
        pass

    def run(self, steps, command):
        assert steps == acquire.SI
        assert np.array_equal(self.action, np.zeros_like(self.action))
        self.intervals += 1
        self.commands.append(command.copy())
        self.action = command.copy()
        phase = 2 * np.pi * self.intervals / acquire.CYCLE_LENGTH
        self.obs[0:2] = [1000 + self.intervals, 2000 + self.intervals]
        self.obs[2:8] = [np.sin(phase + index) for index in range(6)]
        if self.obs.size >= 14:
            self.obs[8:14] = np.arange(6) + self.intervals

    def current_effective_action(self):
        return self.action.copy()

    def solver_clock_state(self):
        return {"solver_absolute_lattice_clock": self.absolute_base + self.intervals * acquire.SI,
                "solver_absolute_control_clock": self.intervals}


def data(n_objects=7):
    phase = 2 * np.pi * np.arange(150) / acquire.CYCLE_LENGTH
    target = np.column_stack([np.sin(phase + index) for index in range(6)]).astype(np.float32)
    result = {
        "config": {"nx": 1280, "ny": 2, "sample_interval": 800,
                   "n_obj_total": n_objects},
        "target_states": target,
    }
    if n_objects == 7:
        result["norm"] = {"save_states": np.zeros((150, 12), dtype=np.float32)}
    return result


def norm():
    return {"force_norm_fact": 1.0, "sens_deviation": np.zeros(6),
            "sens_norm_fact": np.ones(6)}


def capture(flow, nx, ny):
    assert flow.intervals > acquire.WARMUP_INTERVALS
    flow.capture_intervals.append(flow.intervals)
    return {"ux": np.zeros((ny, nx), np.float32), "uy": np.zeros((ny, nx), np.float32)}


def test_target_builder_uses_disturbance_then_three_sensors_without_pinball(monkeypatch):
    class TargetFlow:
        DATA_TYPE = np.float32
        FIELD_SHAPE = (1280, 2, 1)

        def __init__(self, field_cfg, cuda_cfg, device_id):
            self.obs = np.zeros(0, np.float32)
            self.objects = []
            self.interval = 0

        def add_cylinder(self, center, radius):
            self.objects.append(("cylinder", center, radius))
            self.obs = np.zeros(2 * len(self.objects), np.float32)

        def add_sensor(self, center, radius):
            self.objects.append(("sensor", center, radius))
            self.obs = np.zeros(2 * len(self.objects), np.float32)

    class Config:
        def _replace(self, **kwargs):
            return self

    monkeypatch.setattr(builder, "FlowField", TargetFlow)
    monkeypatch.setattr(builder.legacy_utils, "load_cuda_config", lambda path: object())
    monkeypatch.setattr(builder.legacy_utils, "load_flow_field_config", lambda path: Config())
    monkeypatch.setattr(builder, "_stabilize", lambda flow, count: None)
    monkeypatch.setattr(builder, "FIFO_LEN", 2)

    def interval(flow, steps, command):
        flow.interval += 1
        flow.obs[0:2] = [100 + flow.interval, 200 + flow.interval]
        flow.obs[2:8] = np.arange(6) + 10 * flow.interval

    monkeypatch.setattr(builder, "run_historical_interval", interval)
    result = builder.build_karman_target(0, 100.0, sample_interval=800)

    flow = result["flow_field"]
    assert [item[0] for item in flow.objects] == ["cylinder", "sensor", "sensor", "sensor"]
    assert flow.objects[0][1][0] == 10.0 * builder.L0
    assert [item[1][1] for item in flow.objects[1:]] == [40.5, 0.5, -39.5]
    assert result["config"]["n_obj_total"] == 4
    assert np.array_equal(result["target_states"],
                          np.asarray([np.arange(6) + 10, np.arange(6) + 20], np.float32))


def test_zero_role_schedule_actions_rewards_and_fields(monkeypatch, tmp_path):
    flow = RoleFlow(7, 14)
    reward_calls = []

    def rewards(*args):
        reward_calls.append(args)
        return {"reward": 0.4, "reward_cd": 0.5, "reward_cl": 0.6,
                "native_legacy_dtw": 0.7}

    monkeypatch.setattr(acquire, "reward_terms", rewards)
    rows, fields = acquire._collect_role(
        "zero", data(), None, flow, norm(), tmp_path, capture_field=capture,
    )

    assert flow.intervals == len(reward_calls) == 640
    assert flow.capture_intervals == list(range(481, 641))
    assert len(rows) == 160
    assert fields["ux"].shape == fields["uy"].shape == (160, 2, 1280)
    assert not list(tmp_path.glob("boundary_*.npz"))
    counter_bias = np.array([0.0, 0.5, -0.5], np.float32)
    assert np.array_equal(rows[0]["action_normalized"], counter_bias)
    assert all(np.array_equal(command, np.zeros(7, np.float32)) for command in flow.commands)
    assert np.array_equal(rows[0]["commanded_target_omega"], np.zeros(3, np.float32))
    assert np.allclose(rows[0]["effective_smoothed_omega"], 0.0, atol=1e-7)
    assert rows[0]["reward_raw"] == 0.4
    assert rows[0]["native_reward_dtw"] == 0.7
    assert rows[0]["control_index"] == 481 and rows[-1]["control_index"] == 640


def test_zero_role_metadata_names_physical_zero_counter_bias():
    source = Path(acquire.__file__).read_text(encoding="utf-8")
    assert "physical-zero/uncontrolled" in source
    assert "counter-bias normalized action" in source


def test_target_role_ordering_nan_contract_and_schedule(monkeypatch, tmp_path):
    flow = RoleFlow(4, 8)
    monkeypatch.setattr(acquire, "reward_terms", lambda *args: pytest.fail("target has no reward"))
    rows, fields = acquire._collect_role(
        "target", data(n_objects=4), None, flow, None, tmp_path, capture_field=capture,
    )

    assert flow.intervals == 640
    assert all(np.array_equal(command, np.zeros(4, np.float32)) for command in flow.commands)
    assert len(rows) == 160
    assert fields["ux"].shape == (160, 2, 1280)
    first = rows[0]
    phase = 2 * np.pi * 481 / acquire.CYCLE_LENGTH
    expected = np.asarray([np.sin(phase + index) for index in range(6)], np.float32)
    assert np.allclose(first["sensors"], expected)
    assert first["sensors"][0] != flow.obs[0]
    for name in ("forces", "action_normalized", "commanded_target_omega",
                 "effective_smoothed_omega"):
        assert np.isnan(first[name]).all()
    for name in ("reward_raw", "reward_cd", "reward_cl", "reward_sim",
                 "native_reward_dtw"):
        assert np.isnan(first[name])
    assert first["solver_absolute_lattice_step"] == 2_000_000 + 481 * 800


def test_target_sensor_reference_selects_legacy_target_channels():
    six = np.arange(60, dtype=float).reshape(10, 6)
    eight = np.column_stack((np.full((10, 2), -1.0), six))
    assert np.array_equal(acquire._target_sensors(six), six)
    assert np.array_equal(acquire._target_sensors(eight), six)
    with pytest.raises(ValueError, match="six sensor channels"):
        acquire._target_sensors(np.zeros((10, 7)))


def test_target_phase_artifacts_exclude_unavailable_nan_variables():
    columns = {"sensors": np.ones((4, 6)), "forces": np.full((4, 6), np.nan),
               "action_normalized": np.full((4, 3), np.nan),
               "physical_time": np.arange(4), "lattice_step": np.arange(4),
               "solver_absolute_lattice_step": np.arange(4),
               "control_index": np.arange(4), "phase": np.arange(4),
               "cycle_id": np.arange(4)}
    values = acquire._phase_values(columns, "target")
    assert set(values) == {"sensors"}
    assert np.isfinite(values["sensors"]).all()


@pytest.mark.parametrize("role", ["target", "zero"])
def test_role_failure_transaction_preserves_existing(monkeypatch, tmp_path, role):
    root = tmp_path / "out"
    final = root / "legacy" / acquire.CASE_NAME / role
    final.mkdir(parents=True)
    marker = final / "keep.txt"
    marker.write_text("old", encoding="utf-8")
    model = tmp_path / "model.zip"; model.write_bytes(b"model")
    norm_path = tmp_path / "norm.json"; norm_path.write_text("{}", encoding="utf-8")
    bundle = {"case": acquire.get_case(acquire.CASE_NAME), "model_path": model,
              "norm_path": norm_path, "norm": norm(),
              "reference_path": tmp_path / "missing.npz"}
    monkeypatch.setattr(acquire, "_resolve_bundle", lambda repo, *args: bundle)
    monkeypatch.setattr(acquire, "_collect_role", lambda *args, **kwargs: ([], []))

    def storage_validator(**kwargs):
        root.mkdir(exist_ok=True)
        return {"resolved_output_root": root, "device": root.stat().st_dev}

    def runtime_factory(case, current_bundle, device, current_role):
        n_objects = 4 if current_role == "target" else 7
        return data(n_objects), None, RoleFlow(n_objects, 8 if current_role == "target" else 14)

    with pytest.raises(RuntimeError, match="finalize failed"):
        acquire.acquire_role(
            role, output_root=root, overwrite=True, repo_root=tmp_path,
            storage_validator=storage_validator, runtime_factory=runtime_factory,
            finalizer=lambda *args: (_ for _ in ()).throw(RuntimeError("finalize failed")),
        )
    assert marker.read_text(encoding="utf-8") == "old"
    assert not list(final.parent.glob(f".{role}.staging-*"))


def test_cli_exposes_all_roles():
    source = Path(acquire.__file__).read_text(encoding="utf-8")
    assert 'choices=ROLES' in source