import json
from pathlib import Path

import numpy as np
import pytest

from SR_analysis import stage_3_validate as stage3


def formula(path: Path, *, role: str, anchor: str | None = None, expression: str = "0") -> Path:
    feature_names = ["u_a"] if "u_a" in expression else []
    value = {
        "schema_version": "1.0",
        "model_type": "symbolic",
        "feature_names": feature_names,
        "fitted_expression": expression,
        "deployment_expression": expression,
        "role": role,
        "parity": "odd" if role == "front" else None,
        "anchor": anchor,
    }
    path.write_text(json.dumps(value), encoding="utf-8")
    return path


def formula_pair(tmp_path: Path, front_expression: str = "0", rear_expression: str = "1") -> stage3.FormulaPair:
    front = formula(tmp_path / "front.json", role="front", expression=front_expression)
    rear = formula(tmp_path / "rear.json", role="rear_shared", anchor="upper", expression=rear_expression)
    return stage3.load_formula_pair(front, rear)


def karman_cfg() -> dict:
    return {
        "scene_id": "karman",
        "sample_interval": 800,
        "u0": 0.01,
        "mu": 0.02,
        "conv_len": 2,
        "action_layout": stage3.ACTION_ORDER,
        "policy_init_action": (0.0, -4.0, 4.0),
        "action_scale": 8.0,
        "action_bias": (0.0, -4.0, 4.0),
        "s_dim": 12,
    }


def make_plan(tmp_path: Path, *, mode: str = "pysr", pair=None, n_steps: int = 4) -> stage3.ValidationPlan:
    root = tmp_path / "runs"
    stem = "karman_re100__pysr__f-abc__si-800"
    return stage3.ValidationPlan(
        "karman_re100", karman_cfg(), mode, n_steps, "run", root / "run" / "validations",
        root / "run" / "telemetry", stem, root / "run" / "validations" / f"{stem}.json",
        root / "run" / "telemetry" / f"{stem}.npz", pair, None, None, None, None, None,
    )


class FakeEnvironment:
    def __init__(self, *, nonfinite: bool = False):
        self.current_raw = np.arange(12, dtype=float)
        self.target_sensors = np.tile(np.arange(6, dtype=float), (8, 1))
        self.nonfinite = nonfinite
        self.closed = False

    def step(self, omega_native):
        self.current_raw = self.current_raw + 0.01
        if self.nonfinite:
            self.current_raw[0] = np.nan
        return self.current_raw.copy()

    def close(self):
        self.closed = True


def test_formula_roles_hash_and_independent_filename(tmp_path, monkeypatch):
    pair = formula_pair(tmp_path)
    monkeypatch.setattr(stage3, "get_scene", lambda scene: karman_cfg())
    plan = stage3.prepare_plan(
        scene="karman_re100", mode="pysr", n_steps=4, run_id="r1",
        output_root=tmp_path / "out", formula_pair=pair,
    )
    assert pair.front["deployment_expression_hash"]
    assert pair.rear["deployment_expression_hash"]
    assert "karman_re100__pysr__f-" in plan.validation_path.name
    assert "__si-800.json" in plan.validation_path.name
    assert plan.validation_path.parent.name == "validations"
    assert plan.telemetry_path.parent.name == "telemetry"


def test_legacy_top_formula_is_recorded(tmp_path):
    front = tmp_path / "front.json"
    rear = tmp_path / "top.json"
    front.write_text(json.dumps({"channel": "front", "feature_keys": [], "best_sympy": "0"}), encoding="utf-8")
    rear.write_text(json.dumps({"channel": "top", "feature_keys": [], "best_sympy": "0"}), encoding="utf-8")
    pair = stage3.load_formula_pair(front, rear)
    assert "channel=top" in pair.compatibility["rear"]["legacy_role_mapping"]


def test_missing_or_wrong_formula_fails(tmp_path):
    front = formula(tmp_path / "front.json", role="front")
    wrong = formula(tmp_path / "wrong.json", role="rear_shared", anchor="lower")
    with pytest.raises(ValueError, match="anchor='upper'"):
        stage3.load_formula_pair(front, wrong)
    with pytest.raises(FileNotFoundError):
        stage3.load_formula_pair(front, tmp_path / "missing.json")


def test_symbolic_policy_uses_exact_front_odd_projection_and_rear_mapping(tmp_path, monkeypatch):
    pair = formula_pair(tmp_path, rear_expression="u_a")
    policy = stage3.SymbolicPolicy(pair, karman_cfg(), None)
    values = iter([8.0, 2.0, 3.0, 5.0])
    monkeypatch.setattr(policy, "_evaluate", lambda formula, features: next(values))
    omega, _, _ = policy.action(np.arange(12, dtype=float), 0)
    np.testing.assert_allclose(omega / 0.01, [3.0, 3.0, -5.0])


def test_metric_has_single_exact_named_shape(monkeypatch):
    monkeypatch.setattr(
        stage3,
        "_legacy_metric",
        lambda target, sensors, conv_len: {
            "similarity": 0.75,
            "per_channel": {str(index): float(index) for index in range(6)},
            "lag": -2,
        },
    )
    metrics = stage3.compute_versioned_metrics(np.zeros((4, 6)), np.zeros((4, 6)), 2)
    assert set(metrics) == {
        "metric_version",
        "adapter",
        "legacy_reference_cycle_vs_last_recorded_cycle",
    }
    exact = metrics["legacy_reference_cycle_vs_last_recorded_cycle"]
    assert set(exact) == {"similarity", "per_channel", "lag"}
    assert len(exact["per_channel"]) == 6
    assert exact["lag"] == -2


def test_ppo_model_loads_once_for_rollout(tmp_path):
    plan = make_plan(tmp_path, mode="ppo")
    model_path = tmp_path / "model.zip"
    model_path.write_bytes(b"model")
    norm = {"force_norm_fact": 1.0, "sens_deviation": [0.0] * 6, "sens_norm_fact": [1.0] * 6}
    plan = stage3.ValidationPlan(**{**plan.__dict__, "model_path": model_path, "norm": norm})
    calls = []

    class Model:
        def predict(self, observation, deterministic=True):
            return np.zeros(3, dtype=np.float32), None

    def loader(*args, **kwargs):
        calls.append((args, kwargs))
        return Model()

    environment = FakeEnvironment()
    policy = stage3.build_policy(plan, model_loader=loader)
    stage3.run_rollout(plan, environment, policy)
    assert len(calls) == 1
    assert calls[0][1]["device"] == "cpu"


def test_nonfinite_rollout_raises(tmp_path):
    plan = make_plan(tmp_path, mode="uncontrolled")
    with pytest.raises(stage3.RolloutFailure, match="raw observation"):
        stage3.run_rollout(plan, FakeEnvironment(nonfinite=True), stage3.UncontrolledPolicy())


def test_execute_propagates_factory_exception(tmp_path):
    plan = make_plan(tmp_path, mode="uncontrolled")

    def broken(plan, device):
        raise RuntimeError("mock flow failed")

    with pytest.raises(RuntimeError, match="mock flow failed"):
        stage3.execute_plan(plan, 0, environment_factory=broken)


def test_dry_run_requires_no_cuda_and_creates_no_artifacts(tmp_path, monkeypatch, capsys):
    cfg = karman_cfg()
    monkeypatch.setattr(stage3, "get_scene", lambda scene: cfg)
    pair = formula_pair(tmp_path)
    code = stage3.main([
        "--scene", "karman_re100", "--mode", "pysr", "--formula-front", str(pair.front_path),
        "--formula-rear", str(pair.rear_path), "--run-id", "dry", "--output-root", str(tmp_path / "out"),
        "--dry-run",
    ])
    assert code == 0
    assert "validation=" in capsys.readouterr().out
    assert not (tmp_path / "out").exists()


def test_vortex_is_explicit_round1_error(tmp_path, monkeypatch):
    cfg = {**karman_cfg(), "scene_id": "vortex"}
    monkeypatch.setattr(stage3, "get_scene", lambda scene: cfg)
    with pytest.raises(NotImplementedError, match="not supported in round1"):
        stage3.prepare_plan(
            scene="vortex_lamb", mode="uncontrolled", n_steps=4, run_id="r",
            output_root=tmp_path, formula_pair=None,
        )


def test_batch_failure_returns_nonzero(tmp_path, monkeypatch):
    cfg = karman_cfg()
    monkeypatch.setattr(stage3, "get_scene", lambda scene: cfg)
    monkeypatch.setattr(stage3, "execute_plan", lambda plan, device: (_ for _ in ()).throw(RuntimeError("boom")))
    code = stage3.main([
        "--group", "a,b", "--mode", "uncontrolled", "--run-id", "r",
        "--output-root", str(tmp_path), "--steps", "4",
    ])
    assert code != 0
