"""Versioned SR formula loading, evaluation, and validation provenance."""
from __future__ import annotations

import ast
import json
from copy import deepcopy
from numbers import Real
from pathlib import Path
from typing import Any, Dict, Mapping, Optional, Sequence, Tuple, Union

import numpy as np
import sympy

from .provenance import hash_file, hash_json, sha256_text

FORMULA_SCHEMA_VERSION = "1.0"
VALIDATION_SCHEMA_VERSION = "1.0"

FormulaSource = Union[str, Path, Mapping[str, Any]]


def _read_formula(source: FormulaSource) -> Tuple[Dict[str, Any], Optional[Path]]:
    if isinstance(source, Mapping):
        return deepcopy(dict(source)), None
    path = Path(source)
    with path.open("r", encoding="utf-8") as stream:
        value = json.load(stream)
    if not isinstance(value, dict):
        raise ValueError("formula JSON must contain an object")
    return value, path


def expression_hash(expression: str) -> str:
    """Hash an expression after trimming inconsequential outer whitespace."""
    if not isinstance(expression, str) or not expression.strip():
        raise ValueError("expression must be a non-empty string")
    return sha256_text(expression.strip())


def _sympy_expression(expression: str, feature_names: Sequence[str]) -> sympy.Expr:
    if not isinstance(expression, str) or not expression.strip():
        raise ValueError("expression must be a non-empty string")
    _validate_expression_syntax(expression, feature_names)
    symbols = {name: sympy.Symbol(name) for name in feature_names}
    parsed = sympy.sympify(expression, locals=symbols)
    unknown = sorted(str(symbol) for symbol in parsed.free_symbols if str(symbol) not in symbols)
    if unknown:
        raise ValueError(f"expression references undeclared features: {unknown}")
    return parsed


def normalize_symbolic_expression(expression: str, feature_names: Sequence[str]) -> str:
    """Return a deterministic expanded SymPy representation of an expression."""
    return sympy.sstr(sympy.expand(_sympy_expression(expression, feature_names)), order="lex")


def decompose_additive_terms(expression: str, feature_names: Sequence[str]) -> Tuple[str, ...]:
    """Return deterministic top-level terms from the normalized expression."""
    normalized = _sympy_expression(
        normalize_symbolic_expression(expression, feature_names), feature_names
    )
    return tuple(sympy.sstr(term, order="lex") for term in normalized.as_ordered_terms())


def _parent_source_hash(source: FormulaSource, loaded: Mapping[str, Any]) -> str:
    if isinstance(source, Mapping) and isinstance(source.get("source_hash"), str):
        return str(source["source_hash"])
    return str(loaded["source_hash"])


def _term_variant(
    formula_source: FormulaSource,
    term_identity: str,
    *,
    variant_type: str,
    scale: Real,
) -> Dict[str, Any]:
    if isinstance(scale, bool) or not isinstance(scale, Real) or not np.isfinite(float(scale)):
        raise ValueError("scale must be a finite real number")
    parent = load_formula(formula_source)
    deployment = parent.get("deployment_expression")
    if deployment is None:
        raise ValueError("formula variants require a symbolic deployment expression")
    feature_names = parent["feature_names"]
    terms = decompose_additive_terms(deployment, feature_names)
    if term_identity not in terms:
        raise ValueError(f"term identity is not present in deployment expression: {term_identity!r}")

    normalized = _sympy_expression(
        normalize_symbolic_expression(deployment, feature_names), feature_names
    )
    selected = _sympy_expression(term_identity, feature_names)
    scaled = float(scale)
    variant_expression = normalize_symbolic_expression(
        sympy.sstr(normalized - selected + sympy.Float(scaled) * selected, order="lex"),
        feature_names,
    )

    variant = deepcopy(parent)
    variant.pop("source_path", None)
    variant.pop("source_hash", None)
    variant["deployment_expression"] = variant_expression
    variant["deployment_expression_hash"] = expression_hash(variant_expression)
    variant["parent_source_hash"] = _parent_source_hash(formula_source, parent)
    variant["parent_deployment_expression_hash"] = parent["deployment_expression_hash"]
    variant["variant_metadata"] = {
        "type": variant_type,
        "term_identity": term_identity,
        "scale": scaled,
    }
    return load_formula(variant)


def build_term_deletion_variant(
    formula_source: FormulaSource, term_identity: str
) -> Dict[str, Any]:
    """Build an immutable deployment variant with one additive term removed."""
    return _term_variant(
        formula_source, term_identity, variant_type="term_deletion", scale=0.0
    )


def build_coefficient_scaling_variant(
    formula_source: FormulaSource, term_identity: str, scale: Real
) -> Dict[str, Any]:
    """Build an immutable deployment variant with one additive term scaled."""
    return _term_variant(
        formula_source, term_identity, variant_type="coefficient_scaling", scale=scale
    )


def load_formula(source: FormulaSource) -> Dict[str, Any]:
    """Load either a legacy formula or the versioned schema into one shape.

    Legacy aliases are retained while canonical fields are added, allowing old
    consumers and provenance-aware consumers to share the same loaded object.
    """
    formula, path = _read_formula(source)
    raw_source_hash = hash_file(path) if path is not None else hash_json(formula)
    legacy = "schema_version" not in formula

    features = formula.get("feature_names", formula.get("feature_keys"))
    if not isinstance(features, list) or not all(isinstance(item, str) for item in features):
        raise ValueError("formula requires string feature_names or feature_keys")

    fitted = formula.get("fitted_expression")
    if fitted is None:
        fitted = formula.get("best_sympy", formula.get("expression"))
    deployment = formula.get("deployment_expression", fitted)

    model_type = formula.get("model_type", formula.get("schema_type", "symbolic"))
    if fitted is None and model_type != "linear":
        raise ValueError("symbolic formula requires an expression")
    if fitted is not None and (not isinstance(fitted, str) or not fitted.strip()):
        raise ValueError("fitted_expression must be a non-empty string")
    if deployment is not None and (not isinstance(deployment, str) or not deployment.strip()):
        raise ValueError("deployment_expression must be a non-empty string")

    formula["schema_version"] = str(formula.get("schema_version", "legacy"))
    formula["feature_names"] = list(features)
    formula["fitted_expression"] = fitted
    formula["deployment_expression"] = deployment
    formula["model_type"] = model_type
    formula["fit_score"] = formula.get("fit_score", formula.get("best_score"))
    formula["legacy"] = legacy
    if fitted is not None:
        expected = expression_hash(fitted)
        if formula.get("fitted_expression_hash") not in (None, expected):
            raise ValueError("fitted_expression_hash does not match expression")
        formula["fitted_expression_hash"] = expected
    if deployment is not None:
        expected = expression_hash(deployment)
        if formula.get("deployment_expression_hash") not in (None, expected):
            raise ValueError("deployment_expression_hash does not match expression")
        formula["deployment_expression_hash"] = expected
    formula["source_path"] = str(path) if path is not None else None
    formula["source_hash"] = raw_source_hash
    return formula


def build_formula_artifact(
    *,
    artifact_id: str,
    run_id: str,
    role: str,
    anchor: str,
    training_scenes: Sequence[str],
    feature_names: Sequence[str],
    fitted_expression: str,
    deployment_expression: str,
    output: Mapping[str, Any],
    feature_metadata: Mapping[str, Any],
    alignment_metadata: Mapping[str, Any],
    dt_c_metadata: Mapping[str, Any],
    g_metadata: Mapping[str, Any],
    fit_provenance: Mapping[str, Any],
    pysr_provenance: Mapping[str, Any],
    data_provenance: Mapping[str, Any],
    **extra: Any,
) -> Dict[str, Any]:
    """Construct a complete symbolic-formula artifact with expression hashes."""
    required_strings = {
        "artifact_id": artifact_id,
        "run_id": run_id,
        "role": role,
        "anchor": anchor,
    }
    for name, value in required_strings.items():
        if not isinstance(value, str) or not value.strip():
            raise ValueError(f"{name} must be a non-empty string")
    if not training_scenes or not all(isinstance(item, str) and item for item in training_scenes):
        raise ValueError("training_scenes must contain at least one scene")
    if not feature_names or not all(isinstance(item, str) and item for item in feature_names):
        raise ValueError("feature_names must contain at least one feature")

    artifact: Dict[str, Any] = {
        "schema_version": FORMULA_SCHEMA_VERSION,
        **required_strings,
        "training_scenes": list(training_scenes),
        "feature_names": list(feature_names),
        "model_type": "symbolic",
        "fitted_expression": fitted_expression.strip(),
        "deployment_expression": deployment_expression.strip(),
        "fitted_expression_hash": expression_hash(fitted_expression),
        "deployment_expression_hash": expression_hash(deployment_expression),
        "output": dict(output),
        "feature_metadata": dict(feature_metadata),
        "alignment_metadata": dict(alignment_metadata),
        "dt_c_metadata": dict(dt_c_metadata),
        "g_metadata": dict(g_metadata),
        "provenance": {
            "fit": dict(fit_provenance),
            "pysr": dict(pysr_provenance),
            "data": dict(data_provenance),
        },
    }
    artifact.update(extra)
    load_formula(artifact)
    return artifact


def _feature_columns(
    feature_values: Union[Mapping[str, Any], np.ndarray, Sequence[Sequence[float]]],
    feature_names: Sequence[str],
) -> Tuple[Dict[str, np.ndarray], Tuple[int, ...]]:
    if isinstance(feature_values, Mapping):
        missing = [name for name in feature_names if name not in feature_values]
        if missing:
            raise KeyError(f"missing formula features: {missing}")
        raw_columns = [np.asarray(feature_values[name]) for name in feature_names]
        broadcast = np.broadcast_arrays(*raw_columns) if raw_columns else []
        columns = {name: broadcast[index] for index, name in enumerate(feature_names)}
        shape = broadcast[0].shape if broadcast else ()
        return columns, shape
    matrix = np.asarray(feature_values)
    if matrix.ndim == 0 or matrix.shape[-1] != len(feature_names):
        raise ValueError("feature matrix final dimension must match feature_names")
    return {name: matrix[..., index] for index, name in enumerate(feature_names)}, matrix.shape[:-1]


_ALLOWED_EXPRESSION_FUNCTIONS = {"sin", "cos", "tan", "exp", "log", "sqrt", "Abs", "abs"}
_ALLOWED_AST_NODES = (
    ast.Expression,
    ast.BinOp,
    ast.UnaryOp,
    ast.Add,
    ast.Sub,
    ast.Mult,
    ast.Div,
    ast.Pow,
    ast.Mod,
    ast.USub,
    ast.UAdd,
    ast.Name,
    ast.Load,
    ast.Constant,
    ast.Call,
)


def _validate_expression_syntax(expression: str, feature_names: Sequence[str]) -> None:
    try:
        tree = ast.parse(expression, mode="eval")
    except SyntaxError as exc:
        raise ValueError("invalid symbolic expression") from exc
    allowed_names = set(feature_names) | _ALLOWED_EXPRESSION_FUNCTIONS
    for node in ast.walk(tree):
        if not isinstance(node, _ALLOWED_AST_NODES):
            raise ValueError(f"unsupported expression syntax: {type(node).__name__}")
        if isinstance(node, ast.Name) and node.id not in allowed_names:
            raise ValueError(f"expression references undeclared name: {node.id}")
        if isinstance(node, ast.Call):
            if not isinstance(node.func, ast.Name) or node.func.id not in _ALLOWED_EXPRESSION_FUNCTIONS:
                raise ValueError("expression calls an unsupported function")
            if node.keywords:
                raise ValueError("expression function keywords are unsupported")


def formula_safety_report(
    formula_source: FormulaSource,
    feature_ranges: Mapping[str, Sequence[float]],
    *,
    expression_kind: str = "deployment",
    grid_size: int = 9,
    output_range: Optional[Sequence[float]] = None,
) -> Dict[str, Any]:
    """Statically probe a formula over declared feature ranges for deployment safety."""
    formula = load_formula(formula_source)
    if grid_size < 2:
        raise ValueError("grid_size must be at least two")
    missing = [name for name in formula["feature_names"] if name not in feature_ranges]
    if missing:
        raise ValueError(f"missing feature ranges: {missing}")
    probes: Dict[str, np.ndarray] = {}
    for name in formula["feature_names"]:
        bounds = np.asarray(feature_ranges[name], dtype=float)
        if bounds.shape != (2,) or not np.all(np.isfinite(bounds)) or bounds[0] > bounds[1]:
            raise ValueError(f"invalid finite range for feature {name!r}")
        probes[name] = np.linspace(bounds[0], bounds[1], grid_size)
    with np.errstate(all="ignore"):
        values = np.asarray(evaluate_formula(formula, probes, expression_kind=expression_kind), dtype=float)
    finite = bool(np.all(np.isfinite(values)))
    observed = None if not finite or values.size == 0 else [float(values.min()), float(values.max())]
    within_range = True
    expected = None
    if output_range is not None:
        bounds = np.asarray(output_range, dtype=float)
        if bounds.shape != (2,) or not np.all(np.isfinite(bounds)) or bounds[0] > bounds[1]:
            raise ValueError("output_range must be finite ordered bounds")
        expected = bounds.tolist()
        within_range = finite and bool(np.all((values >= bounds[0]) & (values <= bounds[1])))
    return {
        "safe": finite and within_range,
        "finite": finite,
        "within_output_range": within_range,
        "observed_range": observed,
        "required_output_range": expected,
        "grid_size": int(grid_size),
    }


def evaluate_formula(
    formula_source: FormulaSource,
    feature_values: Union[Mapping[str, Any], np.ndarray, Sequence[Sequence[float]]],
    *,
    expression_kind: str = "deployment",
) -> np.ndarray:
    """Evaluate symbolic formulas first; coefficients require explicit linear type."""
    formula = load_formula(formula_source)
    feature_names = formula["feature_names"]
    columns, output_shape = _feature_columns(feature_values, feature_names)
    expression_key = f"{expression_kind}_expression"
    expression = formula.get(expression_key)

    if expression is not None:
        _validate_expression_syntax(expression, feature_names)
        symbols = {name: sympy.Symbol(name) for name in feature_names}
        parsed = sympy.sympify(expression, locals=symbols)
        unknown = sorted(str(symbol) for symbol in parsed.free_symbols if str(symbol) not in symbols)
        if unknown:
            raise ValueError(f"expression references undeclared features: {unknown}")
        function = sympy.lambdify([symbols[name] for name in feature_names], parsed, modules="numpy")
        result = np.asarray(function(*(columns[name] for name in feature_names)), dtype=float)
        if result.shape == () and output_shape:
            result = np.full(output_shape, float(result), dtype=float)
        return result

    if formula.get("model_type") != "linear":
        raise ValueError("formula has no symbolic expression and is not explicit linear schema")
    coefficients = formula.get("coef", formula.get("coefficients"))
    if coefficients is None:
        raise ValueError("explicit linear formula requires coef or coefficients")
    coefficient_array = np.asarray(coefficients, dtype=float)
    if coefficient_array.shape != (len(feature_names),):
        raise ValueError("linear coefficient count must match feature_names")
    matrix = np.stack([np.broadcast_to(columns[name], output_shape) for name in feature_names], axis=-1)
    intercept = formula.get("intercept", 0.0)
    return np.asarray(matrix @ coefficient_array + float(intercept), dtype=float)


def build_validation_record(
    *,
    formula: FormulaSource,
    metric_version: str,
    scene: str,
    mode: str,
    si: bool,
    n_steps: int,
    full_metrics: Mapping[str, Any],
    tail_metrics: Mapping[str, Any],
    per_channel_metrics: Mapping[str, Any],
    action_range: Mapping[str, Any],
    telemetry: Mapping[str, Any],
    provenance: Mapping[str, Any],
    **extra: Any,
) -> Dict[str, Any]:
    """Build a validation record tied to exact formula and metric semantics."""
    loaded = load_formula(formula)
    if not metric_version or not scene or not mode or n_steps < 0:
        raise ValueError("metric_version, scene, mode, and non-negative n_steps are required")
    record: Dict[str, Any] = {
        "schema_version": VALIDATION_SCHEMA_VERSION,
        "formula_hash": loaded["source_hash"],
        "formula_artifact_id": loaded.get("artifact_id"),
        "metric_version": metric_version,
        "scene": scene,
        "mode": mode,
        "si": bool(si),
        "n_steps": int(n_steps),
        "metrics": {
            "full": dict(full_metrics),
            "tail": dict(tail_metrics),
            "per_channel": dict(per_channel_metrics),
            "action_range": dict(action_range),
        },
        "telemetry": dict(telemetry),
        "provenance": dict(provenance),
    }
    record.update(extra)
    record["record_hash"] = hash_json(record)
    return record
