import json
import subprocess
import sys
from pathlib import Path

import numpy as np
import pytest

from SR_analysis import stage_2_fit as stage2


class MockRegressor:
    instances = []

    def __init__(self, **kwargs):
        self.kwargs = kwargs
        self.variable_names = None
        self.target = None
        self.weights = None
        type(self).instances.append(self)

    def fit(self, X, y, variable_names, weights=None):
        self.variable_names = list(variable_names)
        self.target = np.asarray(y).copy()
        self.weights = None if weights is None else np.asarray(weights).copy()
        self.constant = float(np.mean(y))
        return self

    def predict(self, X):
        return np.full(len(X), self.constant)

    def sympy(self):
        return "0.5*u_a"

    equations_ = [
        {"equation": "u_a", "complexity": 1, "loss": 0.2, "score": 0.4},
        {"equation": "0.5*u_a", "complexity": 3, "loss": 0.1, "score": 0.5},
    ]


@pytest.fixture(autouse=True)
def reset_mock():
    MockRegressor.instances.clear()


def write_scene(root: Path, scene: str, *, offset: float = 0.0, target: bool = False) -> Path:
    directory = root / "karman" / scene
    directory.mkdir(parents=True)
    rows = 8
    sensors = np.arange(rows * 6, dtype=np.float64).reshape(rows, 6) * 1e-5
    forces = np.arange(rows * 6, dtype=np.float64).reshape(rows, 6) * 1e-8
    actions = np.column_stack((
        np.arange(rows) + offset,
        np.arange(rows) + 10 + offset,
        np.arange(rows) + 20 + offset,
    )).astype(np.float64) * 1e-1
    payload = {"sensors": sensors, "forces": forces, "actions_alpha_commanded": actions, "dt_c": 0.4}
    if target:
        payload["target_forces"] = np.zeros((rows, 2))
    np.savez_compressed(directory / "controlled.npz", **payload)
    (directory / "config.json").write_text(json.dumps({
        "scene_id": "karman", "mu": 0.02, "sample_interval": 800,
        "u0": 0.01, "action_layout": ["front", "upper", "lower"],
    }), encoding="utf-8")
    return directory


def load_artifacts(paths):
    return [json.loads(Path(path).read_text(encoding="utf-8")) for path in paths]


def test_per_scene_fits_only_front_and_shared_upper_with_seed_and_schema(tmp_path):
    data_root = tmp_path / "data"
    write_scene(data_root, "karman_re100")
    results_root = tmp_path / "results" / "runs"

    paths = stage2.run_fit(
        scenes=["karman_re100"], mode="per-scene", run_id="unit",
        data_root=data_root, results_root=results_root, seed=17,
        niterations=3, smoke=True, regressor_factory=MockRegressor,
        front_deployment_expression="u_a", rear_deployment_expression="-u_a",
    )

    assert [path.name for path in paths] == ["karman_re100_front.json", "karman_re100_rear_shared_upper.json"]
    assert all(path.parent == results_root / "unit" / "formulas" for path in paths)
    assert len(MockRegressor.instances) == 2
    assert all(instance.kwargs["random_state"] == 17 for instance in MockRegressor.instances)
    assert all(instance.kwargs["deterministic"] is True for instance in MockRegressor.instances)
    assert all(instance.kwargs["tournament_selection_n"] == 4 for instance in MockRegressor.instances)
    assert all(instance.kwargs["maxsize"] == 7 for instance in MockRegressor.instances)
    assert all(instance.kwargs["population_size"] == 16 for instance in MockRegressor.instances)
    assert all("procs" not in instance.kwargs for instance in MockRegressor.instances)
    assert MockRegressor.instances[0].target[0] == pytest.approx(0.3)
    assert MockRegressor.instances[0].target[3] == pytest.approx(-0.3)
    assert MockRegressor.instances[1].target[0] == pytest.approx(1.3)
    assert MockRegressor.instances[1].target[3] == pytest.approx(-2.3)

    front, rear = load_artifacts(paths)
    assert front["role"] == "front_odd"
    assert rear["role"] == "rear_shared"
    assert rear["anchor"] == "upper"
    assert front["output"]["native_action_order"] == ["front", "upper", "lower"]
    assert front["fitted_expression"] == "0.5*u_a"
    assert front["deployment_expression"] == "u_a"
    assert front["fitted_expression_hash"] != front["deployment_expression_hash"]
    assert front["g_metadata"]["front_rule"] == "front(x)=0.5*(f(x)-f(Gx)); antisymmetric projection"
    assert rear["g_metadata"]["rear_rule"] == "lower(x)=-upper(Gx)"
    assert front["g_metadata"]["deployment_architecture"] == "mapped_shared"
    assert front["g_metadata"]["fit_augmentation"] == "G"
    assert set(front["provenance"]["fit"]["diagnostic_metrics"]) == {"r2", "mae", "rmse", "max_error"}
    assert front["provenance"]["fit"]["source_n_samples"] == 5
    assert front["provenance"]["fit"]["n_samples"] == 6
    assert front["provenance"]["fit"]["g_augmented"] is True
    assert set(front["provenance"]["fit"]["offline_metrics"]) == {"train", "validation", "blind"}
    assert "manual_validation_required" in front["provenance"]["fit"]["selection_policy"]
    assert front["provenance"]["fit"]["static_safety"]["safe"] is True
    assert len(front["provenance"]["pysr"]["pareto_candidates"]) == 2
    candidate = front["provenance"]["pysr"]["pareto_candidates"][0]
    assert candidate["used_features"] == ["u_a"]
    assert candidate["static_safety"]["safe"] is True
    assert set(candidate["metrics"]) == {"train", "validation", "blind"}
    assert set(candidate["metrics"]["blind"]) == {"r2", "mae", "rmse", "max_error"}


def test_three_head_independent_emits_true_three_head_artifacts(tmp_path):
    data_root = tmp_path / "data"
    write_scene(data_root, "karman_re100")
    paths = stage2.run_fit(
        scenes=["karman_re100"], mode="per-scene", run_id="three-head",
        data_root=data_root, results_root=tmp_path / "results",
        deployment_architecture="three_head_independent", fit_augmentation="none",
        regressor_factory=MockRegressor,
    )
    assert [path.name for path in paths] == [
        "karman_re100_front.json", "karman_re100_upper.json", "karman_re100_lower.json",
    ]
    artifacts = load_artifacts(paths)
    assert [artifact["role"] for artifact in artifacts] == [
        "front_independent", "upper_independent", "lower_independent",
    ]
    assert all(artifact["g_metadata"]["deployment_semantics"] == "three_head_independent" for artifact in artifacts)
    assert all(artifact["provenance"]["fit"]["g_augmented"] is False for artifact in artifacts)


def test_feature_sets_are_complete_and_scene_specific() -> None:
    karman = stage2.get_scene("karman_re100")
    illusion = stage2.get_scene("illusion_1L")
    raw = stage2.feature_names_for_scene(karman, feature_set="raw_complete", joint=False)
    symmetry = stage2.feature_names_for_scene(karman, feature_set="symmetry", joint=False)
    illusion_raw = stage2.feature_names_for_scene(illusion, feature_set="raw_complete", joint=False)
    assert len(raw) == 12
    assert {"v_hat_T", "v_hat_C", "v_hat_B", "Cd_T", "Cd_B", "Cl_T", "Cl_B"} <= set(raw)
    assert {"v_s", "v_a", "v_c", "Cd_rear_s", "Cd_rear_a", "Cl_rear_s", "Cl_rear_a"} <= set(symmetry)
    assert {"target_Cd", "target_Cl", "Cd_err", "Cl_err"} <= set(illusion_raw)
    assert "Cd_err" not in raw and "Cl_err" not in raw

    target_only = stage2.feature_names_for_scene(
        illusion, feature_set="physics_reduced", joint=False, feature_profile="target_only",
    )
    actual_plus_error = stage2.feature_names_for_scene(
        illusion, feature_set="physics_reduced", joint=False, feature_profile="actual_plus_error",
    )
    actual_plus_target = stage2.feature_names_for_scene(
        illusion, feature_set="physics_reduced", joint=False, feature_profile="actual_plus_target",
    )
    assert target_only == ("target_Cd", "target_Cl")
    assert actual_plus_error[-2:] == ("Cd_err", "Cl_err") and "target_Cd" not in actual_plus_error
    assert actual_plus_target[-2:] == ("target_Cd", "target_Cl") and "Cd_err" not in actual_plus_target
    assert stage2.feature_names_for_scene(
        karman, feature_set="physics_reduced", joint=False, feature_profile="actual_only",
    ) == tuple(stage2.FEATURE_SET_KEYS["physics_reduced"])
    with pytest.raises(ValueError, match="only defined for Illusion"):
        stage2.feature_names_for_scene(
            karman, feature_set="physics_reduced", joint=False, feature_profile="actual_plus_error",
        )


def test_joint_merges_both_heads_and_records_scene_counts(tmp_path):
    data_root = tmp_path / "run"
    write_scene(data_root, "karman_re100")
    write_scene(data_root, "karman_re200", offset=100)

    paths = stage2.run_fit(
        scenes=["karman_re100", "karman_re200"], mode="joint", run_id="joint-run",
        data_root=data_root, results_root=tmp_path / "results" / "runs",
        regressor_factory=MockRegressor,
    )
    assert [path.name for path in paths] == ["joint_front.json", "joint_rear_shared_upper.json"]
    front, rear = load_artifacts(paths)
    assert front["training_scenes"] == ["karman_re100", "karman_re200"]
    assert front["provenance"]["fit"]["scene_sample_counts"] == {
        "karman_re100": 5, "karman_re200": 5,
    }
    assert rear["g_metadata"]["rear_rule"] == "lower(x)=-upper(Gx)"
    assert all(instance.weights is not None for instance in MockRegressor.instances)
    for instance in MockRegressor.instances:
        assert instance.weights[:3].sum() == pytest.approx(instance.weights[3:6].sum())


def test_multiple_trajectories_preserve_boundaries_and_provenance(tmp_path):
    first_root = tmp_path / "trajectory-a"
    second_root = tmp_path / "trajectory-b"
    write_scene(first_root, "karman_re100", offset=0)
    write_scene(second_root, "karman_re100", offset=100)

    prepared = stage2.prepare_scene_trajectories(
        "karman_re100", [first_root, second_root], alignment=stage2.CAUSAL_ALIGNMENT,
        joint=False,
    )
    assert prepared.dataset.X.shape[0] == 10
    assert [len(item.dataset.X) for item in prepared.trajectories] == [5, 5]
    assert prepared.dataset.Y[4, 0] == pytest.approx(0.7)
    assert prepared.dataset.Y[5, 0] == pytest.approx(10.3)

    paths = stage2.run_fit(
        scenes=["karman_re100"], mode="per-scene", run_id="multi",
        data_roots=[first_root, second_root], results_root=tmp_path / "results",
        regressor_factory=MockRegressor,
    )
    trajectory_records = load_artifacts(paths)[0]["provenance"]["data"]["trajectories"]["karman_re100"]
    assert [record["id"] for record in trajectory_records] == [
        "karman_re100:trajectory-001", "karman_re100:trajectory-002",
    ]
    assert [record["sample_count"] for record in trajectory_records] == [5, 5]
    assert all(len(record["hash"]) == 64 for record in trajectory_records)
    assert [record["path"] for record in trajectory_records] == [
        str(first_root / "karman" / "karman_re100" / "controlled.npz"),
        str(second_root / "karman" / "karman_re100" / "controlled.npz"),
    ]


def test_refitted_expression_is_canonical_for_sign_changes():
    expression = "-1.0*x + 2.0*y"
    _, tree, names, _ = stage2._topology_evaluator(expression, ("x", "y"))
    rendered = stage2._render_refitted_expression(tree, names, (-2.0, -3.0))
    assert "--" not in rendered
    assert "+ -" not in rendered


def test_fixed_topology_refit_uses_all_rows_two_cases_and_preserves_topology(tmp_path):
    feature_names = ("u_a",)
    scenes = []
    for case_index, name in enumerate(("case_a", "case_b")):
        X = np.arange(1 + case_index, 5 + case_index, dtype=np.float64).reshape(-1, 1)
        front = 2.0 * X[:, 0]
        upper = -1.5 * X[:, 0]
        lower = upper
        dataset = stage2.BatchDataset(
            X, np.column_stack((front, upper, lower)), np.arange(len(X)),
            np.array([], dtype=int),
            {"alignment": stage2.CAUSAL_ALIGNMENT, "warmup": 0, "history_policy": "unit",
             "dt_c": 0.4, "action_mode": "alpha"}, {},
        )
        mirrored = stage2.BatchDataset(
            -X, np.column_stack((-front, -lower, -upper)), np.arange(len(X)),
            np.array([], dtype=int), dict(dataset.metadata), {},
        )
        trajectory = stage2.PreparedTrajectory(
            f"{name}:trajectory-001", tmp_path / f"{name}.npz", None,
            dataset, mirrored, f"hash-{name}",
        )
        scenes.append(stage2.PreparedScene(
            name, {"scene_id": "karman", "u0": 0.01}, trajectory.data_path, None,
            dataset, mirrored, feature_names, ("front", "upper", "lower"),
            trajectory.data_hash, (trajectory,),
        ))

    front = stage2.refit_fixed_topology_head(
        scenes, role="front_odd", channel_index=0, run_id="refit",
        topology_expression="0.5 * u_a", discovery_parent_path=None,
        mode="joint", deployment_architecture="mapped_shared",
        feature_set="physics_reduced", feature_profile="actual_only",
        environment={},
    )
    rear = stage2.refit_fixed_topology_head(
        scenes, role="rear_shared", channel_index=1, run_id="refit",
        topology_expression="0.5 * u_a", discovery_parent_path=None,
        mode="joint", deployment_architecture="mapped_shared",
        feature_set="physics_reduced", feature_profile="actual_only",
        environment={},
    )

    assert front["fitted_expression"].count("u_a") == 1
    assert front["fitted_expression_hash"] == front["deployment_expression_hash"]
    assert front["provenance"]["fit"]["optimizer"]["refitted_constants"] == pytest.approx([2.0])
    assert rear["provenance"]["fit"]["optimizer"]["refitted_constants"] == pytest.approx([-1.5])
    assert front["provenance"]["fit"]["n_samples"] == 16
    assert front["provenance"]["fit"]["sample_weighting"] == "case_equal_trajectory_equal_within_case"
    metrics = front["provenance"]["fit"]["offline_metrics"]
    assert metrics["aggregate"]["rmse"] < 1e-10
    assert set(metrics["per_case"]) == {"case_a", "case_b"}
    assert set(metrics["per_trajectory"]) == {
        "case_a:trajectory-001", "case_b:trajectory-001",
    }
    assert front["g_metadata"]["fit_augmentation"] == "G"


def test_run_fit_fixed_topology_records_discovery_parent_hash(tmp_path, monkeypatch):
    parent = tmp_path / "discovery.json"
    parent.write_text('{"artifact": "discovery"}', encoding="utf-8")
    prepared = []
    for case_index, name in enumerate(("karman_re100", "karman_re200")):
        X = np.arange(1 + case_index, 5 + case_index, dtype=np.float64).reshape(-1, 1)
        front = 2.0 * X[:, 0]
        upper = -1.5 * X[:, 0] + 0.25
        dataset = stage2.BatchDataset(
            X, np.column_stack((front, upper, -upper)), np.arange(len(X)),
            np.array([], dtype=int),
            {"alignment": stage2.CAUSAL_ALIGNMENT, "warmup": 0, "history_policy": "unit",
             "dt_c": 0.4, "action_mode": "alpha"}, {},
        )
        mirrored = stage2.BatchDataset(
            X.copy(), np.column_stack((-front, upper, -upper)), np.arange(len(X)),
            np.array([], dtype=int), dict(dataset.metadata), {},
        )
        trajectory = stage2.PreparedTrajectory(
            f"{name}:trajectory-001", tmp_path / f"{name}.npz", None,
            dataset, mirrored, f"hash-{name}",
        )
        prepared.append(stage2.PreparedScene(
            name, {"scene_id": "karman", "u0": 0.01}, trajectory.data_path, None,
            dataset, mirrored, ("u_a",), ("front", "upper", "lower"),
            trajectory.data_hash, (trajectory,),
        ))
    monkeypatch.setattr(
        stage2, "prepare_scene_trajectories",
        lambda name, *args, **kwargs: prepared[("karman_re100", "karman_re200").index(name)],
    )

    paths = stage2.run_fit(
        scenes=["karman_re100", "karman_re200"], mode="joint", run_id="refit-run",
        results_root=tmp_path / "results", fit_purpose="fixed-topology-refit",
        front_topology_expression="0.5*u_a + 0.1",
        rear_topology_expression="0.5*u_a + 0.1",
        front_discovery_parent_path=parent, rear_discovery_parent_path=parent,
    )

    assert [path.name for path in paths] == ["joint_front.json", "joint_rear_shared_upper.json"]
    artifacts = load_artifacts(paths)
    assert not MockRegressor.instances
    assert all(artifact["discovery_parent"]["path"] == str(parent) for artifact in artifacts)
    assert all(artifact["discovery_parent"]["sha256"] == stage2.hash_file(parent) for artifact in artifacts)
    assert all(artifact["provenance"]["fit"]["fit_purpose"] == "fixed-topology-refit" for artifact in artifacts)


def test_joint_rejects_cross_objective_scene_id_mixing():
    first = stage2.PreparedScene(
        "first", {"scene_id": "karman", "u0": 0.01}, Path("first.npz"), None,
        stage2.BatchDataset(np.zeros((1, 1)), np.zeros((1, 3)), np.array([0]), np.array([], dtype=int),
                            {"action_mode": "alpha"}, {}),
        stage2.BatchDataset(np.zeros((1, 1)), np.zeros((1, 3)), np.array([0]), np.array([], dtype=int),
                            {"action_mode": "alpha"}, {}),
        ("u_a",), ("front", "upper", "lower"), "first-hash",
    )
    second = stage2.PreparedScene(
        "second", {"scene_id": "illusion", "u0": 0.01}, Path("second.npz"), None,
        first.dataset, first.mirrored_dataset, first.feature_names, first.action_layout, "second-hash",
    )
    with pytest.raises(ValueError, match="same scene_id objective"):
        stage2._validate_joint_contract([first, second])


def test_cli_exposes_only_active_architecture_and_comparative_profiles():
    parser = stage2.build_parser()
    architecture = next(action for action in parser._actions if action.dest == "deployment_architecture")
    augmentation = next(action for action in parser._actions if action.dest == "fit_augmentation")
    profile = next(action for action in parser._actions if action.dest == "feature_profile")
    assert tuple(architecture.choices) == ("mapped_shared", "three_head_independent")
    assert tuple(augmentation.choices) == ("none", "G")
    assert tuple(profile.choices) == (
        "actual_only", "target_only", "actual_plus_target", "actual_plus_error",
    )
    assert all(action.dest != "symmetry_mode" for action in parser._actions)
    parsed = parser.parse_args([
        "--scene", "karman_re100", "--data-root", "first", "--data-root", "second",
        "--data-roots", "third,fourth",
    ])
    assert parsed.data_root == ["first", "second"]
    assert parsed.data_roots == "third,fourth"


def test_prepare_only_validates_without_importing_pysr_or_writing(tmp_path):
    data_root = tmp_path / "data"
    write_scene(data_root, "karman_re100")
    results_root = tmp_path / "results" / "runs"
    script = Path(stage2.__file__)
    code = (
        "import runpy,sys; "
        f"sys.argv={[str(script), '--scene', 'karman_re100', '--run-id', 'prepare', '--data-root', str(data_root), '--results-root', str(results_root), '--prepare-only']!r}; "
        f"runpy.run_path({str(script)!r}, run_name='__main__')"
    )
    completed = subprocess.run([sys.executable, "-c", code], cwd=tmp_path, capture_output=True, text=True)
    assert completed.returncode == 0, completed.stderr
    assert "pysr" not in sys.modules
    assert not results_root.exists()


def test_refuses_existing_output_before_fit(tmp_path):
    data_root = tmp_path / "data"
    write_scene(data_root, "karman_re100")
    output = tmp_path / "results" / "runs" / "same" / "formulas"
    output.mkdir(parents=True)
    with pytest.raises(FileExistsError, match="Refusing"):
        stage2.run_fit(
            scenes=["karman_re100"], mode="per-scene", run_id="same",
            data_root=data_root, results_root=tmp_path / "results" / "runs",
            regressor_factory=MockRegressor,
        )
    assert not MockRegressor.instances


def test_illusion_required_target_fails_during_prepare(tmp_path):
    directory = tmp_path / "data" / "illusion" / "illusion_1L"
    directory.mkdir(parents=True)
    np.savez_compressed(
        directory / "controlled.npz", sensors=np.zeros((6, 6)), forces=np.zeros((6, 6)),
        actions_alpha_commanded=np.zeros((6, 3)),
    )
    with pytest.raises(ValueError, match="require target_forces"):
        stage2.prepare_scene("illusion_1L", tmp_path / "data", alignment=stage2.CAUSAL_ALIGNMENT, joint=False)


def test_alignment_option_changes_target_pairing(tmp_path):
    data_root = tmp_path / "data"
    write_scene(data_root, "karman_re100")
    causal = stage2.prepare_scene("karman_re100", data_root, alignment=stage2.CAUSAL_ALIGNMENT, joint=False)
    historical = stage2.prepare_scene(
        "karman_re100", data_root, alignment="historical_post_state_same_action", joint=False,
    )
    assert causal.dataset.Y[0, 0] == pytest.approx(0.3)
    assert historical.dataset.Y[0, 0] == pytest.approx(0.2)
