#!/usr/bin/env python3
"""Stage 2 symbolic fitting with explicit trajectory and artifact contracts.

The module intentionally delays importing PySR so ``--prepare-only`` and unit
 tests can validate all inputs without Julia being installed.
"""
from __future__ import annotations

import argparse
import ast
import json
import os
import inspect
import platform
import re
import subprocess
import sys
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Callable, Mapping, Sequence

import numpy as np
import sympy

REPO_ROOT = Path(__file__).resolve().parents[2]
SRC_ROOT = REPO_ROOT / "src"
SR_ROOT = SRC_ROOT / "SR_analysis"
if str(SRC_ROOT) not in sys.path:
    sys.path.insert(0, str(SRC_ROOT))

from SR_analysis.configs import ACTION_LAYOUT, control_dt, get_scene
from SR_analysis.utils.data_contracts import (
    ALIGNMENTS,
    CAUSAL_ALIGNMENT,
    BatchDataset,
    build_batch_dataset,
    contiguous_block_slices,
)
from SR_analysis.utils.feature_builder import FEATURE_SET_KEYS, ILLUSION_TARGET_KEYS
from SR_analysis.utils.formula_schema import build_formula_artifact, formula_safety_report
from SR_analysis.utils.provenance import atomic_write_json, hash_file, hash_json
from SR_analysis.utils.symmetry_dataset import build_mirrored_batch_dataset

DEFAULT_DATA_ROOT = SR_ROOT / "data"
DEFAULT_RESULTS_ROOT = SR_ROOT / "results" / "runs"
G_VERSION = "legacy-mirror-front-odd-rear-swap-negate-v1"
BINARY_OPERATORS = ("+", "-", "*")
UNARY_OPERATORS = ("square",)
DEPLOYMENT_ARCHITECTURES = (
    "mapped_shared",
    "three_head_independent",
)
FIT_AUGMENTATIONS = ("none", "G")
FIT_PURPOSES = ("discovery", "fixed-topology-refit")
FEATURE_PROFILES = (
    "default",
    "actual_only",
    "target_only",
    "actual_plus_target",
    "actual_plus_error",
)
COMPARATIVE_FEATURE_PROFILES = FEATURE_PROFILES[1:]


@dataclass(frozen=True)
class PreparedTrajectory:
    trajectory_id: str
    data_path: Path
    config_path: Path | None
    dataset: BatchDataset
    mirrored_dataset: BatchDataset
    data_hash: str


@dataclass(frozen=True)
class PreparedScene:
    name: str
    config: dict[str, Any]
    data_path: Path
    config_path: Path | None
    dataset: BatchDataset
    mirrored_dataset: BatchDataset
    feature_names: tuple[str, ...]
    action_layout: tuple[str, ...]
    data_hash: str
    trajectories: tuple[PreparedTrajectory, ...] = ()


def default_run_id() -> str:
    return datetime.now(timezone.utc).strftime("stage2-%Y%m%dT%H%M%SZ")


def resolve_repo_path(value: str | os.PathLike[str] | None, default: Path) -> Path:
    if value is None:
        return default.resolve()
    path = Path(value).expanduser()
    return (path if path.is_absolute() else REPO_ROOT / path).resolve()


def _safe_run_id(run_id: str) -> str:
    if not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9._-]*", run_id):
        raise ValueError("run_id must contain only letters, digits, '.', '_' or '-'")
    return run_id


def formula_output_dir(results_root: Path, run_id: str) -> Path:
    root = results_root.resolve()
    destination = (root / _safe_run_id(run_id) / "formulas").resolve()
    destination.relative_to(root)
    return destination


def refuse_existing_output(path: Path) -> None:
    if path.exists():
        raise FileExistsError(f"Refusing to overwrite existing output directory: {path}")


def _parse_sequence(value: Any) -> tuple[Any, ...]:
    if isinstance(value, str):
        value = ast.literal_eval(value)
    if not isinstance(value, (list, tuple)):
        raise ValueError(f"expected sequence contract, got {value!r}")
    return tuple(value)


def _load_scene_config(scene: str, config_path: Path | None) -> dict[str, Any]:
    config = get_scene(scene)
    if config_path is not None:
        loaded = json.loads(config_path.read_text(encoding="utf-8"))
        if not isinstance(loaded, dict):
            raise ValueError(f"scene config must be a JSON object: {config_path}")
        config.update(loaded)
    for key in ("action_bias", "action_layout", "body_layout"):
        if key in config:
            config[key] = _parse_sequence(config[key])
    config.setdefault("action_layout", tuple(ACTION_LAYOUT))
    return config


def resolve_scene_files(data_root: Path, scene: str, scene_id: str) -> tuple[Path, Path | None]:
    candidates = (
        data_root / scene_id / scene,
        data_root / scene,
        data_root,
    )
    for directory in candidates:
        controlled = directory / "controlled.npz"
        if controlled.is_file():
            config = directory / "config.json"
            return controlled, config if config.is_file() else None
    raise FileNotFoundError(
        f"controlled.npz for {scene!r} not found below {data_root}; expected canonical "
        f"{scene_id}/{scene}/controlled.npz or a run/scene directory"
    )


def feature_names_for_scene(
    config: Mapping[str, Any], *, feature_set: str, joint: bool,
    feature_profile: str = "default",
) -> tuple[str, ...]:
    if feature_set not in FEATURE_SET_KEYS:
        raise ValueError(f"unknown feature set {feature_set!r}; expected one of {tuple(FEATURE_SET_KEYS)}")
    if feature_profile not in FEATURE_PROFILES:
        raise ValueError(f"unknown feature profile {feature_profile!r}; expected one of {FEATURE_PROFILES}")
    base = tuple(FEATURE_SET_KEYS[feature_set])
    if config.get("scene_id") != "illusion":
        if feature_profile not in {"default", "actual_only"}:
            raise ValueError("target/error feature profiles are only defined for Illusion")
        return base
    profile_features = {
        # ``default`` is retained as a non-comparative compatibility profile.
        "default": tuple(ILLUSION_TARGET_KEYS) + ("Cd_err", "Cl_err"),
        "actual_only": (),
        "target_only": tuple(ILLUSION_TARGET_KEYS),
        "actual_plus_target": tuple(ILLUSION_TARGET_KEYS),
        "actual_plus_error": ("Cd_err", "Cl_err"),
    }
    if feature_profile == "target_only":
        return profile_features[feature_profile]
    return base + profile_features[feature_profile]


def _npz_scalar(data: Mapping[str, Any], key: str, default: Any) -> Any:
    if key not in data:
        return default
    value = np.asarray(data[key])
    return value.item() if value.shape == () else default


def prepare_scene(
    scene: str, data_root: Path, *, alignment: str, joint: bool,
    feature_set: str = "physics_reduced", feature_profile: str = "default",
    sample_segment: str = "all", transient_count: int = 50,
) -> PreparedScene:
    baseline = get_scene(scene)
    data_path, config_path = resolve_scene_files(data_root, scene, str(baseline["scene_id"]))
    config = _load_scene_config(scene, config_path)
    feature_names = feature_names_for_scene(
        config, feature_set=feature_set, joint=joint, feature_profile=feature_profile,
    )
    with np.load(data_path, allow_pickle=False) as data:
        sensors_key = "sensors_post" if "sensors_post" in data else "sensors"
        forces_key = "forces_post" if "forces_post" in data else "forces"
        if "actions_alpha_commanded" in data:
            actions = np.asarray(data["actions_alpha_commanded"])
            action_mode = "alpha"
        elif "actions_norm" in data:
            actions = np.asarray(data["actions_norm"])
            action_mode = "normalized"
        elif "actions" in data:
            actions = np.asarray(data["actions"])
            action_mode = "normalized"
        else:
            raise KeyError(f"controlled trajectory missing actions: {data_path}")
        target = np.asarray(data["target_forces"]) if "target_forces" in data else None
        dt_c = float(_npz_scalar(data, "dt_c", config.get("control_dt", control_dt(config["sample_interval"]))))
        common = {
            "mu": float(config["mu"]), "dt_c": dt_c, "alignment": alignment,
            "action_mode": action_mode, "action_scale": float(config.get("action_scale", 1.0)),
            "action_bias": config.get("action_bias", (0.0, 0.0, 0.0)),
            "illusion": config.get("scene_id") == "illusion",
            "u0": float(config.get("u0", 0.01)), "include_cos_sin": False,
            "add_bias": False, "strict": True, "sample_segment": sample_segment,
            "transient_count": transient_count,
        }
        sensors = np.asarray(data[sensors_key])
        forces = np.asarray(data[forces_key])
        dataset = build_batch_dataset(
            sensors, forces, actions, feature_names, target_forces=target, **common,
        )
        mirrored_dataset = build_mirrored_batch_dataset(
            sensors, forces, actions, feature_names, target_forces=target, **common,
        )
    layout = tuple(str(item).lower() for item in config.get("action_layout", ACTION_LAYOUT))
    aliases = {"top": "upper", "bottom": "lower"}
    layout = tuple(aliases.get(item, item) for item in layout)
    if len(layout) != 3 or set(layout) != {"front", "upper", "lower"}:
        raise ValueError(f"invalid native action_layout for {scene}: {layout}")
    if dataset.X.shape[0] == 0:
        raise ValueError(f"no aligned samples remain for {scene}")
    data_hash = hash_json({"controlled_sha256": hash_file(data_path), "config": config})
    trajectory = PreparedTrajectory(
        f"{scene}:trajectory-001", data_path, config_path, dataset, mirrored_dataset, data_hash,
    )
    return PreparedScene(
        scene, config, data_path, config_path, dataset, mirrored_dataset,
        feature_names, layout, data_hash, (trajectory,),
    )


def _stack_batch_datasets(datasets: Sequence[BatchDataset]) -> BatchDataset:
    if not datasets:
        raise ValueError("at least one trajectory dataset is required")
    return BatchDataset(
        X=np.vstack([dataset.X for dataset in datasets]),
        Y=np.vstack([dataset.Y for dataset in datasets]),
        kept_indices=np.concatenate([dataset.kept_indices for dataset in datasets]),
        dropped_indices=np.concatenate([dataset.dropped_indices for dataset in datasets]),
        metadata=dict(datasets[0].metadata),
        features={},
    )


def prepare_scene_trajectories(
    scene: str, data_roots: Sequence[Path], *, alignment: str, joint: bool,
    feature_set: str = "physics_reduced", feature_profile: str = "default",
    sample_segment: str = "all", transient_count: int = 50,
) -> PreparedScene:
    if not data_roots:
        raise ValueError("at least one data root is required")
    prepared = [
        prepare_scene(
            scene, root, alignment=alignment, joint=joint, feature_set=feature_set,
            feature_profile=feature_profile, sample_segment=sample_segment,
            transient_count=transient_count,
        )
        for root in data_roots
    ]
    first = prepared[0]
    for item in prepared[1:]:
        if item.feature_names != first.feature_names or item.action_layout != first.action_layout:
            raise ValueError(f"trajectory contracts differ for scene {scene}")
        if item.config.get("scene_id") != first.config.get("scene_id"):
            raise ValueError(f"trajectory scene_id contracts differ for scene {scene}")
    trajectories = tuple(
        PreparedTrajectory(
            f"{scene}:trajectory-{index:03d}", item.data_path, item.config_path,
            item.dataset, item.mirrored_dataset, item.data_hash,
        )
        for index, item in enumerate(prepared, start=1)
    )
    return PreparedScene(
        scene, first.config, first.data_path, first.config_path,
        _stack_batch_datasets([item.dataset for item in trajectories]),
        _stack_batch_datasets([item.mirrored_dataset for item in trajectories]),
        first.feature_names, first.action_layout,
        hash_json({item.trajectory_id: item.data_hash for item in trajectories}), trajectories,
    )


def _git_sha() -> str | None:
    try:
        return subprocess.run(
            ["git", "-C", str(REPO_ROOT), "rev-parse", "HEAD"], check=True,
            capture_output=True, text=True,
        ).stdout.strip()
    except (OSError, subprocess.CalledProcessError):
        return None


def _environment() -> dict[str, Any]:
    return {
        "python": platform.python_version(),
        "numpy": np.__version__,
        "platform": platform.platform(),
        "pysr": None,
    }


def _regressor_factory() -> Any:
    from pysr import PySRRegressor
    return PySRRegressor


def regressor_parameters(*, seed: int, niterations: int, smoke: bool) -> dict[str, Any]:
    return {
        "binary_operators": list(BINARY_OPERATORS),
        "unary_operators": list(UNARY_OPERATORS),
        "niterations": int(niterations),
        "populations": 2 if smoke else 30,
        "population_size": 16 if smoke else 33,
        "tournament_selection_n": 4 if smoke else 10,
        "maxsize": 7 if smoke else 15,
        "complexity_of_constants": 2,
        "parsimony": 0.01,
        "batching": False,
        "random_state": int(seed),
        "deterministic": True,
        "parallelism": "serial",
    }


def _predict(model: Any, X: np.ndarray) -> np.ndarray:
    prediction = np.asarray(model.predict(X), dtype=np.float64).reshape(-1)
    if prediction.shape != (X.shape[0],):
        raise ValueError(f"regressor prediction has invalid shape {prediction.shape}")
    return prediction


def diagnostic_metrics(y: np.ndarray, prediction: np.ndarray) -> dict[str, float]:
    residual = np.asarray(y) - np.asarray(prediction)
    denominator = float(np.sum((y - np.mean(y)) ** 2))
    r2 = 1.0 - float(np.sum(residual ** 2)) / denominator if denominator > 0 else 0.0
    return {
        "r2": r2,
        "mae": float(np.mean(np.abs(residual))),
        "rmse": float(np.sqrt(np.mean(residual ** 2))),
        "max_error": float(np.max(np.abs(residual))),
    }


def _expression(model: Any) -> str:
    expression = str(model.sympy()).strip()
    if not expression:
        raise ValueError("regressor returned an empty symbolic expression")
    return expression


def _combined_data_hash(scenes: Sequence[PreparedScene]) -> str:
    return hash_json({scene.name: scene.data_hash for scene in scenes})


def validate_deployment_architecture(value: str) -> str:
    if value not in DEPLOYMENT_ARCHITECTURES:
        raise ValueError(f"deployment architecture must be one of {DEPLOYMENT_ARCHITECTURES}")
    return value


def validate_fit_augmentation(value: str) -> str:
    if value not in FIT_AUGMENTATIONS:
        raise ValueError(f"fit augmentation must be one of {FIT_AUGMENTATIONS}")
    return value


def _scene_trajectories(scene: PreparedScene) -> tuple[PreparedTrajectory, ...]:
    if scene.trajectories:
        return scene.trajectories
    return (PreparedTrajectory(
        f"{scene.name}:trajectory-001", scene.data_path, scene.config_path,
        scene.dataset, scene.mirrored_dataset, scene.data_hash,
    ),)


def _head_training_blocks(
    scenes: Sequence[PreparedScene], role: str, channel_index: int,
    *, fit_augmentation: str = "G", split: str | None = None,
) -> list[tuple[str, str, np.ndarray, np.ndarray]]:
    augmentation = validate_fit_augmentation(fit_augmentation)
    blocks = []
    for scene in scenes:
        for trajectory in _scene_trajectories(scene):
            X = trajectory.dataset.X
            GX = trajectory.mirrored_dataset.X
            if GX.shape != X.shape:
                raise ValueError("raw-state mirrored dataset does not align with source dataset")
            y = trajectory.dataset.Y[:, channel_index]
            gy = trajectory.mirrored_dataset.Y[:, channel_index]
            if role == "front_odd":
                np.testing.assert_allclose(gy, -y)
            elif role == "rear_shared":
                lower_index = scene.action_layout.index("lower")
                np.testing.assert_allclose(gy, -trajectory.dataset.Y[:, lower_index])
            elif role not in {"front_independent", "upper_independent", "lower_independent"}:
                raise ValueError(f"unknown fitted head role: {role}")
            block = slice(None) if split is None else contiguous_block_slices(len(X))[split]
            blocks.append((scene.name, trajectory.trajectory_id, X[block], y[block]))
            if augmentation == "G":
                blocks.append((scene.name, trajectory.trajectory_id, GX[block], gy[block]))
    return blocks


def _head_training_data(
    scenes: Sequence[PreparedScene], role: str, channel_index: int,
    *, fit_augmentation: str = "G", split: str | None = None,
) -> tuple[np.ndarray, np.ndarray]:
    blocks = _head_training_blocks(
        scenes, role, channel_index, fit_augmentation=fit_augmentation, split=split,
    )
    return np.vstack([item[2] for item in blocks]), np.concatenate([item[3] for item in blocks])


def _joint_sample_weights(
    blocks: Sequence[tuple[str, str, np.ndarray, np.ndarray]],
) -> np.ndarray:
    cases = tuple(dict.fromkeys(item[0] for item in blocks))
    trajectories = {
        case: tuple(dict.fromkeys(item[1] for item in blocks if item[0] == case))
        for case in cases
    }
    weights = []
    for case, trajectory_id, X, _ in blocks:
        total_for_trajectory = sum(
            len(other_X) for other_case, other_id, other_X, _ in blocks
            if other_case == case and other_id == trajectory_id
        )
        mass = 1.0 / (len(cases) * len(trajectories[case]))
        weights.append(np.full(len(X), mass / total_for_trajectory, dtype=np.float64))
    result = np.concatenate(weights)
    return result * (len(result) / result.sum())


def _parameterize_numeric_constants(
    expression: str, feature_names: Sequence[str],
) -> tuple[ast.Expression, tuple[str, ...], np.ndarray]:
    """Replace every numeric literal with a distinct optimization parameter."""
    tree = ast.parse(expression, mode="eval")
    allowed_names = set(feature_names) | {"sin", "cos", "tan", "exp", "log", "sqrt", "Abs", "abs"}
    parameter_names: list[str] = []
    initial_values: list[float] = []

    class Parameterizer(ast.NodeTransformer):
        def visit_Constant(self, node: ast.Constant) -> ast.AST:
            if isinstance(node.value, bool) or not isinstance(node.value, (int, float)):
                raise ValueError("topology expressions may contain only numeric constants")
            name = f"__constant_{len(parameter_names)}"
            parameter_names.append(name)
            initial_values.append(float(node.value))
            return ast.copy_location(ast.Name(id=name, ctx=ast.Load()), node)

        def visit_Name(self, node: ast.Name) -> ast.AST:
            if node.id not in allowed_names:
                raise ValueError(f"topology expression references undeclared feature {node.id!r}")
            return node

    parameterized = Parameterizer().visit(tree)
    ast.fix_missing_locations(parameterized)
    if not parameter_names:
        raise ValueError("fixed-topology refit requires at least one numeric constant")
    return parameterized, tuple(parameter_names), np.asarray(initial_values, dtype=np.float64)


def _topology_evaluator(
    expression: str, feature_names: Sequence[str],
) -> tuple[Callable[[np.ndarray, np.ndarray], np.ndarray], ast.Expression, tuple[str, ...], np.ndarray]:
    tree, parameter_names, initial = _parameterize_numeric_constants(expression, feature_names)
    symbols = {name: sympy.Symbol(name) for name in (*feature_names, *parameter_names)}
    parsed = sympy.sympify(ast.unparse(tree), locals=symbols, evaluate=False)
    function = sympy.lambdify(
        [symbols[name] for name in (*feature_names, *parameter_names)], parsed, modules="numpy",
    )

    def evaluate(X: np.ndarray, constants: np.ndarray) -> np.ndarray:
        values = function(
            *[X[:, index] for index in range(X.shape[1])], *[float(value) for value in constants],
        )
        prediction = np.asarray(values, dtype=np.float64)
        if prediction.ndim == 0:
            prediction = np.full(len(X), float(prediction))
        prediction = prediction.reshape(-1)
        if prediction.shape != (len(X),) or not np.all(np.isfinite(prediction)):
            raise ValueError("topology expression produced invalid predictions")
        return prediction

    return evaluate, tree, parameter_names, initial


def _render_refitted_expression(
    tree: ast.Expression, parameter_names: Sequence[str], constants: Sequence[float],
) -> str:
    replacements = dict(zip(parameter_names, constants))

    class Renderer(ast.NodeTransformer):
        def visit_Name(self, node: ast.Name) -> ast.AST:
            if node.id not in replacements:
                return node
            return ast.copy_location(ast.Constant(value=float(replacements[node.id])), node)

    rendered = Renderer().visit(ast.fix_missing_locations(ast.parse(ast.unparse(tree), mode="eval")))
    rendered_text = ast.unparse(ast.fix_missing_locations(rendered))
    return str(sympy.simplify(sympy.sympify(rendered_text, evaluate=True)))


def _refit_metrics(
    blocks: Sequence[tuple[str, str, np.ndarray, np.ndarray]],
    evaluate: Callable[[np.ndarray], np.ndarray],
) -> dict[str, Any]:
    aggregate_X = np.vstack([block[2] for block in blocks])
    aggregate_y = np.concatenate([block[3] for block in blocks])
    cases = tuple(dict.fromkeys(block[0] for block in blocks))
    trajectory_keys = tuple(dict.fromkeys((block[0], block[1]) for block in blocks))

    def metrics_for(selected: Sequence[tuple[str, str, np.ndarray, np.ndarray]]) -> dict[str, float]:
        X = np.vstack([block[2] for block in selected])
        y = np.concatenate([block[3] for block in selected])
        return diagnostic_metrics(y, evaluate(X))

    return {
        "aggregate": diagnostic_metrics(aggregate_y, evaluate(aggregate_X)),
        "per_case": {
            case: metrics_for([block for block in blocks if block[0] == case]) for case in cases
        },
        "per_trajectory": {
            trajectory_id: metrics_for([
                block for block in blocks if block[0] == case and block[1] == trajectory_id
            ])
            for case, trajectory_id in trajectory_keys
        },
    }


def _candidate_prediction(
    record: Mapping[str, Any], expression: Any, X: np.ndarray, feature_names: Sequence[str],
) -> np.ndarray:
    callable_expression = record.get("lambda_format")
    if callable(callable_expression):
        try:
            value = callable_expression(X)
        except TypeError:
            value = callable_expression(*[X[:, index] for index in range(X.shape[1])])
    else:
        symbols = [sympy.Symbol(name) for name in feature_names]
        parsed = sympy.sympify(expression, locals=dict(zip(feature_names, symbols)))
        function = sympy.lambdify(symbols, parsed, modules="numpy")
        value = function(*[X[:, index] for index in range(X.shape[1])])
    prediction = np.asarray(value, dtype=np.float64)
    if prediction.ndim == 0:
        prediction = np.full(len(X), float(prediction))
    prediction = prediction.reshape(-1)
    if prediction.shape != (len(X),) or not np.all(np.isfinite(prediction)):
        raise ValueError(f"candidate prediction has invalid values or shape {prediction.shape}")
    return prediction


def _pareto_candidates(
    model: Any, split_data: Mapping[str, tuple[np.ndarray, np.ndarray]],
    feature_names: Sequence[str], feature_ranges: Mapping[str, Sequence[float]], limit: int = 25,
) -> list[dict[str, Any]]:
    equations = getattr(model, "equations_", None)
    if equations is None:
        return []
    records = equations.to_dict("records") if hasattr(equations, "to_dict") else list(equations)
    candidates = []
    symbols = {name: sympy.Symbol(name) for name in feature_names}
    for record in records[:limit]:
        if not isinstance(record, Mapping):
            continue
        expression = record.get("sympy_format", record.get("equation"))
        if expression is None:
            continue
        expression_text = str(expression)
        candidate = {"expression": expression_text}
        for key in ("complexity", "loss", "score"):
            value = record.get(key)
            if value is not None and np.isscalar(value) and np.isfinite(value):
                candidate[key] = float(value)
        try:
            parsed = sympy.sympify(expression, locals=symbols)
            candidate["used_features"] = sorted(str(symbol) for symbol in parsed.free_symbols)
            candidate["static_safety"] = formula_safety_report(
                {"feature_names": list(feature_names), "expression": expression_text}, feature_ranges,
            )
            candidate["metrics"] = {
                name: diagnostic_metrics(target, _candidate_prediction(record, expression, matrix, feature_names))
                for name, (matrix, target) in split_data.items()
            }
        except Exception as exc:
            candidate["evaluation_error"] = f"{type(exc).__name__}: {exc}"
            candidate.setdefault("used_features", [])
            candidate.setdefault("static_safety", {"safe": False, "errors": [candidate["evaluation_error"]]})
            candidate["metrics"] = {name: None for name in split_data}
        candidates.append(candidate)
    return candidates


def refit_fixed_topology_head(
    scenes: Sequence[PreparedScene], *, role: str, channel_index: int, run_id: str,
    topology_expression: str, discovery_parent_path: Path | None,
    mode: str, deployment_architecture: str, feature_set: str, feature_profile: str,
    environment: Mapping[str, Any],
) -> dict[str, Any]:
    try:
        from scipy.optimize import least_squares
    except ImportError as exc:
        raise RuntimeError("fixed-topology refit requires scipy in the active environment") from exc

    architecture = validate_deployment_architecture(deployment_architecture)
    augmentation = "G" if architecture == "mapped_shared" else "none"
    blocks = _head_training_blocks(
        scenes, role, channel_index, fit_augmentation=augmentation, split=None,
    )
    X = np.vstack([block[2] for block in blocks])
    y = np.concatenate([block[3] for block in blocks])
    weights = _joint_sample_weights(blocks)
    evaluate, tree, parameter_names, initial = _topology_evaluator(
        topology_expression, scenes[0].feature_names,
    )
    result = least_squares(
        lambda constants: np.sqrt(weights) * (evaluate(X, constants) - y), initial,
    )
    if not result.success or not np.all(np.isfinite(result.x)):
        raise RuntimeError(f"fixed-topology least-squares refit failed: {result.message}")
    fitted = _render_refitted_expression(tree, parameter_names, result.x)
    predictor = lambda matrix: evaluate(matrix, result.x)
    metrics = _refit_metrics(blocks, predictor)
    feature_ranges = {
        name: [float(X[:, index].min()), float(X[:, index].max())]
        for index, name in enumerate(scenes[0].feature_names)
    }
    safety = formula_safety_report(
        {"feature_names": list(scenes[0].feature_names), "expression": fitted}, feature_ranges,
    )
    counts = {scene.name: int(scene.dataset.X.shape[0]) for scene in scenes}
    parent = None if discovery_parent_path is None else discovery_parent_path.resolve()
    if parent is not None and not parent.is_file():
        raise FileNotFoundError(f"discovery parent artifact not found: {parent}")
    config = scenes[0].config
    alignment = scenes[0].dataset.metadata
    env = dict(environment)
    return build_formula_artifact(
        artifact_id=f"{run_id}-{mode}-{role}-anchor-upper", run_id=run_id, role=role,
        anchor="upper", training_scenes=[scene.name for scene in scenes],
        feature_names=scenes[0].feature_names, fitted_expression=fitted,
        deployment_expression=fitted,
        output={
            "name": "alpha", "definition": "cylinder_surface_tangential_velocity/U0",
            "unit": "dimensionless", "native_action_index": channel_index,
            "native_action_order": list(scenes[0].action_layout),
        },
        feature_metadata={
            "builder": "SR_analysis.utils.data_contracts.build_batch_dataset",
            "definitions_module": "SR_analysis.utils.feature_builder",
            "order": list(scenes[0].feature_names),
        },
        alignment_metadata={
            "alignment": alignment["alignment"], "warmup": alignment["warmup"],
            "sample_segments": alignment.get("sample_segments"),
            "history_policy": alignment["history_policy"],
            "required_target": config.get("scene_id") == "illusion",
        },
        dt_c_metadata={
            "per_scene": {scene.name: float(scene.dataset.metadata["dt_c"]) for scene in scenes},
            "unit": "D/U0",
        },
        g_metadata={
            "version": G_VERSION, "deployment_architecture": architecture,
            "deployment_semantics": architecture, "fit_augmentation": augmentation,
            "front_rule": "front(x)=0.5*(f(x)-f(Gx)); antisymmetric projection" if architecture == "mapped_shared" else None,
            "rear_rule": "lower(x)=-upper(Gx)" if architecture == "mapped_shared" else None,
            "fitted_heads": [role],
        },
        fit_provenance={
            "git_sha": _git_sha(), "environment": env, "command": " ".join(sys.argv),
            "created_at": datetime.now(timezone.utc).isoformat(), "mode": mode,
            "fit_purpose": "fixed-topology-refit", "feature_set": feature_set,
            "feature_profile": feature_profile, "n_samples": int(len(X)),
            "source_n_samples": int(sum(counts.values())), "scene_sample_counts": counts,
            "g_augmented": augmentation == "G", "fit_augmentation": augmentation,
            "deployment_architecture": architecture,
            "sample_weighting": "case_equal_trajectory_equal_within_case",
            "offline_metrics": metrics, "diagnostic_metrics": metrics["aggregate"],
            "optimizer": {
                "name": "scipy.optimize.least_squares", "success": bool(result.success),
                "cost": float(result.cost), "nfev": int(result.nfev),
                "initial_constants": initial.tolist(), "refitted_constants": result.x.tolist(),
            },
            "topology_expression": topology_expression,
            "static_safety": safety,
        },
        pysr_provenance={"used": False, "purpose": "fixed-topology-refit"},
        data_provenance={
            "dataset_hash": _combined_data_hash(scenes),
            "per_scene_hashes": {scene.name: scene.data_hash for scene in scenes},
            "paths": {scene.name: str(scene.data_path) for scene in scenes},
            "trajectories": {
                scene.name: [{
                    "id": trajectory.trajectory_id, "path": str(trajectory.data_path),
                    "hash": trajectory.data_hash,
                    "sample_count": int(trajectory.dataset.X.shape[0]),
                } for trajectory in _scene_trajectories(scene)] for scene in scenes
            },
        },
        discovery_parent={
            "path": None if parent is None else str(parent),
            "sha256": None if parent is None else hash_file(parent),
        },
        git_sha=_git_sha(), environment=env,
    )


def fit_head(
    scenes: Sequence[PreparedScene], *, role: str, channel_index: int, run_id: str,
    seed: int, niterations: int, smoke: bool, deployment_expression: str | None,
    mode: str, deployment_architecture: str, fit_augmentation: str,
    feature_set: str, feature_profile: str,
    regressor_factory: Callable[..., Any], environment: Mapping[str, Any],
) -> dict[str, Any]:
    architecture = validate_deployment_architecture(deployment_architecture)
    augmentation = validate_fit_augmentation(fit_augmentation)
    split_blocks = {
        name: _head_training_blocks(
            scenes, role, channel_index, fit_augmentation=augmentation, split=name,
        )
        for name in ("train", "validation", "blind")
    }
    split_data = {
        name: (np.vstack([item[2] for item in blocks]), np.concatenate([item[3] for item in blocks]))
        for name, blocks in split_blocks.items()
    }
    X, y = split_data["train"]
    parameters = regressor_parameters(seed=seed, niterations=niterations, smoke=smoke)
    model = regressor_factory(**parameters)
    fit_kwargs: dict[str, Any] = {"variable_names": list(scenes[0].feature_names)}
    weights = _joint_sample_weights(split_blocks["train"]) if mode == "joint" else None
    if weights is not None:
        try:
            signature = inspect.signature(model.fit)
        except (TypeError, ValueError):
            signature = None
        supports_weights = signature is None or "weights" in signature.parameters or any(
            parameter.kind == inspect.Parameter.VAR_KEYWORD
            for parameter in signature.parameters.values()
        )
        if not supports_weights:
            raise TypeError(
                "joint fitting requires PySR model.fit(X, y, weights=...); "
                "the configured regressor does not support explicit sample weights"
            )
        fit_kwargs["weights"] = weights
    model.fit(X, y, **fit_kwargs)
    fitted = _expression(model)
    deployment = fitted if deployment_expression is None else deployment_expression
    counts = {scene.name: int(scene.dataset.X.shape[0]) for scene in scenes}
    augmented_counts = {name: 2 * count for name, count in counts.items()}
    metrics = {
        name: diagnostic_metrics(target, _predict(model, matrix))
        for name, (matrix, target) in split_data.items()
    }
    all_X = np.vstack([matrix for matrix, _ in split_data.values()])
    feature_ranges = {
        name: [float(all_X[:, index].min()), float(all_X[:, index].max())]
        for index, name in enumerate(scenes[0].feature_names)
    }
    safety = formula_safety_report(
        {"feature_names": list(scenes[0].feature_names), "expression": deployment},
        feature_ranges,
    )
    config = scenes[0].config
    alignment = scenes[0].dataset.metadata
    artifact_id = f"{run_id}-{mode}-{role}-anchor-upper"
    env = dict(environment)
    env["pysr"] = getattr(sys.modules.get("pysr"), "__version__", env.get("pysr"))
    return build_formula_artifact(
        artifact_id=artifact_id, run_id=run_id, role=role, anchor="upper",
        training_scenes=[scene.name for scene in scenes], feature_names=scenes[0].feature_names,
        fitted_expression=fitted, deployment_expression=deployment,
        output={
            "name": "alpha",
            "definition": "cylinder_surface_tangential_velocity/U0",
            "unit": "dimensionless",
            "native_action_index": channel_index,
            "native_action_order": list(scenes[0].action_layout),
        },
        feature_metadata={
            "builder": "SR_analysis.utils.data_contracts.build_batch_dataset",
            "definitions_module": "SR_analysis.utils.feature_builder",
            "order": list(scenes[0].feature_names),
        },
        alignment_metadata={
            "alignment": alignment["alignment"], "warmup": alignment["warmup"],
            "sample_segments": alignment.get("sample_segments"),
            "history_policy": alignment["history_policy"], "required_target": config.get("scene_id") == "illusion",
        },
        dt_c_metadata={"per_scene": {scene.name: float(scene.dataset.metadata["dt_c"]) for scene in scenes}, "unit": "D/U0"},
        g_metadata={
            "version": G_VERSION,
            "deployment_architecture": architecture,
            "deployment_semantics": architecture,
            "fit_augmentation": augmentation,
            "front_rule": (
                "front(x)=0.5*(f(x)-f(Gx)); antisymmetric projection"
                if architecture == "mapped_shared" else None
            ),
            "rear_rule": "lower(x)=-upper(Gx)" if architecture == "mapped_shared" else None,
            "fitted_heads": [role],
        },
        fit_provenance={
            "git_sha": _git_sha(),
            "environment": env,
            "command": " ".join(sys.argv),
            "created_at": datetime.now(timezone.utc).isoformat(),
            "mode": mode, "seed": int(seed), "niterations": int(niterations), "smoke": bool(smoke),
            "feature_set": feature_set, "feature_profile": feature_profile,
            "n_samples": int(X.shape[0]), "source_n_samples": int(sum(counts.values())),
            "scene_sample_counts": counts, "scene_augmented_sample_counts": (
                augmented_counts if augmentation == "G" else counts
            ),
            "g_augmented": augmentation == "G",
            "fit_augmentation": augmentation,
            "deployment_architecture": architecture,
            "deployment_semantics": architecture,
            "sample_weighting": (
                "case_equal_trajectory_equal_within_case" if weights is not None else "uniform"
            ),
            "offline_metrics": metrics,
            "diagnostic_metrics": metrics["train"],
            "selection_policy": "manual_validation_required_r2_never_auto_promotes",
            "static_safety": safety,
        },
        pysr_provenance={
            "parameters": parameters, "binary_operators": list(BINARY_OPERATORS),
            "unary_operators": list(UNARY_OPERATORS), "complexity": {"maxsize": parameters["maxsize"], "parsimony": parameters["parsimony"]},
            "pareto_candidates": _pareto_candidates(
                model, split_data, scenes[0].feature_names, feature_ranges,
            ),
            "candidate_limit": 25,
        },
        data_provenance={
            "dataset_hash": _combined_data_hash(scenes),
            "per_scene_hashes": {scene.name: scene.data_hash for scene in scenes},
            "paths": {scene.name: str(scene.data_path) for scene in scenes},
            "trajectories": {
                scene.name: [
                    {
                        "id": trajectory.trajectory_id,
                        "path": str(trajectory.data_path),
                        "hash": trajectory.data_hash,
                        "sample_count": int(trajectory.dataset.X.shape[0]),
                    }
                    for trajectory in _scene_trajectories(scene)
                ]
                for scene in scenes
            },
        },
        git_sha=_git_sha(), environment=env,
    )


def _validate_joint_contract(scenes: Sequence[PreparedScene]) -> None:
    first = scenes[0]
    objective = first.config.get("scene_id")
    for scene in scenes[1:]:
        if scene.config.get("scene_id") != objective:
            raise ValueError("joint scenes must belong to the same scene_id objective")
        if scene.feature_names != first.feature_names:
            raise ValueError("joint scenes must have identical feature definitions/order")
        if scene.action_layout != first.action_layout:
            raise ValueError("joint scenes must have identical native action order")
        action_contract = (
            scene.dataset.metadata["action_mode"],
            float(scene.config.get("action_scale", 1.0)),
            tuple(scene.config.get("action_bias", (0.0, 0.0, 0.0))),
            float(scene.config.get("u0", 0.01)),
        )
        first_action_contract = (
            first.dataset.metadata["action_mode"],
            float(first.config.get("action_scale", 1.0)),
            tuple(first.config.get("action_bias", (0.0, 0.0, 0.0))),
            float(first.config.get("u0", 0.01)),
        )
        if action_contract != first_action_contract:
            raise ValueError("joint scenes must have identical action conversion contracts")


def run_fit(
    *, scenes: Sequence[str], mode: str, run_id: str, data_root: Path | None = None,
    data_roots: Sequence[Path] | None = None,
    results_root: Path = DEFAULT_RESULTS_ROOT, alignment: str = CAUSAL_ALIGNMENT,
    feature_set: str = "physics_reduced", feature_profile: str = "default",
    sample_segment: str = "all", transient_count: int = 50,
    deployment_architecture: str = "mapped_shared", fit_augmentation: str = "G",
    seed: int = 0, niterations: int = 40,
    smoke: bool = False, prepare_only: bool = False,
    fit_purpose: str = "discovery",
    front_topology_expression: str | None = None,
    rear_topology_expression: str | None = None,
    front_discovery_parent_path: Path | None = None,
    rear_discovery_parent_path: Path | None = None,
    front_deployment_expression: str | None = None,
    rear_deployment_expression: str | None = None,
    regressor_factory: Callable[..., Any] | None = None,
) -> list[Path]:
    if not scenes:
        raise ValueError("at least one scene is required")
    if mode == "joint" and len(scenes) < 2:
        raise ValueError("joint mode requires at least two scenes")
    if fit_purpose not in FIT_PURPOSES:
        raise ValueError(f"fit purpose must be one of {FIT_PURPOSES}")
    if fit_purpose == "fixed-topology-refit":
        if not front_topology_expression or not rear_topology_expression:
            raise ValueError("fixed-topology refit requires front and rear topology expressions")
        if deployment_architecture != "mapped_shared":
            raise ValueError("fixed-topology refit currently requires mapped_shared architecture")
    output_dir = formula_output_dir(results_root, run_id)
    refuse_existing_output(output_dir)
    roots = tuple(data_roots or (() if data_root is None else (data_root,)))
    if not roots:
        roots = (DEFAULT_DATA_ROOT,)
    prepared = [
        prepare_scene_trajectories(
            name, roots, alignment=alignment, joint=mode == "joint",
            feature_set=feature_set, feature_profile=feature_profile,
            sample_segment=sample_segment, transient_count=transient_count,
        )
        for name in scenes
    ]
    if mode == "joint":
        _validate_joint_contract(prepared)
    if prepare_only:
        return []
    factory = None if fit_purpose == "fixed-topology-refit" else (regressor_factory or _regressor_factory())
    environment = _environment()
    groups = [prepared] if mode == "joint" else [[scene] for scene in prepared]
    artifacts: list[tuple[str, dict[str, Any]]] = []
    for group in groups:
        layout = group[0].action_layout
        front_index = layout.index("front")
        upper_index = layout.index("upper")
        label = "joint" if mode == "joint" else group[0].name
        architecture = validate_deployment_architecture(deployment_architecture)
        augmentation = validate_fit_augmentation(fit_augmentation)
        if architecture == "three_head_independent":
            head_specs = (
                ("front", "front_independent", front_index, front_deployment_expression),
                ("upper", "upper_independent", upper_index, rear_deployment_expression),
                ("lower", "lower_independent", layout.index("lower"), None),
            )
        else:
            head_specs = (
                ("front", "front_odd", front_index, front_deployment_expression),
                ("rear_shared_upper", "rear_shared", upper_index, rear_deployment_expression),
            )
        for filename_role, role, channel_index, deployment_expression in head_specs:
            if fit_purpose == "fixed-topology-refit":
                is_front = filename_role == "front"
                topology_expression = front_topology_expression if is_front else rear_topology_expression
                parent_path = front_discovery_parent_path if is_front else rear_discovery_parent_path
                artifact = refit_fixed_topology_head(
                    group, role=role, channel_index=channel_index, run_id=run_id,
                    topology_expression=str(topology_expression), discovery_parent_path=parent_path,
                    mode=mode, deployment_architecture=architecture, feature_set=feature_set,
                    feature_profile=feature_profile, environment=environment,
                )
            else:
                artifact = fit_head(
                    group, role=role, channel_index=channel_index, run_id=run_id, seed=seed,
                    niterations=niterations, smoke=smoke,
                    deployment_expression=deployment_expression, mode=mode,
                    deployment_architecture=architecture, fit_augmentation=augmentation,
                    feature_set=feature_set, feature_profile=feature_profile,
                    regressor_factory=factory, environment=environment,
                )
            artifacts.append((f"{label}_{filename_role}.json", artifact))
    output_dir.mkdir(parents=True, exist_ok=False)
    paths = []
    for filename, artifact in artifacts:
        paths.append(atomic_write_json(output_dir / filename, artifact))
    return paths


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Stage 2 contract-aware PySR fitting")
    selection = parser.add_mutually_exclusive_group(required=True)
    selection.add_argument("--scene")
    selection.add_argument("--scenes", help="comma-separated scenes")
    parser.add_argument("--mode", choices=("per-scene", "joint"), default="per-scene")
    parser.add_argument("--run-id", default=None, help="isolated results run ID; defaults to a UTC ID")
    parser.add_argument(
        "--data-root", action="append", default=None,
        help="repeatable canonical data root, data/runs/<id>, or scene directory",
    )
    parser.add_argument(
        "--data-roots", default=None,
        help="comma-separated independent trajectory roots",
    )
    parser.add_argument("--results-root", default=None, help=argparse.SUPPRESS)
    parser.add_argument("--alignment", choices=ALIGNMENTS, default=CAUSAL_ALIGNMENT)
    parser.add_argument("--feature-set", choices=tuple(FEATURE_SET_KEYS), default="physics_reduced")
    parser.add_argument("--feature-profile", choices=COMPARATIVE_FEATURE_PROFILES, default="default")
    parser.add_argument("--sample-segment", choices=("all", "transient", "steady"), default="all")
    parser.add_argument("--transient-count", type=int, default=50)
    parser.add_argument("--deployment-architecture", choices=DEPLOYMENT_ARCHITECTURES, default="mapped_shared")
    parser.add_argument("--fit-augmentation", choices=FIT_AUGMENTATIONS, default="G")
    parser.add_argument("--fit-purpose", choices=FIT_PURPOSES, default="discovery")
    parser.add_argument("--front-topology-expression")
    parser.add_argument("--rear-topology-expression")
    parser.add_argument("--front-discovery-parent-path")
    parser.add_argument("--rear-discovery-parent-path")
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--niterations", type=int, default=40)
    parser.add_argument("--smoke", action="store_true")
    parser.add_argument("--prepare-only", action="store_true")
    parser.add_argument("--front-deployment-expression")
    parser.add_argument("--rear-deployment-expression")
    return parser


def main(argv: Sequence[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)
    if args.niterations < 1:
        parser.error("--niterations must be positive")
    scenes = [args.scene] if args.scene else [item.strip() for item in args.scenes.split(",") if item.strip()]
    root_values = list(args.data_root or [])
    if args.data_roots:
        root_values.extend(item.strip() for item in args.data_roots.split(",") if item.strip())
    data_roots = [resolve_repo_path(value, DEFAULT_DATA_ROOT) for value in root_values]
    if not data_roots:
        data_roots = [DEFAULT_DATA_ROOT.resolve()]
    try:
        paths = run_fit(
            scenes=scenes, mode=args.mode, run_id=args.run_id or default_run_id(),
            data_roots=data_roots,
            results_root=resolve_repo_path(args.results_root, DEFAULT_RESULTS_ROOT),
            alignment=args.alignment, feature_set=args.feature_set,
            feature_profile=args.feature_profile, sample_segment=args.sample_segment,
            transient_count=args.transient_count,
            deployment_architecture=args.deployment_architecture,
            fit_augmentation=args.fit_augmentation, seed=args.seed, niterations=args.niterations,
            smoke=args.smoke, prepare_only=args.prepare_only, fit_purpose=args.fit_purpose,
            front_topology_expression=args.front_topology_expression,
            rear_topology_expression=args.rear_topology_expression,
            front_discovery_parent_path=(
                resolve_repo_path(args.front_discovery_parent_path, REPO_ROOT)
                if args.front_discovery_parent_path else None
            ),
            rear_discovery_parent_path=(
                resolve_repo_path(args.rear_discovery_parent_path, REPO_ROOT)
                if args.rear_discovery_parent_path else None
            ),
            front_deployment_expression=args.front_deployment_expression,
            rear_deployment_expression=args.rear_deployment_expression,
        )
    except Exception as exc:
        print(f"stage_2_fit: error: {exc}", file=sys.stderr)
        return 1
    if args.prepare_only:
        print(f"Prepared {len(scenes)} scene(s); inputs and output path validated.")
    else:
        for path in paths:
            print(path)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
