import json
from pathlib import Path

import numpy as np
import pytest

from SR_analysis.utils.formula_schema import (
    build_coefficient_scaling_variant,
    build_formula_artifact,
    build_term_deletion_variant,
    build_validation_record,
    decompose_additive_terms,
    evaluate_formula,
    formula_safety_report,
    load_formula,
    normalize_symbolic_expression,
)
from SR_analysis.utils.provenance import atomic_write_json_below, hash_json, safe_path


def test_expression_alias_is_supported():
    formula = {"feature_names": ["x"], "expression": "2*x + 1", "best_score": 0.9}
    np.testing.assert_allclose(evaluate_formula(formula, {"x": [1.0, 3.0]}), [3.0, 7.0])


def test_formula_static_safety_detects_finite_and_range_failures():
    safe = formula_safety_report(
        {"feature_names": ["x"], "expression": "2*x"}, {"x": [-1, 1]}, output_range=[-2, 2],
    )
    assert safe["safe"] is True
    out_of_range = formula_safety_report(
        {"feature_names": ["x"], "expression": "10*x"}, {"x": [-1, 1]}, output_range=[-2, 2],
    )
    assert out_of_range["finite"] is True
    assert out_of_range["safe"] is False
    nonfinite = formula_safety_report(
        {"feature_names": ["x"], "expression": "1/x"}, {"x": [-1, 1]}, grid_size=3,
    )
    assert nonfinite["finite"] is False


def test_missing_coef_never_silently_produces_zero():
    formula = {
        "schema_version": "1.0",
        "model_type": "linear",
        "feature_names": ["x"],
        "fitted_expression": None,
        "deployment_expression": None,
    }

    with pytest.raises(ValueError, match="requires coef"):
        evaluate_formula(formula, {"x": np.array([2.0, 4.0])})


def test_symbolic_expression_wins_over_incidental_coefficients():
    formula = {
        "schema_version": "1.0",
        "model_type": "linear",
        "feature_names": ["x"],
        "deployment_expression": "3*x",
        "fitted_expression": "3*x",
        "coef": [0.0],
    }
    np.testing.assert_allclose(evaluate_formula(formula, {"x": [2.0]}), [6.0])


def test_build_versioned_artifact_and_validation_record():
    artifact = build_formula_artifact(
        artifact_id="formula-001",
        run_id="run-001",
        role="controller",
        anchor="front",
        training_scenes=["scene-a"],
        feature_names=["x"],
        fitted_expression="2*x",
        deployment_expression="2*x",
        output={"name": "alpha", "unit": "rad"},
        feature_metadata={"version": "features-v1"},
        alignment_metadata={"policy": "same-step"},
        dt_c_metadata={"value": 0.2, "unit": "s"},
        g_metadata={"operator": "identity"},
        fit_provenance={"seed": 7},
        pysr_provenance={"version": "test"},
        data_provenance={"dataset_hash": "abc"},
    )
    record = build_validation_record(
        formula=artifact,
        metric_version="metrics-v1",
        scene="scene-a",
        mode="closed-loop",
        si=True,
        n_steps=100,
        full_metrics={"rmse": 0.1},
        tail_metrics={"rmse": 0.05},
        per_channel_metrics={"front": {"rmse": 0.1}},
        action_range={"min": -1.0, "max": 1.0},
        telemetry={"complete": True},
        provenance={"validator": "pytest"},
    )

    assert artifact["fitted_expression_hash"] == artifact["deployment_expression_hash"]
    assert record["formula_hash"] == hash_json(artifact)
    assert record["metrics"]["tail"]["rmse"] == 0.05
    assert record["record_hash"] == hash_json({k: v for k, v in record.items() if k != "record_hash"})


def _variant_parent():
    return load_formula(
        build_formula_artifact(
            artifact_id="formula-variant-parent",
            run_id="run-variant",
            role="controller",
            anchor="front",
            training_scenes=["scene-a"],
            feature_names=["x", "y"],
            fitted_expression="x*(y + 2) + 3*x - 4",
            deployment_expression="x*(y + 2) + 3*x - 4",
            output={"name": "alpha", "unit": "rad"},
            feature_metadata={"version": "features-v1"},
            alignment_metadata={"policy": "same-step"},
            dt_c_metadata={"value": 0.2, "unit": "s"},
            g_metadata={"operator": "identity"},
            fit_provenance={"seed": 7},
            pysr_provenance={"version": "test"},
            data_provenance={"dataset_hash": "abc"},
        )
    )


def test_symbolic_normalization_and_term_decomposition_are_deterministic():
    expected = "x*y + 5*x - 4"
    assert normalize_symbolic_expression("3*x - 4 + x*(2 + y)", ["x", "y"]) == expected
    assert normalize_symbolic_expression("x*y - 4 + 5*x", ["x", "y"]) == expected
    assert decompose_additive_terms("3*x - 4 + x*(2 + y)", ["x", "y"]) == (
        "x*y",
        "5*x",
        "-4",
    )


def test_term_variants_are_immutable_validated_and_preserve_provenance():
    parent = _variant_parent()
    original = json.loads(json.dumps(parent))
    term = "5*x"

    deleted = build_term_deletion_variant(parent, term)
    scaled = build_coefficient_scaling_variant(parent, term, 0.5)

    assert parent == original
    assert deleted["fitted_expression"] == parent["fitted_expression"]
    assert deleted["fitted_expression_hash"] == parent["fitted_expression_hash"]
    assert deleted["deployment_expression"] == "x*y - 4"
    assert scaled["deployment_expression"] == "x*y + 2.5*x - 4"
    assert deleted["deployment_expression_hash"] != parent["deployment_expression_hash"]
    assert scaled["deployment_expression_hash"] != parent["deployment_expression_hash"]
    assert deleted["parent_source_hash"] == parent["source_hash"]
    assert deleted["parent_deployment_expression_hash"] == parent["deployment_expression_hash"]
    assert deleted["variant_metadata"] == {
        "type": "term_deletion",
        "term_identity": term,
        "scale": 0.0,
    }
    assert scaled["variant_metadata"] == {
        "type": "coefficient_scaling",
        "term_identity": term,
        "scale": 0.5,
    }
    for key in (
        "role",
        "g_metadata",
        "feature_metadata",
        "alignment_metadata",
        "dt_c_metadata",
        "provenance",
    ):
        assert deleted[key] == parent[key]
        assert scaled[key] == parent[key]
    assert load_formula(deleted)["deployment_expression_hash"] == deleted[
        "deployment_expression_hash"
    ]


def test_zero_scaling_is_symbolically_and_numerically_equivalent_to_deletion():
    parent = _variant_parent()
    term = "5*x"
    deleted = build_term_deletion_variant(parent, term)
    zero_scaled = build_coefficient_scaling_variant(parent, term, 0)

    assert normalize_symbolic_expression(
        deleted["deployment_expression"], parent["feature_names"]
    ) == normalize_symbolic_expression(
        zero_scaled["deployment_expression"], parent["feature_names"]
    )
    values = {"x": np.array([-2.0, 0.5, 3.0]), "y": np.array([1.0, -4.0, 2.0])}
    np.testing.assert_allclose(
        evaluate_formula(deleted, values), evaluate_formula(zero_scaled, values)
    )


def test_scaling_accepts_caller_grid_values_and_rejects_unknown_terms():
    parent = _variant_parent()
    variants = [
        build_coefficient_scaling_variant(parent, "5*x", scale)
        for scale in (-1.25, 0.25, 1.75)
    ]
    assert [variant["variant_metadata"]["scale"] for variant in variants] == [
        -1.25,
        0.25,
        1.75,
    ]
    with pytest.raises(ValueError, match="term identity is not present"):
        build_term_deletion_variant(parent, "x")


def test_stage3_three_head_policy_deploys_each_formula_directly():
    from SR_analysis.stage_3_validate import FormulaPair, SymbolicPolicy

    def formula(role, expression):
        return load_formula({
            "schema_version": "1.0", "feature_names": ["u_m"],
            "role": role, "anchor": role, "fitted_expression": expression,
            "deployment_expression": expression,
        })

    front = formula("front_independent", "1")
    upper = formula("upper_independent", "2")
    lower = formula("lower_independent", "3")
    pair = FormulaPair(
        front, upper, Path("front"), Path("upper"), "hash", {},
        lower=lower, lower_path=Path("lower"), deployment_semantics="three_head_independent",
    )
    cfg = {
        "mu": 0.02, "sample_interval": 800, "u0": 0.01, "scene_id": "karman",
        "policy_init_action": [0, 0, 0],
    }
    omega, _, _ = SymbolicPolicy(pair, cfg, None).action(np.zeros(12), 0)
    np.testing.assert_allclose(omega, [0.01, 0.02, 0.03])


def test_stage3_rollout_failure_preserves_partial_telemetry():
    from SR_analysis.stage_3_validate import RolloutFailure, run_rollout

    class Environment:
        current_raw = np.zeros(12)
        target_sensors = np.zeros((72, 6))

        def step(self, action):
            return np.full(12, np.nan)

    class Policy:
        def action(self, raw, step):
            return np.zeros(3), None, None

        def observe(self, action):
            raise AssertionError("failed observations must not update policy history")

    class Plan:
        n_steps = 3
        cfg = {"u0": 0.01}

    with pytest.raises(RolloutFailure) as captured:
        run_rollout(Plan(), Environment(), Policy())

    assert captured.value.step == 0
    assert captured.value.telemetry["sensors"].shape == (0, 6)
    assert "non-finite" in str(captured.value)


def test_stage3_karman_trained_group_matches_stage1_contract():
    from SR_analysis.stage_3_validate import GROUPS

    assert list(GROUPS["karman_trained"]()) == [
        "karman_re50",
        "karman_re100",
        "karman_re200",
        "karman_re400",
    ]


def test_safe_atomic_json_write_rejects_escape(tmp_path):
    destination = atomic_write_json_below(tmp_path, "nested/artifact.json", {"b": 2, "a": 1})
    assert destination == safe_path(tmp_path, "nested/artifact.json")
    assert json.loads(destination.read_text(encoding="utf-8")) == {"a": 1, "b": 2}

    with pytest.raises(ValueError, match="escapes"):
        safe_path(tmp_path, "../outside.json")
    with pytest.raises(ValueError, match="relative"):
        safe_path(tmp_path, tmp_path / "absolute.json")
