from __future__ import annotations

import json
import logging
import uuid
from typing import Any

import numpy as np

from metrics_lie.artifacts.plots import (
    plot_calibration_curve,
    plot_metric_distribution,
    plot_subgroup_bars,
    plot_threshold_curve,
)
from metrics_lie.datasets.loaders import load_binary_csv, load_dataset
from metrics_lie.metrics.applicability import (
    ApplicableMetricSet,
    DatasetProperties,
    MetricResolver,
)
from metrics_lie.model.surface import SurfaceType
from metrics_lie.task_types import TaskType
from metrics_lie.surface_compat import (
    DEFAULT_THRESHOLD,
    SURFACE_TYPE_MAP,
    filter_compatible_scenarios,
    filter_compatible_scenarios_by_task,
)
from metrics_lie.analysis import (
    analyze_metric_disagreements,
    build_dashboard_summary,
    locate_failure_modes,
    run_sensitivity_analysis,
    run_threshold_sweep,
)
from metrics_lie.diagnostics.calibration import brier_score, expected_calibration_error
from metrics_lie.metrics.core import METRICS, compute_metric
from metrics_lie.schema import Artifact, MetricSummary, ResultBundle, ScenarioResult
from metrics_lie.spec import load_experiment_spec
from metrics_lie.utils.paths import get_run_dir
from metrics_lie.experiments.datasets import dataset_fingerprint_csv
from metrics_lie.experiments.definition import ExperimentDefinition
from metrics_lie.experiments.registry import (
    upsert_experiment as upsert_experiment_jsonl,
    log_run as log_run_jsonl,
)
from metrics_lie.experiments.runs import RunRecord
from metrics_lie.db.session import get_session
from metrics_lie.db.crud import (
    upsert_experiment,
    insert_run,
    update_run,
    insert_artifacts,
    get_experiment_spec_json,
    get_experiment_id_for_run,
)
from metrics_lie.experiments.identity import canonical_json

# Ensure scenario registration occurs (import-time registration).
import metrics_lie.scenarios.class_imbalance  # noqa: F401
import metrics_lie.scenarios.label_noise  # noqa: F401
import metrics_lie.scenarios.score_noise  # noqa: F401
import metrics_lie.scenarios.threshold_gaming  # noqa: F401
from metrics_lie.runner import RunConfig, run_scenarios
from metrics_lie.scenarios.base import ScenarioContext
from metrics_lie.scenarios.registry import create_scenario

logger = logging.getLogger(__name__)


def _summary_from_single_value(v: float) -> MetricSummary:
    # Baseline is a single deterministic value (Phase 1.3/1.4)
    return MetricSummary(mean=v, std=0.0, q05=v, q50=v, q95=v, n=1)


def run_from_spec_dict(
    spec_dict: dict,
    *,
    spec_path_for_notes: str | None = None,
    rerun_of: str | None = None,
) -> str:
    """
    Execute an experiment run given a parsed spec dictionary.

    This is the canonical execution path used both for CLI `run` (from file)
    and for `rerun` (from a stored spec_json snapshot in the DB).
    """
    spec = load_experiment_spec(spec_dict)
    task_type = TaskType(spec.task)

    dataset_fp = dataset_fingerprint_csv(spec.dataset.path)
    exp_def = ExperimentDefinition.from_spec(spec, dataset_fingerprint=dataset_fp)

    # Canonical JSON snapshot of the original spec for deterministic reruns.
    spec_json_str = canonical_json(spec_dict)

    # Phase 2.2: Write to DB
    with get_session() as session:
        upsert_experiment(session, exp_def, spec_json_str)

    # Phase 2.1: Keep JSONL logging optional
    upsert_experiment_jsonl(exp_def)

    if spec.metric not in METRICS:
        raise ValueError(
            f"Unknown metric '{spec.metric}'. Supported: {sorted(METRICS.keys())}"
        )

    # Load dataset. If model_source is provided, require features.
    require_features = spec.model_source is not None
    allow_missing_score = spec.model_source is not None
    score_validation = (
        spec.surface_source.surface_type if spec.surface_source else "probability"
    )

    if task_type.is_binary:
        ds = load_binary_csv(
            path=spec.dataset.path,
            y_true_col=spec.dataset.y_true_col,
            y_score_col=spec.dataset.y_score_col,
            subgroup_col=spec.dataset.subgroup_col,
            feature_cols=spec.dataset.feature_cols,
            require_features=require_features,
            allow_missing_score=allow_missing_score,
            score_validation=score_validation,
        )
    else:
        ds = load_dataset(
            path=spec.dataset.path,
            task_type=spec.task,
            y_true_col=spec.dataset.y_true_col,
            y_score_col=spec.dataset.y_score_col,
            subgroup_col=spec.dataset.subgroup_col,
            feature_cols=spec.dataset.feature_cols,
            require_features=require_features,
            allow_missing_score=allow_missing_score,
        )

    y_true = ds.y_true.to_numpy(dtype=float if task_type.is_regression else int)
    y_score = ds.y_score.to_numpy(dtype=float)
    subgroup = None
    if ds.subgroup is not None:
        subgroup = ds.subgroup.to_numpy()

    # If model source is provided, run inference to produce surface.
    prediction_surface = None
    surface_type = SurfaceType.PROBABILITY
    if spec.model_source is not None:
        from metrics_lie.model.surface import CalibrationState

        if ds.X is None:
            raise ValueError("Model inference requires feature columns in dataset.")

        kind = spec.model_source.kind

        if spec.model_source.path:
            from metrics_lie.model.security import check_trust_policy, scan_model_file

            check_trust_policy(
                spec.model_source.path,
                trust_pickle=spec.model_source.trust_pickle,
            )
            scan_result = scan_model_file(spec.model_source.path)
            if not scan_result.is_safe:
                raise ValueError(
                    f"Model file failed security scan: {scan_result.issues}"
                )

        if kind in ("pickle", "import"):
            # Legacy sklearn path
            from metrics_lie.model.adapter import SklearnAdapter
            from metrics_lie.model.sources import ModelSourceImport, ModelSourcePickle

            if kind == "pickle":
                if not spec.model_source.path:
                    raise ValueError("model_source.path is required for kind=pickle")
                source = ModelSourcePickle(path=spec.model_source.path)
            else:
                if not spec.model_source.import_path:
                    raise ValueError("model_source.import_path is required for kind=import")
                source = ModelSourceImport(import_path=spec.model_source.import_path)

            adapter = SklearnAdapter(
                source,
                threshold=spec.model_source.threshold or DEFAULT_THRESHOLD,
                positive_label=spec.model_source.positive_label or 1,
                calibration_state=CalibrationState.UNKNOWN,
                task_type=task_type,
            )
        elif kind == "http":
            from metrics_lie.model.adapters.http_adapter import HTTPAdapter

            adapter = HTTPAdapter(
                endpoint=spec.model_source.endpoint,
                task_type=TaskType(spec.task),
                threshold=spec.model_source.threshold or DEFAULT_THRESHOLD,
                positive_label=spec.model_source.positive_label or 1,
            )
        elif kind == "onnx":
            from metrics_lie.model.adapters.onnx_adapter import ONNXAdapter

            if not spec.model_source.path:
                raise ValueError("model_source.path is required for kind=onnx")
            adapter = ONNXAdapter(
                path=spec.model_source.path,
                task_type=TaskType(spec.task),
                threshold=spec.model_source.threshold or DEFAULT_THRESHOLD,
                positive_label=spec.model_source.positive_label or 1,
            )
        elif kind in ("xgboost", "lightgbm", "catboost"):
            from metrics_lie.model.adapters.boosting_adapter import BoostingAdapter

            if not spec.model_source.path:
                raise ValueError(f"model_source.path is required for kind={kind}")
            adapter = BoostingAdapter(
                path=spec.model_source.path,
                framework=kind,
                task_type=TaskType(spec.task),
                threshold=spec.model_source.threshold or DEFAULT_THRESHOLD,
                positive_label=spec.model_source.positive_label or 1,
            )
        elif kind == "pytorch":
            from metrics_lie.model.adapters.pytorch_adapter import PyTorchAdapter

            if not spec.model_source.path:
                raise ValueError("model_source.path is required for kind=pytorch")
            adapter = PyTorchAdapter(
                path=spec.model_source.path,
                task_type=TaskType(spec.task),
                threshold=spec.model_source.threshold or DEFAULT_THRESHOLD,
                positive_label=spec.model_source.positive_label or 1,
            )
        elif kind == "tensorflow":
            from metrics_lie.model.adapters.tensorflow_adapter import TensorFlowAdapter

            if not spec.model_source.path:
                raise ValueError("model_source.path is required for kind=tensorflow")
            adapter = TensorFlowAdapter(
                path=spec.model_source.path,
                task_type=TaskType(spec.task),
                threshold=spec.model_source.threshold or DEFAULT_THRESHOLD,
                positive_label=spec.model_source.positive_label or 1,
            )
        elif kind == "huggingface":
            from metrics_lie.model.adapters.huggingface_adapter import HuggingFaceAdapter

            model_path = spec.model_source.path or ""
            adapter = HuggingFaceAdapter.from_model_path(
                model_path,
                task_type=TaskType(spec.task),
            )
        else:
            raise ValueError(f"Unsupported model_source kind: {kind}")

        surfaces = adapter.get_all_surfaces(ds.X.to_numpy())
        if task_type.is_regression and SurfaceType.CONTINUOUS in surfaces:
            prediction_surface = surfaces[SurfaceType.CONTINUOUS]
        elif SurfaceType.PROBABILITY in surfaces:
            prediction_surface = surfaces[SurfaceType.PROBABILITY]
        elif SurfaceType.SCORE in surfaces:
            prediction_surface = surfaces[SurfaceType.SCORE]
        elif SurfaceType.LABEL in surfaces:
            prediction_surface = surfaces[SurfaceType.LABEL]
        elif SurfaceType.CONTINUOUS in surfaces:
            prediction_surface = surfaces[SurfaceType.CONTINUOUS]
        else:
            raise ValueError("Model adapter produced no usable surfaces.")
        surface_type = prediction_surface.surface_type
        y_score = prediction_surface.values.astype(float)

    # Phase 8: Direct surface ingestion from CSV (no model needed).
    elif spec.surface_source is not None:
        from metrics_lie.model.surface import (
            CalibrationState,
            PredictionSurface,
            validate_surface,
        )

        # Map spec surface_type to enum (Phase 9A: score; 9B: label).
        surface_type_enum = SURFACE_TYPE_MAP[spec.surface_source.surface_type]
        validated_values = validate_surface(
            surface_type=surface_type_enum,
            values=y_score,
            expected_n_samples=len(y_true),
            threshold=spec.surface_source.threshold,
        )
        prediction_surface = PredictionSurface(
            surface_type=surface_type_enum,
            values=validated_values,
            dtype=validated_values.dtype,
            n_samples=len(y_true),
            class_names=("0", "1"),
            positive_label=spec.surface_source.positive_label,
            threshold=spec.surface_source.threshold,
            calibration_state=CalibrationState.UNKNOWN,
            model_hash=None,
            is_deterministic=True,
        )
        surface_type = surface_type_enum

    # Resolve applicable metrics based on surface type and dataset properties.
    # Triggered when model_source OR surface_source is present (multi-metric mode).
    if spec.model_source is not None or spec.surface_source is not None:
        if task_type.is_regression:
            dataset_props = DatasetProperties(
                n_samples=int(len(y_true)),
                n_positive=0,
                n_negative=0,
                has_subgroups=subgroup is not None,
                positive_rate=0.0,
            )
        elif task_type.is_binary:
            dataset_props = DatasetProperties(
                n_samples=int(len(y_true)),
                n_positive=int(np.sum(y_true == 1)),
                n_negative=int(np.sum(y_true == 0)),
                has_subgroups=subgroup is not None,
                positive_rate=float(np.mean(y_true)) if len(y_true) > 0 else 0.0,
            )
        else:
            # Multiclass / multilabel
            dataset_props = DatasetProperties(
                n_samples=int(len(y_true)),
                n_positive=0,
                n_negative=0,
                has_subgroups=subgroup is not None,
                positive_rate=0.0,
                n_classes=int(len(np.unique(y_true))),
            )
        resolver = MetricResolver()
        applicable = resolver.resolve(
            task_type=spec.task,
            surface_type=surface_type,
            dataset_props=dataset_props,
        )
    else:
        applicable = ApplicableMetricSet(
            task_type=spec.task,
            surface_type=surface_type,
            metrics=[spec.metric],
            excluded=[],
            reasoning_trace=["manual_metric_selection"],
            warnings=[],
        )

    # Primary metric selection: respect spec.metric if applicable.
    primary_metric = spec.metric if spec.metric in applicable.metrics else None
    if primary_metric is None and applicable.metrics:
        primary_metric = applicable.metrics[0]
    if primary_metric is None:
        raise ValueError("No applicable metrics available for this surface.")

    # Phase 9B: filter scenarios by surface compatibility when using surface_source.
    if spec.surface_source is not None:
        effective_scenarios, skipped = filter_compatible_scenarios(
            spec.scenarios, surface_type
        )
        if skipped:
            logger.warning(
                "Skipping scenarios not compatible with surface_type=%s: %s",
                surface_type.value,
                skipped,
            )
    else:
        effective_scenarios = spec.scenarios

    # Filter scenarios by task-type compatibility.
    effective_scenarios, task_skipped = filter_compatible_scenarios_by_task(
        effective_scenarios, spec.task
    )
    if task_skipped:
        logger.warning(
            "Skipping scenarios not compatible with task_type=%s: %s",
            spec.task,
            task_skipped,
        )

    metric_results: dict[str, MetricSummary] = {}
    scenario_results_by_metric: dict[str, list[ScenarioResult]] = {}

    for metric_id in applicable.metrics:
        metric_fn = METRICS[metric_id]
        baseline_value = compute_metric(
            metric_id, metric_fn, y_true, y_score, threshold=DEFAULT_THRESHOLD
        )
        metric_results[metric_id] = _summary_from_single_value(baseline_value)

        # --- Phase 1.4: run scenario stress tests (Monte Carlo) ---
        scenario_results = run_scenarios(
            y_true=y_true,
            y_score=y_score,
            metric_name=metric_id,
            metric_fn=metric_fn,
            scenario_specs=[s.model_dump() for s in effective_scenarios],
            cfg=RunConfig(n_trials=spec.n_trials, seed=spec.seed),
            ctx=ScenarioContext(task=spec.task, surface_type=surface_type.value),
            subgroup=subgroup,
        )

        # --- Phase 1.5: add sensitivity_abs diagnostic ---
        baseline_mean = baseline_value
        scenario_results_with_diag = []
        for sr in scenario_results:
            sensitivity_abs = abs(sr.metric.mean - baseline_mean)
            diag = sr.diagnostics.copy()
            diag["sensitivity_abs"] = sensitivity_abs
            scenario_results_with_diag.append(
                ScenarioResult(
                    scenario_id=sr.scenario_id,
                    params=sr.params,
                    metric=sr.metric,
                    diagnostics=diag,
                    artifacts=sr.artifacts,
                )
            )
        scenario_results_by_metric[metric_id] = scenario_results_with_diag

    baseline_cal = {}
    if task_type.is_binary and surface_type == SurfaceType.PROBABILITY:
        baseline_cal = {
            "brier": brier_score(y_true, y_score),
            "ece": expected_calibration_error(y_true, y_score, n_bins=10),
        }
    elif (
        task_type == TaskType.MULTICLASS_CLASSIFICATION
        and surface_type == SurfaceType.PROBABILITY
        and y_score.ndim == 2
    ):
        from metrics_lie.diagnostics.calibration import (
            multiclass_brier_score,
            multiclass_ece,
        )

        baseline_cal = {
            "multiclass_brier": multiclass_brier_score(y_true, y_score),
            "multiclass_ece": multiclass_ece(y_true, y_score),
        }

    # Use primary metric for headline baseline/scenarios.
    baseline_value = metric_results[primary_metric].mean
    scenario_results_with_diag = scenario_results_by_metric.get(primary_metric, [])

    run_id = uuid.uuid4().hex[:10].upper()
    paths = get_run_dir(run_id)
    paths.ensure()

    run_record = RunRecord(
        run_id=run_id,
        experiment_id=exp_def.experiment_id,
        results_path=str(paths.results_json),
        artifacts_dir=str(paths.artifacts_dir),
        seed_used=spec.seed,
        rerun_of=rerun_of,
    )

    # Phase 2.2: Write to DB (queued)
    with get_session() as session:
        insert_run(session, run_record)

    # Phase 2.1: Keep JSONL logging optional
    log_run_jsonl(run_record)

    try:
        run_record.mark_running()

        # Phase 2.2: Update DB (running)
        with get_session() as session:
            update_run(session, run_record)

        # Phase 2.1: Keep JSONL logging optional
        log_run_jsonl(run_record)

        # --- Phase 1.7B: generate artifacts (plots) ---
        rng_artifacts = np.random.default_rng(spec.seed)
        scenario_results_with_artifacts: list[ScenarioResult] = []
        for sr in scenario_results_with_diag:
            artifacts_list: list[Artifact] = []
            scenario_id = sr.scenario_id

            # 1. Metric distribution plot
            try:
                metric_dist_path = (
                    paths.artifacts_dir / f"metric_dist_{scenario_id}.png"
                )
                plot_metric_distribution(
                    metric_summary=sr.metric.model_dump(),
                    metric_name=spec.metric,
                    scenario_id=scenario_id,
                    out_path=metric_dist_path,
                )
                artifacts_list.append(
                    Artifact(
                        kind="plot",
                        path=f"artifacts/metric_dist_{scenario_id}.png",
                        meta={"type": "metric_distribution"},
                    )
                )
            except Exception:
                pass  # Skip if plot generation fails

            # 2. Calibration curve (only for binary classification)
            if task_type.is_binary:
                try:
                    scenario = create_scenario(scenario_id, sr.params)
                    y_p_rep, s_p_rep = scenario.apply(
                        y_true,
                        y_score,
                        rng_artifacts,
                        ScenarioContext(task=spec.task, surface_type=surface_type.value),
                    )
                    if len(y_p_rep) > 0 and len(s_p_rep) > 0:
                        cal_path = paths.artifacts_dir / f"calibration_{scenario_id}.png"
                        plot_calibration_curve(
                            y_true=y_p_rep,
                            y_score=s_p_rep,
                            scenario_id=scenario_id,
                            out_path=cal_path,
                        )
                        artifacts_list.append(
                            Artifact(
                                kind="plot",
                                path=f"artifacts/calibration_{scenario_id}.png",
                                meta={"type": "calibration_curve"},
                            )
                        )
                except Exception:
                    pass  # Skip if plot generation fails

            # 3. Subgroup metric bars (if subgroup diagnostics exist)
            try:
                subgroup_metric = sr.diagnostics.get("subgroup_metric")
                if subgroup_metric:
                    group_means = {k: v["mean"] for k, v in subgroup_metric.items()}
                    if group_means:
                        subgroup_path = (
                            paths.artifacts_dir / f"subgroup_metric_{scenario_id}.png"
                        )
                        plot_subgroup_bars(
                            group_means=group_means,
                            scenario_id=scenario_id,
                            out_path=subgroup_path,
                        )
                        artifacts_list.append(
                            Artifact(
                                kind="plot",
                                path=f"artifacts/subgroup_metric_{scenario_id}.png",
                                meta={"type": "subgroup_comparison"},
                            )
                        )
            except Exception:
                pass  # Skip if plot generation fails

            # 4. Threshold curve (only for accuracy with metric_inflation)
            if spec.metric == "accuracy":
                try:
                    metric_inflation = sr.diagnostics.get("metric_inflation")
                    if metric_inflation:
                        scenario = create_scenario(scenario_id, sr.params)
                        y_p_rep, s_p_rep = scenario.apply(
                            y_true,
                            y_score,
                            rng_artifacts,
                            ScenarioContext(
                                task=spec.task, surface_type=surface_type.value
                            ),
                        )
                        if len(y_p_rep) > 0 and len(s_p_rep) > 0:
                            # Get mean optimal threshold from diagnostics (approximate)
                            # Use a representative threshold from the inflation data
                            baseline_thresh = DEFAULT_THRESHOLD
                            # Estimate optimized threshold from delta (use DEFAULT_THRESHOLD + small adjustment as proxy)
                            # Actually, we need to recompute or store it - let's use a simple heuristic
                            # For now, use DEFAULT_THRESHOLD as baseline and compute optimal from representative trial
                            from metrics_lie.diagnostics.metric_gaming import (
                                find_optimal_threshold,
                            )

                            thresholds = np.linspace(0.05, 0.95, 19)
                            opt_thresh, _ = find_optimal_threshold(
                                y_p_rep, s_p_rep, thresholds
                            )

                            threshold_path = (
                                paths.artifacts_dir
                                / f"threshold_curve_{scenario_id}.png"
                            )
                            plot_threshold_curve(
                                y_true=y_p_rep,
                                y_score=s_p_rep,
                                baseline_threshold=baseline_thresh,
                                optimized_threshold=opt_thresh,
                                scenario_id=scenario_id,
                                out_path=threshold_path,
                            )
                            artifacts_list.append(
                                Artifact(
                                    kind="plot",
                                    path=f"artifacts/threshold_curve_{scenario_id}.png",
                                    meta={"type": "threshold_optimization"},
                                )
                            )
                except Exception:
                    pass  # Skip if plot generation fails

            scenario_results_with_artifacts.append(
                ScenarioResult(
                    scenario_id=sr.scenario_id,
                    params=sr.params,
                    metric=sr.metric,
                    diagnostics=sr.diagnostics,
                    artifacts=artifacts_list,
                )
            )

            if artifacts_list:
                print(
                    f"[PLOT] Saved {len(artifacts_list)} artifacts for scenario {scenario_id}"
                )

        # Task-specific summary artifacts for web display
        task_specific: dict[str, Any] = {}

        if task_type.is_classification and prediction_surface is not None:
            from sklearn.metrics import confusion_matrix as _confusion_matrix

            if prediction_surface.surface_type == SurfaceType.LABEL:
                y_pred_labels = prediction_surface.values.astype(int)
            elif prediction_surface.surface_type == SurfaceType.PROBABILITY:
                if prediction_surface.values.ndim == 2:
                    y_pred_labels = np.argmax(prediction_surface.values, axis=1)
                else:
                    threshold = prediction_surface.threshold or 0.5
                    y_pred_labels = (prediction_surface.values >= threshold).astype(int)
            else:
                y_pred_labels = None

            if y_pred_labels is not None:
                cm = _confusion_matrix(y_true, y_pred_labels)
                task_specific["confusion_matrix"] = cm.tolist()

                unique_classes = sorted(set(y_true.tolist()) | set(y_pred_labels.tolist()))
                task_specific["class_names"] = [str(c) for c in unique_classes]

                if not task_type.is_binary:
                    from sklearn.metrics import precision_recall_fscore_support
                    prec, rec, f1s, sup = precision_recall_fscore_support(
                        y_true, y_pred_labels, labels=unique_classes, zero_division=0
                    )
                    per_class = {}
                    for i, cls in enumerate(unique_classes):
                        per_class[str(cls)] = {
                            "precision": float(prec[i]),
                            "recall": float(rec[i]),
                            "f1": float(f1s[i]),
                            "support": int(sup[i]),
                        }
                    task_specific["per_class_metrics"] = per_class

        elif task_type.is_regression and prediction_surface is not None:
            residuals = y_true - prediction_surface.values
            task_specific["residual_stats"] = {
                "mean": float(np.mean(residuals)),
                "std": float(np.std(residuals)),
                "min": float(np.min(residuals)),
                "max": float(np.max(residuals)),
                "median": float(np.median(residuals)),
                "mae": float(np.mean(np.abs(residuals))),
                "rmse": float(np.sqrt(np.mean(residuals ** 2))),
            }

        notes = {
            "phase": "1.7B",
            "spec_path": spec_path_for_notes,
            "baseline_diagnostics": baseline_cal,
            "applicable_metrics": applicable.metrics,
            "metric_resolution": {
                "excluded": applicable.excluded,
                "warnings": applicable.warnings,
                "reasoning_trace": applicable.reasoning_trace,
            },
            "task_specific": task_specific,
        }

        analysis_artifacts: dict[str, Any] = {}

        # Binary-only analysis: threshold sweep, sensitivity, disagreement
        if (
            task_type.is_binary
            and prediction_surface is not None
            and prediction_surface.surface_type
            in (SurfaceType.PROBABILITY, SurfaceType.SCORE)
        ):
            sweep = run_threshold_sweep(
                y_true=y_true,
                surface=prediction_surface,
                metrics=applicable.metrics,
                n_points=101,
            )
            analysis_artifacts["threshold_sweep"] = sweep.to_jsonable()
            sensitivity = run_sensitivity_analysis(
                y_true=y_true,
                surface=prediction_surface,
                metrics=applicable.metrics,
                perturbation_type="score_noise",
                magnitudes=[0.01, 0.02, 0.05, 0.1, 0.2],
                n_trials=50,
                seed=spec.seed,
            )
            analysis_artifacts["sensitivity"] = sensitivity.to_jsonable()
            disagreements = analyze_metric_disagreements(
                y_true=y_true,
                surface=prediction_surface,
                thresholds=sweep.optimal_thresholds,
                metrics=applicable.metrics,
            )
            analysis_artifacts["metric_disagreements"] = [
                d.to_jsonable() for d in disagreements
            ]

        # Sensitivity for non-binary: works on CONTINUOUS and LABEL surfaces
        elif (
            not task_type.is_binary
            and prediction_surface is not None
            and prediction_surface.surface_type
            in (SurfaceType.CONTINUOUS, SurfaceType.LABEL)
        ):
            sensitivity_metrics = [
                m for m in applicable.metrics
                if m not in ("macro_auc", "top_k_accuracy")
            ]
            if sensitivity_metrics:
                sensitivity = run_sensitivity_analysis(
                    y_true=y_true,
                    surface=prediction_surface,
                    metrics=sensitivity_metrics,
                    perturbation_type="score_noise",
                    magnitudes=[0.01, 0.02, 0.05, 0.1, 0.2],
                    n_trials=50,
                    seed=spec.seed,
                )
                analysis_artifacts["sensitivity"] = sensitivity.to_jsonable()

        # Failure modes: works for all task types and surface types
        if prediction_surface is not None:
            failures = locate_failure_modes(
                y_true=y_true,
                surface=prediction_surface,
                metrics=applicable.metrics,
                subgroup=subgroup,
                top_k=20,
            )
            analysis_artifacts["failure_modes"] = failures.to_jsonable()

        # Dashboard: works for all task types with 2+ metrics
        if len(applicable.metrics) > 1:
            from metrics_lie.metrics.registry import METRIC_DIRECTION

            dashboard = build_dashboard_summary(
                primary_metric=primary_metric,
                surface_type=surface_type.value,
                metric_results={k: v.model_dump() for k, v in metric_results.items()},
                scenario_results_by_metric={
                    k: [sr.model_dump() for sr in v]
                    for k, v in scenario_results_by_metric.items()
                },
                metric_directions=METRIC_DIRECTION,
            )
            analysis_artifacts["dashboard_summary"] = dashboard.to_jsonable()

        # Fairness analysis (optional — requires fairlearn and sensitive_feature)
        sensitive_feature = spec_dict.get("sensitive_feature")
        if sensitive_feature and prediction_surface is not None:
            try:
                from metrics_lie.diagnostics.fairness import compute_fairness_report
                import pandas as _pd_fair

                raw_df = _pd_fair.read_csv(spec.dataset.path)
                if sensitive_feature in raw_df.columns:
                    sensitive_features = raw_df[sensitive_feature].values
                    if prediction_surface.values.ndim == 2:
                        y_pred_fairness = np.argmax(prediction_surface.values, axis=1)
                    elif prediction_surface.surface_type == SurfaceType.LABEL:
                        y_pred_fairness = prediction_surface.values.astype(int)
                    elif prediction_surface.surface_type == SurfaceType.PROBABILITY:
                        threshold = prediction_surface.threshold or 0.5
                        y_pred_fairness = (prediction_surface.values >= threshold).astype(int)
                    else:
                        y_pred_fairness = prediction_surface.values

                    fairness_report = compute_fairness_report(
                        y_true=y_true,
                        y_pred=y_pred_fairness,
                        sensitive_features=sensitive_features,
                        metric_fns={
                            "accuracy": lambda yt, yp: float(np.mean(yt == yp)),
                        },
                    )
                    analysis_artifacts["fairness"] = fairness_report
            except ImportError:
                pass  # fairlearn not installed — skip

        # Drift detection (optional — requires evidently and reference_dataset)
        reference_dataset = spec_dict.get("reference_dataset")
        if reference_dataset:
            try:
                from metrics_lie.diagnostics.drift import compute_drift_report
                import pandas as _pd_drift

                ref_df = _pd_drift.read_csv(reference_dataset)
                eval_df = _pd_drift.read_csv(spec.dataset.path)
                feature_cols = [
                    c for c in eval_df.columns
                    if c not in ("y_true", "y_score", "subgroup")
                    and not c.startswith("y_score_")
                ]
                shared_cols = sorted(set(feature_cols) & set(ref_df.columns))
                if shared_cols:
                    drift_report = compute_drift_report(
                        reference=ref_df,
                        current=eval_df,
                        feature_columns=shared_cols,
                    )
                    analysis_artifacts["drift"] = drift_report
            except ImportError:
                pass  # evidently not installed — skip

        bundle = ResultBundle(
            run_id=run_id,
            experiment_name=spec.name,
            metric_name=primary_metric,
            task_type=spec.task,
            baseline=_summary_from_single_value(baseline_value),
            scenarios=scenario_results_with_artifacts,
            prediction_surface=prediction_surface.to_jsonable()
            if prediction_surface
            else None,
            applicable_metrics=applicable.metrics,
            metric_results=metric_results,
            scenario_results_by_metric=scenario_results_by_metric,
            analysis_artifacts=analysis_artifacts,
            notes=notes,
        )

        paths.results_json.write_text(bundle.to_pretty_json(), encoding="utf-8")
        print(f"[OK] Wrote results: {paths.results_json}")
        print(f"Baseline {spec.metric} = {baseline_value:.6f}")
        if spec.scenarios:
            print(
                f"[OK] Ran {len(spec.scenarios)} scenario(s) with n_trials={spec.n_trials}"
            )

        # Phase 2.2: Insert artifacts into DB
        all_artifacts: list[Artifact] = []
        for sr in scenario_results_with_artifacts:
            all_artifacts.extend(sr.artifacts)
        if all_artifacts:
            with get_session() as session:
                insert_artifacts(session, run_id, all_artifacts)

        run_record.mark_completed()

        # Phase 2.2: Update DB (completed)
        with get_session() as session:
            update_run(session, run_record)

        # Phase 2.1: Keep JSONL logging optional
        log_run_jsonl(run_record)
    except Exception as exc:  # pragma: no cover - simple logging wrapper
        run_record.mark_failed(str(exc))

        # Phase 2.2: Update DB (failed)
        with get_session() as session:
            update_run(session, run_record)

        # Phase 2.1: Keep JSONL logging optional
        log_run_jsonl(run_record)
        raise

    return run_id


def rerun(run_id: str) -> str:
    """
    Deterministically rerun a completed experiment using the stored spec_json snapshot.

    This creates a new run (with a new run_id) linked to the same experiment_id,
    and optionally records the original run_id in the DB.
    """
    with get_session() as session:
        experiment_id = get_experiment_id_for_run(session, run_id)
        spec_json_str = get_experiment_spec_json(session, experiment_id)

    if not spec_json_str:
        raise ValueError(
            "No stored spec_json found for experiment. "
            "Please run database migrations with 'alembic upgrade head' before using 'rerun'."
        )

    spec_dict = json.loads(spec_json_str)
    # Use a descriptive marker in notes so results.json keeps the same schema.
    spec_path_for_notes = f"<rerun_of:{run_id}>"
    new_run_id = run_from_spec_dict(
        spec_dict, spec_path_for_notes=spec_path_for_notes, rerun_of=run_id
    )
    print(new_run_id)
    return new_run_id
