gam-models 0.3.151

//! Tests for the survival marginal-slope family (relocated verbatim).

use super::*;
use crate::custom_family::{CustomFamily, ExactOuterDerivativeOrder};
use approx::assert_relative_eq;
use faer::sparse::{SparseColMat, Triplet};
use gam_linalg::matrix::{DenseDesignMatrix, SymmetricMatrix};
// Trait import for the `.value()` calls in the linux-gated #932 pullback test —
// on other targets that test is compiled out and the import would be dead.
#[cfg(target_os = "linux")]
use gam_math::nested_dual::JetField;
use ndarray::array;

/// Local scalar closeness assertion used throughout this module's exactness
/// gates. Asserts `|lhs - rhs| <= tol`, reporting both operands on failure.
fn assert_close(lhs: f64, rhs: f64, tol: f64, label: &str) {
    assert!(
        (lhs - rhs).abs() <= tol,
        "{label} mismatch: lhs={lhs:.12e}, rhs={rhs:.12e}, tol={tol:.3e}"
    );
}

/// Pin the survival cross-block W metric to the survival row Hessian
/// formula `u2_eta1 = (1-d)·w·k2(-η, w·(1-d)) + d·w` rather than the
/// Bernoulli probit proxy `φ²/(Φ·(1-Φ))`. Three rows exercise the
/// censored-only, event-only, and weighted-censored branches; the
/// expected values are recomputed in-test from
/// `signed_probit_neglog_derivatives_up_to_fourth` to bind the metric
/// to its source-of-truth derivative routine.
#[test]
fn survival_pilot_irls_row_metric_matches_u2_eta1_formula() {
    let eta = ndarray::Array1::from(vec![0.3_f64, -0.7, 1.2]);
    let weights = ndarray::Array1::from(vec![1.0_f64, 2.5, 0.5]);
    let event = ndarray::Array1::from(vec![0.0_f64, 1.0, 0.0]);
    let computed = super::survival_pilot_irls_row_metric_at_eta(&eta, &weights, &event)
        .expect("survival pilot W metric");
    assert_eq!(computed.len(), 3);
    for i in 0..3 {
        let d = event[i];
        let w = weights[i];
        let e = eta[i];
        let (_, k2, _, _) = super::signed_probit_neglog_derivatives_up_to_fourth(-e, w * (1.0 - d))
            .expect("probit k2");
        let expected = k2 + w * d;
        approx::assert_relative_eq!(computed[i], expected, max_relative = 1e-12);
    }
    // Censored row (d=0): metric must be strictly positive at finite η —
    // the Mills-ratio second derivative of -log Φ(-η) is positive on ℝ.
    assert!(
        computed[0] > 0.0,
        "censored row metric must be > 0, got {}",
        computed[0],
    );
    // Event row (d=1): the formula collapses to w·d exactly; the censored
    // branch contributes zero because the k2 weight `w·(1-d)` is zero.
    approx::assert_relative_eq!(computed[1], weights[1], max_relative = 1e-12);
}

/// Length-mismatch contract: the metric helper must reject misaligned
/// inputs instead of producing a truncated W vector that silently
/// passes the cross-block routine but breaks the W-inner product.
#[test]
fn survival_pilot_irls_row_metric_rejects_length_mismatch() {
    let eta = ndarray::Array1::from(vec![0.0_f64, 1.0]);
    let weights = ndarray::Array1::from(vec![1.0_f64]); // wrong length
    let event = ndarray::Array1::from(vec![0.0_f64, 0.0]);
    let result = super::survival_pilot_irls_row_metric_at_eta(&eta, &weights, &event);
    match result {
        Ok(_) => panic!("expected length-mismatch error, got Ok"),
        Err(msg) => assert!(
            msg.contains("length mismatch"),
            "expected 'length mismatch' in error, got: {msg}",
        ),
    }
}

fn empty_termspec() -> TermCollectionSpec {
    TermCollectionSpec {
        linear_terms: vec![],
        random_effect_terms: vec![],
        smooth_terms: vec![],
    }
}

fn unit_score_covariance() -> MarginalSlopeCovariance {
    MarginalSlopeCovariance::diagonal(array![1.0]).unwrap()
}

fn base_time_block() -> TimeBlockInput {
    TimeBlockInput {
        design_entry: DesignMatrix::from(Array2::zeros((1, 1))),
        design_exit: DesignMatrix::from(Array2::zeros((1, 1))),
        design_derivative_exit: DesignMatrix::from(Array2::ones((1, 1))),
        offset_entry: Array1::zeros(1),
        offset_exit: Array1::zeros(1),
        derivative_offset_exit: Array1::from_elem(
            1,
            DEFAULT_SURVIVAL_MARGINAL_SLOPE_DERIVATIVE_GUARD,
        ),
        time_monotonicity:
            crate::survival::location_scale::TimeBlockMonotonicity::EnforcedByRowConstraint,
        penalties: Vec::new(),
        nullspace_dims: Vec::new(),
        initial_log_lambdas: None,
        initial_beta: Some(Array1::zeros(1)),
    }
}

/// Endpoint evaluations whose average empirical function Gram is exactly
/// `(2/3) I`.  Both endpoint charts span the two coefficient directions, so
/// they exercise function-space shrinkage without the singular all-zero
/// design that the production generalized eigensolve correctly rejects.
fn full_span_time_endpoint_designs() -> (DesignMatrix, DesignMatrix) {
    (
        DesignMatrix::from(array![[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]]),
        DesignMatrix::from(array![[1.0, 0.0], [0.0, 1.0], [1.0, -1.0]]),
    )
}

#[test]
fn time_nullspace_shrinkage_adds_precision_for_uncontrolled_time_direction() {
    let (design_entry, design_exit) = full_span_time_endpoint_designs();
    let mut block = TimeBlockInput {
        design_entry,
        design_exit,
        design_derivative_exit: DesignMatrix::from(Array2::ones((3, 2))),
        offset_entry: Array1::zeros(3),
        offset_exit: Array1::zeros(3),
        derivative_offset_exit: Array1::from_elem(
            3,
            DEFAULT_SURVIVAL_MARGINAL_SLOPE_DERIVATIVE_GUARD,
        ),
        penalties: vec![array![[1.0, 0.0], [0.0, 0.0]]],
        nullspace_dims: vec![1],
        initial_beta: Some(Array1::zeros(2)),
        ..base_time_block()
    };

    assert!(
        install_time_nullspace_shrinkage_penalty(&mut block)
            .expect("time nullspace shrinkage should build"),
        "expected a shrinkage penalty to be appended",
    );
    assert_eq!(block.penalties.len(), 2);
    assert_eq!(block.nullspace_dims, vec![1, 0]);
    let expected = array![[0.0, 0.0], [0.0, 2.0 / 3.0]];
    for i in 0..2 {
        for j in 0..2 {
            assert_close(
                block.penalties[1][[i, j]],
                expected[[i, j]],
                1e-12,
                &format!("time nullspace function-metric ridge ({i},{j})"),
            );
        }
    }
}

#[test]
fn time_nullspace_shrinkage_is_noop_for_full_rank_time_penalty() {
    let (design_entry, design_exit) = full_span_time_endpoint_designs();
    let mut block = TimeBlockInput {
        design_entry,
        design_exit,
        design_derivative_exit: DesignMatrix::from(Array2::ones((3, 2))),
        offset_entry: Array1::zeros(3),
        offset_exit: Array1::zeros(3),
        derivative_offset_exit: Array1::from_elem(
            3,
            DEFAULT_SURVIVAL_MARGINAL_SLOPE_DERIVATIVE_GUARD,
        ),
        penalties: vec![Array2::<f64>::eye(2)],
        nullspace_dims: vec![0],
        initial_beta: Some(Array1::zeros(2)),
        ..base_time_block()
    };

    assert!(
        !install_time_nullspace_shrinkage_penalty(&mut block)
            .expect("full-rank time penalty should be accepted"),
        "full-rank time penalties should not get another penalty",
    );
    assert_eq!(block.penalties.len(), 1);
    assert_eq!(block.nullspace_dims, vec![0]);
}

fn sparse_design(dense: &Array2<f64>) -> DesignMatrix {
    let mut triplets = Vec::<Triplet<usize, usize, f64>>::new();
    for i in 0..dense.nrows() {
        for j in 0..dense.ncols() {
            let value = dense[[i, j]];
            if value != 0.0 {
                triplets.push(Triplet::new(i, j, value));
            }
        }
    }
    let sparse = SparseColMat::try_new_from_triplets(dense.nrows(), dense.ncols(), &triplets)
        .expect("assemble sparse design");
    DesignMatrix::Sparse(gam_linalg::matrix::SparseDesignMatrix::new(sparse))
}

/// Build an n-row closed-form survival family with empty time/marginal/
/// logslope blocks (so q0/q1/qd1 come from offsets only). No flex
/// deviations are configured, so `log_likelihood_only` takes the
/// closed-form fast path.
fn make_closed_form_test_family(n: usize) -> SurvivalMarginalSlopeFamily {
    // Pseudo-random rows uncorrelated with row parity, so an even-only
    // subsample is representative for the Horvitz-Thompson rescaling
    // check.
    let event: Array1<f64> =
        Array1::from_iter((0..n).map(|i| if (i * 31 + 7) % 5 >= 3 { 1.0 } else { 0.0 }));
    let weights: Array1<f64> =
        Array1::from_iter((0..n).map(|i| 0.5 + ((i * 13 + 4) % 5) as f64 * 0.1));
    let z: Array1<f64> = Array1::from_iter(
        (0..n).map(|i| -1.0 + 2.0 * (((i * 17 + 5) % n) as f64 + 0.5) / (n as f64)),
    );
    let offset_entry: Array1<f64> = Array1::from_iter(
        (0..n).map(|i| -0.4 + 0.7 * (((i * 11 + 3) % n) as f64 + 0.5) / (n as f64)),
    );
    let offset_exit: Array1<f64> = Array1::from_iter(
        (0..n).map(|i| 0.1 + 0.6 * (((i * 19 + 7) % n) as f64 + 0.5) / (n as f64)),
    );
    // qd1 must remain strictly above the derivative guard.
    let derivative_offset_exit: Array1<f64> =
        Array1::from_iter((0..n).map(|i| 0.5 + 0.05 * ((i * 23 + 1) % 3) as f64));
    SurvivalMarginalSlopeFamily {
        n,
        event: Arc::new(event),
        weights: Arc::new(weights),
        z: Arc::new(z.insert_axis(Axis(1))),
        score_covariance: unit_score_covariance(),
        gaussian_frailty_sd: None,
        family_hyper: SurvivalMarginalSlopeFamilyHyperState::default(),
        derivative_guard: 1e-6,
        // Empty time/marginal/logslope designs: `n_rows × 0` so the
        // closed-form q geometry is driven entirely by offsets.
        design_entry: DesignMatrix::from(Array2::zeros((n, 0))),
        design_exit: DesignMatrix::from(Array2::zeros((n, 0))),
        design_derivative_exit: DesignMatrix::from(Array2::zeros((n, 0))),
        offset_entry: Arc::new(offset_entry),
        offset_exit: Arc::new(offset_exit),
        derivative_offset_exit: Arc::new(derivative_offset_exit),
        marginal_design: DesignMatrix::from(Array2::zeros((n, 0))),
        logslope_layout: (DesignMatrix::from(Array2::zeros((n, 0)))).into(),
        score_warp: None,
        link_dev: None,
        influence_absorber: None,
        time_linear_constraints: None,
        time_wiggle_knots: None,
        time_wiggle_degree: None,
        time_wiggle_ncols: 0,
        intercept_warm_starts: None,
        auto_subsample_phase_counter: Arc::new(AtomicUsize::new(0)),
        auto_subsample_last_rho: Arc::new(Mutex::new(None)),
    }
}

#[test]
fn k1_shared_logslope_uses_cached_arbitrary_variance_932() {
    let mut family = make_closed_form_test_family(3);
    family.score_covariance = MarginalSlopeCovariance::diagonal(array![3.75]).unwrap();
    let cached = family.score_covariance.ones_quadratic_form();
    assert!((cached - 3.75).abs() <= f64::EPSILON);
    assert_eq!(family.shared_logslope_covariance_scale(), cached);
}

fn closed_form_block_states(
    family: &SurvivalMarginalSlopeFamily,
    g: f64,
) -> Vec<ParameterBlockState> {
    let n = family.n;
    vec![
        // Time block: empty beta; per-row eta entries unused (designs
        // are zero-column).
        ParameterBlockState {
            beta: Array1::zeros(0),
            eta: Array1::zeros(n),
        },
        // Marginal block: empty beta.
        ParameterBlockState {
            beta: Array1::zeros(0),
            eta: Array1::zeros(n),
        },
        // Log-slope block: empty beta with per-row eta = g (constant
        // log-slope across rows).
        ParameterBlockState {
            beta: Array1::zeros(0),
            eta: Array1::from_elem(n, g),
        },
    ]
}

#[test]
fn converged_identifiability_scalars_use_current_vector_geometry_932() {
    let n = 2;
    let mut family = make_closed_form_test_family(n);
    let design = array![[1.0], [2.0]];
    let offset = array![0.2, -0.1];
    family.logslope_layout = LogslopeTopology::shared()
        .materialize_identity(DesignMatrix::from(design.clone()), &offset)
        .unwrap();
    let beta_logslope = array![0.5];
    let eta_logslope = design.dot(&beta_logslope) + &offset;
    let states = vec![
        ParameterBlockState {
            beta: Array1::zeros(0),
            eta: Array1::zeros(n),
        },
        ParameterBlockState {
            beta: Array1::zeros(0),
            eta: Array1::zeros(n),
        },
        ParameterBlockState {
            beta: beta_logslope,
            eta: eta_logslope.clone(),
        },
    ];

    let erased =
        <SurvivalMarginalSlopeFamily as CustomFamily>::current_identifiability_family_scalars(
            &family, &states,
        )
        .unwrap()
        .expect("survival must expose converged scalars");
    let scalars = erased
        .downcast_ref::<SurvivalMarginalSlopeFamilyScalars>()
        .expect("survival scalar type");
    for row in 0..n {
        assert_eq!(scalars.q0_i[row], family.offset_entry[row]);
        assert_eq!(scalars.q1_i[row], family.offset_exit[row]);
        assert_eq!(scalars.qd1_i[row], family.derivative_offset_exit[row]);
        assert_eq!(
            scalars.c_i[row],
            (1.0 + eta_logslope[row] * eta_logslope[row]).sqrt(),
        );
    }
}

#[test]
fn survival_primary_g_fourth_cell_partials_are_zero() {
    let family = make_closed_form_test_family(1);
    let primary = flex_primary_slices(&family);
    let score_span = exact_kernel::LocalSpanCubic {
        left: -1.0,
        right: 1.0,
        c0: 0.2,
        c1: -0.1,
        c2: 0.05,
        c3: -0.02,
    };
    let link_span = exact_kernel::LocalSpanCubic {
        left: -0.5,
        right: 0.5,
        c0: 0.1,
        c1: 0.3,
        c2: -0.2,
        c3: 0.4,
    };
    let fixed = family
        .denested_cell_primary_fixed_partials(&primary, 0.2, 0.7, score_span, link_span, 0.0, 0.0)
        .expect("primary fixed partials");
    let (_, dc_daab, dc_dabb, dc_dbbb) = exact_kernel::denested_cell_third_partials(link_span);

    assert_eq!(fixed.coeff_aau[primary.g], dc_daab);
    assert_eq!(fixed.coeff_abu[primary.g], dc_dabb);
    assert_eq!(fixed.coeff_bbu[primary.g], dc_dbbb);
    assert!(dc_daab.iter().any(|value| *value != 0.0));
    assert!(dc_dabb.iter().any(|value| *value != 0.0));
    assert!(dc_dbbb.iter().any(|value| *value != 0.0));
    assert_eq!(fixed.coeff_aaau[primary.g], [0.0; 4]);
    assert_eq!(fixed.coeff_aabu[primary.g], [0.0; 4]);
    assert_eq!(fixed.coeff_abbu[primary.g], [0.0; 4]);
    assert_eq!(fixed.coeff_bbbu[primary.g], [0.0; 4]);
}

#[test]
fn survival_log_likelihood_subsample_full_equals_unsampled() {
    use crate::outer_subsample::OuterScoreSubsample;
    let n = 200usize;
    let family = make_closed_form_test_family(n);
    let states = closed_form_block_states(&family, 0.25);

    let baseline = family
        .log_likelihood_only(&states)
        .expect("baseline ll (no subsample)");

    let mut opts_full = BlockwiseFitOptions::default();
    opts_full.outer_score_subsample = Some(Arc::new(
        OuterScoreSubsample::from_uniform_inclusion_mask((0..n).collect(), n, 0xDEADBEEF),
    ));
    let with_full_mask = family
        .log_likelihood_only_with_options(&states, &opts_full)
        .expect("ll with mask=full");

    let rel = ((with_full_mask - baseline) / baseline.abs().max(1.0)).abs();
    assert!(
        rel < 1e-12,
        "subsample(mask=full) {} differs from baseline {} by rel {}",
        with_full_mask,
        baseline,
        rel
    );
}

#[test]
fn survival_log_likelihood_subsample_half_scales_correctly() {
    use crate::outer_subsample::OuterScoreSubsample;
    let n = 200usize;
    let family = make_closed_form_test_family(n);
    let states = closed_form_block_states(&family, 0.25);

    let even_mask: Vec<usize> = (0..n).filter(|i| i % 2 == 0).collect();
    let m = even_mask.len();

    let mut opts_half = BlockwiseFitOptions::default();
    opts_half.outer_score_subsample = Some(Arc::new(
        OuterScoreSubsample::from_uniform_inclusion_mask(even_mask.clone(), n, 0xCAFE),
    ));
    let scaled = family
        .log_likelihood_only_with_options(&states, &opts_half)
        .expect("ll with mask=even");

    // Raw even-row sum: same mask but weight_scale = 1.0.
    let mut opts_even_unscaled = BlockwiseFitOptions::default();
    opts_even_unscaled.outer_score_subsample = Some(Arc::new(
        OuterScoreSubsample::with_uniform_weight(even_mask, m, 0, 1.0),
    ));
    let raw_even_sum = family
        .log_likelihood_only_with_options(&states, &opts_even_unscaled)
        .expect("raw even-row ll sum");

    let expected_scaled = (n as f64 / m as f64) * raw_even_sum;
    let rel = ((scaled - expected_scaled) / expected_scaled.abs().max(1.0)).abs();
    assert!(
        rel < 1e-12,
        "scaled {} != 2*even_sum {} (rel {})",
        scaled,
        expected_scaled,
        rel
    );

    // Horvitz-Thompson check: 2 * Σ_even ≈ full-data sum.
    let baseline = family.log_likelihood_only(&states).expect("baseline ll");
    let ht_rel = ((scaled - baseline) / baseline.abs().max(1.0)).abs();
    assert!(
        ht_rel < 0.05,
        "Horvitz-Thompson scaled {} not near baseline {} (rel {})",
        scaled,
        baseline,
        ht_rel
    );
}

fn dummy_blockspec(cols: usize) -> ParameterBlockSpec {
    // `validate_blockspecs` enforces unique block names so coefficient
    // labels stay unambiguous. A monotonic per-process counter keeps
    // each call's name distinct even when multiple specs are stacked
    // into one `Vec`.
    use std::sync::atomic::{AtomicUsize, Ordering};
    static SEQ: AtomicUsize = AtomicUsize::new(0);
    let idx = SEQ.fetch_add(1, Ordering::Relaxed);
    ParameterBlockSpec {
        name: format!("dummy_{idx}"),
        design: DesignMatrix::Dense(DenseDesignMatrix::from(Array2::zeros((1, cols)))),
        offset: Array1::zeros(1),
        penalties: Vec::new(),
        nullspace_dims: Vec::new(),
        initial_log_lambdas: Array1::zeros(0),
        initial_beta: Some(Array1::zeros(cols)),
        gauge_priority: 100,
        jacobian_callback: None,
        stacked_design: None,
        stacked_offset: None,
    }
}

fn dummy_penalized_blockspec(cols: usize, penalties: usize) -> ParameterBlockSpec {
    let mut spec = dummy_blockspec(cols);
    spec.penalties = (0..penalties)
        .map(|_| PenaltyMatrix::Dense(Array2::eye(cols)))
        .collect();
    spec.nullspace_dims = vec![0; penalties];
    spec.initial_log_lambdas = Array1::zeros(penalties);
    spec
}

fn test_deviation_runtime() -> DeviationRuntime {
    build_score_warp_deviation_block_from_seed(
        &array![-1.0, 0.0, 1.0],
        &DeviationBlockConfig {
            degree: 3,
            num_internal_knots: 1,
            penalty_order: 2,
            penalty_orders: vec![1, 2, 3],
            double_penalty: false,
            monotonicity_eps: 1e-4,
        },
    )
    .expect("build test deviation runtime")
    .runtime
}

fn max_abs_diff_vec(lhs: &Array1<f64>, rhs: &Array1<f64>) -> f64 {
    lhs.iter()
        .zip(rhs.iter())
        .map(|(left, right)| (left - right).abs())
        .fold(0.0_f64, f64::max)
}

fn max_abs_diff_mat(lhs: &Array2<f64>, rhs: &Array2<f64>) -> f64 {
    lhs.iter()
        .zip(rhs.iter())
        .map(|(left, right)| (left - right).abs())
        .fold(0.0_f64, f64::max)
}

fn assert_blockwise_matches_joint_principal_blocks(
    family: &SurvivalMarginalSlopeFamily,
    block_states: &[ParameterBlockState],
) {
    let eval = family
        .evaluate_blockwise_exact_newton(block_states)
        .expect("blockwise exact-newton evaluation");
    let (joint_ll, joint_gradient, joint_hessian) = family
        .evaluate_exact_newton_joint_dense(block_states)
        .expect("joint dense exact-newton evaluation");
    let slices = block_slices(family, block_states);
    let mut block_ranges = vec![
        slices.time.clone(),
        slices.marginal.clone(),
        slices.logslope.clone(),
    ];
    if let Some(range) = slices.score_warp.clone() {
        block_ranges.push(range);
    }
    if let Some(range) = slices.link_dev.clone() {
        block_ranges.push(range);
    }

    assert!((eval.log_likelihood - joint_ll).abs() <= 1e-10);
    assert_eq!(eval.blockworking_sets.len(), block_ranges.len());
    for (work, range) in eval.blockworking_sets.iter().zip(block_ranges.iter()) {
        let BlockWorkingSet::ExactNewton { gradient, hessian } = work else {
            panic!("expected exact-newton block working set");
        };
        let expected_gradient = joint_gradient.slice(s![range.clone()]).to_owned();
        let expected_hessian = joint_hessian
            .slice(s![range.clone(), range.clone()])
            .to_owned();
        assert!(
            max_abs_diff_vec(gradient, &expected_gradient) <= 1e-10,
            "gradient block mismatch"
        );
        assert!(
            max_abs_diff_mat(&hessian.to_dense(), &expected_hessian) <= 1e-10,
            "hessian block mismatch"
        );
    }
}

fn test_family(
    score_warp: Option<DeviationRuntime>,
    link_dev: Option<DeviationRuntime>,
) -> SurvivalMarginalSlopeFamily {
    SurvivalMarginalSlopeFamily {
        n: 1,
        event: Arc::new(array![0.0]),
        weights: Arc::new(array![1.0]),
        z: Arc::new(array![0.0].insert_axis(Axis(1))),
        score_covariance: unit_score_covariance(),
        gaussian_frailty_sd: None,
        family_hyper: SurvivalMarginalSlopeFamilyHyperState::default(),
        derivative_guard: 1e-6,
        design_entry: DesignMatrix::from(Array2::zeros((1, 1))),
        design_exit: DesignMatrix::from(Array2::zeros((1, 1))),
        design_derivative_exit: DesignMatrix::from(Array2::ones((1, 1))),
        offset_entry: Arc::new(Array1::zeros(1)),
        offset_exit: Arc::new(Array1::zeros(1)),
        derivative_offset_exit: Arc::new(Array1::from_elem(1, 1e-6)),
        marginal_design: DesignMatrix::from(Array2::zeros((1, 2))),
        logslope_layout: (DesignMatrix::from(Array2::zeros((1, 3)))).into(),
        score_warp,
        link_dev,
        influence_absorber: None,
        time_linear_constraints: None,
        time_wiggle_knots: None,
        time_wiggle_degree: None,
        time_wiggle_ncols: 0,
        intercept_warm_starts: None,
        auto_subsample_phase_counter: Arc::new(AtomicUsize::new(0)),
        auto_subsample_last_rho: Arc::new(Mutex::new(None)),
    }
}

#[test]
fn validate_spec_rejects_coordinate_cone_without_guard_offset() {
    let spec = SurvivalMarginalSlopeTermSpec {
        age_entry: array![0.0, 0.0],
        age_exit: array![1.0, 1.0],
        event_target: array![0.0, 1.0],
        weights: array![1.0, 1.0],
        z: array![-1.0, 1.0].insert_axis(Axis(1)),
        base_link: InverseLink::Standard(StandardLink::Probit),
        marginalspec: empty_termspec(),
        marginal_offset: Array1::zeros(2),
        frailty: FrailtySpec::None,
        derivative_guard: 1e-4,
        baseline_hyper: SurvivalMarginalSlopeBaselineHyperSpec::Linear {
            config: crate::survival::construction::SurvivalBaselineConfig {
                target: crate::survival::construction::SurvivalBaselineTarget::Linear,
                scale: None,
                shape: None,
                rate: None,
                makeham: None,
            },
        },
        time_block: TimeBlockInput {
            design_entry: DesignMatrix::from(Array2::zeros((2, 1))),
            design_exit: DesignMatrix::from(Array2::zeros((2, 1))),
            design_derivative_exit: DesignMatrix::from(Array2::ones((2, 1))),
            offset_entry: Array1::zeros(2),
            offset_exit: Array1::zeros(2),
            derivative_offset_exit: Array1::zeros(2),
            time_monotonicity:
                crate::survival::location_scale::TimeBlockMonotonicity::EnforcedByCoordinateCone,
            ..base_time_block()
        },
        timewiggle_block: None,
        logslopespec: empty_termspec(),
        logslopespecs: None,
        logslope_offset: Array1::zeros(2),
        score_warp: None,
        link_dev: None,
        score_influence_jacobian: None,
        latent_z_policy: LatentZPolicy::default(),
    };

    let err = validate_spec(&spec).expect_err("coordinate cone without guard offset should fail");
    assert!(
        err.contains("coordinate-cone time block requires derivative offset >= guard"),
        "unexpected error: {err}"
    );
}

#[test]
fn validate_spec_accepts_learned_gaussian_shift_sigma() {
    let spec = SurvivalMarginalSlopeTermSpec {
        age_entry: array![0.0, 0.0],
        age_exit: array![1.0, 1.0],
        event_target: array![0.0, 1.0],
        weights: array![1.0, 1.0],
        z: array![-1.0, 1.0].insert_axis(Axis(1)),
        base_link: InverseLink::Standard(StandardLink::Probit),
        marginalspec: empty_termspec(),
        marginal_offset: Array1::zeros(2),
        frailty: FrailtySpec::GaussianShift {
            scale: FrailtyScale::Learned { initial_sigma: 0.5 },
        },
        derivative_guard: 1e-4,
        baseline_hyper: SurvivalMarginalSlopeBaselineHyperSpec::Linear {
            config: crate::survival::construction::SurvivalBaselineConfig {
                target: crate::survival::construction::SurvivalBaselineTarget::Linear,
                scale: None,
                shape: None,
                rate: None,
                makeham: None,
            },
        },
        time_block: base_time_block(),
        timewiggle_block: None,
        logslopespec: empty_termspec(),
        logslopespecs: None,
        logslope_offset: Array1::zeros(2),
        score_warp: None,
        link_dev: None,
        score_influence_jacobian: None,
        latent_z_policy: LatentZPolicy::default(),
    };

    let validation = validate_spec(&spec);
    assert!(
        validation.is_ok(),
        "learned GaussianShift scale must be an explicit family coordinate: {validation:?}"
    );
}

#[test]
fn block_slices_handles_link_only_survival_flex_layout() {
    let link_runtime = test_deviation_runtime();
    let family = test_family(None, Some(link_runtime.clone()));
    let block_states = vec![
        ParameterBlockState {
            beta: Array1::zeros(1),
            eta: Array1::zeros(1),
        },
        ParameterBlockState {
            beta: Array1::zeros(2),
            eta: Array1::zeros(1),
        },
        ParameterBlockState {
            beta: Array1::zeros(3),
            eta: Array1::zeros(1),
        },
        ParameterBlockState {
            beta: Array1::zeros(link_runtime.basis_dim()),
            eta: Array1::zeros(1),
        },
    ];

    let slices = block_slices(&family, &block_states);
    assert!(slices.score_warp.is_none());
    assert_eq!(
        slices.link_dev.as_ref().expect("link-only slice").len(),
        link_runtime.basis_dim()
    );
    assert_eq!(slices.total, 1 + 2 + 3 + link_runtime.basis_dim());
}

#[test]
fn exact_survival_callbacks_lock_the_family_owned_coefficient_chart() {
    use crate::custom_family::BlockEffectiveJacobian;

    let one = Arc::new(Array2::<f64>::ones((2, 1)));
    let time = TimeBlockJacobian::new(Arc::clone(&one), Arc::clone(&one), Arc::clone(&one));
    let marginal = MarginalBlockJacobian::new(Arc::clone(&one));
    let family = test_family(None, None);
    let logslope = LogslopeBlockJacobian::new(
        family.logslope_layout.clone(),
        Arc::clone(&family.z),
        family.score_covariance.clone(),
    )
    .expect("test logslope callback");

    assert!(time.locks_raw_width_reduction());
    assert!(marginal.locks_raw_width_reduction());
    assert!(logslope.locks_raw_width_reduction());
}

// ── Single-source block-layout parity (#428) ─────────────────────────
//
// `HessBlock` + `BlockHessianAccumulator::block_view` are the one place
// the 5×5 block layout and its transpose relationships live. Every
// assembler (`to_dense`, `diagonal`, operator `mul_vec`, operator
// `bilinear`) is driven by them. These tests pin that machinery against a
// *separately hand-written* fifteen-block scatter so a layout or
// transpose regression in the shared helpers cannot hide behind itself.

/// Build contiguous block slices for the given per-block widths. A flex
/// block with width 0 is absent (`None`) and consumes no coordinates,
/// exactly mirroring how `block_slices` lays out optional deviation blocks.
fn parity_make_slices(
    pt: usize,
    pm: usize,
    pg: usize,
    ph: usize,
    pw: usize,
    pi: usize,
) -> BlockSlices {
    let mut cursor = 0usize;
    let mut take = |n: usize| {
        let r = cursor..cursor + n;
        cursor += n;
        r
    };
    let time = take(pt);
    let marginal = take(pm);
    let logslope = take(pg);
    let score_warp = (ph > 0).then(|| take(ph));
    let link_dev = (pw > 0).then(|| take(pw));
    let influence = (pi > 0).then(|| take(pi));
    let total = cursor;
    BlockSlices {
        time,
        marginal,
        logslope,
        score_warp,
        link_dev,
        influence,
        total,
    }
}

/// A genuinely-symmetric block (diagonal blocks of a Hessian are symmetric
/// by construction), with a per-block `tag` so cross-block contamination
/// is detectable.
fn parity_sym(n: usize, tag: f64) -> Array2<f64> {
    Array2::from_shape_fn((n, n), |(i, j)| {
        tag + (i as f64 + j as f64) * 0.5 + (i as f64) * (j as f64) * 0.0625
    })
}

/// A general (non-symmetric) off-diagonal block, distinct per `tag` and
/// deliberately asymmetric in (i, j) so a missing/extra transpose shows up.
fn parity_gen(rows: usize, cols: usize, tag: f64) -> Array2<f64> {
    Array2::from_shape_fn((rows, cols), |(i, j)| {
        tag + (i as f64) * 0.5 + (j as f64) * 0.125
    })
}

fn parity_filled_accumulator(
    pt: usize,
    pm: usize,
    pg: usize,
    ph: usize,
    pw: usize,
    pi: usize,
) -> BlockHessianAccumulator {
    BlockHessianAccumulator {
        h_tt: parity_sym(pt, 1.0),
        h_mm: parity_sym(pm, 2.0),
        h_gg: parity_sym(pg, 3.0),
        h_hh: parity_sym(ph, 4.0),
        h_ww: parity_sym(pw, 5.0),
        h_ii: parity_sym(pi, 6.0),
        h_tm: parity_gen(pt, pm, 10.0),
        h_tg: parity_gen(pt, pg, 11.0),
        h_th: parity_gen(pt, ph, 12.0),
        h_tw: parity_gen(pt, pw, 13.0),
        h_ti: parity_gen(pt, pi, 20.0),
        h_mg: parity_gen(pm, pg, 14.0),
        h_mh: parity_gen(pm, ph, 15.0),
        h_mw: parity_gen(pm, pw, 16.0),
        h_mi: parity_gen(pm, pi, 21.0),
        h_gh: parity_gen(pg, ph, 17.0),
        h_gw: parity_gen(pg, pw, 18.0),
        h_gi: parity_gen(pg, pi, 22.0),
        h_hw: parity_gen(ph, pw, 19.0),
        h_hi: parity_gen(ph, pi, 23.0),
        h_wi: parity_gen(pw, pi, 24.0),
    }
}

/// Independent dense oracle: scatter the fifteen stored blocks by hand,
/// placing each off-diagonal block in its upper position and its explicit
/// transpose in the mirror position. Deliberately avoids `block_view`,
/// `range_of`, and `for_each_offdiagonal_pair`.
fn parity_reference_dense(acc: &BlockHessianAccumulator, sl: &BlockSlices) -> Array2<f64> {
    let mut out = Array2::zeros((sl.total, sl.total));
    out.slice_mut(s![sl.time.clone(), sl.time.clone()])
        .assign(&acc.h_tt);
    out.slice_mut(s![sl.marginal.clone(), sl.marginal.clone()])
        .assign(&acc.h_mm);
    out.slice_mut(s![sl.logslope.clone(), sl.logslope.clone()])
        .assign(&acc.h_gg);
    if let Some(h) = &sl.score_warp {
        out.slice_mut(s![h.clone(), h.clone()]).assign(&acc.h_hh);
    }
    if let Some(w) = &sl.link_dev {
        out.slice_mut(s![w.clone(), w.clone()]).assign(&acc.h_ww);
    }
    if let Some(i) = &sl.influence {
        out.slice_mut(s![i.clone(), i.clone()]).assign(&acc.h_ii);
    }
    let mut place =
        |r: std::ops::Range<usize>, c: std::ops::Range<usize>, m: ArrayView2<'_, f64>| {
            out.slice_mut(s![r.clone(), c.clone()]).assign(&m);
            out.slice_mut(s![c, r]).assign(&m.t());
        };
    place(sl.time.clone(), sl.marginal.clone(), acc.h_tm.view());
    place(sl.time.clone(), sl.logslope.clone(), acc.h_tg.view());
    place(sl.marginal.clone(), sl.logslope.clone(), acc.h_mg.view());
    if let Some(h) = &sl.score_warp {
        place(sl.time.clone(), h.clone(), acc.h_th.view());
        place(sl.marginal.clone(), h.clone(), acc.h_mh.view());
        place(sl.logslope.clone(), h.clone(), acc.h_gh.view());
    }
    if let Some(w) = &sl.link_dev {
        place(sl.time.clone(), w.clone(), acc.h_tw.view());
        place(sl.marginal.clone(), w.clone(), acc.h_mw.view());
        place(sl.logslope.clone(), w.clone(), acc.h_gw.view());
    }
    if let (Some(h), Some(w)) = (&sl.score_warp, &sl.link_dev) {
        place(h.clone(), w.clone(), acc.h_hw.view());
    }
    if let Some(i) = &sl.influence {
        place(sl.time.clone(), i.clone(), acc.h_ti.view());
        place(sl.marginal.clone(), i.clone(), acc.h_mi.view());
        place(sl.logslope.clone(), i.clone(), acc.h_gi.view());
        if let Some(h) = &sl.score_warp {
            place(h.clone(), i.clone(), acc.h_hi.view());
        }
        if let Some(w) = &sl.link_dev {
            place(w.clone(), i.clone(), acc.h_wi.view());
        }
    }
    out
}

const PARITY_LAYOUTS: [(usize, usize, usize, usize, usize, usize); 7] = [
    (2, 3, 2, 4, 3, 0), // full flex, no absorber
    (2, 3, 2, 0, 0, 0), // rigid: no flex blocks
    (2, 3, 2, 4, 0, 0), // score-warp only
    (2, 3, 2, 0, 3, 0), // link-deviation only
    (2, 3, 2, 0, 0, 5), // absorber only (no flex)
    (2, 3, 2, 4, 3, 5), // full flex + absorber (#461)
    (2, 3, 2, 4, 0, 5), // score-warp + absorber
];

#[test]
fn block_to_dense_matches_hand_scatter_bit_exact() {
    for (pt, pm, pg, ph, pw, pi) in PARITY_LAYOUTS {
        let sl = parity_make_slices(pt, pm, pg, ph, pw, pi);
        let acc = parity_filled_accumulator(pt, pm, pg, ph, pw, pi);
        let got = acc.to_dense(&sl);
        let want = parity_reference_dense(&acc, &sl);
        assert_eq!(
            got,
            want,
            "to_dense diverged from hand scatter for layout {:?}",
            (pt, pm, pg, ph, pw)
        );
    }
}

#[test]
fn block_diagonal_matches_dense_diagonal_bit_exact() {
    for (pt, pm, pg, ph, pw, pi) in PARITY_LAYOUTS {
        let sl = parity_make_slices(pt, pm, pg, ph, pw, pi);
        let acc = parity_filled_accumulator(pt, pm, pg, ph, pw, pi);
        let got = acc.diagonal(&sl);
        let want = parity_reference_dense(&acc, &sl).diag().to_owned();
        assert_eq!(
            got,
            want,
            "diagonal diverged from dense diagonal for layout {:?}",
            (pt, pm, pg, ph, pw)
        );
    }
}

#[test]
fn block_operator_matvec_matches_dense_gemv() {
    for (pt, pm, pg, ph, pw, pi) in PARITY_LAYOUTS {
        let sl = parity_make_slices(pt, pm, pg, ph, pw, pi);
        let acc = parity_filled_accumulator(pt, pm, pg, ph, pw, pi);
        let dense = parity_reference_dense(&acc, &sl);
        let v = Array1::from_shape_fn(sl.total, |i| (i as f64 * 0.37).sin());
        let op = acc.into_operator(sl.clone());
        let got = op.mul_vec(&v);
        let want = dense.dot(&v);
        assert_relative_eq!(
            got.as_slice().unwrap(),
            want.as_slice().unwrap(),
            max_relative = 1e-12,
            epsilon = 1e-12
        );
    }
}

#[test]
fn block_operator_bilinear_matches_dense_quadratic_form() {
    for (pt, pm, pg, ph, pw, pi) in PARITY_LAYOUTS {
        let sl = parity_make_slices(pt, pm, pg, ph, pw, pi);
        let acc = parity_filled_accumulator(pt, pm, pg, ph, pw, pi);
        let dense = parity_reference_dense(&acc, &sl);
        let v = Array1::from_shape_fn(sl.total, |i| (i as f64 * 0.37).sin());
        let u = Array1::from_shape_fn(sl.total, |i| (i as f64 * 0.53).cos());
        let want = v.dot(&dense.dot(&u));
        let op = acc.into_operator(sl.clone());
        let got = op.bilinear(&v, &u);
        assert_relative_eq!(got, want, max_relative = 1e-12, epsilon = 1e-12);
    }
}

#[test]
fn block_operator_dense_matches_accumulator_dense_bit_exact() {
    for (pt, pm, pg, ph, pw, pi) in PARITY_LAYOUTS {
        let sl = parity_make_slices(pt, pm, pg, ph, pw, pi);
        let acc = parity_filled_accumulator(pt, pm, pg, ph, pw, pi);
        let direct = acc.to_dense(&sl);
        let op = acc.into_operator(sl.clone());
        let via_op = op.to_dense();
        assert_eq!(
            direct,
            via_op,
            "operator to_dense diverged from accumulator to_dense for layout {:?}",
            (pt, pm, pg, ph, pw)
        );
    }
}

#[test]
fn block_view_is_transpose_symmetric_across_present_pairs() {
    for (pt, pm, pg, ph, pw, pi) in PARITY_LAYOUTS {
        let sl = parity_make_slices(pt, pm, pg, ph, pw, pi);
        let acc = parity_filled_accumulator(pt, pm, pg, ph, pw, pi);
        for a in HessBlock::ALL {
            if sl.range_of(a).is_none() {
                continue;
            }
            for b in HessBlock::ALL {
                if sl.range_of(b).is_none() {
                    continue;
                }
                let ab = acc.block_view(a, b).to_owned();
                let ba_t = acc.block_view(b, a).t().to_owned();
                assert_eq!(
                    ab, ba_t,
                    "block_view({a:?},{b:?}) != block_view({b:?},{a:?})^T"
                );
            }
        }
    }
}

#[test]
fn exact_flex_row_matches_rigid_closed_form_without_deviations() {
    let family = SurvivalMarginalSlopeFamily {
        n: 1,
        event: Arc::new(array![1.0]),
        weights: Arc::new(array![1.7]),
        z: Arc::new(array![0.25].insert_axis(Axis(1))),
        score_covariance: unit_score_covariance(),
        gaussian_frailty_sd: None,
        family_hyper: SurvivalMarginalSlopeFamilyHyperState::default(),
        derivative_guard: 1e-6,
        design_entry: DesignMatrix::from(Array2::zeros((1, 1))),
        design_exit: DesignMatrix::from(Array2::zeros((1, 1))),
        design_derivative_exit: DesignMatrix::from(Array2::zeros((1, 1))),
        offset_entry: Arc::new(array![0.2]),
        offset_exit: Arc::new(array![0.4]),
        derivative_offset_exit: Arc::new(array![0.8]),
        marginal_design: DesignMatrix::from(Array2::zeros((1, 0))),
        logslope_layout: (DesignMatrix::from(Array2::zeros((1, 0)))).into(),
        score_warp: None,
        link_dev: None,
        influence_absorber: None,
        time_linear_constraints: None,
        time_wiggle_knots: None,
        time_wiggle_degree: None,
        time_wiggle_ncols: 0,
        intercept_warm_starts: None,
        auto_subsample_phase_counter: Arc::new(AtomicUsize::new(0)),
        auto_subsample_last_rho: Arc::new(Mutex::new(None)),
    };
    let block_states = vec![
        ParameterBlockState {
            beta: Array1::zeros(1),
            eta: Array1::zeros(1),
        },
        ParameterBlockState {
            beta: Array1::zeros(0),
            eta: Array1::zeros(1),
        },
        ParameterBlockState {
            beta: Array1::zeros(0),
            eta: array![0.6],
        },
    ];
    let q_geom = family
        .row_dynamic_q_geometry(0, &block_states)
        .expect("row geometry");
    let primary = flex_primary_slices(&family);
    let (nll_exact, grad_exact, hess_exact) = family
        .compute_row_flex_primary_gradient_hessian_exact(0, &block_states, &q_geom, &primary)
        .expect("exact flex row");
    let (nll_rigid, grad_rigid, hess_rigid) = row_primary_closed_form(
        q_geom.q0,
        q_geom.q1,
        q_geom.qd1,
        block_states[2].eta[0],
        family.z[[0, 0]],
        family.weights[0],
        family.event[0],
        family.derivative_guard,
        family.probit_frailty_scale(),
    )
    .expect("rigid row");

    assert!((nll_exact - nll_rigid).abs() < 1e-10);
    for idx in 0..N_PRIMARY {
        assert!((grad_exact[idx] - grad_rigid[idx]).abs() < 1e-8);
    }
    for i in 0..N_PRIMARY {
        for j in 0..N_PRIMARY {
            assert!((hess_exact[[i, j]] - hess_rigid[i][j]).abs() < 1e-7);
        }
    }
}

#[test]
fn row_primary_closed_form_rejects_negative_infinite_signed_margin() {
    let err = row_primary_closed_form(f64::INFINITY, 0.0, 1.0, 0.0, 0.0, 1.0, 1.0, 1e-6, 1.0)
        .expect_err("exact closed-form row should reject -inf signed margins");
    assert!(err.contains("non-finite signed margin"));
}

/// Mechanism-of-ρ=2 proof for the inner-PIRLS pathology on large-scale
/// saturated probit fits.
///
/// `add_pullback_primary_hessian` (line 9753) sums `h[0,0] + h[1,1]`
/// over rows into the marginal-block joint Hessian (β_marg enters both
/// `q0` and `q1` with the SAME `marginal_design` row, so the Jacobian
/// pullback adds the q0 and q1 second derivatives at the same slot).
///
/// Mathematical facts pinned here:
///   - Censored rows: `h[0,0] + h[1,1] = 0` to ULP at every η ≥ 0.
///     The entry term `+w·log Φ(−η₀)` (concave, curvature −w·c²) and
///     the exit term `−w·log Φ(−η₁)` (convex, curvature +w·c²) have
///     equal-and-opposite second derivatives when q₀ = q₁.
///   - Event rows: residual `+w·c²·(1/η² + O(1/η⁴))` from the Mills
///     asymptotic. The event-density term `+w·η₁²/2` contributes
///     exactly `+w·c²`; the entry-survival term contributes
///     `−w·c²·(1 − 1/η²)`. Sum = `+w·c²/η²`.
///
/// At large-scale saturation (η ~ 988), the marginal-block joint Hessian
/// collapses to `O(w/η²) = O(1e-6)` per event row from the likelihood
/// side; censored contributions are 0 to ULP. The Newton step in that
/// block is then dominated by the smoothing penalty `S_marg`. When
/// the saturating direction lies in the null space of `S_marg`
/// (typical for the duchon-smooth's polynomial null space), the
/// effective curvature drops to the f64 ridge floor, the inner
/// Newton step is set by ridge alone, and `actual = rhs·δ` while
/// `predicted = ½·rhs·δ` — yielding ρ ≡ 2 to floating-point precision
/// as observed.
#[test]
fn marginal_block_hessian_cancels_in_saturated_regime() {
    let probit_scale = 1.0_f64;
    let w = 1.0_f64;
    let derivative_guard = 1e-6;
    let qd1 = 1.0_f64;
    let g = 0.0_f64;
    let z = 0.0_f64;

    // Censored rows, q0 = q1 = η, at a wide range of saturations:
    // cancellation must be ULP-exact for every η.
    for &eta in &[0.5_f64, 1.0, 2.0, 5.0, 10.0, 40.0, 100.0, 500.0, 988.0] {
        let (_nll, _grad, hess) =
            row_primary_closed_form(eta, eta, qd1, g, z, w, 0.0, derivative_guard, probit_scale)
                .expect("rigid censored row");
        let sum = hess[0][0] + hess[1][1];
        assert!(
            sum.abs() <= 1e-12 * (hess[0][0].abs() + hess[1][1].abs()).max(1.0),
            "censored cancellation broke at η={eta}: h[0,0]={:.3e} h[1,1]={:.3e} sum={:.3e}",
            hess[0][0],
            hess[1][1],
            sum,
        );
    }

    // Event rows, q0 = q1 = η, deep saturation: residual scales as
    // 1/η² by Mills asymptotic M(−η) = η + 1/η + O(1/η³).
    for &eta in &[40.0_f64, 100.0, 500.0, 988.0] {
        let (_nll, _grad, hess) =
            row_primary_closed_form(eta, eta, qd1, g, z, w, 1.0, derivative_guard, probit_scale)
                .expect("rigid event row");
        let sum = hess[0][0] + hess[1][1];
        let bound = 2.0 / (eta * eta);
        assert!(
            sum > 0.0 && sum <= bound,
            "event cancellation residual at η={eta}: sum={:.3e} expected (0, {:.3e}]",
            sum,
            bound,
        );
    }

    // Cross-check at η = 988 (the user's large-scale saturation):
    // both kinds of rows hit the predicted floor exactly.
    let (_, _, ev) =
        row_primary_closed_form(988.0, 988.0, 1.0, 0.0, 0.0, 1.0, 1.0, 1e-6, 1.0).unwrap();
    let (_, _, ce) =
        row_primary_closed_form(988.0, 988.0, 1.0, 0.0, 0.0, 1.0, 0.0, 1e-6, 1.0).unwrap();
    let ev_sum = ev[0][0] + ev[1][1];
    let ce_sum = ce[0][0] + ce[1][1];
    assert!(
        ev_sum > 0.0 && ev_sum < 2.0e-6,
        "event saturated h[0,0]+h[1,1] = {ev_sum:.3e}, expected ~1/988² ≈ 1e-6",
    );
    assert_eq!(
        ce_sum, 0.0,
        "censored saturated h[0,0]+h[1,1] must be EXACTLY 0, got {ce_sum:.3e}",
    );
}

#[test]
fn row_primary_closed_form_rejects_nan_signed_margin() {
    let err = row_primary_closed_form(f64::NAN, 0.0, 1.0, 0.0, 0.0, 1.0, 1.0, 1e-6, 1.0)
        .expect_err("exact closed-form row should reject NaN signed margins");
    assert!(err.contains("non-finite signed margin"));
}

#[test]
fn rigid_row_kernel_propagates_invalid_nonfinite_signed_margin_errors() {
    let mut family = test_family(None, None);
    family.offset_entry = Arc::new(array![f64::INFINITY]);
    family.offset_exit = Arc::new(array![0.0]);
    family.derivative_offset_exit = Arc::new(array![1.0]);
    family.event = Arc::new(array![1.0]);

    let kernel = SurvivalMarginalSlopeRowKernel::new(
        family,
        vec![
            ParameterBlockState {
                beta: Array1::zeros(1),
                eta: Array1::zeros(1),
            },
            ParameterBlockState {
                beta: Array1::zeros(2),
                eta: Array1::zeros(1),
            },
            ParameterBlockState {
                beta: Array1::zeros(3),
                eta: Array1::zeros(1),
            },
        ],
    );

    let err =
        <SurvivalMarginalSlopeRowKernel as crate::row_kernel::RowKernel<4>>::row_kernel(&kernel, 0)
            .expect_err("row kernel should propagate exact probit boundary failures");
    assert!(err.contains("non-finite signed margin"));
}

#[test]
fn rigid_row_kernel_propagates_nan_signed_margin_errors() {
    let mut family = test_family(None, None);
    family.offset_entry = Arc::new(array![f64::NAN]);
    family.offset_exit = Arc::new(array![0.0]);
    family.derivative_offset_exit = Arc::new(array![1.0]);
    family.event = Arc::new(array![1.0]);

    let kernel = SurvivalMarginalSlopeRowKernel::new(
        family,
        vec![
            ParameterBlockState {
                beta: Array1::zeros(1),
                eta: Array1::zeros(1),
            },
            ParameterBlockState {
                beta: Array1::zeros(2),
                eta: Array1::zeros(1),
            },
            ParameterBlockState {
                beta: Array1::zeros(3),
                eta: Array1::zeros(1),
            },
        ],
    );

    let err =
        <SurvivalMarginalSlopeRowKernel as crate::row_kernel::RowKernel<4>>::row_kernel(&kernel, 0)
            .expect_err("row kernel should propagate NaN probit boundary failures");
    assert!(err.contains("non-finite signed margin"));
}

/// The single-expression Taylor-jet tower (#932) of the rigid K=1
/// survival marginal-slope row NLL, written ONCE over generic jet
/// primaries `(q0, q1, qd1, g)`. It reuses the family's OWN hand-certified
/// `[f64; 5]` special-function derivative stacks (`unary_derivatives_sqrt`
/// / `_neglog_phi` / `_log_normal_pdf` / `_log`) through
/// `JetScalar::compose_unary`, so no probit/log primitive is re-derived here:
/// the tower mechanizes only the Leibniz / Faà di Bruno composition that
/// `row_primary_closed_form` previously coded by hand (where #736's
/// cross-block sign flip lived). The value channel of the returned tower
/// is the row NLL expression, so every derivative channel is exact by
/// construction.
struct SurvivalMarginalSlopeRigidNllProgram {
    primaries: Vec<[f64; 4]>,
    z: Vec<f64>,
    w: Vec<f64>,
    d: Vec<f64>,
    probit_scale: f64,
}

impl gam_math::jet_tower::RowProgram<4> for SurvivalMarginalSlopeRigidNllProgram {
    fn n_rows(&self) -> usize {
        self.primaries.len()
    }

    fn primaries(&self, row: usize) -> Result<[f64; 4], String> {
        self.primaries
            .get(row)
            .copied()
            .ok_or_else(|| format!("rigid nll program: row {row} out of range"))
    }

    fn eval<S: gam_math::jet_scalar::JetScalar<4>>(
        &self,
        row: usize,
        p: &[S; 4],
    ) -> Result<S, String> {
        let z = *self
            .z
            .get(row)
            .ok_or_else(|| format!("rigid nll program: z row {row} out of range"))?;
        let w = self.w[row];
        let d = self.d[row];
        let s_f = self.probit_scale;
        let q0 = p[0];
        let q1 = p[1];
        let qd1 = p[2];
        let g = p[3];

        // c(g) = sqrt(1 + (s_f g)^2)  — K=1 covariance_ones = 1, exactly the
        // shared MultiDirJet `one_plus_b2 -> sqrt` composition.
        let observed_g = g.scale(s_f);
        let one_plus_b2 = observed_g.mul(&observed_g).add(&S::constant(1.0));
        let c = one_plus_b2.compose_unary(unary_derivatives_sqrt(one_plus_b2.value()));

        let eta0 = q0.mul(&c).add(&observed_g.scale(z));
        let eta1 = q1.mul(&c).add(&observed_g.scale(z));
        let ad1 = qd1.mul(&c);

        // Entry survival: +w logΦ(-η0) = -1 * (-w logΦ(-η0)).
        let neg_eta0 = eta0.neg();
        let entry = neg_eta0
            .compose_unary(unary_derivatives_neglog_phi(neg_eta0.value(), w))
            .scale(-1.0);
        // Exit survival: (1-d) * (-w logΦ(-η1)) carried with weight w(1-d).
        let neg_eta1 = eta1.neg();
        let exit = neg_eta1.compose_unary(unary_derivatives_neglog_phi(
            neg_eta1.value(),
            w * (1.0 - d),
        ));
        // Event density: -w d logφ(η1).
        let event_density = if d > 0.0 {
            eta1.compose_unary(unary_derivatives_log_normal_pdf(eta1.value()))
                .scale(-w * d)
        } else {
            S::constant(0.0)
        };
        // Time derivative: -w d log(ad1).
        let time_deriv = if d > 0.0 {
            ad1.compose_unary(unary_derivatives_log(ad1.value()))
                .scale(-w * d)
        } else {
            S::constant(0.0)
        };

        Ok(exit.add(&entry).add(&event_density).add(&time_deriv))
    }
}

/// Build a rigid K=1 survival marginal-slope family whose entry/exit/
/// derivative/marginal/logslope designs are all nontrivial dense blocks,
/// so every one of the four primaries `(q0, q1, qd1, g)` is exercised
/// when the kernel reads its designs. `n` rows, `event` per row.
fn oracle_rigid_family(
    n: usize,
    z: &[f64],
    weights: &[f64],
    event: &[f64],
    gaussian_frailty_sd: Option<f64>,
) -> SurvivalMarginalSlopeFamily {
    let z_col = Array2::from_shape_fn((n, 1), |(r, _)| z[r]);
    // Distinct entry/exit/derivative rows so q0 != q1 != qd1.
    let design_entry = Array2::from_shape_fn((n, 1), |(r, _)| {
        0.4 + 0.13 * (r as f64) - 0.05 * (r as f64).cos()
    });
    let design_exit = Array2::from_shape_fn((n, 1), |(r, _)| {
        0.9 + 0.07 * (r as f64) + 0.04 * (r as f64).sin()
    });
    // Strictly positive derivative-exit design so qd1 > 0 (monotone).
    let design_deriv = Array2::from_shape_fn((n, 1), |(r, _)| 1.2 + 0.21 * (r as f64).abs().sqrt());
    SurvivalMarginalSlopeFamily {
        n,
        event: Arc::new(Array1::from(event.to_vec())),
        weights: Arc::new(Array1::from(weights.to_vec())),
        z: Arc::new(z_col),
        score_covariance: unit_score_covariance(),
        gaussian_frailty_sd,
        family_hyper: SurvivalMarginalSlopeFamilyHyperState::default(),
        derivative_guard: 1e-8,
        design_entry: DesignMatrix::from(design_entry),
        design_exit: DesignMatrix::from(design_exit),
        design_derivative_exit: DesignMatrix::from(design_deriv),
        offset_entry: Arc::new(Array1::from_shape_fn(n, |r| 0.05 * (r as f64) - 0.2)),
        offset_exit: Arc::new(Array1::from_shape_fn(n, |r| 0.15 - 0.03 * (r as f64))),
        derivative_offset_exit: Arc::new(Array1::from_elem(n, 0.0)),
        marginal_design: DesignMatrix::from(Array2::zeros((n, 0))),
        logslope_layout: (DesignMatrix::from(Array2::zeros((n, 0)))).into(),
        score_warp: None,
        link_dev: None,
        influence_absorber: None,
        time_linear_constraints: None,
        time_wiggle_knots: None,
        time_wiggle_degree: None,
        time_wiggle_ncols: 0,
        intercept_warm_starts: None,
        auto_subsample_phase_counter: Arc::new(AtomicUsize::new(0)),
        auto_subsample_last_rho: Arc::new(Mutex::new(None)),
    }
}

/// #932 dedicated correctness oracle for the rigid family-direction
/// lowerings (`rigid_family_direction_terms` /
/// `rigid_family_direction_beta_drift` in `eval_family.rs`) — previously the
/// only jet-derived production lowerings of this family covered solely
/// transitively (through the ψ-terms integration tests).
///
/// The Dual2-over-`Order2` output must equal a central finite-difference
/// witness of the direct symbolic order-2 lowering (`rigid_row_order2`, a
/// different code path from the nested-dual instantiation) along the
/// quadratic primary path `p(t) = p + t·first + (t²/2)·second`:
/// the first channel is `d/dt|₀` and the second channel is `d²/dt²|₀` of
/// every (objective, gradient, Hessian) entry. The beta-drift lowering must
/// equal a central FD over realizable block-state shifts (marginal-η,
/// logslope-η, and time-β directions) of the first-channel terms themselves,
/// pinning the `OneSeed::eps` wiring end-to-end through the production
/// helper.
#[test]
fn rigid_family_direction_terms_match_fd_witness_932() {
    use super::row_kernel::{rigid_row_inputs, rigid_row_kernel_primaries, rigid_row_order2};

    let n = 5;
    let z = [0.4, -1.1, 0.0, 0.7, -0.3];
    let weights = [1.0, 0.8, 1.3, 0.9, 1.1];
    let event = [1.0, 0.0, 0.0, 1.0, 1.0];
    let g_eta = array![0.2, -0.5, 0.35, -0.15, 0.6];
    let marginal_eta = array![0.1, -0.2, 0.05, 0.12, -0.08];

    let first = [0.3_f64, -0.6, 0.4, 0.7];
    let second = [-0.2_f64, 0.5, 0.1, -0.3];

    let close = |label: &str, got: f64, want: f64, band_scale: f64| {
        let band = band_scale * got.abs().max(want.abs()).max(1.0);
        assert!(
            got.is_finite() && want.is_finite() && (got - want).abs() <= band,
            "{label}: production={got:+.12e} fd={want:+.12e} band={band:.3e}"
        );
    };

    for frailty in [None, Some(0.6_f64)] {
        let family = oracle_rigid_family(n, &z, &weights, &event, frailty);
        let beta_time = array![0.85];
        let block_states = vec![
            ParameterBlockState {
                beta: beta_time.clone(),
                eta: Array1::zeros(n),
            },
            ParameterBlockState {
                beta: Array1::zeros(0),
                eta: marginal_eta.clone(),
            },
            ParameterBlockState {
                beta: Array1::zeros(0),
                eta: g_eta.clone(),
            },
        ];

        for row in 0..n {
            let (terms_first, terms_second) = family
                .rigid_family_direction_terms(row, &block_states, first, second)
                .expect("family-direction terms");

            let p = rigid_row_kernel_primaries(&family, &block_states, row).expect("primaries");
            let inputs = rigid_row_inputs(&family, &block_states, row, "family-direction oracle")
                .expect("row inputs");
            let at = |t: f64| -> (f64, [f64; 4], [[f64; 4]; 4]) {
                let shifted: [f64; 4] =
                    std::array::from_fn(|a| p[a] + t * first[a] + 0.5 * t * t * second[a]);
                rigid_row_order2(&shifted, &inputs).expect("shifted symbolic row")
            };

            // First channel: d/dt at 0 (central, h=1e-6).
            let h1 = 1.0e-6;
            let (plus_v, plus_g, plus_h) = at(h1);
            let (minus_v, minus_g, minus_h) = at(-h1);
            close(
                &format!("row {row} first objective"),
                terms_first.objective,
                (plus_v - minus_v) / (2.0 * h1),
                1.0e-7,
            );
            for a in 0..4 {
                close(
                    &format!("row {row} first gradient[{a}]"),
                    terms_first.gradient[a],
                    (plus_g[a] - minus_g[a]) / (2.0 * h1),
                    1.0e-7,
                );
                for b in 0..4 {
                    close(
                        &format!("row {row} first hessian[{a},{b}]"),
                        terms_first.hessian[[a, b]],
                        (plus_h[a][b] - minus_h[a][b]) / (2.0 * h1),
                        1.0e-7,
                    );
                }
            }

            // Second channel: d²/dt² at 0 (central, h=1e-4).
            let h2 = 1.0e-4;
            let (p2v, p2g, p2h) = at(h2);
            let (m2v, m2g, m2h) = at(-h2);
            let (zv, zg, zh) = at(0.0);
            close(
                &format!("row {row} second objective"),
                terms_second.objective,
                (p2v - 2.0 * zv + m2v) / (h2 * h2),
                5.0e-5,
            );
            for a in 0..4 {
                close(
                    &format!("row {row} second gradient[{a}]"),
                    terms_second.gradient[a],
                    (p2g[a] - 2.0 * zg[a] + m2g[a]) / (h2 * h2),
                    5.0e-5,
                );
                for b in 0..4 {
                    close(
                        &format!("row {row} second hessian[{a},{b}]"),
                        terms_second.hessian[[a, b]],
                        (p2h[a][b] - 2.0 * zh[a][b] + m2h[a][b]) / (h2 * h2),
                        5.0e-5,
                    );
                }
            }

            // Beta-drift channel: for each realizable block-state direction,
            // production `OneSeed::eps` must equal the central FD of the
            // first-channel terms across the shifted states.
            let probes: [(&str, usize, [f64; 4]); 3] = [
                ("marginal-eta", 1, [1.0, 1.0, 0.0, 0.0]),
                ("logslope-eta", 2, [0.0, 0.0, 0.0, 1.0]),
                (
                    "time-beta",
                    0,
                    [
                        family.design_entry.dot_row(row, &array![1.0]),
                        family.design_exit.dot_row(row, &array![1.0]),
                        family.design_derivative_exit.dot_row(row, &array![1.0]),
                        0.0,
                    ],
                ),
            ];
            for (label, block, beta_dir) in probes {
                let drift = family
                    .rigid_family_direction_beta_drift(row, &block_states, first, beta_dir)
                    .expect("family-direction beta drift");
                let hs = 1.0e-5;
                let shifted_terms = |s: f64| {
                    let mut states = block_states.clone();
                    if block == 0 {
                        states[0].beta[0] += s;
                    } else {
                        states[block].eta[row] += s;
                    }
                    family
                        .rigid_family_direction_terms(row, &states, first, [0.0; 4])
                        .expect("shifted family-direction terms")
                        .0
                };
                let plus = shifted_terms(hs);
                let minus = shifted_terms(-hs);
                close(
                    &format!("row {row} {label} drift objective"),
                    drift.objective,
                    (plus.objective - minus.objective) / (2.0 * hs),
                    1.0e-6,
                );
                for a in 0..4 {
                    close(
                        &format!("row {row} {label} drift gradient[{a}]"),
                        drift.gradient[a],
                        (plus.gradient[a] - minus.gradient[a]) / (2.0 * hs),
                        1.0e-6,
                    );
                    for b in 0..4 {
                        close(
                            &format!("row {row} {label} drift hessian[{a},{b}]"),
                            drift.hessian[[a, b]],
                            (plus.hessian[[a, b]] - minus.hessian[[a, b]]) / (2.0 * hs),
                            1.0e-6,
                        );
                    }
                }
            }
        }
    }
}

/// #932 universal oracle on the ONLY production `RowKernel` impl.
///
/// Audits every channel the production `SurvivalMarginalSlopeRowKernel`
/// emits — value / gradient / Hessian / `row_third_contracted(dir)` /
/// `row_fourth_contracted(u, v)` — against the single-expression
/// `RowProgram<4>`-derived tower truth, over several fixture rows
/// (mixed event/censored, with and without Gaussian frailty so the probit
/// scale ≠ 1) and several random direction vectors. The cross blocks that
/// #736's sign flip corrupted are contracted explicitly. Agreement here is
/// the proof the production kernel is correct; the planted-flip test in
/// `jet_tower` proves the same harness is loud on disagreement.
#[test]
fn rigid_row_kernel_agrees_with_jet_tower_program_all_channels() {
    use crate::row_kernel::RowKernel;
    use gam_math::jet_tower::{KernelChannels, program_full_tower, verify_kernel_channels};

    let n = 7;
    let z = [0.4, -1.1, 0.0, 0.7, -0.3, 1.6, -1.4];
    let weights = [1.0, 0.8, 1.3, 0.9, 1.1, 0.7, 1.4];
    // Mix exact events and right-censored rows; the last two rows push
    // eta into opposite normal-tail regimes while retaining finite truth.
    let event = [1.0, 0.0, 0.0, 1.0, 1.0, 0.0, 1.0];
    // logslope eta (g) per row, fed through block 2.
    let g_eta = array![0.2, -0.5, 0.35, -0.15, 0.6, 0.45, -0.55];
    // marginal eta (block 1) — additive index shared by η0 and η1.
    let marginal_eta = array![0.1, -0.2, 0.05, 0.12, -0.08, 6.2, -7.4];

    // Deterministic pseudo-random direction vectors (no RNG dependency).
    let dirs: [[f64; 4]; 3] = [
        [0.7, -1.3, 0.5, 0.9],
        [-0.4, 0.6, -1.1, 0.3],
        [1.2, 0.2, -0.7, -0.5],
    ];

    for frailty in [None, Some(0.6_f64)] {
        let family = oracle_rigid_family(n, &z, &weights, &event, frailty);
        let probit_scale = family.probit_frailty_scale();
        let beta_time = array![0.85]; // single time coefficient
        let block_states = vec![
            ParameterBlockState {
                beta: beta_time.clone(),
                eta: Array1::zeros(n),
            },
            ParameterBlockState {
                beta: Array1::zeros(0),
                eta: marginal_eta.clone(),
            },
            ParameterBlockState {
                beta: Array1::zeros(0),
                eta: g_eta.clone(),
            },
        ];

        let kernel = SurvivalMarginalSlopeRowKernel::new(family.clone(), block_states.clone());

        // Build the tower program's primaries exactly as `row_kernel` reads
        // them (designs · beta_time + offsets + shared marginal/logslope eta).
        let mut primaries = Vec::with_capacity(n);
        for row in 0..n {
            let q0 = family.design_entry.dot_row(row, &beta_time)
                + family.offset_entry[row]
                + marginal_eta[row];
            let q1 = family.design_exit.dot_row(row, &beta_time)
                + family.offset_exit[row]
                + marginal_eta[row];
            let qd1 = family.design_derivative_exit.dot_row(row, &beta_time)
                + family.derivative_offset_exit[row];
            let g = g_eta[row];
            primaries.push([q0, q1, qd1, g]);
        }

        let program = SurvivalMarginalSlopeRigidNllProgram {
            primaries,
            z: z.to_vec(),
            w: weights.to_vec(),
            d: event.to_vec(),
            probit_scale,
        };

        for row in 0..n {
            let tower = program_full_tower(&program, row).expect("tower evaluation");

            let (value, gradient, hessian) =
                RowKernel::row_kernel(&kernel, row).expect("production kernel value/grad/hess");

            let third: Vec<([f64; 4], [[f64; 4]; 4])> = dirs
                .iter()
                .map(|dir| {
                    let claim = RowKernel::row_third_contracted(&kernel, row, dir)
                        .expect("production kernel third");
                    (*dir, claim)
                })
                .collect();

            let fourth: Vec<([f64; 4], [f64; 4], [[f64; 4]; 4])> = dirs
                .iter()
                .enumerate()
                .map(|(i, u)| {
                    let v = dirs[(i + 1) % dirs.len()];
                    let claim = RowKernel::row_fourth_contracted(&kernel, row, u, &v)
                        .expect("production kernel fourth");
                    (*u, v, claim)
                })
                .collect();

            let claims = KernelChannels {
                value,
                gradient,
                hessian,
                third,
                fourth,
            };

            verify_kernel_channels(&tower, &claims, 1e-9).unwrap_or_else(|e| {
                    panic!(
                        "frailty {frailty:?} row {row}: production RowKernel disagrees with #932 jet-tower truth: {e}"
                    )
                });
        }
    }
}

/// #932 unified-feature oracle: the scalar/shared 5->4 order-two pullback must
/// agree with the generic `Order2<4>` feature-map evaluation on every channel,
/// and its admission wrapper must agree with the dependency-sliced witness
/// surface. The random grid covers both event branches, non-unit covariance,
/// and non-unit frailty scale from the sole `rigid_feature_program` declaration.
// The witness half of this fixture calls `rigid_row_admission_witnesses`, which
// is `cfg(target_os = "linux")` production code, so off-Linux the test target
// does not compile at all. Gate the fixture with the surface it exercises.
#[cfg(target_os = "linux")]
#[test]
fn rigid_feature_program_scalar_pullback_matches_generic_and_witnesses_932() {
    use gam_math::jet_scalar::{JetScalar, Order2};

    // Deterministic xorshift grid (no RNG dependency).
    let mut s: u64 = 0x9E3779B97F4A7C15;
    let mut nx = || {
        s ^= s << 13;
        s ^= s >> 7;
        s ^= s << 17;
        ((s >> 11) as f64) / ((1u64 << 53) as f64) * 2.0 - 1.0
    };

    let mut max_rel = 0.0_f64;
    for _ in 0..4000 {
        let p = [nx() * 1.5, nx() * 1.5, 0.5 + nx().abs() * 2.0, nx() * 1.2];
        let inputs = RigidRowInputs {
            row: 0,
            wi: 0.5 + nx().abs(),
            di: if nx() > 0.0 { 1.0 } else { 0.0 },
            z_sum: nx() * 1.2,
            covariance_ones: 0.7 + nx().abs(),
            probit_scale: 0.6 + nx().abs(),
            qd1_lower: -1.0,
        };

        let dense_vars: [Order2<4>; 4] = std::array::from_fn(|a| Order2::variable(p[a], a));
        let dense = rigid_row_nll(&dense_vars, &inputs).expect("generic scalar feature map");
        let (value, gradient, hessian) =
            rigid_row_order2(&p, &inputs).expect("scalar feature pullback");
        let observed_g = inputs.probit_scale * p[3];
        let (_, _, _, semantic_witnesses) = rigid_feature_program_order2(
            p[0],
            p[1],
            p[2],
            observed_g * inputs.z_sum,
            (p[3] * p[3]) * inputs.covariance_ones,
            inputs.wi,
            inputs.di,
            inputs.probit_scale,
        );
        let sliced_witnesses = rigid_row_admission_witnesses(&p, &inputs);

        let mut check = |a: f64, b: f64| {
            let rel = (a - b).abs() / (1.0 + a.abs().max(b.abs()));
            if rel > max_rel {
                max_rel = rel;
            }
            assert!(
                (a - b).abs() <= 1e-12 + 1e-12 * a.abs().max(b.abs()),
                "generated vs generic channel disagreement: {a:+.16e} vs {b:+.16e}"
            );
        };
        check(value, dense.value());
        for a in 0..4 {
            check(gradient[a], dense.g()[a]);
            for b in 0..4 {
                check(hessian[a][b], dense.h()[a][b]);
            }
        }
        for witness in 0..3 {
            check(semantic_witnesses[witness], sliced_witnesses[witness]);
        }
    }
    assert!(
        max_rel <= 1e-12,
        "generated vs generic max relative error {max_rel:.3e} exceeds 1e-12"
    );
}

#[test]
fn exact_flex_row_value_matches_rigid_with_zero_score_and_link_coefficients() {
    let score_runtime = test_deviation_runtime();
    let link_runtime = test_deviation_runtime();
    let family = SurvivalMarginalSlopeFamily {
        n: 1,
        event: Arc::new(array![0.0]),
        weights: Arc::new(array![0.9]),
        z: Arc::new(array![-0.35].insert_axis(Axis(1))),
        score_covariance: unit_score_covariance(),
        gaussian_frailty_sd: None,
        family_hyper: SurvivalMarginalSlopeFamilyHyperState::default(),
        derivative_guard: 1e-6,
        design_entry: DesignMatrix::from(Array2::zeros((1, 1))),
        design_exit: DesignMatrix::from(Array2::zeros((1, 1))),
        design_derivative_exit: DesignMatrix::from(Array2::zeros((1, 1))),
        offset_entry: Arc::new(array![-0.1]),
        offset_exit: Arc::new(array![0.15]),
        derivative_offset_exit: Arc::new(array![0.6]),
        marginal_design: DesignMatrix::from(Array2::zeros((1, 0))),
        logslope_layout: (DesignMatrix::from(Array2::zeros((1, 0)))).into(),
        score_warp: Some(score_runtime.clone()),
        link_dev: Some(link_runtime.clone()),
        influence_absorber: None,
        time_linear_constraints: None,
        time_wiggle_knots: None,
        time_wiggle_degree: None,
        time_wiggle_ncols: 0,
        intercept_warm_starts: None,
        auto_subsample_phase_counter: Arc::new(AtomicUsize::new(0)),
        auto_subsample_last_rho: Arc::new(Mutex::new(None)),
    };
    let block_states = vec![
        ParameterBlockState {
            beta: Array1::zeros(1),
            eta: Array1::zeros(1),
        },
        ParameterBlockState {
            beta: Array1::zeros(0),
            eta: Array1::zeros(1),
        },
        ParameterBlockState {
            beta: Array1::zeros(0),
            eta: array![0.45],
        },
        ParameterBlockState {
            beta: Array1::zeros(score_runtime.basis_dim()),
            eta: Array1::zeros(1),
        },
        ParameterBlockState {
            beta: Array1::zeros(link_runtime.basis_dim()),
            eta: Array1::zeros(1),
        },
    ];
    let q_geom = family
        .row_dynamic_q_geometry(0, &block_states)
        .expect("row geometry");
    let primary = flex_primary_slices(&family);
    let (nll_exact, grad_exact, hess_exact) = family
        .compute_row_flex_primary_gradient_hessian_exact(0, &block_states, &q_geom, &primary)
        .expect("exact flex row");
    let (nll_rigid, grad_rigid, hess_rigid) = row_primary_closed_form(
        q_geom.q0,
        q_geom.q1,
        q_geom.qd1,
        block_states[2].eta[0],
        family.z[[0, 0]],
        family.weights[0],
        family.event[0],
        family.derivative_guard,
        family.probit_frailty_scale(),
    )
    .expect("rigid row");

    assert!((nll_exact - nll_rigid).abs() < 1e-10);
    assert!((grad_exact[primary.q0] - grad_rigid[0]).abs() < 1e-8);
    assert!((grad_exact[primary.q1] - grad_rigid[1]).abs() < 1e-8);
    assert!((grad_exact[primary.qd1] - grad_rigid[2]).abs() < 1e-8);
    assert!((grad_exact[primary.g] - grad_rigid[3]).abs() < 1e-8);
    assert!((hess_exact[[primary.q0, primary.q0]] - hess_rigid[0][0]).abs() < 1e-7);
    assert!((hess_exact[[primary.q0, primary.g]] - hess_rigid[0][3]).abs() < 1e-7);
    assert!((hess_exact[[primary.q1, primary.q1]] - hess_rigid[1][1]).abs() < 1e-7);
    assert!((hess_exact[[primary.q1, primary.g]] - hess_rigid[1][3]).abs() < 1e-7);
    assert!((hess_exact[[primary.qd1, primary.qd1]] - hess_rigid[2][2]).abs() < 1e-7);
    assert!((hess_exact[[primary.g, primary.g]] - hess_rigid[3][3]).abs() < 1e-7);
}

/// gam#932/#979: INDEPENDENT witness for the survival marginal-slope FLEX
/// higher-order tower (`row_flex_primary_{third,fourth}_contracted_exact`).
///
/// The rigid K=1 path is independently guarded by
/// `SurvivalMarginalSlopeRigidNllProgram` (a single-expression `Tower4` algebra
/// re-derivation, no shared jet code); the flex path was not.
///
/// This closes that gap on the `(q0, q1, qd1, g)` primary block: at ZERO
/// deviation coefficients the flex de-nested calibration / cell-moment / intercept
/// machinery is fully exercised (it does NOT short-circuit to the rigid kernel —
/// it runs the partition with zero-coefficient cubics), and its third/fourth
/// directional contractions over `(q0, q1, qd1, g)` MUST equal the independent
/// rigid `Tower4` truth. A planted cross-block sign flip must leave the band,
/// proving resolving power (the #736 genus the shared-input parity cannot catch).
#[test]
fn flex_contracted_tower_matches_independent_rigid_tower_and_catches_sign_flip() {
    use gam_math::jet_tower::{program_fourth_contracted, program_third_contracted};

    let score_runtime = test_deviation_runtime();
    let link_runtime = test_deviation_runtime();
    // Several fixture rows: events + censored, distinct q-geometry, frailty off
    // (probit scale = 1; the rigid program and the flex path share the closed
    // form there) and a non-zero logslope g so the c(g) coupling is live.
    struct Fix {
        event: f64,
        weight: f64,
        z: f64,
        q0: f64,
        q1: f64,
        qd1: f64,
        g: f64,
    }
    let fixtures = [
        Fix {
            event: 1.0,
            weight: 0.75,
            z: -0.2,
            q0: -0.4,
            q1: 0.6,
            qd1: 0.85,
            g: 0.32,
        },
        Fix {
            event: 0.0,
            weight: 1.35,
            z: -1.15,
            q0: -1.35,
            q1: -0.9,
            qd1: 0.42,
            g: -0.55,
        },
        Fix {
            event: 1.0,
            weight: 0.9,
            z: 0.7,
            q0: 0.15,
            q1: 1.05,
            qd1: 0.6,
            g: 0.45,
        },
    ];

    for fix in &fixtures {
        let family = SurvivalMarginalSlopeFamily {
            n: 1,
            event: Arc::new(array![fix.event]),
            weights: Arc::new(array![fix.weight]),
            z: Arc::new(array![fix.z].insert_axis(Axis(1))),
            score_covariance: unit_score_covariance(),
            gaussian_frailty_sd: None,
            family_hyper: SurvivalMarginalSlopeFamilyHyperState::default(),
            derivative_guard: 1e-6,
            design_entry: DesignMatrix::from(Array2::zeros((1, 1))),
            design_exit: DesignMatrix::from(Array2::zeros((1, 1))),
            design_derivative_exit: DesignMatrix::from(Array2::zeros((1, 1))),
            offset_entry: Arc::new(array![fix.q0]),
            offset_exit: Arc::new(array![fix.q1]),
            derivative_offset_exit: Arc::new(array![fix.qd1]),
            marginal_design: DesignMatrix::from(Array2::zeros((1, 0))),
            logslope_layout: (DesignMatrix::from(Array2::zeros((1, 0)))).into(),
            score_warp: Some(score_runtime.clone()),
            link_dev: Some(link_runtime.clone()),
            influence_absorber: None,
            time_linear_constraints: None,
            time_wiggle_knots: None,
            time_wiggle_degree: None,
            time_wiggle_ncols: 0,
            intercept_warm_starts: None,
            auto_subsample_phase_counter: Arc::new(AtomicUsize::new(0)),
            auto_subsample_last_rho: Arc::new(Mutex::new(None)),
        };
        // ZERO deviation coefficients: the flex calculus runs in full, but the
        // primary NLL reduces to the rigid closed form so the rigid Tower4 is the
        // exact independent truth on the (q0,q1,qd1,g) block.
        let block_states = vec![
            ParameterBlockState {
                beta: Array1::zeros(1),
                eta: Array1::zeros(1),
            },
            ParameterBlockState {
                beta: Array1::zeros(0),
                eta: Array1::zeros(1),
            },
            ParameterBlockState {
                beta: Array1::zeros(0),
                eta: array![fix.g],
            },
            ParameterBlockState {
                beta: Array1::zeros(score_runtime.basis_dim()),
                eta: Array1::zeros(1),
            },
            ParameterBlockState {
                beta: Array1::zeros(link_runtime.basis_dim()),
                eta: Array1::zeros(1),
            },
        ];

        let primary = flex_primary_slices(&family);
        let p = primary.total;
        // The four time/marginal/logslope primaries occupy the leading slots.
        let block_idx = [primary.q0, primary.q1, primary.qd1, primary.g];

        // Independent rigid Tower4 program at the SAME primaries.
        let program = SurvivalMarginalSlopeRigidNllProgram {
            primaries: vec![[fix.q0, fix.q1, fix.qd1, fix.g]],
            z: vec![fix.z],
            w: vec![fix.weight],
            d: vec![fix.event],
            probit_scale: family.probit_frailty_scale(),
        };

        // Distinct 4-vector directions confined to (q0,q1,qd1,g); embed each into
        // the full p-vector at the leading slots for the flex call.
        let dirs4: [[f64; 4]; 3] = [
            [0.7, -1.3, 0.5, 0.9],
            [-0.4, 0.6, -1.1, 0.3],
            [1.2, 0.2, -0.7, -0.5],
        ];
        let embed = |d4: &[f64; 4]| -> Array1<f64> {
            let mut full = Array1::zeros(p);
            for (k, &slot) in block_idx.iter().enumerate() {
                full[slot] = d4[k];
            }
            full
        };

        // ── Third-order: D_dir H over the (q0,q1,qd1,g) block ───────────────
        for d4 in &dirs4 {
            let flex_full = family
                .row_flex_primary_third_contracted_exact(0, &block_states, &embed(d4))
                .expect("flex third contracted at zero deviation");
            let rigid = program_third_contracted(&program, 0, d4).expect("rigid third");
            let scale = rigid
                .iter()
                .flatten()
                .fold(0.0_f64, |m, v| m.max(v.abs()))
                .max(1.0);
            for (u, &bu) in block_idx.iter().enumerate() {
                for (v, &bv) in block_idx.iter().enumerate() {
                    let got = flex_full[[bu, bv]];
                    let want = rigid[u][v];
                    assert!(
                        (got - want).abs() <= 1e-7 * scale,
                        "third[{u},{v}] flex {got:+.9e} != independent rigid tower {want:+.9e} (z={}, event={})",
                        fix.z,
                        fix.event
                    );
                }
            }
        }

        // Planted sign-flip tripwire on a representative third-order cross block.
        {
            let d4 = &dirs4[0];
            let flex_full = family
                .row_flex_primary_third_contracted_exact(0, &block_states, &embed(d4))
                .expect("flex third contracted (tripwire)");
            let rigid = program_third_contracted(&program, 0, d4).expect("rigid third (tripwire)");
            // (q0, g) cross block — the marginal↔logslope coupling.
            let want = rigid[0][3];
            let scale = want.abs().max(1.0);
            if want.abs() > 1e-6 {
                let flipped = -flex_full[[primary.q0, primary.g]];
                assert!(
                    (flipped - want).abs() > 1e-7 * scale,
                    "independent rigid tower failed to reject a planted (q0,g) sign flip: flipped {flipped:+.9e} vs truth {want:+.9e}"
                );
            }
        }

        // ── Fourth-order: D_u D_v H over the (q0,q1,qd1,g) block ─────────────
        let quad_pairs = [(0usize, 1usize), (1, 2), (2, 0)];
        for &(iu, iv) in &quad_pairs {
            let du = &dirs4[iu];
            let dv = &dirs4[iv];
            let flex_full = family
                .row_flex_primary_fourth_contracted_exact(0, &block_states, &embed(du), &embed(dv))
                .expect("flex fourth contracted at zero deviation");
            let rigid = program_fourth_contracted(&program, 0, du, dv).expect("rigid fourth");
            let scale = rigid
                .iter()
                .flatten()
                .fold(0.0_f64, |m, v| m.max(v.abs()))
                .max(1.0);
            for (u, &bu) in block_idx.iter().enumerate() {
                for (v, &bv) in block_idx.iter().enumerate() {
                    let got = flex_full[[bu, bv]];
                    let want = rigid[u][v];
                    assert!(
                        (got - want).abs() <= 1e-6 * scale,
                        "fourth[{u},{v}] flex {got:+.9e} != independent rigid tower {want:+.9e} (z={}, event={})",
                        fix.z,
                        fix.event
                    );
                }
            }
        }
    }
}

/// gam#932/#979: INDEPENDENT finite-difference witness for the survival
/// marginal-slope FLEX higher-order tower with NON-ZERO deviation coefficients —
/// the part Arm A (`flex_contracted_tower_matches_independent_rigid_tower_*`)
/// cannot reach, and the part most likely to harbor a shared-input bug.
///
/// The witness re-derives the scalar survival flex row NLL FROM SCRATCH over the
/// primary vector `p = [q0, q1, qd1, g, β_h..., β_w...]`:
///   * the deviation index `index(a, g, z) = a + g·warp(z) + linkdev(a + g·z)`
///     is reconstructed from the RAW basis matrices (`DeviationRuntime::design`),
///     dotted with the coefficients — no production jet / cell-moment code;
///   * the per-timepoint intercept `a(q)` is the calibration root
///     `∫ Φ(−index(a,g,z))·φ(z) dz = Φ(−q)`, solved by an independent secant on
///     a fine composite-Simpson quadrature of the latent normal, with the density
///     normalization `D = |∂F/∂a|` from the same quadrature;
///   * the NLL is assembled from the closed-form survival pieces
///     `w·[logΦ(−η0) − (1−d)logΦ(−η1) − d·logφ(η1) − d·log χ1 − d·logφ(q1)
///        + d·log D1 − d·log qd1]`.
///
/// To DE-RISK the witness (a witness-side re-derivation error would masquerade as
/// a production disagreement), the witness scalar NLL is first self-validated
/// against the production `row_neglog_flex_value` at the SAME non-zero β; only
/// then is it finite-differenced (Richardson O(h⁴)) and compared to the
/// production `row_flex_primary_{third,fourth}_contracted_exact`. A planted sign
/// flip on a deviation-touching cross block must leave the band.
#[test]
fn flex_contracted_tower_matches_independent_fd_witness_nonzero_deviation() {
    let score_runtime = test_deviation_runtime();
    let link_runtime = test_deviation_runtime();
    let h_dim = score_runtime.basis_dim();
    let w_dim = link_runtime.basis_dim();

    // No frailty ⇒ probit scale = 1 ⇒ the calibration measure is the standard
    // normal latent and the index is unscaled — the regime the witness derives.
    let z_row = 0.3_f64;
    let q0v = -0.25_f64;
    let q1v = 0.7_f64;
    let qd1v = 0.9_f64;
    let gv = 0.4_f64;
    let weight = 0.85_f64;
    let event = 1.0_f64;

    let family = SurvivalMarginalSlopeFamily {
        n: 1,
        event: Arc::new(array![event]),
        weights: Arc::new(array![weight]),
        z: Arc::new(array![z_row].insert_axis(Axis(1))),
        score_covariance: unit_score_covariance(),
        gaussian_frailty_sd: None,
        family_hyper: SurvivalMarginalSlopeFamilyHyperState::default(),
        derivative_guard: 1e-6,
        design_entry: DesignMatrix::from(Array2::zeros((1, 1))),
        design_exit: DesignMatrix::from(Array2::zeros((1, 1))),
        design_derivative_exit: DesignMatrix::from(Array2::zeros((1, 1))),
        offset_entry: Arc::new(array![q0v]),
        offset_exit: Arc::new(array![q1v]),
        derivative_offset_exit: Arc::new(array![qd1v]),
        marginal_design: DesignMatrix::from(Array2::zeros((1, 0))),
        logslope_layout: (DesignMatrix::from(Array2::zeros((1, 0)))).into(),
        score_warp: Some(score_runtime.clone()),
        link_dev: Some(link_runtime.clone()),
        influence_absorber: None,
        time_linear_constraints: None,
        time_wiggle_knots: None,
        time_wiggle_degree: None,
        time_wiggle_ncols: 0,
        intercept_warm_starts: None,
        auto_subsample_phase_counter: Arc::new(AtomicUsize::new(0)),
        auto_subsample_last_rho: Arc::new(Mutex::new(None)),
    };
    let primary = flex_primary_slices(&family);
    let p = primary.total;
    let h_range = primary.h.clone().expect("score-warp primary range");
    let w_range = primary.w.clone().expect("link-dev primary range");

    // Small, distinct non-zero deviation coefficients so every basis column
    // carries signal into the derivative chain.
    let beta_h0: Vec<f64> = (0..h_dim)
        .map(|k| 0.04 * ((k as f64 + 1.3).sin()))
        .collect();
    let beta_w0: Vec<f64> = (0..w_dim)
        .map(|k| 0.035 * ((k as f64 + 0.7).cos()))
        .collect();

    // ── Independent basis evaluations (raw design rows · β) ──────────────────
    let warp_eval = |beta_h: &[f64], z: f64| -> f64 {
        let row = score_runtime
            .design(&array![z])
            .expect("score-warp basis row");
        row.row(0).iter().zip(beta_h).map(|(b, c)| b * c).sum()
    };
    let linkdev_eval = |beta_w: &[f64], u: f64| -> f64 {
        let row = link_runtime.design(&array![u]).expect("link-dev basis row");
        row.row(0).iter().zip(beta_w).map(|(b, c)| b * c).sum()
    };
    // Survival deviation index inside Φ: a + g·z + g·warp(z) + linkdev(a + g·z).
    // The rigid slope term `g·z` is explicit: the score-warp value basis is a
    // pure DEVIATION (identity excluded — deviation_runtime.rs documents "zero
    // coefficients mean the identity map"), and production adds the rigid `b·z`
    // separately in the denested cell coefficient. Omitting it here left an
    // odd-in-z residual that the φ(z)-weighted calibration mostly absorbed into
    // the intercept, surfacing as a ~8.4e-4 self-validation gap.
    let index = |a: f64, g: f64, beta_h: &[f64], beta_w: &[f64], z: f64| -> f64 {
        a + g * z + g * warp_eval(beta_h, z) + linkdev_eval(beta_w, a + g * z)
    };
    // Witness-exact standard-normal primitives (`libm::erfc`, no piecewise
    // rational approximation). The intercept density `d1 = |∂F/∂a|` is taken by
    // a central FD with ε = 1e-6, which divides any error in `F` by 2ε = 2e-6;
    // production's `normal_cdf` carries an ~5e-12 oscillating approximation
    // error, so `F(a±ε)` inherit independent ~5e-12 perturbations whose FD is
    // ~4e-6 — exactly the residual `d1` gap (and hence the `ln d1` term left a
    // ~1.3e-5 self-validation gap, #979). Routing the calibration and the
    // survival log-CDF through ulp-accurate `erfc` removes that amplified noise.
    fn wnorm_cdf(x: f64) -> f64 {
        0.5 * libm::erfc(-x / std::f64::consts::SQRT_2)
    }
    fn wnorm_pdf(x: f64) -> f64 {
        (-0.5 * x * x).exp() / (2.0 * std::f64::consts::PI).sqrt()
    }
    // ── Like-for-like model representation: production's own denested cells ──
    // Production does NOT Simpson-integrate the symbolic composed `index` over a
    // uniform latent grid. It partitions z into cells at every score-warp knot and
    // every link-knot crossing `z=(τ−a)/g`, and on each cell represents the index
    // as the EXACT piecewise cubic `cell.eta(z)=c0+c1 z+c2 z²+c3 z³`. Because the
    // deviation bases are piecewise-cubic, `cell.eta(z)` equals the symbolic
    // `index(a,g,z)` pointwise — but a single uniform Simpson grid that straddles
    // those interior breakpoints integrates a piecewise-cubic integrand with nodes
    // off the breaks, leaving a ~4e-6 aliasing error in `d1` (the #979 value gap).
    // The witness here re-derives production's SAME moment integral by an
    // INDEPENDENT quadrature: composite Simpson applied PER production cell (so every
    // breakpoint is an exact node, and each subinterval integrand is a smooth product
    // of a single cubic with φ). Infinite tail cells are clamped to ±8 (φ decay).
    //
    // `index(a,g,z)` is retained above and used for the observed-node eta/χ, where a
    // single cell covers z_row and the symbolic form coincides with `cell.eta`.
    let denested_cells = |a: f64, g: f64, beta_h: &[f64], beta_w: &[f64]| {
        let beta_h_arr = Array1::from(beta_h.to_vec());
        let beta_w_arr = Array1::from(beta_w.to_vec());
        family
            .denested_partition_cells(a, g, Some(&beta_h_arr), Some(&beta_w_arr))
            .expect("production denested partition cells")
    };
    // Composite Simpson of `f` over a single (possibly infinite) cell, clamped to
    // [-8, 8]; `sub` even subintervals so every cell is integrated to round-off for
    // the smooth single-cubic·φ integrand.
    let cell_simpson = |left: f64, right: f64, sub: usize, f: &dyn Fn(f64) -> f64| -> f64 {
        let lo = left.max(-8.0_f64);
        let hi = right.min(8.0_f64);
        if !(hi > lo) {
            return 0.0;
        }
        let m = sub; // even
        let h = (hi - lo) / m as f64;
        let mut acc = 0.0_f64;
        for k in 0..=m {
            let z = lo + h * k as f64;
            let coef = if k == 0 || k == m {
                1.0
            } else if k % 2 == 1 {
                4.0
            } else {
                2.0
            };
            acc += coef * f(z);
        }
        acc * h / 3.0
    };
    // Calibration F(a) = ∑cells ∫_cell Φ(−eta_cell(z)) φ(z) dz − Φ(−q), Simpson per
    // production cell on its exact representation of the index.
    let calibration = |a: f64, q: f64, g: f64, beta_h: &[f64], beta_w: &[f64]| -> f64 {
        let cells = denested_cells(a, g, beta_h, beta_w);
        let mut acc = 0.0_f64;
        for partition_cell in &cells {
            let cell = partition_cell.cell;
            acc += cell_simpson(cell.left, cell.right, 1024, &|z| {
                wnorm_cdf(-cell.eta(z)) * wnorm_pdf(z)
            });
        }
        acc - wnorm_cdf(-q)
    };
    // Intercept root + density normalization D = |∂F/∂a|.
    let solve_intercept = |q: f64, g: f64, beta_h: &[f64], beta_w: &[f64]| -> (f64, f64) {
        let f = |a: f64| calibration(a, q, g, beta_h, beta_w);
        // Monotone in a; secant from two seeds around the rigid closed form.
        let c = (1.0 + g * g).sqrt();
        let mut a0 = q * c - 0.5;
        let mut a1 = q * c + 0.5;
        let mut f0 = f(a0);
        for _ in 0..200 {
            let f1 = f(a1);
            if (f1 - f0).abs() <= f64::MIN_POSITIVE {
                break;
            }
            let a2 = a1 - f1 * (a1 - a0) / (f1 - f0);
            a0 = a1;
            f0 = f1;
            a1 = a2;
            if (a1 - a0).abs() <= 1e-13 {
                break;
            }
        }
        let a = a1;
        // Density D = |∂F/∂a| = ∑cells ∫_cell φ(eta_cell(z))·(∂eta_cell/∂a)(z)·φ(z) dz,
        // Simpson per production cell. On each cell `∂eta/∂a` is the EXACT cubic in z
        // whose coefficients production derives from the same score/link spans via
        // `denested_cell_coefficient_partials` (this equals `1 + linkdev'(a+g·z)` on
        // the cell, but taken from production's own cell algebra so the witness
        // re-derives production's density with no symbolic-vs-cell aliasing — the
        // ~4e-6 `d1` value gap of #979). Integrating per cell puts every breakpoint on
        // a Simpson node, so the single-cubic·φ integrand resolves to round-off.
        let cells = denested_cells(a, g, beta_h, beta_w);
        let mut acc = 0.0_f64;
        for partition_cell in &cells {
            let cell = partition_cell.cell;
            let (dc_da, _) = exact_kernel::denested_cell_coefficient_partials(
                partition_cell.score_span,
                partition_cell.link_span,
                a,
                g,
            );
            acc += cell_simpson(cell.left, cell.right, 1024, &|z| {
                let deta_da = dc_da[0] + dc_da[1] * z + dc_da[2] * z * z + dc_da[3] * z * z * z;
                wnorm_pdf(cell.eta(z)) * deta_da * wnorm_pdf(z)
            });
        }
        let d = acc.abs();
        (a, d)
    };
    // linkdev'(u) from the link runtime's analytic first-derivative design —
    // the EXACT slope of the link deviation, no finite difference.
    let linkdev_prime_eval = |beta_w: &[f64], u: f64| -> f64 {
        let row = link_runtime
            .first_derivative_design(&array![u])
            .expect("link-dev first-derivative basis row");
        row.row(0).iter().zip(beta_w).map(|(b, c)| b * c).sum()
    };
    // χ1 = ∂η1/∂a at the observed node. The index is
    // a + g·z + g·warp(z) + linkdev(a + g·z), so ∂/∂a = 1 + linkdev'(a + g·z)
    // exactly. An earlier eps=1e-6 central FD here inherited the same amplified
    // ~5e-11 cancellation noise #979 removed from the intercept density: each
    // index value carries ~1e-16 absolute round-off, the FD divides the
    // difference by 2eps=2e-6, and the third-order witness stencils then
    // multiply that floor by ~1/h³≈5e7. Taking the slope analytically removes
    // that amplified noise so the witness third derivatives are limited only by
    // the (Richardson-extrapolated) truncation of the scalar NLL itself.
    let observed_eta_chi = |a: f64, g: f64, beta_h: &[f64], beta_w: &[f64]| -> (f64, f64) {
        let eta = index(a, g, beta_h, beta_w, z_row);
        let chi = 1.0 + linkdev_prime_eval(beta_w, a + g * z_row);
        (eta, chi)
    };
    // Independent scalar survival flex row NLL over the primary vector.
    let witness_nll = |pv: &[f64]| -> f64 {
        let q0 = pv[primary.q0];
        let q1 = pv[primary.q1];
        let qd1 = pv[primary.qd1];
        let g = pv[primary.g];
        let beta_h: Vec<f64> = h_range.clone().map(|i| pv[i]).collect();
        let beta_w: Vec<f64> = w_range.clone().map(|i| pv[i]).collect();
        let (a0, _) = solve_intercept(q0, g, &beta_h, &beta_w);
        let (a1, d1) = solve_intercept(q1, g, &beta_h, &beta_w);
        let (eta0, _) = observed_eta_chi(a0, g, &beta_h, &beta_w);
        let (eta1, chi1) = observed_eta_chi(a1, g, &beta_h, &beta_w);
        let log_surv0 = wnorm_cdf(-eta0).ln();
        let log_surv1 = wnorm_cdf(-eta1).ln();
        let tau_ln = std::f64::consts::TAU.ln();
        let log_phi_eta1 = -0.5 * (eta1 * eta1 + tau_ln);
        let log_phi_q1 = -0.5 * (q1 * q1 + tau_ln);
        weight
            * (log_surv0
                - (1.0 - event) * log_surv1
                - event * log_phi_eta1
                - event * chi1.ln()
                - event * log_phi_q1
                + event * d1.ln()
                - event * qd1.ln())
    };

    // Base primary point with NON-ZERO deviation coefficients.
    let mut p0 = vec![0.0_f64; p];
    p0[primary.q0] = q0v;
    p0[primary.q1] = q1v;
    p0[primary.qd1] = qd1v;
    p0[primary.g] = gv;
    for (k, i) in h_range.clone().enumerate() {
        p0[i] = beta_h0[k];
    }
    for (k, i) in w_range.clone().enumerate() {
        p0[i] = beta_w0[k];
    }

    // ── De-risk: witness scalar NLL must match the production scalar value ───
    let block_states = vec![
        ParameterBlockState {
            beta: Array1::zeros(1),
            eta: array![0.0],
        },
        ParameterBlockState {
            beta: Array1::zeros(0),
            eta: array![0.0],
        },
        ParameterBlockState {
            beta: Array1::zeros(0),
            eta: array![gv],
        },
        ParameterBlockState {
            beta: Array1::from(beta_h0.clone()),
            eta: array![0.0],
        },
        ParameterBlockState {
            beta: Array1::from(beta_w0.clone()),
            eta: array![0.0],
        },
    ];
    let prod_value = family
        .row_neglog_flex_value(0, &block_states)
        .expect("production flex row value");
    let wit_value = witness_nll(&p0);
    assert!(
        (prod_value - wit_value).abs() <= 1e-7 * prod_value.abs().max(1.0),
        "witness re-derivation disagrees with production scalar NLL: witness {wit_value:+.10e} vs production {prod_value:+.10e} \
         (this is a witness-side error to fix BEFORE trusting the FD jets, NOT a production bug)"
    );

    // ── Richardson central differences of the witness scalar NLL ────────────
    let central = |axes: &[(usize, usize)], h: f64| -> f64 {
        fn stencil(order: usize) -> &'static [(i64, f64)] {
            match order {
                1 => &[(-1, -0.5), (1, 0.5)],
                2 => &[(-1, 1.0), (0, -2.0), (1, 1.0)],
                3 => &[(-2, -0.5), (-1, 1.0), (1, -1.0), (2, 0.5)],
                4 => &[(-2, 1.0), (-1, -4.0), (0, 6.0), (1, -4.0), (2, 1.0)],
                _ => &[(0, 1.0)],
            }
        }
        fn walk(
            stencils: &[(usize, &'static [(i64, f64)])],
            h: f64,
            coeff_acc: f64,
            point: &mut Vec<f64>,
            f: &dyn Fn(&[f64]) -> f64,
        ) -> f64 {
            match stencils.split_first() {
                None => coeff_acc * f(point),
                Some((&(idx, st), rest)) => {
                    let mut acc = 0.0;
                    let saved = point[idx];
                    for &(off, c) in st {
                        point[idx] = saved + (off as f64) * h;
                        acc += walk(rest, h, coeff_acc * c, point, f);
                    }
                    point[idx] = saved;
                    acc
                }
            }
        }
        // Coalesce repeated axes before building stencils. A coordinate that
        // appears as two separate order-1 entries would be differenced by
        // COMPOSING two independent ±h shifts, i.e. 0.25·[f(x+2h) − 2f(x) +
        // f(x−2h)] — a second difference at the DOUBLED step 2h, whose leading
        // O(h²) truncation is ~16× that of the compact 3-point order-2 stencil
        // [(-1,1),(0,-2),(1,1)] at step h. Richardson extrapolation cancels the
        // leading term in both cases, but the residual O(h⁴) constant inherits
        // the same blow-up, so a cross-derivative that hits one axis twice
        // (e.g. the ∂³/∂g²∂β_w block, axes [(g,1),(β_w,1),(g,1)]) drifts far
        // enough to break the tight third-order gate while every all-distinct
        // block stays well inside it. Summing the orders of repeated indices
        // keeps each coordinate on its tightest single stencil at step h, so
        // the witness is a faithful ground truth for repeated-axis blocks too.
        let mut merged: Vec<(usize, usize)> = Vec::with_capacity(axes.len());
        for &(idx, ord) in axes {
            if let Some(slot) = merged.iter_mut().find(|(i, _)| *i == idx) {
                slot.1 += ord;
            } else {
                merged.push((idx, ord));
            }
        }
        let mut total_order = 0usize;
        let stencils: Vec<(usize, &'static [(i64, f64)])> = merged
            .iter()
            .map(|&(idx, ord)| {
                total_order += ord;
                (idx, stencil(ord))
            })
            .collect();
        let mut point = p0.clone();
        let raw = walk(&stencils, h, 1.0, &mut point, &witness_nll);
        raw / h.powi(total_order as i32)
    };
    let central_rich = |axes: &[(usize, usize)], h: f64| -> f64 {
        let coarse = central(axes, h);
        let fine = central(axes, h * 0.5);
        (4.0 * fine - coarse) / 3.0
    };

    // The deviation axes the witness must witness (q, g, and a β_h / β_w coord).
    let q0i = primary.q0;
    let gi = primary.g;
    let hi0 = h_range.start;
    let wi0 = w_range.start;

    // Embed a unit direction along a single primary axis.
    let unit = |idx: usize| -> Array1<f64> {
        let mut d = Array1::zeros(p);
        d[idx] = 1.0;
        d
    };

    let q_geom = family
        .row_dynamic_q_geometry(0, &block_states)
        .expect("diagnostic q geometry");
    let (_, _, prod_hess) = family
        .compute_row_flex_primary_gradient_hessian_exact(0, &block_states, &q_geom, &primary)
        .expect("diagnostic production flex hessian");
    let witness_h_gw = central_rich(&[(gi, 1), (wi0, 1)], 6e-3);
    eprintln!(
        "#932 diagnostic base hess[g,w0]: production {:+.6e} witness {:+.6e}",
        prod_hess[[gi, wi0]],
        witness_h_gw
    );
    // gam#1454 regression: the BASE intercept Hessian must match an independent
    // gradient-FD witness at [g, w0] BEFORE the directional third is even asked
    // — the owner's decisive localizer (a base mismatch means the directional
    // third merely inherits it). The base f_aa/f_au moving-boundary a-axis flux
    // (first_full.rs `moving_density_boundary_flux_a`) closes this; if it
    // regresses, this fires pointing at the base Hessian rather than the
    // directional chain.
    {
        let want = witness_h_gw;
        let got = prod_hess[[gi, wi0]];
        let scale = want.abs().max(1.0);
        assert!(
            (got - want).abs() <= 5e-3 * scale + 1e-6,
            "gam#1454 base hess[g,w0] production {got:+.6e} != gradient-FD witness {want:+.6e} \
             (base intercept moving-boundary flux missing/incorrect; the directional third \
             cannot be correct while the base Hessian it differentiates is off)"
        );
    }
    // ── Third order: production D_dir H[u,v] vs witness ∂³ along (u,v,dir) ───
    // Contract along the logslope axis g; check cross blocks touching the
    // deviation coordinates (the channels Arm A cannot reach).
    let third = family
        .row_flex_primary_third_contracted_exact(0, &block_states, &unit(gi))
        .expect("production third contracted (nonzero deviation)");
    let third_checks = [(q0i, hi0), (gi, wi0), (hi0, wi0), (q0i, wi0)];
    for &(u, v) in &third_checks {
        let want = central_rich(&[(u, 1), (v, 1), (gi, 1)], 6e-3);
        let got = third[[u, v]];
        let scale = want.abs().max(1.0);
        assert!(
            (got - want).abs() <= 5e-3 * scale + 1e-6,
            "third[{u},{v}] (contract g) production {got:+.6e} != independent FD witness {want:+.6e}"
        );
    }
    // Planted sign-flip tripwire on a deviation-touching cross block.
    {
        let want = central_rich(&[(q0i, 1), (wi0, 1), (gi, 1)], 6e-3);
        if want.abs() > 1e-5 {
            let flipped = -third[[q0i, wi0]];
            assert!(
                (flipped - want).abs() > 5e-3 * want.abs().max(1.0) + 1e-6,
                "independent FD witness failed to reject a planted (q0, β_w) sign flip: flipped {flipped:+.6e} vs witness {want:+.6e}"
            );
        }
    }

    // ── Fourth order: production D_u D_v H[p,q] vs witness ∂⁴ ────────────────
    let fourth = family
        .row_flex_primary_fourth_contracted_exact(0, &block_states, &unit(gi), &unit(wi0))
        .expect("production fourth contracted (nonzero deviation)");
    let fourth_checks = [(q0i, gi), (q0i, hi0), (gi, hi0)];
    for &(u, v) in &fourth_checks {
        let want = central_rich(&[(u, 1), (v, 1), (gi, 1), (wi0, 1)], 9e-3);
        let got = fourth[[u, v]];
        let scale = want.abs().max(1.0);
        assert!(
            (got - want).abs() <= 2e-2 * scale + 1e-5,
            "fourth[{u},{v}] (contract g,β_w) production {got:+.6e} != independent FD witness {want:+.6e}"
        );
    }
}

#[test]
fn link_flex_family_supports_second_order_exact_outer_path() {
    let score_runtime = test_deviation_runtime();
    let link_runtime = test_deviation_runtime();
    let family = SurvivalMarginalSlopeFamily {
        n: 1,
        event: Arc::new(array![0.0]),
        weights: Arc::new(array![1.0]),
        z: Arc::new(array![0.0].insert_axis(Axis(1))),
        score_covariance: unit_score_covariance(),
        gaussian_frailty_sd: None,
        family_hyper: SurvivalMarginalSlopeFamilyHyperState::default(),
        derivative_guard: 1e-6,
        design_entry: DesignMatrix::from(Array2::zeros((1, 1))),
        design_exit: DesignMatrix::from(Array2::zeros((1, 1))),
        design_derivative_exit: DesignMatrix::from(Array2::ones((1, 1))),
        offset_entry: Arc::new(Array1::zeros(1)),
        offset_exit: Arc::new(Array1::zeros(1)),
        derivative_offset_exit: Arc::new(Array1::ones(1)),
        marginal_design: DesignMatrix::from(Array2::zeros((1, 0))),
        logslope_layout: (DesignMatrix::from(Array2::zeros((1, 0)))).into(),
        score_warp: Some(score_runtime.clone()),
        link_dev: Some(link_runtime.clone()),
        influence_absorber: None,
        time_linear_constraints: None,
        time_wiggle_knots: None,
        time_wiggle_degree: None,
        time_wiggle_ncols: 0,
        intercept_warm_starts: None,
        auto_subsample_phase_counter: Arc::new(AtomicUsize::new(0)),
        auto_subsample_last_rho: Arc::new(Mutex::new(None)),
    };
    let specs = vec![
        dummy_blockspec(1),
        dummy_blockspec(0),
        dummy_blockspec(score_runtime.basis_dim()),
        dummy_blockspec(link_runtime.basis_dim()),
    ];
    assert_eq!(
        family.exact_outer_derivative_order(&specs, &BlockwiseFitOptions::default()),
        ExactOuterDerivativeOrder::Second
    );
}

#[test]
fn timewiggle_scorewarp_family_supports_second_order_exact_outer_path() {
    let score_runtime = test_deviation_runtime();
    let (time_wiggle_knots, time_wiggle_degree, time_wiggle_ncols) = standard_test_time_wiggle();
    let family = SurvivalMarginalSlopeFamily {
        n: 1,
        event: Arc::new(array![0.0]),
        weights: Arc::new(array![1.0]),
        z: Arc::new(array![0.0].insert_axis(Axis(1))),
        score_covariance: unit_score_covariance(),
        gaussian_frailty_sd: None,
        family_hyper: SurvivalMarginalSlopeFamilyHyperState::default(),
        derivative_guard: 1e-6,
        design_entry: DesignMatrix::from(Array2::zeros((1, 5))),
        design_exit: DesignMatrix::from(Array2::zeros((1, 5))),
        design_derivative_exit: DesignMatrix::from(Array2::ones((1, 5))),
        offset_entry: Arc::new(Array1::zeros(1)),
        offset_exit: Arc::new(Array1::zeros(1)),
        derivative_offset_exit: Arc::new(Array1::ones(1)),
        marginal_design: DesignMatrix::from(Array2::zeros((1, 0))),
        logslope_layout: (DesignMatrix::from(Array2::zeros((1, 0)))).into(),
        score_warp: Some(score_runtime.clone()),
        link_dev: None,
        influence_absorber: None,
        time_linear_constraints: None,
        time_wiggle_knots: Some(time_wiggle_knots),
        time_wiggle_degree: Some(time_wiggle_degree),
        time_wiggle_ncols,
        intercept_warm_starts: None,
        auto_subsample_phase_counter: Arc::new(AtomicUsize::new(0)),
        auto_subsample_last_rho: Arc::new(Mutex::new(None)),
    };
    let specs = vec![
        dummy_blockspec(5),
        dummy_blockspec(0),
        dummy_blockspec(score_runtime.basis_dim()),
    ];
    assert_eq!(
        family.exact_outer_derivative_order(&specs, &BlockwiseFitOptions::default()),
        ExactOuterDerivativeOrder::Second
    );
}

#[test]
fn timewiggle_time_jacobian_nonzero_at_zero_beta_linearization() {
    // Regression: when `timewiggle(...)` disables the base time basis the
    // time block's coefficient slots are zero placeholder columns. The
    // identifiability canonicaliser linearises every block at β=0 by
    // calling `effective_jacobian_at` with `beta = &[]`. Previously the
    // timewiggle callback inferred block existence from the (empty)
    // coefficient slice, drove `beta_tw` empty, and returned an all-zero
    // time Jacobian — so the compiler reported "block 0 fully aliased:
    // structural residual Gram has no positive eigenspace". The true
    // derivative ∂q/∂β_tw[j] = B_j(h) at β=0 is the wiggle basis value and
    // is nonzero.
    let (knots, degree, p_tw) = standard_test_time_wiggle();
    assert!(p_tw > 0);
    let n = 3usize;
    // Base time basis disabled: p_base = 0, every column is a wiggle slot,
    // densified as zeros (the placeholder tail the workflow appends).
    let zeros = Arc::new(Array2::<f64>::zeros((n, p_tw)));
    // Pilot coordinates in the interior of the knot span [0, 1].
    let offset_entry = Arc::new(array![0.2, 0.4, 0.6]);
    let offset_exit = Arc::new(array![0.5, 0.7, 0.9]);
    let offset_deriv = Arc::new(array![1.0, 1.0, 1.0]);
    let jac_cb = SmsTimewiggleTimeJacobian::new(
        Arc::clone(&zeros),
        Arc::clone(&zeros),
        Arc::clone(&zeros),
        Arc::new(Array2::<f64>::zeros((n, 0))), // p_m = 0
        Arc::clone(&offset_entry),
        Arc::clone(&offset_exit),
        Arc::clone(&offset_deriv),
        Arc::new(Array1::<f64>::zeros(n)), // marginal_offset = 0
        knots.clone(),
        degree,
        p_tw,
        0,
    );
    let empty: Vec<f64> = Vec::new();
    let state = crate::custom_family::FamilyLinearizationState {
        beta: &empty,
        family_scalars: None,
        channel_hessian: None,
        probit_frailty_scale: 1.0,
    };
    let jac = crate::custom_family::BlockEffectiveJacobian::effective_jacobian_at(&jac_cb, &state)
        .expect("timewiggle time jacobian at beta=0");
    assert_eq!(jac.dim(), (3 * n, p_tw));

    // q0 rows (0..n) must equal the wiggle basis at the entry pilot
    // coordinate (c_i = 1 at β_g = 0), not the broken all-zero placeholder.
    let basis_entry =
        monotone_wiggle_basis_with_derivative_order(offset_entry.view(), &knots, degree, 0)
            .expect("entry basis");
    for i in 0..n {
        for j in 0..p_tw {
            assert_close(
                jac[[i, j]],
                basis_entry[[i, j]],
                1e-12,
                &format!("q0 wiggle col ({i},{j})"),
            );
        }
    }
    // The time block must carry a positive structural eigenspace: its Gram
    // Jᵀ J is not the zero matrix.
    let gram = jac.t().dot(&jac);
    let trace: f64 = (0..p_tw).map(|j| gram[[j, j]]).sum();
    assert!(
        trace > 1e-6,
        "time block Gram is structurally zero (fully-aliased artifact): trace={trace}"
    );
}

#[test]
fn survival_marginal_slope_advertises_outer_hvp_at_large_psi_dim() {
    // `dummy_penalized_blockspec` materializes single-row designs, so the
    // family row count must be 1 to satisfy the HVP availability row guard
    // (`parameter_block_specs_match_rows`, added in 28a1c035f). The
    // "large psi dim" under test is the 32-column block below, not `n`.
    let n = 1usize;
    let family = make_block_psi_test_family(n);
    let specs = vec![
        dummy_penalized_blockspec(0, 0),
        dummy_penalized_blockspec(1, 31),
        dummy_penalized_blockspec(1, 1),
    ];
    let options = BlockwiseFitOptions {
        use_remlobjective: true,
        use_outer_hessian: true,
        ..BlockwiseFitOptions::default()
    };

    let (gradient, hessian) = custom_family_outer_derivatives(&family, &specs, &options);

    assert!(family.inner_coefficient_hessian_hvp_available(&specs));
    assert!(family.outer_hyper_hessian_hvp_available(&specs));
    assert_eq!(
        family.exact_outer_derivative_order(&specs, &options),
        ExactOuterDerivativeOrder::Second
    );
    assert!(
        gam_solve::estimate::reml::reml_outer_engine::prefer_outer_hessian_operator(50_001, 2, 32,)
    );
    assert_eq!(gradient, gam_problem::Derivative::Analytic);
    assert_eq!(hessian, gam_problem::DeclaredHessianForm::Either);
}

#[test]
fn survival_marginal_slope_coefficient_cost_uses_joint_coupled_formula() {
    // Rigid three-block shape: time p=12, marginal p=20, log-slope p=8.
    // The row kernel couples all three blocks, so the joint Hessian is
    // dense over (12+20+8)²=1600 entries per row. The override must
    // return n·(Σ p_b)², not the block-diagonal Σ n·p_b².
    let n = 200usize;
    let p_time = 12usize;
    let p_marg = 20usize;
    let p_log = 8usize;
    let family = SurvivalMarginalSlopeFamily {
        n,
        event: Arc::new(Array1::zeros(n)),
        weights: Arc::new(Array1::from_elem(n, 1.0)),
        z: Arc::new(Array1::zeros(n).insert_axis(Axis(1))),
        score_covariance: unit_score_covariance(),
        gaussian_frailty_sd: None,
        family_hyper: SurvivalMarginalSlopeFamilyHyperState::default(),
        derivative_guard: 1e-6,
        design_entry: DesignMatrix::from(Array2::zeros((n, p_time))),
        design_exit: DesignMatrix::from(Array2::zeros((n, p_time))),
        design_derivative_exit: DesignMatrix::from(Array2::ones((n, p_time))),
        offset_entry: Arc::new(Array1::zeros(n)),
        offset_exit: Arc::new(Array1::zeros(n)),
        derivative_offset_exit: Arc::new(Array1::ones(n)),
        marginal_design: DesignMatrix::from(Array2::zeros((n, p_marg))),
        logslope_layout: (DesignMatrix::from(Array2::zeros((n, p_log)))).into(),
        score_warp: None,
        link_dev: None,
        influence_absorber: None,
        time_linear_constraints: None,
        time_wiggle_knots: None,
        time_wiggle_degree: None,
        time_wiggle_ncols: 0,
        intercept_warm_starts: None,
        auto_subsample_phase_counter: Arc::new(AtomicUsize::new(0)),
        auto_subsample_last_rho: Arc::new(Mutex::new(None)),
    };
    let specs = vec![
        dummy_penalized_blockspec(p_time, 1),
        dummy_penalized_blockspec(p_marg, 1),
        dummy_penalized_blockspec(p_log, 1),
    ];
    let p_total = (p_time + p_marg + p_log) as u64;
    let expected_joint = (n as u64) * p_total * p_total;
    let expected_block_diag_at_full_n =
        (n as u64) * ((p_time * p_time + p_marg * p_marg + p_log * p_log) as u64);
    assert_eq!(family.coefficient_hessian_cost(&specs), expected_joint);
    // Joint coupling exceeds block-diagonal by the cross-block fill
    // 2·n·(p_t·p_m + p_t·p_l + p_m·p_l).
    assert!(expected_joint > expected_block_diag_at_full_n);
}

#[test]
fn exact_outer_row_work_gate_keeps_large_timewiggle_link_models_under_linear_flex_budget() {
    let link_runtime = test_deviation_runtime();
    let (time_wiggle_knots, time_wiggle_degree, time_wiggle_ncols) = standard_test_time_wiggle();
    let family = SurvivalMarginalSlopeFamily {
        n: 80,
        event: Arc::new(array![0.0]),
        weights: Arc::new(array![1.0]),
        z: Arc::new(array![0.0].insert_axis(Axis(1))),
        score_covariance: unit_score_covariance(),
        gaussian_frailty_sd: None,
        family_hyper: SurvivalMarginalSlopeFamilyHyperState::default(),
        derivative_guard: 1e-6,
        design_entry: DesignMatrix::from(Array2::zeros((1, 12))),
        design_exit: DesignMatrix::from(Array2::zeros((1, 12))),
        design_derivative_exit: DesignMatrix::from(Array2::ones((1, 12))),
        offset_entry: Arc::new(Array1::zeros(1)),
        offset_exit: Arc::new(Array1::zeros(1)),
        derivative_offset_exit: Arc::new(Array1::ones(1)),
        marginal_design: DesignMatrix::from(Array2::zeros((1, 20))),
        logslope_layout: (DesignMatrix::from(Array2::zeros((1, 20)))).into(),
        score_warp: None,
        link_dev: Some(link_runtime.clone()),
        influence_absorber: None,
        time_linear_constraints: None,
        time_wiggle_knots: Some(time_wiggle_knots),
        time_wiggle_degree: Some(time_wiggle_degree),
        time_wiggle_ncols,
        intercept_warm_starts: None,
        auto_subsample_phase_counter: Arc::new(AtomicUsize::new(0)),
        auto_subsample_last_rho: Arc::new(Mutex::new(None)),
    };
    let specs = vec![
        dummy_penalized_blockspec(12, 2),
        dummy_penalized_blockspec(20, 2),
        dummy_penalized_blockspec(link_runtime.basis_dim(), 2),
    ];
    assert_eq!(
        family.exact_outer_derivative_order(&specs, &BlockwiseFitOptions::default()),
        ExactOuterDerivativeOrder::Second
    );
}

#[test]
fn timewiggle_scorewarp_beta_hessian_directional_derivative_returns_finite_matrix() {
    let score_runtime = test_deviation_runtime();
    let marginal_design = array![[0.7, -0.2]];
    let marginal_beta = array![0.35, -0.1];
    let (time_wiggle_knots, time_wiggle_degree, time_wiggle_ncols) = standard_test_time_wiggle();
    let family = SurvivalMarginalSlopeFamily {
        n: 1,
        event: Arc::new(array![1.0]),
        weights: Arc::new(array![1.0]),
        z: Arc::new(array![0.15].insert_axis(Axis(1))),
        score_covariance: unit_score_covariance(),
        gaussian_frailty_sd: None,
        family_hyper: SurvivalMarginalSlopeFamilyHyperState::default(),
        derivative_guard: 1e-6,
        design_entry: DesignMatrix::from(array![[0.0, 0.0, 0.0, 0.0, 0.0]]),
        design_exit: DesignMatrix::from(array![[0.0, 0.0, 0.0, 0.0, 0.0]]),
        design_derivative_exit: DesignMatrix::from(array![[1.0, 0.0, 0.0, 0.0, 0.0]]),
        offset_entry: Arc::new(array![0.05]),
        offset_exit: Arc::new(array![0.15]),
        derivative_offset_exit: Arc::new(array![0.9]),
        marginal_design: DesignMatrix::from(marginal_design.clone()),
        logslope_layout: (DesignMatrix::from(array![[1.0]])).into(),
        score_warp: Some(score_runtime.clone()),
        link_dev: None,
        influence_absorber: None,
        time_linear_constraints: None,
        time_wiggle_knots: Some(time_wiggle_knots),
        time_wiggle_degree: Some(time_wiggle_degree),
        time_wiggle_ncols,
        intercept_warm_starts: None,
        auto_subsample_phase_counter: Arc::new(AtomicUsize::new(0)),
        auto_subsample_last_rho: Arc::new(Mutex::new(None)),
    };
    let block_states = vec![
        ParameterBlockState {
            beta: array![0.0, 0.08, -0.03, 0.02, -0.01],
            eta: array![0.0],
        },
        ParameterBlockState {
            beta: marginal_beta.clone(),
            eta: marginal_design.dot(&marginal_beta),
        },
        ParameterBlockState {
            beta: array![0.2],
            eta: array![0.2],
        },
        ParameterBlockState {
            beta: Array1::zeros(score_runtime.basis_dim()),
            eta: Array1::zeros(1),
        },
    ];
    let slices = block_slices(&family, &block_states);
    let mut d_beta_flat = Array1::zeros(slices.total);
    d_beta_flat[slices.time.start] = 0.07;
    d_beta_flat[slices.time.start + 1] = -0.03;
    d_beta_flat[slices.marginal.start] = 0.05;
    d_beta_flat[slices.logslope.start] = -0.04;
    if let Some(h_range) = slices.score_warp.as_ref() {
        d_beta_flat[h_range.start] = 0.02;
    }

    let directional = family
        .exact_newton_joint_hessian_directional_derivative(&block_states, &d_beta_flat)
        .expect("timewiggle flex beta-Hessian directional derivative should evaluate")
        .expect("directional derivative should exist");
    assert_eq!(directional.nrows(), slices.total);
    assert_eq!(directional.ncols(), slices.total);
    assert!(directional.iter().all(|value| value.is_finite()));
}

#[test]
fn timewiggle_scorewarp_beta_hessian_second_directional_derivative_returns_finite_matrix() {
    let score_runtime = test_deviation_runtime();
    let marginal_design = array![[0.7, -0.2]];
    let marginal_beta = array![0.35, -0.1];
    let (time_wiggle_knots, time_wiggle_degree, time_wiggle_ncols) = standard_test_time_wiggle();
    let family = SurvivalMarginalSlopeFamily {
        n: 1,
        event: Arc::new(array![1.0]),
        weights: Arc::new(array![1.0]),
        z: Arc::new(array![0.15].insert_axis(Axis(1))),
        score_covariance: unit_score_covariance(),
        gaussian_frailty_sd: None,
        family_hyper: SurvivalMarginalSlopeFamilyHyperState::default(),
        derivative_guard: 1e-6,
        design_entry: DesignMatrix::from(array![[0.0, 0.0, 0.0, 0.0, 0.0]]),
        design_exit: DesignMatrix::from(array![[0.0, 0.0, 0.0, 0.0, 0.0]]),
        design_derivative_exit: DesignMatrix::from(array![[1.0, 0.0, 0.0, 0.0, 0.0]]),
        offset_entry: Arc::new(array![0.05]),
        offset_exit: Arc::new(array![0.15]),
        derivative_offset_exit: Arc::new(array![0.9]),
        marginal_design: DesignMatrix::from(marginal_design.clone()),
        logslope_layout: (DesignMatrix::from(array![[1.0]])).into(),
        score_warp: Some(score_runtime.clone()),
        link_dev: None,
        influence_absorber: None,
        time_linear_constraints: None,
        time_wiggle_knots: Some(time_wiggle_knots),
        time_wiggle_degree: Some(time_wiggle_degree),
        time_wiggle_ncols,
        intercept_warm_starts: None,
        auto_subsample_phase_counter: Arc::new(AtomicUsize::new(0)),
        auto_subsample_last_rho: Arc::new(Mutex::new(None)),
    };
    let block_states = vec![
        ParameterBlockState {
            beta: array![0.0, 0.08, -0.03, 0.02, -0.01],
            eta: array![0.0],
        },
        ParameterBlockState {
            beta: marginal_beta.clone(),
            eta: marginal_design.dot(&marginal_beta),
        },
        ParameterBlockState {
            beta: array![0.2],
            eta: array![0.2],
        },
        ParameterBlockState {
            beta: Array1::zeros(score_runtime.basis_dim()),
            eta: Array1::zeros(1),
        },
    ];
    let slices = block_slices(&family, &block_states);
    let mut d_beta_u = Array1::zeros(slices.total);
    let mut d_beta_v = Array1::zeros(slices.total);
    d_beta_u[slices.time.start] = 0.07;
    d_beta_u[slices.time.start + 1] = -0.03;
    d_beta_u[slices.marginal.start] = 0.05;
    d_beta_u[slices.logslope.start] = -0.04;
    d_beta_v[slices.time.start + 2] = 0.06;
    d_beta_v[slices.marginal.start + 1] = -0.02;
    d_beta_v[slices.logslope.start] = 0.03;
    if let Some(h_range) = slices.score_warp.as_ref() {
        d_beta_u[h_range.start] = 0.02;
        d_beta_v[h_range.start] = -0.01;
    }

    let second = family
        .exact_newton_joint_hessiansecond_directional_derivative(
            &block_states,
            &d_beta_u,
            &d_beta_v,
        )
        .expect("timewiggle flex beta-Hessian second directional derivative should evaluate")
        .expect("second directional derivative should exist");
    assert_eq!(second.nrows(), slices.total);
    assert_eq!(second.ncols(), slices.total);
    assert!(second.iter().all(|value| value.is_finite()));
}

#[test]
fn link_flex_blockwise_exact_newton_matches_joint_principal_blocks() {
    let score_runtime = test_deviation_runtime();
    let link_runtime = test_deviation_runtime();
    let marginal_design = array![[0.7, -0.2], [0.1, 0.6]];
    let logslope_design = array![[1.0], [0.5]];
    let family = SurvivalMarginalSlopeFamily {
        n: 2,
        event: Arc::new(array![1.0, 0.0]),
        weights: Arc::new(array![1.0, 0.8]),
        z: Arc::new(array![0.15, -0.25].insert_axis(Axis(1))),
        score_covariance: unit_score_covariance(),
        gaussian_frailty_sd: None,
        family_hyper: SurvivalMarginalSlopeFamilyHyperState::default(),
        derivative_guard: 1e-6,
        design_entry: DesignMatrix::from(Array2::zeros((2, 1))),
        design_exit: DesignMatrix::from(Array2::zeros((2, 1))),
        design_derivative_exit: DesignMatrix::from(Array2::ones((2, 1))),
        offset_entry: Arc::new(array![0.05, -0.02]),
        offset_exit: Arc::new(array![0.15, 0.08]),
        derivative_offset_exit: Arc::new(array![0.9, 1.1]),
        marginal_design: DesignMatrix::from(marginal_design.clone()),
        logslope_layout: (DesignMatrix::from(logslope_design.clone())).into(),
        score_warp: Some(score_runtime.clone()),
        link_dev: Some(link_runtime.clone()),
        influence_absorber: None,
        time_linear_constraints: None,
        time_wiggle_knots: None,
        time_wiggle_degree: None,
        time_wiggle_ncols: 0,
        intercept_warm_starts: None,
        auto_subsample_phase_counter: Arc::new(AtomicUsize::new(0)),
        auto_subsample_last_rho: Arc::new(Mutex::new(None)),
    };
    let marginal_beta = array![0.35, -0.1];
    let logslope_beta = array![0.2];
    let block_states = vec![
        ParameterBlockState {
            beta: array![0.0],
            eta: array![0.0, 0.0],
        },
        ParameterBlockState {
            beta: marginal_beta.clone(),
            eta: marginal_design.dot(&marginal_beta),
        },
        ParameterBlockState {
            beta: logslope_beta.clone(),
            eta: logslope_design.dot(&logslope_beta),
        },
        ParameterBlockState {
            beta: Array1::zeros(score_runtime.basis_dim()),
            eta: Array1::zeros(2),
        },
        ParameterBlockState {
            beta: Array1::zeros(link_runtime.basis_dim()),
            eta: Array1::zeros(2),
        },
    ];

    assert_blockwise_matches_joint_principal_blocks(&family, &block_states);
}

#[test]
fn link_flex_marginal_psi_terms_return_finite_joint_terms() {
    let score_runtime = test_deviation_runtime();
    let link_runtime = test_deviation_runtime();
    let marginal_design = array![[0.7, -0.2]];
    let marginal_beta = array![0.35, -0.1];
    let logslope_beta = array![0.2];
    let family = SurvivalMarginalSlopeFamily {
        n: 1,
        event: Arc::new(array![1.0]),
        weights: Arc::new(array![1.0]),
        z: Arc::new(array![0.15].insert_axis(Axis(1))),
        score_covariance: unit_score_covariance(),
        gaussian_frailty_sd: None,
        family_hyper: SurvivalMarginalSlopeFamilyHyperState::default(),
        derivative_guard: 1e-6,
        design_entry: DesignMatrix::from(Array2::zeros((1, 1))),
        design_exit: DesignMatrix::from(Array2::zeros((1, 1))),
        design_derivative_exit: DesignMatrix::from(Array2::ones((1, 1))),
        offset_entry: Arc::new(array![0.05]),
        offset_exit: Arc::new(array![0.15]),
        derivative_offset_exit: Arc::new(array![0.9]),
        marginal_design: DesignMatrix::from(marginal_design.clone()),
        logslope_layout: (DesignMatrix::from(array![[1.0]])).into(),
        score_warp: Some(score_runtime.clone()),
        link_dev: Some(link_runtime.clone()),
        influence_absorber: None,
        time_linear_constraints: None,
        time_wiggle_knots: None,
        time_wiggle_degree: None,
        time_wiggle_ncols: 0,
        intercept_warm_starts: None,
        auto_subsample_phase_counter: Arc::new(AtomicUsize::new(0)),
        auto_subsample_last_rho: Arc::new(Mutex::new(None)),
    };
    let block_states = vec![
        ParameterBlockState {
            beta: array![0.0],
            eta: array![0.0],
        },
        ParameterBlockState {
            beta: marginal_beta.clone(),
            eta: marginal_design.dot(&marginal_beta),
        },
        ParameterBlockState {
            beta: logslope_beta.clone(),
            eta: logslope_beta.clone(),
        },
        ParameterBlockState {
            beta: Array1::zeros(score_runtime.basis_dim()),
            eta: Array1::zeros(score_runtime.basis_dim()),
        },
        ParameterBlockState {
            beta: Array1::zeros(link_runtime.basis_dim()),
            eta: Array1::zeros(link_runtime.basis_dim()),
        },
    ];
    let derivative_blocks = vec![
        Vec::new(),
        vec![crate::custom_family::CustomFamilyBlockPsiDerivative::new(
            None,
            array![[1.0, -0.4]],
            Array2::zeros((2, 2)),
            None,
            None,
            None,
            None,
        )],
        Vec::new(),
    ];

    let slices = block_slices(&family, &block_states);
    let terms = family
        .psi_terms(&block_states, &derivative_blocks, 0)
        .expect("link flex psi terms should evaluate")
        .expect("psi terms should exist");
    assert!(terms.objective_psi.is_finite());
    assert_eq!(terms.score_psi.len(), slices.total);
    assert!(terms.score_psi.iter().all(|value| value.is_finite()));
    assert!(terms.hessian_psi_operator.is_some());
}

#[test]
fn link_flex_marginal_psi_second_order_returns_finite_joint_terms() {
    let score_runtime = test_deviation_runtime();
    let link_runtime = test_deviation_runtime();
    let marginal_design = array![[0.7, -0.2]];
    let marginal_beta = array![0.35, -0.1];
    let logslope_design = array![[1.2, -0.3]];
    let logslope_beta = array![0.2, -0.05];
    let family = SurvivalMarginalSlopeFamily {
        n: 1,
        event: Arc::new(array![1.0]),
        weights: Arc::new(array![1.0]),
        z: Arc::new(array![0.15].insert_axis(Axis(1))),
        score_covariance: unit_score_covariance(),
        gaussian_frailty_sd: None,
        family_hyper: SurvivalMarginalSlopeFamilyHyperState::default(),
        derivative_guard: 1e-6,
        design_entry: DesignMatrix::from(Array2::zeros((1, 1))),
        design_exit: DesignMatrix::from(Array2::zeros((1, 1))),
        design_derivative_exit: DesignMatrix::from(Array2::ones((1, 1))),
        offset_entry: Arc::new(array![0.05]),
        offset_exit: Arc::new(array![0.15]),
        derivative_offset_exit: Arc::new(array![0.9]),
        marginal_design: DesignMatrix::from(marginal_design.clone()),
        logslope_layout: (DesignMatrix::from(logslope_design.clone())).into(),
        score_warp: Some(score_runtime.clone()),
        link_dev: Some(link_runtime.clone()),
        influence_absorber: None,
        time_linear_constraints: None,
        time_wiggle_knots: None,
        time_wiggle_degree: None,
        time_wiggle_ncols: 0,
        intercept_warm_starts: None,
        auto_subsample_phase_counter: Arc::new(AtomicUsize::new(0)),
        auto_subsample_last_rho: Arc::new(Mutex::new(None)),
    };
    let block_states = vec![
        ParameterBlockState {
            beta: array![0.0],
            eta: array![0.0],
        },
        ParameterBlockState {
            beta: marginal_beta.clone(),
            eta: marginal_design.dot(&marginal_beta),
        },
        ParameterBlockState {
            beta: logslope_beta.clone(),
            eta: logslope_design.dot(&logslope_beta),
        },
        ParameterBlockState {
            beta: Array1::zeros(score_runtime.basis_dim()),
            eta: Array1::zeros(1),
        },
        ParameterBlockState {
            beta: Array1::zeros(link_runtime.basis_dim()),
            eta: Array1::zeros(1),
        },
    ];
    let derivative_blocks = vec![
        Vec::new(),
        vec![crate::custom_family::CustomFamilyBlockPsiDerivative::new(
            None,
            array![[1.0, -0.4]],
            Array2::zeros((2, 2)),
            None,
            None,
            None,
            None,
        )],
        vec![crate::custom_family::CustomFamilyBlockPsiDerivative::new(
            None,
            array![[0.3, 0.8]],
            Array2::zeros((2, 2)),
            None,
            None,
            None,
            None,
        )],
        Vec::new(),
    ];

    let slices = block_slices(&family, &block_states);
    let terms = family
        .psi_second_order_terms(&block_states, &derivative_blocks, 0, 1)
        .expect("link flex psi second-order path should evaluate")
        .expect("psi second-order terms should exist");
    assert!(terms.objective_psi_psi.is_finite());
    assert_eq!(terms.score_psi_psi.len(), slices.total);
    assert!(terms.score_psi_psi.iter().all(|value| value.is_finite()));
    assert!(terms.hessian_psi_psi_operator.is_some());
}

#[test]
fn link_flex_marginal_psi_hessian_directional_returns_finite_matrix() {
    let score_runtime = test_deviation_runtime();
    let link_runtime = test_deviation_runtime();
    let marginal_design = array![[0.7, -0.2]];
    let marginal_beta = array![0.35, -0.1];
    let family = SurvivalMarginalSlopeFamily {
        n: 1,
        event: Arc::new(array![1.0]),
        weights: Arc::new(array![1.0]),
        z: Arc::new(array![0.15].insert_axis(Axis(1))),
        score_covariance: unit_score_covariance(),
        gaussian_frailty_sd: None,
        family_hyper: SurvivalMarginalSlopeFamilyHyperState::default(),
        derivative_guard: 1e-6,
        design_entry: DesignMatrix::from(Array2::zeros((1, 1))),
        design_exit: DesignMatrix::from(Array2::zeros((1, 1))),
        design_derivative_exit: DesignMatrix::from(Array2::ones((1, 1))),
        offset_entry: Arc::new(array![0.05]),
        offset_exit: Arc::new(array![0.15]),
        derivative_offset_exit: Arc::new(array![0.9]),
        marginal_design: DesignMatrix::from(marginal_design.clone()),
        logslope_layout: (DesignMatrix::from(array![[1.0]])).into(),
        score_warp: Some(score_runtime.clone()),
        link_dev: Some(link_runtime.clone()),
        influence_absorber: None,
        time_linear_constraints: None,
        time_wiggle_knots: None,
        time_wiggle_degree: None,
        time_wiggle_ncols: 0,
        intercept_warm_starts: None,
        auto_subsample_phase_counter: Arc::new(AtomicUsize::new(0)),
        auto_subsample_last_rho: Arc::new(Mutex::new(None)),
    };
    let block_states = vec![
        ParameterBlockState {
            beta: array![0.0],
            eta: array![0.0],
        },
        ParameterBlockState {
            beta: marginal_beta.clone(),
            eta: marginal_design.dot(&marginal_beta),
        },
        ParameterBlockState {
            beta: array![0.2],
            eta: array![0.2],
        },
        ParameterBlockState {
            beta: Array1::zeros(score_runtime.basis_dim()),
            eta: Array1::zeros(1),
        },
        ParameterBlockState {
            beta: Array1::zeros(link_runtime.basis_dim()),
            eta: Array1::zeros(1),
        },
    ];
    let derivative_blocks = vec![
        Vec::new(),
        vec![crate::custom_family::CustomFamilyBlockPsiDerivative::new(
            None,
            array![[1.0, -0.4]],
            Array2::zeros((2, 2)),
            None,
            None,
            None,
            None,
        )],
        Vec::new(),
    ];
    let slices = block_slices(&family, &block_states);
    let mut d_beta_flat = Array1::zeros(slices.total);
    d_beta_flat[slices.time.start] = 0.07;
    d_beta_flat[slices.marginal.start] = 0.05;
    d_beta_flat[slices.logslope.start] = -0.04;
    if let Some(h_range) = slices.score_warp.as_ref() {
        d_beta_flat[h_range.start] = 0.02;
    }
    if let Some(w_range) = slices.link_dev.as_ref() {
        d_beta_flat[w_range.start] = -0.03;
    }

    let hess_dir = family
        .psi_hessian_directional_derivative(&block_states, &derivative_blocks, 0, &d_beta_flat)
        .expect("link flex psi-Hessian directional path should evaluate")
        .expect("psi-Hessian directional derivative should exist");
    assert_eq!(hess_dir.nrows(), slices.total);
    assert_eq!(hess_dir.ncols(), slices.total);
    assert!(hess_dir.iter().all(|value| value.is_finite()));
}

#[test]
fn timewiggle_marginal_psi_terms_return_finite_joint_terms() {
    let score_runtime = test_deviation_runtime();
    let marginal_design = array![[0.7, -0.2]];
    let marginal_beta = array![0.35, -0.1];
    let (time_wiggle_knots, time_wiggle_degree, time_wiggle_ncols) = standard_test_time_wiggle();
    let family = SurvivalMarginalSlopeFamily {
        n: 1,
        event: Arc::new(array![1.0]),
        weights: Arc::new(array![1.0]),
        z: Arc::new(array![0.15].insert_axis(Axis(1))),
        score_covariance: unit_score_covariance(),
        gaussian_frailty_sd: None,
        family_hyper: SurvivalMarginalSlopeFamilyHyperState::default(),
        derivative_guard: 1e-6,
        design_entry: DesignMatrix::from(array![[0.0, 0.0, 0.0, 0.0, 0.0]]),
        design_exit: DesignMatrix::from(array![[0.0, 0.0, 0.0, 0.0, 0.0]]),
        design_derivative_exit: DesignMatrix::from(array![[1.0, 0.0, 0.0, 0.0, 0.0]]),
        offset_entry: Arc::new(array![0.05]),
        offset_exit: Arc::new(array![0.15]),
        derivative_offset_exit: Arc::new(array![0.9]),
        marginal_design: DesignMatrix::from(marginal_design.clone()),
        logslope_layout: (DesignMatrix::from(array![[1.0]])).into(),
        score_warp: Some(score_runtime.clone()),
        link_dev: None,
        influence_absorber: None,
        time_linear_constraints: None,
        time_wiggle_knots: Some(time_wiggle_knots),
        time_wiggle_degree: Some(time_wiggle_degree),
        time_wiggle_ncols,
        intercept_warm_starts: None,
        auto_subsample_phase_counter: Arc::new(AtomicUsize::new(0)),
        auto_subsample_last_rho: Arc::new(Mutex::new(None)),
    };
    let block_states = vec![
        ParameterBlockState {
            beta: array![0.0, 0.08, -0.03, 0.02, -0.01],
            eta: array![0.0],
        },
        ParameterBlockState {
            beta: marginal_beta.clone(),
            eta: marginal_design.dot(&marginal_beta),
        },
        ParameterBlockState {
            beta: array![0.2],
            eta: array![0.2],
        },
        ParameterBlockState {
            beta: Array1::zeros(score_runtime.basis_dim()),
            eta: Array1::zeros(1),
        },
    ];
    let derivative_blocks = vec![
        Vec::new(),
        vec![crate::custom_family::CustomFamilyBlockPsiDerivative::new(
            None,
            array![[1.0, -0.4]],
            Array2::zeros((2, 2)),
            None,
            None,
            None,
            None,
        )],
        Vec::new(),
    ];

    let slices = block_slices(&family, &block_states);
    let terms = family
        .psi_terms(&block_states, &derivative_blocks, 0)
        .expect("timewiggle psi terms should evaluate")
        .expect("psi terms should exist");
    assert!(terms.objective_psi.is_finite());
    assert_eq!(terms.score_psi.len(), slices.total);
    assert!(terms.score_psi.iter().all(|value| value.is_finite()));
    assert!(terms.hessian_psi_operator.is_some());
}

#[test]
fn timewiggle_blockwise_exact_newton_matches_joint_principal_blocks() {
    let score_runtime = test_deviation_runtime();
    let marginal_design = array![[0.7, -0.2], [0.1, 0.6]];
    let marginal_beta = array![0.35, -0.1];
    let logslope_beta = array![0.2];
    let (time_wiggle_knots, time_wiggle_degree, time_wiggle_ncols) = standard_test_time_wiggle();
    let family = SurvivalMarginalSlopeFamily {
        n: 2,
        event: Arc::new(array![1.0, 0.0]),
        weights: Arc::new(array![1.0, 0.8]),
        z: Arc::new(array![0.15, -0.25].insert_axis(Axis(1))),
        score_covariance: unit_score_covariance(),
        gaussian_frailty_sd: None,
        family_hyper: SurvivalMarginalSlopeFamilyHyperState::default(),
        derivative_guard: 1e-6,
        design_entry: DesignMatrix::from(array![
            [0.0, 0.0, 0.0, 0.0, 0.0],
            [0.0, 0.0, 0.0, 0.0, 0.0]
        ]),
        design_exit: DesignMatrix::from(array![
            [0.0, 0.0, 0.0, 0.0, 0.0],
            [0.0, 0.0, 0.0, 0.0, 0.0]
        ]),
        design_derivative_exit: DesignMatrix::from(array![
            [1.0, 0.0, 0.0, 0.0, 0.0],
            [1.0, 0.0, 0.0, 0.0, 0.0]
        ]),
        offset_entry: Arc::new(array![0.05, -0.02]),
        offset_exit: Arc::new(array![0.15, 0.08]),
        derivative_offset_exit: Arc::new(array![0.9, 1.1]),
        marginal_design: DesignMatrix::from(marginal_design.clone()),
        logslope_layout: (DesignMatrix::from(array![[1.0], [0.5]])).into(),
        score_warp: Some(score_runtime.clone()),
        link_dev: None,
        influence_absorber: None,
        time_linear_constraints: None,
        time_wiggle_knots: Some(time_wiggle_knots),
        time_wiggle_degree: Some(time_wiggle_degree),
        time_wiggle_ncols,
        intercept_warm_starts: None,
        auto_subsample_phase_counter: Arc::new(AtomicUsize::new(0)),
        auto_subsample_last_rho: Arc::new(Mutex::new(None)),
    };
    let block_states = vec![
        ParameterBlockState {
            beta: array![0.0, 0.08, -0.03, 0.02, -0.01],
            eta: array![0.0, 0.0],
        },
        ParameterBlockState {
            beta: marginal_beta.clone(),
            eta: marginal_design.dot(&marginal_beta),
        },
        ParameterBlockState {
            beta: logslope_beta.clone(),
            eta: array![0.2, 0.1],
        },
        ParameterBlockState {
            beta: Array1::zeros(score_runtime.basis_dim()),
            eta: Array1::zeros(2),
        },
    ];

    assert_blockwise_matches_joint_principal_blocks(&family, &block_states);
}

#[test]
fn flex_timewiggle_fast_gradient_matches_dense_joint_gradient() {
    let score_runtime = test_deviation_runtime();
    let link_runtime = test_deviation_runtime();
    let marginal_design = array![[0.7, -0.2]];
    let marginal_beta = array![0.35, -0.1];
    let logslope_design = array![[1.0]];
    let logslope_beta = array![0.2];
    let (time_wiggle_knots, time_wiggle_degree, time_wiggle_ncols) = standard_test_time_wiggle();
    let family = SurvivalMarginalSlopeFamily {
        n: 1,
        event: Arc::new(array![1.0]),
        weights: Arc::new(array![1.0]),
        z: Arc::new(array![0.15].insert_axis(Axis(1))),
        score_covariance: unit_score_covariance(),
        gaussian_frailty_sd: None,
        family_hyper: SurvivalMarginalSlopeFamilyHyperState::default(),
        derivative_guard: 1e-6,
        design_entry: DesignMatrix::from(array![[0.0, 0.0, 0.0, 0.0, 0.0]]),
        design_exit: DesignMatrix::from(array![[0.0, 0.0, 0.0, 0.0, 0.0]]),
        design_derivative_exit: DesignMatrix::from(array![[1.0, 0.0, 0.0, 0.0, 0.0]]),
        offset_entry: Arc::new(array![0.05]),
        offset_exit: Arc::new(array![0.15]),
        derivative_offset_exit: Arc::new(array![0.9]),
        marginal_design: DesignMatrix::from(marginal_design.clone()),
        logslope_layout: (DesignMatrix::from(logslope_design.clone())).into(),
        score_warp: Some(score_runtime.clone()),
        link_dev: Some(link_runtime.clone()),
        influence_absorber: None,
        time_linear_constraints: None,
        time_wiggle_knots: Some(time_wiggle_knots),
        time_wiggle_degree: Some(time_wiggle_degree),
        time_wiggle_ncols,
        intercept_warm_starts: None,
        auto_subsample_phase_counter: Arc::new(AtomicUsize::new(0)),
        auto_subsample_last_rho: Arc::new(Mutex::new(None)),
    };
    let block_states = vec![
        ParameterBlockState {
            beta: array![0.0, 0.08, -0.03, 0.02, -0.01],
            eta: array![0.0],
        },
        ParameterBlockState {
            beta: marginal_beta.clone(),
            eta: marginal_design.dot(&marginal_beta),
        },
        ParameterBlockState {
            beta: logslope_beta.clone(),
            eta: logslope_design.dot(&logslope_beta),
        },
        ParameterBlockState {
            beta: Array1::zeros(score_runtime.basis_dim()),
            eta: Array1::zeros(1),
        },
        ParameterBlockState {
            beta: Array1::zeros(link_runtime.basis_dim()),
            eta: Array1::zeros(1),
        },
    ];

    let (fast_ll, fast_grad) = family
        .evaluate_exact_newton_joint_gradient_dynamic_q(&block_states)
        .expect("fast gradient should evaluate");
    let (dense_ll, dense_grad, _) = family
        .evaluate_exact_newton_joint_dynamic_q_dense(&block_states)
        .expect("dense joint derivatives should evaluate");

    assert_close(fast_ll, dense_ll, 1e-10, "log-likelihood");
    assert_eq!(fast_grad.len(), dense_grad.len());
    for idx in 0..fast_grad.len() {
        assert_close(
            fast_grad[idx],
            dense_grad[idx],
            1e-8,
            &format!("gradient[{idx}]"),
        );
    }
}

/// #932 genus / sibling-of-flex-Hessian witness: the survival marginal-slope
/// dynamic-q JOINT Hessian is a HAND-assembled chain-rule pullback of the
/// per-row primary (4×4) Hessian through the `dq`/`d2q` time-wiggle geometry
/// (`accumulate_dynamic_q_core_hessian` + `row_dynamic_q_geometry_into`). The
/// flex primary-Hessian bug (a dropped/mis-signed 2nd-order coupling that
/// shipped because its FD oracle was committed UNRUN) proved that an
/// unvalidated bespoke pullback can be silently wrong. The ONLY pre-existing
/// joint-Hessian test here checks `blockwise == joint-dense` — both sides share
/// THIS pullback, so it cannot catch a wrong pullback. This oracle closes that
/// gap directly: it central-differences the full joint GRADIENT and compares to
/// the analytic joint HESSIAN on the TIME-WIGGLE-active path (the only path
/// where `q` is nonlinear in the coefficients, hence the only path with a
/// nontrivial `d2q` that a hand-derivation can get wrong). A real wrong
/// derivative is h-independent and orders above the FD floor; FD truncation
/// shrinks with h.
#[test]
fn timewiggle_joint_hessian_matches_central_fd_of_joint_gradient() {
    let score_runtime = test_deviation_runtime();
    let link_runtime = test_deviation_runtime();
    let marginal_design = array![[0.7, -0.2]];
    let marginal_beta = array![0.35, -0.1];
    let logslope_design = array![[1.0]];
    let logslope_beta = array![0.2];
    let (time_wiggle_knots, time_wiggle_degree, time_wiggle_ncols) = standard_test_time_wiggle();
    let family = SurvivalMarginalSlopeFamily {
        n: 1,
        event: Arc::new(array![1.0]),
        weights: Arc::new(array![1.0]),
        z: Arc::new(array![0.15].insert_axis(Axis(1))),
        score_covariance: unit_score_covariance(),
        gaussian_frailty_sd: None,
        family_hyper: SurvivalMarginalSlopeFamilyHyperState::default(),
        derivative_guard: 1e-6,
        design_entry: DesignMatrix::from(array![[0.4, -0.1, 0.2, 0.0, 0.0]]),
        design_exit: DesignMatrix::from(array![[0.6, 0.3, -0.15, 0.0, 0.0]]),
        design_derivative_exit: DesignMatrix::from(array![[1.0, 0.2, -0.1, 0.0, 0.0]]),
        offset_entry: Arc::new(array![0.05]),
        offset_exit: Arc::new(array![0.15]),
        derivative_offset_exit: Arc::new(array![0.9]),
        marginal_design: DesignMatrix::from(marginal_design.clone()),
        logslope_layout: (DesignMatrix::from(logslope_design.clone())).into(),
        score_warp: Some(score_runtime.clone()),
        link_dev: Some(link_runtime.clone()),
        influence_absorber: None,
        time_linear_constraints: None,
        time_wiggle_knots: Some(time_wiggle_knots),
        time_wiggle_degree: Some(time_wiggle_degree),
        time_wiggle_ncols,
        intercept_warm_starts: None,
        auto_subsample_phase_counter: Arc::new(AtomicUsize::new(0)),
        auto_subsample_last_rho: Arc::new(Mutex::new(None)),
    };
    assert!(
        family.flex_timewiggle_active(),
        "fixture must engage the nonlinear time-wiggle q geometry"
    );

    // Base block states. Per-block coefficient consumption (verified against the
    // geometry / dense entry): the TIME block (0) feeds `q` through its design
    // directly (its `eta` is unused); the MARGINAL (1) and LOGSLOPE (2) blocks
    // feed through their `eta = design·beta`; the SCORE-WARP (3) and LINK-DEV
    // (4) deviation blocks feed through their `beta` directly. A faithful
    // perturbation therefore rebuilds `eta = design·beta` for blocks 1 and 2.
    let base_states = vec![
        ParameterBlockState {
            beta: array![0.0, 0.08, -0.03, 0.02, -0.01],
            eta: array![0.0],
        },
        ParameterBlockState {
            beta: marginal_beta.clone(),
            eta: marginal_design.dot(&marginal_beta),
        },
        ParameterBlockState {
            beta: logslope_beta.clone(),
            eta: logslope_design.dot(&logslope_beta),
        },
        ParameterBlockState {
            // score-warp block width is `basis_dim * score_dim`; with a
            // single score column here that is `basis_dim`. Seed it with a
            // small smoothly-varying nonzero so the deviation block carries
            // genuine curvature into the joint Hessian.
            beta: Array1::from_iter(
                (0..score_runtime.basis_dim() * family.score_dim())
                    .map(|j| 0.03 * (-0.4_f64).powi(j as i32)),
            ),
            eta: Array1::zeros(1),
        },
        ParameterBlockState {
            beta: Array1::from_iter(
                (0..link_runtime.basis_dim()).map(|j| -0.02 * (0.5_f64).powi(j as i32)),
            ),
            eta: Array1::zeros(1),
        },
    ];
    assert_eq!(
        base_states[3].beta.len(),
        score_runtime.basis_dim() * family.score_dim()
    );
    assert_eq!(base_states[4].beta.len(), link_runtime.basis_dim());

    let slices = block_slices(&family, &base_states);
    let p_total = slices.total;

    // Rebuild block states from a flat coefficient vector, keeping each block's
    // `eta` consistent with its `beta` exactly as the production solver does.
    let states_from_flat = |flat: &Array1<f64>| -> Vec<ParameterBlockState> {
        let mut states = base_states.clone();
        states[0].beta = flat.slice(s![slices.time.clone()]).to_owned();
        states[1].beta = flat.slice(s![slices.marginal.clone()]).to_owned();
        states[1].eta = marginal_design.dot(&states[1].beta);
        states[2].beta = flat.slice(s![slices.logslope.clone()]).to_owned();
        states[2].eta = logslope_design.dot(&states[2].beta);
        if let Some(range) = slices.score_warp.clone() {
            states[3].beta = flat.slice(s![range]).to_owned();
        }
        if let Some(range) = slices.link_dev.clone() {
            let block_index = 3 + usize::from(family.score_warp.is_some());
            states[block_index].beta = flat.slice(s![range]).to_owned();
        }
        states
    };

    // Flatten the base states into the joint coefficient layout.
    let mut base_flat = Array1::<f64>::zeros(p_total);
    base_flat
        .slice_mut(s![slices.time.clone()])
        .assign(&base_states[0].beta);
    base_flat
        .slice_mut(s![slices.marginal.clone()])
        .assign(&base_states[1].beta);
    base_flat
        .slice_mut(s![slices.logslope.clone()])
        .assign(&base_states[2].beta);
    if let Some(range) = slices.score_warp.clone() {
        base_flat.slice_mut(s![range]).assign(&base_states[3].beta);
    }
    if let Some(range) = slices.link_dev.clone() {
        let block_index = 3 + usize::from(family.score_warp.is_some());
        base_flat
            .slice_mut(s![range])
            .assign(&base_states[block_index].beta);
    }

    let joint_gradient_at = |flat: &Array1<f64>| -> Array1<f64> {
        let states = states_from_flat(flat);
        let (_ll, grad) = family
            .evaluate_exact_newton_joint_gradient_dynamic_q(&states)
            .expect("perturbed joint gradient");
        grad
    };

    let (_ll, _grad, analytic_hessian) = family
        .evaluate_exact_newton_joint_dynamic_q_dense(&base_states)
        .expect("analytic joint dense gradient + hessian");
    assert_eq!(analytic_hessian.shape(), &[p_total, p_total]);

    // Central-difference step. The joint gradient is smooth in every joint
    // coefficient; 1e-5 balances O(h^2) truncation against the FP cancellation
    // floor of the per-perturbation primary intercept re-solve.
    let h = 1e-5;
    let mut max_rel = 0.0_f64;
    let mut worst = (0usize, 0usize, 0.0_f64, 0.0_f64);
    for u in 0..p_total {
        let mut plus = base_flat.clone();
        plus[u] += h;
        let mut minus = base_flat.clone();
        minus[u] -= h;
        let grad_plus = joint_gradient_at(&plus);
        let grad_minus = joint_gradient_at(&minus);
        for v in 0..p_total {
            // The joint "gradient" returned by the exact-Newton path is the
            // score (∇ of the log-likelihood = −∇nll), while the joint dense
            // Hessian is the observed information (+∇²nll = −∇(score)). The
            // analytic Hessian therefore equals the NEGATED Jacobian of the
            // joint gradient, so the central FD must carry the leading minus —
            // exactly as every sibling FD-vs-Hessian oracle in the codebase
            // does (cf. survival/location_scale `fd_neggrad_jac`). Omitting it
            // produced an exact sign flip on every entry.
            let fd = -(grad_plus[v] - grad_minus[v]) / (2.0 * h);
            let analytic = analytic_hessian[[v, u]];
            let denom = 1.0 + analytic.abs().max(fd.abs());
            let rel = (analytic - fd).abs() / denom;
            if rel > max_rel {
                max_rel = rel;
                worst = (v, u, analytic, fd);
            }
        }
    }

    assert!(
        max_rel <= 1e-6,
        "timewiggle joint Hessian disagrees with central FD of the joint gradient: \
         worst entry H[{}][{}] = {:.6e} vs FD {:.6e} (rel {max_rel:.3e}); a chain-rule \
         pullback term (dq/d2q) is dropped or mis-signed",
        worst.0,
        worst.1,
        worst.2,
        worst.3,
    );
}

#[test]
fn row_dynamic_q_geometry_into_pooled_matches_fresh_allocation_bitwise() {
    // Regression: `row_dynamic_q_geometry_into` reuses a caller-owned
    // `SurvivalMarginalSlopeDynamicRow` workspace (resized + zero-filled
    // in place) instead of allocating nine fresh Array2/Array1 buffers
    // per row. Both code paths must return bit-for-bit identical
    // contents, with the workspace path additionally verified to leave
    // the same answer when re-entered against an already-populated
    // buffer (so the in-place `reset` correctly wipes stale state).
    let score_runtime = test_deviation_runtime();
    let link_runtime = test_deviation_runtime();
    let marginal_design = array![[0.7, -0.2], [0.1, 0.6]];
    let marginal_beta = array![0.35, -0.1];
    let logslope_design = array![[1.0], [0.5]];
    let logslope_beta = array![0.2];
    let (time_wiggle_knots, time_wiggle_degree, time_wiggle_ncols) = standard_test_time_wiggle();
    let family = SurvivalMarginalSlopeFamily {
        n: 2,
        event: Arc::new(array![1.0, 0.0]),
        weights: Arc::new(array![1.0, 0.8]),
        z: Arc::new(array![0.15, -0.25].insert_axis(Axis(1))),
        score_covariance: unit_score_covariance(),
        gaussian_frailty_sd: None,
        family_hyper: SurvivalMarginalSlopeFamilyHyperState::default(),
        derivative_guard: 1e-6,
        design_entry: DesignMatrix::from(array![
            [0.0, 0.0, 0.0, 0.0, 0.0],
            [0.0, 0.0, 0.0, 0.0, 0.0]
        ]),
        design_exit: DesignMatrix::from(array![
            [0.0, 0.0, 0.0, 0.0, 0.0],
            [0.0, 0.0, 0.0, 0.0, 0.0]
        ]),
        design_derivative_exit: DesignMatrix::from(array![
            [1.0, 0.0, 0.0, 0.0, 0.0],
            [1.0, 0.0, 0.0, 0.0, 0.0]
        ]),
        offset_entry: Arc::new(array![0.05, -0.02]),
        offset_exit: Arc::new(array![0.15, 0.08]),
        derivative_offset_exit: Arc::new(array![0.9, 1.1]),
        marginal_design: DesignMatrix::from(marginal_design.clone()),
        logslope_layout: (DesignMatrix::from(logslope_design.clone())).into(),
        score_warp: Some(score_runtime.clone()),
        link_dev: Some(link_runtime.clone()),
        influence_absorber: None,
        time_linear_constraints: None,
        time_wiggle_knots: Some(time_wiggle_knots),
        time_wiggle_degree: Some(time_wiggle_degree),
        time_wiggle_ncols,
        intercept_warm_starts: None,
        auto_subsample_phase_counter: Arc::new(AtomicUsize::new(0)),
        auto_subsample_last_rho: Arc::new(Mutex::new(None)),
    };
    let block_states = vec![
        ParameterBlockState {
            beta: array![0.0, 0.08, -0.03, 0.02, -0.01],
            eta: array![0.0, 0.0],
        },
        ParameterBlockState {
            beta: marginal_beta.clone(),
            eta: marginal_design.dot(&marginal_beta),
        },
        ParameterBlockState {
            beta: logslope_beta.clone(),
            eta: logslope_design.dot(&logslope_beta),
        },
        ParameterBlockState {
            beta: Array1::zeros(score_runtime.basis_dim()),
            eta: Array1::zeros(2),
        },
        ParameterBlockState {
            beta: Array1::zeros(link_runtime.basis_dim()),
            eta: Array1::zeros(2),
        },
    ];

    // Compare each row, twice — first into a virgin workspace, then
    // into the same workspace re-used after a different row has
    // already populated it. Both must equal the fresh-allocation path
    // bit-for-bit.
    let mut workspace = SurvivalMarginalSlopeDynamicRow::empty_workspace();
    // Pre-load the workspace with row 1 so the row 0 comparison
    // exercises the buffer-reuse zeroing logic on a non-trivial state.
    family
        .row_dynamic_q_geometry_into(1, &block_states, &mut workspace)
        .expect("preload workspace with row 1");
    for row in [0usize, 1usize, 0usize] {
        let fresh = family
            .row_dynamic_q_geometry(row, &block_states)
            .expect("fresh-allocation row geometry");
        family
            .row_dynamic_q_geometry_into(row, &block_states, &mut workspace)
            .expect("pooled-workspace row geometry");
        assert_eq!(workspace.q0.to_bits(), fresh.q0.to_bits(), "row {row} q0");
        assert_eq!(workspace.q1.to_bits(), fresh.q1.to_bits(), "row {row} q1");
        assert_eq!(
            workspace.qd1.to_bits(),
            fresh.qd1.to_bits(),
            "row {row} qd1"
        );
        let array1_pairs: [(&Array1<f64>, &Array1<f64>, &str); 6] = [
            (&workspace.dq0_time, &fresh.dq0_time, "dq0_time"),
            (&workspace.dq1_time, &fresh.dq1_time, "dq1_time"),
            (&workspace.dqd1_time, &fresh.dqd1_time, "dqd1_time"),
            (&workspace.dq0_marginal, &fresh.dq0_marginal, "dq0_marginal"),
            (&workspace.dq1_marginal, &fresh.dq1_marginal, "dq1_marginal"),
            (
                &workspace.dqd1_marginal,
                &fresh.dqd1_marginal,
                "dqd1_marginal",
            ),
        ];
        for (lhs, rhs, label) in array1_pairs {
            assert_eq!(lhs.shape(), rhs.shape(), "row {row} {label} shape");
            for (i, (l, r)) in lhs.iter().zip(rhs.iter()).enumerate() {
                assert_eq!(
                    l.to_bits(),
                    r.to_bits(),
                    "row {row} {label}[{i}] lhs={l:.17e} rhs={r:.17e}",
                );
            }
        }
        let array2_pairs: [(&Array2<f64>, &Array2<f64>, &str); 9] = [
            (
                &workspace.d2q0_time_time,
                &fresh.d2q0_time_time,
                "d2q0_time_time",
            ),
            (
                &workspace.d2q1_time_time,
                &fresh.d2q1_time_time,
                "d2q1_time_time",
            ),
            (
                &workspace.d2qd1_time_time,
                &fresh.d2qd1_time_time,
                "d2qd1_time_time",
            ),
            (
                &workspace.d2q0_time_marginal,
                &fresh.d2q0_time_marginal,
                "d2q0_time_marginal",
            ),
            (
                &workspace.d2q1_time_marginal,
                &fresh.d2q1_time_marginal,
                "d2q1_time_marginal",
            ),
            (
                &workspace.d2qd1_time_marginal,
                &fresh.d2qd1_time_marginal,
                "d2qd1_time_marginal",
            ),
            (
                &workspace.d2q0_marginal_marginal,
                &fresh.d2q0_marginal_marginal,
                "d2q0_marginal_marginal",
            ),
            (
                &workspace.d2q1_marginal_marginal,
                &fresh.d2q1_marginal_marginal,
                "d2q1_marginal_marginal",
            ),
            (
                &workspace.d2qd1_marginal_marginal,
                &fresh.d2qd1_marginal_marginal,
                "d2qd1_marginal_marginal",
            ),
        ];
        for (lhs, rhs, label) in array2_pairs {
            assert_eq!(lhs.shape(), rhs.shape(), "row {row} {label} shape");
            for ((idx, l), r) in lhs.indexed_iter().zip(rhs.iter()) {
                assert_eq!(
                    l.to_bits(),
                    r.to_bits(),
                    "row {row} {label}[{idx:?}] lhs={l:.17e} rhs={r:.17e}",
                );
            }
        }
    }
}

#[test]
fn flex_timewiggle_operator_to_dense_matches_evaluate_dense_joint_hessian() {
    // Regression: SurvivalMarginalSlopeExactNewtonJointHessianWorkspace::
    // hessian_dense() now returns the already-built operator's
    // to_dense() instead of re-running
    // evaluate_exact_newton_joint_dynamic_q_dense (a second full n-row
    // sweep). Both code paths must agree on the joint p×p Hessian.
    let score_runtime = test_deviation_runtime();
    let link_runtime = test_deviation_runtime();
    let marginal_design = array![[0.7, -0.2], [0.1, 0.6]];
    let marginal_beta = array![0.35, -0.1];
    let logslope_design = array![[1.0], [0.5]];
    let logslope_beta = array![0.2];
    let (time_wiggle_knots, time_wiggle_degree, time_wiggle_ncols) = standard_test_time_wiggle();
    let family = SurvivalMarginalSlopeFamily {
        n: 2,
        event: Arc::new(array![1.0, 0.0]),
        weights: Arc::new(array![1.0, 0.8]),
        z: Arc::new(array![0.15, -0.25].insert_axis(Axis(1))),
        score_covariance: unit_score_covariance(),
        gaussian_frailty_sd: None,
        family_hyper: SurvivalMarginalSlopeFamilyHyperState::default(),
        derivative_guard: 1e-6,
        design_entry: DesignMatrix::from(array![
            [0.0, 0.0, 0.0, 0.0, 0.0],
            [0.0, 0.0, 0.0, 0.0, 0.0]
        ]),
        design_exit: DesignMatrix::from(array![
            [0.0, 0.0, 0.0, 0.0, 0.0],
            [0.0, 0.0, 0.0, 0.0, 0.0]
        ]),
        design_derivative_exit: DesignMatrix::from(array![
            [1.0, 0.0, 0.0, 0.0, 0.0],
            [1.0, 0.0, 0.0, 0.0, 0.0]
        ]),
        offset_entry: Arc::new(array![0.05, -0.02]),
        offset_exit: Arc::new(array![0.15, 0.08]),
        derivative_offset_exit: Arc::new(array![0.9, 1.1]),
        marginal_design: DesignMatrix::from(marginal_design.clone()),
        logslope_layout: (DesignMatrix::from(logslope_design.clone())).into(),
        score_warp: Some(score_runtime.clone()),
        link_dev: Some(link_runtime.clone()),
        influence_absorber: None,
        time_linear_constraints: None,
        time_wiggle_knots: Some(time_wiggle_knots),
        time_wiggle_degree: Some(time_wiggle_degree),
        time_wiggle_ncols,
        intercept_warm_starts: None,
        auto_subsample_phase_counter: Arc::new(AtomicUsize::new(0)),
        auto_subsample_last_rho: Arc::new(Mutex::new(None)),
    };
    let block_states = vec![
        ParameterBlockState {
            beta: array![0.0, 0.08, -0.03, 0.02, -0.01],
            eta: array![0.0, 0.0],
        },
        ParameterBlockState {
            beta: marginal_beta.clone(),
            eta: marginal_design.dot(&marginal_beta),
        },
        ParameterBlockState {
            beta: logslope_beta.clone(),
            eta: logslope_design.dot(&logslope_beta),
        },
        ParameterBlockState {
            beta: Array1::zeros(score_runtime.basis_dim()),
            eta: Array1::zeros(2),
        },
        ParameterBlockState {
            beta: Array1::zeros(link_runtime.basis_dim()),
            eta: Array1::zeros(2),
        },
    ];

    let (_, _, dense) = family
        .evaluate_exact_newton_joint_dynamic_q_dense(&block_states)
        .expect("dense joint Hessian");
    let (operator, _, _, _) = family
        .exact_newton_joint_hessian_operator(&block_states, &BlockwiseFitOptions::default())
        .expect("joint Hessian operator");
    let op_dense = operator.to_dense();

    assert_eq!(op_dense.shape(), dense.shape());
    let diff = max_abs_diff_mat(&op_dense, &dense);
    assert!(
        diff <= 1e-10,
        "operator.to_dense() differs from evaluate_dense by {diff:.3e}",
    );
}

#[test]
fn timewiggle_marginal_logslope_psi_second_order_returns_finite_joint_terms() {
    let score_runtime = test_deviation_runtime();
    let marginal_design = array![[0.7, -0.2]];
    let marginal_beta = array![0.35, -0.1];
    let logslope_design = array![[1.2, -0.3]];
    let logslope_beta = array![0.2, -0.05];
    let (time_wiggle_knots, time_wiggle_degree, time_wiggle_ncols) = standard_test_time_wiggle();
    let family = SurvivalMarginalSlopeFamily {
        n: 1,
        event: Arc::new(array![1.0]),
        weights: Arc::new(array![1.0]),
        z: Arc::new(array![0.15].insert_axis(Axis(1))),
        score_covariance: unit_score_covariance(),
        gaussian_frailty_sd: None,
        family_hyper: SurvivalMarginalSlopeFamilyHyperState::default(),
        derivative_guard: 1e-6,
        design_entry: DesignMatrix::from(array![[0.0, 0.0, 0.0, 0.0, 0.0]]),
        design_exit: DesignMatrix::from(array![[0.0, 0.0, 0.0, 0.0, 0.0]]),
        design_derivative_exit: DesignMatrix::from(array![[1.0, 0.0, 0.0, 0.0, 0.0]]),
        offset_entry: Arc::new(array![0.05]),
        offset_exit: Arc::new(array![0.15]),
        derivative_offset_exit: Arc::new(array![0.9]),
        marginal_design: DesignMatrix::from(marginal_design.clone()),
        logslope_layout: (DesignMatrix::from(logslope_design.clone())).into(),
        score_warp: Some(score_runtime.clone()),
        link_dev: None,
        influence_absorber: None,
        time_linear_constraints: None,
        time_wiggle_knots: Some(time_wiggle_knots),
        time_wiggle_degree: Some(time_wiggle_degree),
        time_wiggle_ncols,
        intercept_warm_starts: None,
        auto_subsample_phase_counter: Arc::new(AtomicUsize::new(0)),
        auto_subsample_last_rho: Arc::new(Mutex::new(None)),
    };
    let block_states = vec![
        ParameterBlockState {
            beta: array![0.0, 0.08, -0.03, 0.02, -0.01],
            eta: array![0.0],
        },
        ParameterBlockState {
            beta: marginal_beta.clone(),
            eta: marginal_design.dot(&marginal_beta),
        },
        ParameterBlockState {
            beta: logslope_beta.clone(),
            eta: logslope_design.dot(&logslope_beta),
        },
        ParameterBlockState {
            beta: Array1::zeros(score_runtime.basis_dim()),
            eta: Array1::zeros(1),
        },
    ];
    let derivative_blocks = vec![
        Vec::new(),
        vec![crate::custom_family::CustomFamilyBlockPsiDerivative::new(
            None,
            array![[1.0, -0.4]],
            Array2::zeros((2, 2)),
            None,
            None,
            None,
            None,
        )],
        vec![crate::custom_family::CustomFamilyBlockPsiDerivative::new(
            None,
            array![[0.3, 0.8]],
            Array2::zeros((2, 2)),
            None,
            None,
            None,
            None,
        )],
        Vec::new(),
    ];

    let slices = block_slices(&family, &block_states);
    let terms = family
        .psi_second_order_terms(&block_states, &derivative_blocks, 0, 1)
        .expect("timewiggle scorewarp psi second-order path should evaluate")
        .expect("psi second-order terms should exist");
    assert!(terms.objective_psi_psi.is_finite());
    assert_eq!(terms.score_psi_psi.len(), slices.total);
    assert!(terms.score_psi_psi.iter().all(|value| value.is_finite()));
    assert!(terms.hessian_psi_psi_operator.is_some());
}

#[test]
fn timewiggle_marginal_psi_hessian_directional_returns_finite_matrix() {
    let score_runtime = test_deviation_runtime();
    let marginal_design = array![[0.7, -0.2]];
    let marginal_beta = array![0.35, -0.1];
    let (time_wiggle_knots, time_wiggle_degree, time_wiggle_ncols) = standard_test_time_wiggle();
    let family = SurvivalMarginalSlopeFamily {
        n: 1,
        event: Arc::new(array![1.0]),
        weights: Arc::new(array![1.0]),
        z: Arc::new(array![0.15].insert_axis(Axis(1))),
        score_covariance: unit_score_covariance(),
        gaussian_frailty_sd: None,
        family_hyper: SurvivalMarginalSlopeFamilyHyperState::default(),
        derivative_guard: 1e-6,
        design_entry: DesignMatrix::from(array![[0.0, 0.0, 0.0, 0.0, 0.0]]),
        design_exit: DesignMatrix::from(array![[0.0, 0.0, 0.0, 0.0, 0.0]]),
        design_derivative_exit: DesignMatrix::from(array![[1.0, 0.0, 0.0, 0.0, 0.0]]),
        offset_entry: Arc::new(array![0.05]),
        offset_exit: Arc::new(array![0.15]),
        derivative_offset_exit: Arc::new(array![0.9]),
        marginal_design: DesignMatrix::from(marginal_design.clone()),
        logslope_layout: (DesignMatrix::from(array![[1.0]])).into(),
        score_warp: Some(score_runtime.clone()),
        link_dev: None,
        influence_absorber: None,
        time_linear_constraints: None,
        time_wiggle_knots: Some(time_wiggle_knots),
        time_wiggle_degree: Some(time_wiggle_degree),
        time_wiggle_ncols,
        intercept_warm_starts: None,
        auto_subsample_phase_counter: Arc::new(AtomicUsize::new(0)),
        auto_subsample_last_rho: Arc::new(Mutex::new(None)),
    };
    let block_states = vec![
        ParameterBlockState {
            beta: array![0.0, 0.08, -0.03, 0.02, -0.01],
            eta: array![0.0],
        },
        ParameterBlockState {
            beta: marginal_beta.clone(),
            eta: marginal_design.dot(&marginal_beta),
        },
        ParameterBlockState {
            beta: array![0.2],
            eta: array![0.2],
        },
        ParameterBlockState {
            beta: Array1::zeros(score_runtime.basis_dim()),
            eta: Array1::zeros(1),
        },
    ];
    let derivative_blocks = vec![
        Vec::new(),
        vec![crate::custom_family::CustomFamilyBlockPsiDerivative::new(
            None,
            array![[1.0, -0.4]],
            Array2::zeros((2, 2)),
            None,
            None,
            None,
            None,
        )],
        Vec::new(),
    ];
    let slices = block_slices(&family, &block_states);
    let mut d_beta_flat = Array1::zeros(slices.total);
    d_beta_flat[slices.time.start] = 0.07;
    d_beta_flat[slices.time.start + 1] = -0.03;
    d_beta_flat[slices.marginal.start] = 0.05;
    d_beta_flat[slices.logslope.start] = -0.04;
    if let Some(h_range) = slices.score_warp.as_ref() {
        d_beta_flat[h_range.start] = 0.02;
    }

    let slices = block_slices(&family, &block_states);
    let hess_dir = family
        .psi_hessian_directional_derivative(&block_states, &derivative_blocks, 0, &d_beta_flat)
        .expect("timewiggle scorewarp psi-Hessian directional path should evaluate")
        .expect("psi-Hessian directional derivative should exist");
    assert_eq!(hess_dir.nrows(), slices.total);
    assert_eq!(hess_dir.ncols(), slices.total);
    assert!(hess_dir.iter().all(|value| value.is_finite()));
}

#[test]
fn sigma_exact_joint_psi_terms_returns_analytic_terms() {
    let marginal_design = array![[0.7, -0.2]];
    let marginal_beta = array![0.35, -0.1];
    let logslope_beta = array![0.2];
    let sigma = 0.65;
    let family = SurvivalMarginalSlopeFamily {
        n: 1,
        event: Arc::new(array![1.0]),
        weights: Arc::new(array![1.0]),
        z: Arc::new(array![0.15].insert_axis(Axis(1))),
        score_covariance: unit_score_covariance(),
        gaussian_frailty_sd: Some(sigma),
        family_hyper: SurvivalMarginalSlopeFamilyHyperState::default(),
        derivative_guard: 1e-6,
        design_entry: DesignMatrix::from(Array2::zeros((1, 1))),
        design_exit: DesignMatrix::from(Array2::zeros((1, 1))),
        design_derivative_exit: DesignMatrix::from(Array2::ones((1, 1))),
        offset_entry: Arc::new(array![0.05]),
        offset_exit: Arc::new(array![0.15]),
        derivative_offset_exit: Arc::new(array![0.9]),
        marginal_design: DesignMatrix::from(marginal_design.clone()),
        logslope_layout: (DesignMatrix::from(array![[1.0]])).into(),
        score_warp: None,
        link_dev: None,
        influence_absorber: None,
        time_linear_constraints: None,
        time_wiggle_knots: None,
        time_wiggle_degree: None,
        time_wiggle_ncols: 0,
        intercept_warm_starts: None,
        auto_subsample_phase_counter: Arc::new(AtomicUsize::new(0)),
        auto_subsample_last_rho: Arc::new(Mutex::new(None)),
    };
    let block_states = vec![
        ParameterBlockState {
            beta: array![0.0],
            eta: array![0.0],
        },
        ParameterBlockState {
            beta: marginal_beta.clone(),
            eta: marginal_design.dot(&marginal_beta),
        },
        ParameterBlockState {
            beta: logslope_beta.clone(),
            eta: logslope_beta.clone(),
        },
    ];
    let specs = vec![
        dummy_blockspec(1),
        dummy_blockspec(marginal_design.ncols()),
        dummy_blockspec(1),
    ];

    let terms = family
        .sigma_exact_joint_psi_terms(&block_states, &specs)
        .expect("sigma psi terms should evaluate analytically")
        .expect("sigma psi terms should be present");
    assert!(terms.objective_psi.is_finite());
    assert_eq!(
        terms.score_psi.len(),
        block_slices(&family, &block_states).total
    );
    assert!(terms.score_psi.iter().all(|value| value.is_finite()));
    assert!(terms.hessian_psi_operator.is_some());
}

#[test]
fn censored_rows_still_reject_invalid_time_derivative() {
    let family = SurvivalMarginalSlopeFamily {
        n: 1,
        event: Arc::new(array![0.0]),
        weights: Arc::new(array![1.0]),
        z: Arc::new(array![0.0].insert_axis(Axis(1))),
        score_covariance: unit_score_covariance(),
        gaussian_frailty_sd: None,
        family_hyper: SurvivalMarginalSlopeFamilyHyperState::default(),
        derivative_guard: 1e-4,
        design_entry: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
            Array2::zeros((1, 1)),
        )),
        design_exit: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
            Array2::zeros((1, 1)),
        )),
        design_derivative_exit: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
            Array2::ones((1, 1)),
        )),
        offset_entry: Arc::new(Array1::zeros(1)),
        offset_exit: Arc::new(Array1::zeros(1)),
        derivative_offset_exit: Arc::new(Array1::from_elem(1, 1e-6)),
        marginal_design: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
            Array2::zeros((1, 0)),
        )),
        logslope_layout: (DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
            Array2::zeros((1, 0)),
        )))
        .into(),
        score_warp: None,
        link_dev: None,
        influence_absorber: None,
        time_linear_constraints: None,
        time_wiggle_knots: None,
        time_wiggle_degree: None,
        time_wiggle_ncols: 0,
        intercept_warm_starts: None,
        auto_subsample_phase_counter: Arc::new(AtomicUsize::new(0)),
        auto_subsample_last_rho: Arc::new(Mutex::new(None)),
    };
    let block_states = vec![
        ParameterBlockState {
            beta: array![0.0],
            eta: array![0.0],
        },
        ParameterBlockState {
            beta: Array1::zeros(0),
            eta: array![0.0],
        },
        ParameterBlockState {
            beta: Array1::zeros(0),
            eta: array![0.0],
        },
    ];

    let err = family
        .compute_row_primary_gradient_hessian_uncached(0, &block_states)
        .expect_err("censored rows must still enforce the time-derivative domain");
    assert!(
        err.contains("monotonicity violated at row 0"),
        "unexpected error: {err}"
    );
}

fn standard_test_time_wiggle() -> (Array1<f64>, usize, usize) {
    let knots = array![0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 1.0];
    let degree = 3usize;
    let ncols = time_wiggle_basis_ncols(&knots, degree).expect("timewiggle basis width");
    (knots, degree, ncols)
}

#[test]
fn exact_newton_evaluation_propagates_invalid_rows() {
    let family = SurvivalMarginalSlopeFamily {
        n: 1,
        event: Arc::new(array![1.0]),
        weights: Arc::new(array![1.0]),
        z: Arc::new(array![0.0].insert_axis(Axis(1))),
        score_covariance: unit_score_covariance(),
        gaussian_frailty_sd: None,
        family_hyper: SurvivalMarginalSlopeFamilyHyperState::default(),
        derivative_guard: 1e-4,
        design_entry: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(array![[
            0.0
        ]])),
        design_exit: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(array![[
            0.0
        ]])),
        design_derivative_exit: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
            array![[1.0]],
        )),
        offset_entry: Arc::new(array![0.0]),
        offset_exit: Arc::new(array![0.0]),
        derivative_offset_exit: Arc::new(array![1e-6]),
        marginal_design: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
            Array2::zeros((1, 0)),
        )),
        logslope_layout: (DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
            Array2::zeros((1, 0)),
        )))
        .into(),
        score_warp: None,
        link_dev: None,
        influence_absorber: None,
        time_linear_constraints: None,
        time_wiggle_knots: None,
        time_wiggle_degree: None,
        time_wiggle_ncols: 0,
        intercept_warm_starts: None,
        auto_subsample_phase_counter: Arc::new(AtomicUsize::new(0)),
        auto_subsample_last_rho: Arc::new(Mutex::new(None)),
    };
    let block_states = vec![
        ParameterBlockState {
            beta: array![0.0],
            eta: array![0.0],
        },
        ParameterBlockState {
            beta: Array1::zeros(0),
            eta: array![0.0],
        },
        ParameterBlockState {
            beta: Array1::zeros(0),
            eta: array![0.0],
        },
    ];

    let err = family
        .evaluate(&block_states)
        .expect_err("invalid rows must abort exact-newton evaluation");
    assert!(
        err.contains("monotonicity violated"),
        "unexpected error: {err}"
    );
}

#[test]
fn time_constraints_use_exact_derivative_guard_rows() {
    let family = SurvivalMarginalSlopeFamily {
        n: 2,
        event: Arc::new(array![0.0, 1.0]),
        weights: Arc::new(array![1.0, 1.0]),
        z: Arc::new(array![0.0, 0.0].insert_axis(Axis(1))),
        score_covariance: unit_score_covariance(),
        gaussian_frailty_sd: None,
        family_hyper: SurvivalMarginalSlopeFamilyHyperState::default(),
        derivative_guard: 1e-4,
        design_entry: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
            Array2::zeros((2, 2)),
        )),
        design_exit: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
            Array2::zeros((2, 2)),
        )),
        design_derivative_exit: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
            array![[1.0, 2.0], [3.0, 4.0]],
        )),
        offset_entry: Arc::new(Array1::zeros(2)),
        offset_exit: Arc::new(Array1::zeros(2)),
        derivative_offset_exit: Arc::new(array![0.25, 0.5]),
        marginal_design: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
            Array2::zeros((2, 0)),
        )),
        logslope_layout: (DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
            Array2::zeros((2, 0)),
        )))
        .into(),
        time_linear_constraints: time_derivative_guard_constraints(
            &DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(array![
                [1.0, 1.0],
                [1.0, -1.0]
            ])),
            &array![0.0, 0.25],
            1.0,
        )
        .expect("time derivative guard constraints"),
        score_warp: None,
        link_dev: None,
        influence_absorber: None,
        time_wiggle_knots: None,
        time_wiggle_degree: None,
        time_wiggle_ncols: 0,
        intercept_warm_starts: None,
        auto_subsample_phase_counter: Arc::new(AtomicUsize::new(0)),
        auto_subsample_last_rho: Arc::new(Mutex::new(None)),
    };
    let spec = ParameterBlockSpec {
        name: "time_surface".to_string(),
        design: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(array![
            [0.0, 0.0],
            [0.0, 0.0]
        ])),
        offset: Array1::zeros(2),
        penalties: Vec::new(),
        nullspace_dims: Vec::new(),
        initial_log_lambdas: Array1::zeros(0),
        initial_beta: None,
        gauge_priority: 100,
        jacobian_callback: None,
        stacked_design: None,
        stacked_offset: None,
    };
    let constraints = match family
        .block_linear_constraints(&[], 0, &spec)
        .expect("constraint lookup")
        .expect("time constraints")
    {
        gam_problem::ConstraintSet::Dense(dense) => dense,
        other => panic!("time-block constraints must be dense rows, got {other:?}"),
    };
    let row_scale = 2.0_f64.sqrt();
    let expected_a = array![
        [1.0 / row_scale, 1.0 / row_scale],
        [1.0 / row_scale, -1.0 / row_scale]
    ];
    assert_eq!(constraints.a.dim(), expected_a.dim());
    for (got, want) in constraints.a.iter().zip(expected_a.iter()) {
        assert_relative_eq!(*got, *want, epsilon = 1e-12);
    }
    let expected_b = array![1.0 / row_scale, 0.75 / row_scale];
    assert_eq!(constraints.b.dim(), expected_b.dim());
    for (got, want) in constraints.b.iter().zip(expected_b.iter()) {
        assert_relative_eq!(*got, *want, epsilon = 1e-12);
    }
}

#[test]
fn time_block_constraints_synthesize_qd1_rows_when_stored_constraints_missing() {
    let family = SurvivalMarginalSlopeFamily {
        n: 1,
        event: Arc::new(array![1.0]),
        weights: Arc::new(array![1.0]),
        z: Arc::new(array![0.0].insert_axis(Axis(1))),
        score_covariance: unit_score_covariance(),
        gaussian_frailty_sd: None,
        family_hyper: SurvivalMarginalSlopeFamilyHyperState::default(),
        derivative_guard: 1e-6,
        design_entry: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(array![[
            0.0, 0.0
        ]])),
        design_exit: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(array![[
            0.0, 0.0
        ]])),
        design_derivative_exit: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
            array![[1.0, 0.0]],
        )),
        offset_entry: Arc::new(array![0.0]),
        offset_exit: Arc::new(array![0.0]),
        derivative_offset_exit: Arc::new(array![1e-6]),
        marginal_design: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
            Array2::zeros((1, 0)),
        )),
        logslope_layout: (DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
            Array2::zeros((1, 0)),
        )))
        .into(),
        score_warp: None,
        link_dev: None,
        influence_absorber: None,
        time_linear_constraints: None,
        time_wiggle_knots: None,
        time_wiggle_degree: None,
        time_wiggle_ncols: 0,
        intercept_warm_starts: None,
        auto_subsample_phase_counter: Arc::new(AtomicUsize::new(0)),
        auto_subsample_last_rho: Arc::new(Mutex::new(None)),
    };
    let spec = ParameterBlockSpec {
        name: "time_surface".to_string(),
        design: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(array![[
            0.0, 0.0
        ]])),
        offset: Array1::zeros(1),
        penalties: Vec::new(),
        nullspace_dims: Vec::new(),
        initial_log_lambdas: Array1::zeros(0),
        initial_beta: None,
        gauge_priority: 100,
        jacobian_callback: None,
        stacked_design: None,
        stacked_offset: None,
    };

    let constraints = match family
        .block_linear_constraints(&[], 0, &spec)
        .expect("synthesized constraints")
        .expect("qd1 row")
    {
        gam_problem::ConstraintSet::Dense(dense) => dense,
        other => panic!("synthesized qd1 constraints must be dense rows, got {other:?}"),
    };
    assert_eq!(constraints.a, array![[1.0, 0.0]]);
    assert_eq!(constraints.b, array![0.0]);
}

#[test]
fn time_block_max_feasible_step_uses_synthesized_qd1_rows() {
    let family = SurvivalMarginalSlopeFamily {
        n: 1,
        event: Arc::new(array![1.0]),
        weights: Arc::new(array![1.0]),
        z: Arc::new(array![0.0].insert_axis(Axis(1))),
        score_covariance: unit_score_covariance(),
        gaussian_frailty_sd: None,
        family_hyper: SurvivalMarginalSlopeFamilyHyperState::default(),
        derivative_guard: 1e-6,
        design_entry: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(array![[
            0.0, 0.0
        ]])),
        design_exit: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(array![[
            0.0, 0.0
        ]])),
        design_derivative_exit: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
            array![[1.0, 0.0]],
        )),
        offset_entry: Arc::new(array![0.0]),
        offset_exit: Arc::new(array![0.0]),
        derivative_offset_exit: Arc::new(array![1e-6]),
        marginal_design: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
            Array2::zeros((1, 0)),
        )),
        logslope_layout: (DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
            Array2::zeros((1, 0)),
        )))
        .into(),
        score_warp: None,
        link_dev: None,
        influence_absorber: None,
        time_linear_constraints: None,
        time_wiggle_knots: None,
        time_wiggle_degree: None,
        time_wiggle_ncols: 0,
        intercept_warm_starts: None,
        auto_subsample_phase_counter: Arc::new(AtomicUsize::new(0)),
        auto_subsample_last_rho: Arc::new(Mutex::new(None)),
    };
    let states = vec![ParameterBlockState {
        beta: array![0.4, 7.0],
        eta: array![0.0],
    }];
    let alpha = family
        .max_feasible_step_size(&states, 0, &array![-1.0, -10.0])
        .expect("synthesized qd1 step ceiling")
        .expect("binding synthesized qd1 row");

    assert_relative_eq!(alpha, 0.398, epsilon = 1e-12);
}

#[test]
fn coupled_qd1_guard_limits_time_step_before_post_update_projection() {
    let constraints = time_derivative_guard_constraints(
        &DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(array![[
            1.0, 1.0
        ]])),
        &array![0.0],
        1.0,
    )
    .expect("time derivative guard constraints")
    .expect("coupled row");
    let family = SurvivalMarginalSlopeFamily {
        n: 1,
        event: Arc::new(array![0.0]),
        weights: Arc::new(array![1.0]),
        z: Arc::new(array![0.0].insert_axis(Axis(1))),
        score_covariance: unit_score_covariance(),
        gaussian_frailty_sd: None,
        family_hyper: SurvivalMarginalSlopeFamilyHyperState::default(),
        derivative_guard: 1.0,
        design_entry: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(array![[
            0.0, 0.0
        ]])),
        design_exit: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(array![[
            0.0, 0.0
        ]])),
        design_derivative_exit: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
            array![[1.0, 1.0]],
        )),
        offset_entry: Arc::new(array![0.0]),
        offset_exit: Arc::new(array![0.0]),
        derivative_offset_exit: Arc::new(array![0.0]),
        marginal_design: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
            Array2::zeros((1, 0)),
        )),
        logslope_layout: (DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
            Array2::zeros((1, 0)),
        )))
        .into(),
        time_linear_constraints: Some(constraints),
        score_warp: None,
        link_dev: None,
        influence_absorber: None,
        time_wiggle_knots: None,
        time_wiggle_degree: None,
        time_wiggle_ncols: 0,
        intercept_warm_starts: None,
        auto_subsample_phase_counter: Arc::new(AtomicUsize::new(0)),
        auto_subsample_last_rho: Arc::new(Mutex::new(None)),
    };
    let states = vec![ParameterBlockState {
        beta: array![0.6, 0.6],
        eta: array![0.0],
    }];
    let alpha = family
        .max_feasible_step_size(&states, 0, &array![-1.0, 0.0])
        .expect("coupled qd1 step limit")
        .expect("binding coupled row");

    assert_relative_eq!(alpha, 0.199, epsilon = 1e-12);
}

#[test]
fn timewiggle_tail_constraints_are_part_of_time_block_feasibility() {
    let structural = time_derivative_guard_constraints(
        &DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(array![[
            1.0, 0.0, 0.0
        ]])),
        &array![1e-6],
        1e-6,
    )
    .expect("time derivative guard constraints");
    let constraints = append_timewiggle_tail_nonnegative_constraints(structural, 3, 2)
        .expect("combined constraints")
        .expect("time constraints");

    assert_eq!(constraints.a, Array2::<f64>::eye(3));
    assert_eq!(constraints.b, Array1::<f64>::zeros(3));
}

#[test]
fn timewiggle_tail_step_is_clipped_before_it_can_flip_derivative() {
    let constraints = append_timewiggle_tail_nonnegative_constraints(None, 2, 1)
        .expect("tail constraints")
        .expect("time constraints");
    let family = SurvivalMarginalSlopeFamily {
        n: 1,
        event: Arc::new(array![0.0]),
        weights: Arc::new(array![1.0]),
        z: Arc::new(array![0.0].insert_axis(Axis(1))),
        score_covariance: unit_score_covariance(),
        gaussian_frailty_sd: None,
        family_hyper: SurvivalMarginalSlopeFamilyHyperState::default(),
        derivative_guard: 1e-6,
        design_entry: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(array![[
            0.0, 0.0
        ]])),
        design_exit: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(array![[
            0.0, 0.0
        ]])),
        design_derivative_exit: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
            array![[0.0, 0.0]],
        )),
        offset_entry: Arc::new(array![0.0]),
        offset_exit: Arc::new(array![0.0]),
        derivative_offset_exit: Arc::new(array![1e-6]),
        marginal_design: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
            Array2::zeros((1, 0)),
        )),
        logslope_layout: (DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
            Array2::zeros((1, 0)),
        )))
        .into(),
        time_linear_constraints: Some(constraints),
        score_warp: None,
        link_dev: None,
        influence_absorber: None,
        time_wiggle_knots: None,
        time_wiggle_degree: None,
        time_wiggle_ncols: 1,
        intercept_warm_starts: None,
        auto_subsample_phase_counter: Arc::new(AtomicUsize::new(0)),
        auto_subsample_last_rho: Arc::new(Mutex::new(None)),
    };
    let states = vec![ParameterBlockState {
        beta: array![0.0, 0.5],
        eta: array![0.0],
    }];
    let alpha = family
        .max_feasible_step_size(&states, 0, &array![0.0, -1.0])
        .expect("timewiggle tail step ceiling")
        .expect("negative tail step should be bounded");
    assert_relative_eq!(alpha, 0.4975, epsilon = 1e-12);
}

#[test]
fn time_block_post_update_rejects_infeasible_beta_instead_of_projecting() {
    let family = SurvivalMarginalSlopeFamily {
        n: 1,
        event: Arc::new(array![0.0]),
        weights: Arc::new(array![1.0]),
        z: Arc::new(array![0.0].insert_axis(Axis(1))),
        score_covariance: unit_score_covariance(),
        gaussian_frailty_sd: None,
        family_hyper: SurvivalMarginalSlopeFamilyHyperState::default(),
        derivative_guard: 1e-6,
        design_entry: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(array![[
            1.0, 0.0
        ]])),
        design_exit: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(array![[
            1.0, 0.0
        ]])),
        design_derivative_exit: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
            array![[1.0, 0.0]],
        )),
        offset_entry: Arc::new(array![0.0]),
        offset_exit: Arc::new(array![0.0]),
        derivative_offset_exit: Arc::new(array![1e-6]),
        marginal_design: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
            Array2::zeros((1, 0)),
        )),
        logslope_layout: (DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
            Array2::zeros((1, 0)),
        )))
        .into(),
        time_linear_constraints: append_timewiggle_tail_nonnegative_constraints(
            time_derivative_guard_constraints(
                &DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(array![[
                    1.0, 0.0
                ]])),
                &array![1e-6],
                1e-6,
            )
            .expect("time derivative guard constraints"),
            2,
            1,
        )
        .expect("combined time constraints"),
        score_warp: None,
        link_dev: None,
        influence_absorber: None,
        time_wiggle_knots: None,
        time_wiggle_degree: None,
        time_wiggle_ncols: 1,
        intercept_warm_starts: None,
        auto_subsample_phase_counter: Arc::new(AtomicUsize::new(0)),
        auto_subsample_last_rho: Arc::new(Mutex::new(None)),
    };
    let spec = ParameterBlockSpec {
        name: "time_surface".to_string(),
        design: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(array![[
            1.0, 0.0
        ]])),
        offset: Array1::zeros(1),
        penalties: Vec::new(),
        nullspace_dims: Vec::new(),
        initial_log_lambdas: Array1::zeros(0),
        initial_beta: None,
        gauge_priority: 100,
        jacobian_callback: None,
        stacked_design: None,
        stacked_offset: None,
    };
    let err = family
        .post_update_block_beta(
            &[ParameterBlockState {
                beta: array![0.0, 0.0],
                eta: array![0.0],
            }],
            0,
            &spec,
            array![-0.3, -0.2],
        )
        .expect_err("post-update must not project an infeasible time beta");
    assert!(
        err.contains("violates monotonicity") && err.contains("proposed"),
        "unexpected error message: {err}"
    );
}

/// Regression guard for the phantom-multiplier failure mode: if a
/// hand-constructed family omits the derivative-guard rows from
/// `time_linear_constraints`, post-update must not repair the proposed
/// time beta by a hidden projection. The KKT system needs those rows; a
/// missing-row violation is an error, not a convergence mechanism.
#[test]
fn time_block_post_update_rejects_qd1_when_no_linear_constraints() {
    let family = SurvivalMarginalSlopeFamily {
        n: 1,
        event: Arc::new(array![1.0]),
        weights: Arc::new(array![1.0]),
        z: Arc::new(array![0.0].insert_axis(Axis(1))),
        score_covariance: unit_score_covariance(),
        gaussian_frailty_sd: None,
        family_hyper: SurvivalMarginalSlopeFamilyHyperState::default(),
        derivative_guard: 1e-6,
        design_entry: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(array![[
            0.0
        ]])),
        design_exit: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(array![[
            0.0
        ]])),
        // qd1 = 1.0 · β[0] + 0.0 · β[1] + offset
        design_derivative_exit: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
            array![[1.0, 0.0]],
        )),
        offset_entry: Arc::new(array![0.0]),
        offset_exit: Arc::new(array![0.0]),
        // offset = derivative_guard exactly (the production setup).
        derivative_offset_exit: Arc::new(array![1e-6]),
        marginal_design: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
            Array2::zeros((1, 0)),
        )),
        logslope_layout: (DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
            Array2::zeros((1, 0)),
        )))
        .into(),
        score_warp: None,
        link_dev: None,
        influence_absorber: None,
        time_linear_constraints: None,
        time_wiggle_knots: None,
        time_wiggle_degree: None,
        time_wiggle_ncols: 0,
        intercept_warm_starts: None,
        auto_subsample_phase_counter: Arc::new(AtomicUsize::new(0)),
        auto_subsample_last_rho: Arc::new(Mutex::new(None)),
    };
    let spec = ParameterBlockSpec {
        name: "time_surface".to_string(),
        design: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(array![[
            0.0, 0.0
        ]])),
        offset: Array1::zeros(1),
        penalties: Vec::new(),
        nullspace_dims: Vec::new(),
        initial_log_lambdas: Array1::zeros(0),
        initial_beta: None,
        gauge_priority: 100,
        jacobian_callback: None,
        stacked_design: None,
        stacked_offset: None,
    };
    let current = array![0.4, 7.0];
    // qd1 at current = 1.0·0.4 + 0.0·7.0 + 1e-6 ≈ 0.4 (feasible)
    // qd1 at proposed = 1.0·-0.6 + 0.0·-3.0 + 1e-6 ≈ -0.6 (infeasible)
    let proposed = array![-0.6, -3.0];
    let err = family
        .post_update_block_beta(
            &[ParameterBlockState {
                beta: current.clone(),
                eta: array![0.0],
            }],
            0,
            &spec,
            proposed.clone(),
        )
        .expect_err("missing qd1 constraints must not be repaired by projection");
    assert!(
        err.contains("violates monotonicity")
            && err.contains("proposed")
            && err.contains("time_linear_constraints"),
        "unexpected error message: {err}"
    );
}

/// When `current` already violates the qd1 monotonicity, the projection
/// surfaces a structured error (with the row index and qd1 value)
/// instead of silently returning a still-infeasible β. This is the
/// invariant the score_warp / link_dev projection enforces; the time
/// block now matches.
#[test]
fn time_block_post_update_errors_when_current_violates_qd1() {
    let family = SurvivalMarginalSlopeFamily {
        n: 1,
        event: Arc::new(array![1.0]),
        weights: Arc::new(array![1.0]),
        z: Arc::new(array![0.0].insert_axis(Axis(1))),
        score_covariance: unit_score_covariance(),
        gaussian_frailty_sd: None,
        family_hyper: SurvivalMarginalSlopeFamilyHyperState::default(),
        derivative_guard: 1e-6,
        design_entry: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(array![[
            0.0
        ]])),
        design_exit: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(array![[
            0.0
        ]])),
        design_derivative_exit: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
            array![[1.0]],
        )),
        offset_entry: Arc::new(array![0.0]),
        offset_exit: Arc::new(array![0.0]),
        derivative_offset_exit: Arc::new(array![1e-6]),
        marginal_design: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
            Array2::zeros((1, 0)),
        )),
        logslope_layout: (DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
            Array2::zeros((1, 0)),
        )))
        .into(),
        score_warp: None,
        link_dev: None,
        influence_absorber: None,
        time_linear_constraints: None,
        time_wiggle_knots: None,
        time_wiggle_degree: None,
        time_wiggle_ncols: 0,
        intercept_warm_starts: None,
        auto_subsample_phase_counter: Arc::new(AtomicUsize::new(0)),
        auto_subsample_last_rho: Arc::new(Mutex::new(None)),
    };
    let spec = ParameterBlockSpec {
        name: "time_surface".to_string(),
        design: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(array![[0.0]])),
        offset: Array1::zeros(1),
        penalties: Vec::new(),
        nullspace_dims: Vec::new(),
        initial_log_lambdas: Array1::zeros(0),
        initial_beta: None,
        gauge_priority: 100,
        jacobian_callback: None,
        stacked_design: None,
        stacked_offset: None,
    };
    // current qd1 = -1.0 + 1e-6 < guard → infeasible.
    let err = family
        .post_update_block_beta(
            &[ParameterBlockState {
                beta: array![-1.0],
                eta: array![0.0],
            }],
            0,
            &spec,
            array![0.5],
        )
        .expect_err("infeasible current must surface an error");
    assert!(
        err.contains("violates monotonicity") && err.contains("row 0"),
        "unexpected error message: {err}"
    );
}

#[test]
fn time_block_feasible_step_stays_inside_derivative_guard() {
    let family = SurvivalMarginalSlopeFamily {
        n: 1,
        event: Arc::new(array![0.0]),
        weights: Arc::new(array![1.0]),
        z: Arc::new(array![0.0].insert_axis(Axis(1))),
        score_covariance: unit_score_covariance(),
        gaussian_frailty_sd: None,
        family_hyper: SurvivalMarginalSlopeFamilyHyperState::default(),
        derivative_guard: 1e-4,
        design_entry: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(array![[
            0.0, 0.0
        ]])),
        design_exit: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(array![[
            0.0, 0.0
        ]])),
        design_derivative_exit: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
            array![[1.0, 0.0]],
        )),
        offset_entry: Arc::new(array![0.0]),
        offset_exit: Arc::new(array![0.0]),
        derivative_offset_exit: Arc::new(array![0.2]),
        marginal_design: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
            Array2::zeros((1, 0)),
        )),
        logslope_layout: (DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
            Array2::zeros((1, 0)),
        )))
        .into(),
        score_warp: None,
        link_dev: None,
        influence_absorber: None,
        time_linear_constraints: time_derivative_guard_constraints(
            &DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(array![[
                1.0, 0.0
            ]])),
            &array![0.2],
            1e-4,
        )
        .expect("time derivative guard constraints"),
        time_wiggle_knots: None,
        time_wiggle_degree: None,
        time_wiggle_ncols: 0,
        intercept_warm_starts: None,
        auto_subsample_phase_counter: Arc::new(AtomicUsize::new(0)),
        auto_subsample_last_rho: Arc::new(Mutex::new(None)),
    };
    let states = vec![
        ParameterBlockState {
            beta: array![0.0, 0.0],
            eta: array![0.0],
        },
        ParameterBlockState {
            beta: Array1::zeros(0),
            eta: array![0.0],
        },
        ParameterBlockState {
            beta: Array1::zeros(0),
            eta: array![0.0],
        },
    ];
    let alpha = family
        .max_feasible_step_size(&states, 0, &array![-1.0, 0.0])
        .expect("time step ceiling")
        .expect("time step should be bounded");
    // Starting at beta=[0,0] with derivative q' = design·beta + offset = 0.2,
    // far above the 1e-4 guard. Stepping along [-1, 0] drives q' toward the
    // guard; the largest feasible α satisfies -α + 0.2 = 1e-4, i.e. α ≈ 0.1999,
    // and the shared `feasible_step_fraction` applies a 0.995 boundary backoff
    // so the post-step iterate stays *strictly* interior (slack > 0).
    assert!(
        alpha > 0.0 && alpha < 1.0,
        "expected an interior step, got {alpha}"
    );
    let feasible = &states[0].beta + &(array![-1.0, 0.0] * alpha);
    let slack = family
        .time_linear_constraints
        .as_ref()
        .expect("constraints")
        .a
        .row(0)
        .dot(&feasible)
        - family
            .time_linear_constraints
            .as_ref()
            .expect("constraints")
            .b[0];
    assert!(
        slack > 0.0,
        "boundary-backed-off step must stay strictly interior; slack={slack}"
    );
}

#[test]
fn mixed_blockwise_exact_newton_preserves_sparse_block_hessians() {
    let family = SurvivalMarginalSlopeFamily {
        n: 2,
        event: Arc::new(array![1.0, 0.0]),
        weights: Arc::new(array![1.0, 0.8]),
        z: Arc::new(array![0.1, -0.2].insert_axis(Axis(1))),
        score_covariance: unit_score_covariance(),
        gaussian_frailty_sd: None,
        family_hyper: SurvivalMarginalSlopeFamilyHyperState::default(),
        derivative_guard: 1e-6,
        design_entry: DesignMatrix::Dense(DenseDesignMatrix::from(array![[1.0], [0.6]])),
        design_exit: DesignMatrix::Dense(DenseDesignMatrix::from(array![[0.9], [0.5]])),
        design_derivative_exit: DesignMatrix::Dense(DenseDesignMatrix::from(array![[1.0], [1.0]])),
        offset_entry: Arc::new(array![0.0, 0.0]),
        offset_exit: Arc::new(array![0.0, 0.0]),
        derivative_offset_exit: Arc::new(array![0.05, 0.05]),
        marginal_design: sparse_design(&array![[1.0, 0.0], [0.0, 1.0]]),
        logslope_layout: (DesignMatrix::Dense(DenseDesignMatrix::from(array![[1.0], [0.5]])))
            .into(),
        score_warp: None,
        link_dev: None,
        influence_absorber: None,
        time_linear_constraints: None,
        time_wiggle_knots: None,
        time_wiggle_degree: None,
        time_wiggle_ncols: 0,
        intercept_warm_starts: None,
        auto_subsample_phase_counter: Arc::new(AtomicUsize::new(0)),
        auto_subsample_last_rho: Arc::new(Mutex::new(None)),
    };
    let block_states = vec![
        ParameterBlockState {
            beta: array![0.4],
            eta: array![0.0, 0.0],
        },
        ParameterBlockState {
            beta: array![0.2, -0.1],
            eta: array![0.0, 0.0],
        },
        ParameterBlockState {
            beta: array![0.3],
            eta: array![0.3, 0.3],
        },
    ];

    let eval = family
        .evaluate_blockwise_exact_newton(&block_states)
        .expect("mixed exact-newton evaluation");

    assert!(matches!(
        &eval.blockworking_sets[0],
        BlockWorkingSet::ExactNewton {
            hessian: SymmetricMatrix::Dense(_),
            ..
        }
    ));
    assert!(matches!(
        &eval.blockworking_sets[1],
        BlockWorkingSet::ExactNewton {
            hessian: SymmetricMatrix::Dense(_) | SymmetricMatrix::Sparse(_),
            ..
        }
    ));
    assert!(matches!(
        &eval.blockworking_sets[2],
        BlockWorkingSet::ExactNewton {
            hessian: SymmetricMatrix::Dense(_),
            ..
        }
    ));
}

/// Closed-form test family with `gaussian_frailty_sd` set so the sigma-aware
/// joint psi paths fire. Same pseudo-random row layout as
/// `make_closed_form_test_family` so half-row subsamples remain
/// representative.
fn make_sigma_aware_closed_form_test_family(n: usize) -> SurvivalMarginalSlopeFamily {
    let mut family = make_closed_form_test_family(n);
    family.gaussian_frailty_sd = Some(0.6);
    family
}

fn rel_diff_array1_survival(a: &Array1<f64>, b: &Array1<f64>) -> f64 {
    let mut max = 0.0f64;
    for i in 0..a.len() {
        let d = (a[i] - b[i]).abs() / b[i].abs().max(1.0);
        if d > max {
            max = d;
        }
    }
    max
}

fn rel_diff_array2_survival(a: &Array2<f64>, b: &Array2<f64>) -> f64 {
    let mut max = 0.0f64;
    for ((i, j), &av) in a.indexed_iter() {
        let bv = b[[i, j]];
        let d = (av - bv).abs() / bv.abs().max(1.0);
        if d > max {
            max = d;
        }
    }
    max
}

#[test]
fn survival_sigma_psi_terms_subsample_full_equals_unsampled() {
    use crate::outer_subsample::OuterScoreSubsample;
    let n = 200usize;
    let family = make_sigma_aware_closed_form_test_family(n);
    let states = closed_form_block_states(&family, 0.25);
    let specs = vec![dummy_blockspec(0), dummy_blockspec(0), dummy_blockspec(0)];

    let baseline = family
        .sigma_exact_joint_psi_terms(&states, &specs)
        .expect("baseline psi terms")
        .expect("some");

    let mut opts_full = BlockwiseFitOptions::default();
    opts_full.outer_score_subsample = Some(Arc::new(
        OuterScoreSubsample::from_uniform_inclusion_mask((0..n).collect(), n, 0xDEADBEEF),
    ));
    let with_full = family
        .sigma_exact_joint_psi_terms_with_options(&states, &specs, &opts_full)
        .expect("with full mask")
        .expect("some");

    let obj_rel = ((with_full.objective_psi - baseline.objective_psi)
        / baseline.objective_psi.abs().max(1.0))
    .abs();
    assert!(obj_rel < 1e-12, "objective_psi rel {}", obj_rel);
    let score_rel = rel_diff_array1_survival(&with_full.score_psi, &baseline.score_psi);
    assert!(score_rel < 1e-12, "score_psi rel {}", score_rel);
}

#[test]
fn survival_sigma_psi_terms_subsample_half_scales_correctly() {
    use crate::outer_subsample::OuterScoreSubsample;
    let n = 200usize;
    let family = make_sigma_aware_closed_form_test_family(n);
    let states = closed_form_block_states(&family, 0.25);
    let specs = vec![dummy_blockspec(0), dummy_blockspec(0), dummy_blockspec(0)];

    let even_mask: Vec<usize> = (0..n).filter(|i| i % 2 == 0).collect();
    let m = even_mask.len();

    let mut opts_half = BlockwiseFitOptions::default();
    opts_half.outer_score_subsample = Some(Arc::new(
        OuterScoreSubsample::from_uniform_inclusion_mask(even_mask.clone(), n, 0xCAFE),
    ));
    let scaled = family
        .sigma_exact_joint_psi_terms_with_options(&states, &specs, &opts_half)
        .expect("scaled")
        .expect("some");

    let mut opts_raw = BlockwiseFitOptions::default();
    opts_raw.outer_score_subsample = Some(Arc::new(OuterScoreSubsample::with_uniform_weight(
        even_mask, m, 0, 1.0,
    )));
    let raw = family
        .sigma_exact_joint_psi_terms_with_options(&states, &specs, &opts_raw)
        .expect("raw")
        .expect("some");

    let factor = n as f64 / m as f64;
    let exp_obj = factor * raw.objective_psi;
    let obj_rel = ((scaled.objective_psi - exp_obj) / exp_obj.abs().max(1.0)).abs();
    assert!(obj_rel < 1e-12, "objective_psi rel {}", obj_rel);
    let exp_score = &raw.score_psi * factor;
    let score_rel = rel_diff_array1_survival(&scaled.score_psi, &exp_score);
    assert!(score_rel < 1e-12, "score_psi rel {}", score_rel);
}

#[test]
fn survival_sigma_psi_second_order_subsample_full_equals_unsampled() {
    use crate::outer_subsample::OuterScoreSubsample;
    let n = 200usize;
    let family = make_sigma_aware_closed_form_test_family(n);
    let states = closed_form_block_states(&family, 0.25);

    let baseline = family
        .sigma_exact_joint_psisecond_order_terms(&states)
        .expect("baseline")
        .expect("some");

    let mut opts_full = BlockwiseFitOptions::default();
    opts_full.outer_score_subsample = Some(Arc::new(
        OuterScoreSubsample::from_uniform_inclusion_mask((0..n).collect(), n, 0xDEADBEEF),
    ));
    let with_full = family
        .sigma_exact_joint_psisecond_order_terms_with_options(&states, &opts_full)
        .expect("with full mask")
        .expect("some");

    let obj_rel = ((with_full.objective_psi_psi - baseline.objective_psi_psi)
        / baseline.objective_psi_psi.abs().max(1.0))
    .abs();
    assert!(obj_rel < 1e-12, "objective rel {}", obj_rel);
    let score_rel = rel_diff_array1_survival(&with_full.score_psi_psi, &baseline.score_psi_psi);
    assert!(score_rel < 1e-12, "score rel {}", score_rel);
}

#[test]
fn survival_sigma_psi_second_order_subsample_half_scales_correctly() {
    use crate::outer_subsample::OuterScoreSubsample;
    let n = 200usize;
    let family = make_sigma_aware_closed_form_test_family(n);
    let states = closed_form_block_states(&family, 0.25);

    let even_mask: Vec<usize> = (0..n).filter(|i| i % 2 == 0).collect();
    let m = even_mask.len();

    let mut opts_half = BlockwiseFitOptions::default();
    opts_half.outer_score_subsample = Some(Arc::new(
        OuterScoreSubsample::from_uniform_inclusion_mask(even_mask.clone(), n, 0xCAFE),
    ));
    let scaled = family
        .sigma_exact_joint_psisecond_order_terms_with_options(&states, &opts_half)
        .expect("scaled")
        .expect("some");

    let mut opts_raw = BlockwiseFitOptions::default();
    opts_raw.outer_score_subsample = Some(Arc::new(OuterScoreSubsample::with_uniform_weight(
        even_mask, m, 0, 1.0,
    )));
    let raw = family
        .sigma_exact_joint_psisecond_order_terms_with_options(&states, &opts_raw)
        .expect("raw")
        .expect("some");

    let factor = n as f64 / m as f64;
    let exp_obj = factor * raw.objective_psi_psi;
    let obj_rel = ((scaled.objective_psi_psi - exp_obj) / exp_obj.abs().max(1.0)).abs();
    assert!(obj_rel < 1e-12, "objective rel {}", obj_rel);
    let exp_score = &raw.score_psi_psi * factor;
    let score_rel = rel_diff_array1_survival(&scaled.score_psi_psi, &exp_score);
    assert!(score_rel < 1e-12, "score rel {}", score_rel);
}

#[test]
fn survival_sigma_psihessian_directional_derivative_subsample_full_equals_unsampled() {
    use crate::outer_subsample::OuterScoreSubsample;
    let n = 200usize;
    let family = make_sigma_aware_closed_form_test_family(n);
    let states = closed_form_block_states(&family, 0.25);
    let slices = block_slices(&family, &states);
    let d_beta_flat = Array1::<f64>::zeros(slices.total);

    let baseline = family
        .sigma_exact_joint_psihessian_directional_derivative(&states, &d_beta_flat)
        .expect("baseline")
        .expect("some");

    let mut opts_full = BlockwiseFitOptions::default();
    opts_full.outer_score_subsample = Some(Arc::new(
        OuterScoreSubsample::from_uniform_inclusion_mask((0..n).collect(), n, 0xDEADBEEF),
    ));
    let with_full = family
        .sigma_exact_joint_psihessian_directional_derivative_with_options(
            &states,
            &d_beta_flat,
            &opts_full,
        )
        .expect("with full")
        .expect("some");

    let rel = rel_diff_array2_survival(&with_full, &baseline);
    assert!(rel < 1e-12, "drift rel {}", rel);
}

#[test]
fn survival_sigma_psihessian_directional_derivative_subsample_half_scales_correctly() {
    use crate::outer_subsample::OuterScoreSubsample;
    let n = 200usize;
    let family = make_sigma_aware_closed_form_test_family(n);
    let states = closed_form_block_states(&family, 0.25);
    let slices = block_slices(&family, &states);
    let d_beta_flat = Array1::<f64>::zeros(slices.total);

    let even_mask: Vec<usize> = (0..n).filter(|i| i % 2 == 0).collect();
    let m = even_mask.len();

    let mut opts_half = BlockwiseFitOptions::default();
    opts_half.outer_score_subsample = Some(Arc::new(
        OuterScoreSubsample::from_uniform_inclusion_mask(even_mask.clone(), n, 0xCAFE),
    ));
    let scaled = family
        .sigma_exact_joint_psihessian_directional_derivative_with_options(
            &states,
            &d_beta_flat,
            &opts_half,
        )
        .expect("scaled")
        .expect("some");

    let mut opts_raw = BlockwiseFitOptions::default();
    opts_raw.outer_score_subsample = Some(Arc::new(OuterScoreSubsample::with_uniform_weight(
        even_mask, m, 0, 1.0,
    )));
    let raw = family
        .sigma_exact_joint_psihessian_directional_derivative_with_options(
            &states,
            &d_beta_flat,
            &opts_raw,
        )
        .expect("raw")
        .expect("some");

    let factor = n as f64 / m as f64;
    let exp = &raw * factor;
    let rel = rel_diff_array2_survival(&scaled, &exp);
    assert!(rel < 1e-12, "drift rel {}", rel);
}

/// Multi-row test family with non-empty marginal/logslope designs but no
/// score_warp / link_dev / time_wiggle. Drives the rigid block path of
/// `psi_terms_inner` so we can subsample-check Horvitz-Thompson scaling.
fn make_block_psi_test_family(n: usize) -> SurvivalMarginalSlopeFamily {
    let event: Array1<f64> =
        Array1::from_iter((0..n).map(|i| if (i * 31 + 7) % 5 >= 3 { 1.0 } else { 0.0 }));
    let weights: Array1<f64> =
        Array1::from_iter((0..n).map(|i| 0.5 + ((i * 13 + 4) % 5) as f64 * 0.1));
    let z: Array1<f64> = Array1::from_iter(
        (0..n).map(|i| -1.0 + 2.0 * (((i * 17 + 5) % n) as f64 + 0.5) / (n as f64)),
    );
    let offset_entry: Array1<f64> = Array1::from_iter(
        (0..n).map(|i| -0.4 + 0.7 * (((i * 11 + 3) % n) as f64 + 0.5) / (n as f64)),
    );
    let offset_exit: Array1<f64> = Array1::from_iter(
        (0..n).map(|i| 0.1 + 0.6 * (((i * 19 + 7) % n) as f64 + 0.5) / (n as f64)),
    );
    let derivative_offset_exit: Array1<f64> =
        Array1::from_iter((0..n).map(|i| 0.5 + 0.05 * ((i * 23 + 1) % 3) as f64));
    // Single-column marginal/logslope designs with row-varying entries.
    let marginal_design = Array2::from_shape_fn((n, 1), |(i, _)| {
        0.3 + 0.4 * (((i * 29 + 11) % n) as f64) / (n as f64)
    });
    let logslope_design = Array2::from_shape_fn((n, 1), |(i, _)| {
        0.2 + 0.5 * (((i * 37 + 9) % n) as f64) / (n as f64)
    });
    SurvivalMarginalSlopeFamily {
        n,
        event: Arc::new(event),
        weights: Arc::new(weights),
        z: Arc::new(z.insert_axis(Axis(1))),
        score_covariance: unit_score_covariance(),
        gaussian_frailty_sd: None,
        family_hyper: SurvivalMarginalSlopeFamilyHyperState::default(),
        derivative_guard: 1e-6,
        design_entry: DesignMatrix::from(Array2::zeros((n, 0))),
        design_exit: DesignMatrix::from(Array2::zeros((n, 0))),
        design_derivative_exit: DesignMatrix::from(Array2::zeros((n, 0))),
        offset_entry: Arc::new(offset_entry),
        offset_exit: Arc::new(offset_exit),
        derivative_offset_exit: Arc::new(derivative_offset_exit),
        marginal_design: DesignMatrix::from(marginal_design),
        logslope_layout: (DesignMatrix::from(logslope_design)).into(),
        score_warp: None,
        link_dev: None,
        influence_absorber: None,
        time_linear_constraints: None,
        time_wiggle_knots: None,
        time_wiggle_degree: None,
        time_wiggle_ncols: 0,
        intercept_warm_starts: None,
        auto_subsample_phase_counter: Arc::new(AtomicUsize::new(0)),
        auto_subsample_last_rho: Arc::new(Mutex::new(None)),
    }
}

fn block_psi_test_block_states(
    family: &SurvivalMarginalSlopeFamily,
    m_beta: f64,
    g_beta: f64,
) -> Vec<ParameterBlockState> {
    let n = family.n;
    let m_design = family.marginal_design.to_dense().to_owned();
    let g_design = family
        .logslope_layout
        .coefficient_design()
        .to_dense()
        .to_owned();
    let m_eta = m_design.dot(&array![m_beta]);
    let g_eta = g_design.dot(&array![g_beta]);
    vec![
        ParameterBlockState {
            beta: Array1::zeros(0),
            eta: Array1::zeros(n),
        },
        ParameterBlockState {
            beta: array![m_beta],
            eta: m_eta,
        },
        ParameterBlockState {
            beta: array![g_beta],
            eta: g_eta,
        },
    ]
}

fn make_rigid_baseline_psi_test_family(
    n: usize,
) -> (
    SurvivalMarginalSlopeFamily,
    crate::custom_family::CustomFamilyHyperLayout,
) {
    let mut family = make_block_psi_test_family(n);
    let age_entry = Array1::from_shape_fn(n, |row| 0.25 + 0.1 * row as f64);
    let age_exit = Array1::from_shape_fn(n, |row| age_entry[row] + 0.75 + 0.03 * row as f64);
    let baseline_config = crate::survival::construction::SurvivalBaselineConfig {
        target: crate::survival::construction::SurvivalBaselineTarget::GompertzMakeham,
        scale: None,
        shape: Some(0.08),
        rate: Some(0.22),
        makeham: Some(0.04),
    };
    let geometry = Arc::new(
        crate::survival::construction::build_survival_marginal_slope_baseline_geometry(
            &age_entry,
            &age_exit,
            &baseline_config,
        )
        .expect("build rigid baseline geometry")
        .expect("Gompertz-Makeham has a nonlinear baseline chart"),
    );
    family.offset_entry = Arc::new(geometry.offset_entry.clone());
    family.offset_exit = Arc::new(geometry.offset_exit.clone());
    family.derivative_offset_exit = Arc::new(geometry.derivative_offset_exit.clone());
    family.family_hyper =
        SurvivalMarginalSlopeFamilyHyperState::new(Some(Arc::clone(&geometry)), None)
            .expect("install rigid baseline family coordinates");

    let family_axes = (0..geometry.theta.len()).collect::<Vec<_>>();
    let layout = crate::custom_family::CustomFamilyHyperLayout::new(
        vec![Vec::new(), Vec::new(), Vec::new()],
        family_axes,
        geometry.theta.clone(),
    )
    .expect("build typed baseline hyper layout");
    (family, layout)
}

fn make_flex_baseline_psi_test_fixture() -> (
    SurvivalMarginalSlopeFamily,
    Vec<ParameterBlockState>,
    Vec<ParameterBlockSpec>,
    crate::custom_family::CustomFamilyHyperLayout,
) {
    let score_runtime = test_deviation_runtime();
    let link_runtime = test_deviation_runtime();
    let (knots, degree, wiggle_width) = standard_test_time_wiggle();
    let time_width = 1 + wiggle_width;
    let age_entry = array![0.25];
    let age_exit = array![1.15];
    let baseline_config = crate::survival::construction::SurvivalBaselineConfig {
        target: crate::survival::construction::SurvivalBaselineTarget::GompertzMakeham,
        scale: None,
        shape: Some(0.08),
        rate: Some(0.22),
        makeham: Some(0.04),
    };
    let geometry = Arc::new(
        crate::survival::construction::build_survival_marginal_slope_baseline_geometry(
            &age_entry,
            &age_exit,
            &baseline_config,
        )
        .expect("build FLEX baseline geometry")
        .expect("Gompertz-Makeham has a nonlinear baseline chart"),
    );
    let mut entry_design = Array2::zeros((1, time_width));
    let mut exit_design = Array2::zeros((1, time_width));
    let mut derivative_design = Array2::zeros((1, time_width));
    entry_design[[0, 0]] = 0.25;
    exit_design[[0, 0]] = 0.45;
    derivative_design[[0, 0]] = 0.15;
    let marginal_design = array![[0.30]];
    let logslope_design = array![[0.40]];
    let family_hyper =
        SurvivalMarginalSlopeFamilyHyperState::new(Some(Arc::clone(&geometry)), None)
            .expect("install FLEX baseline family coordinates");
    let family = SurvivalMarginalSlopeFamily {
        n: 1,
        event: Arc::new(array![1.0]),
        weights: Arc::new(array![0.9]),
        z: Arc::new(array![0.2].insert_axis(Axis(1))),
        score_covariance: unit_score_covariance(),
        gaussian_frailty_sd: None,
        family_hyper,
        derivative_guard: 1e-6,
        design_entry: DesignMatrix::from(entry_design),
        design_exit: DesignMatrix::from(exit_design.clone()),
        design_derivative_exit: DesignMatrix::from(derivative_design),
        offset_entry: Arc::new(geometry.offset_entry.clone()),
        offset_exit: Arc::new(geometry.offset_exit.clone()),
        derivative_offset_exit: Arc::new(geometry.derivative_offset_exit.clone()),
        marginal_design: DesignMatrix::from(marginal_design.clone()),
        logslope_layout: DesignMatrix::from(logslope_design.clone()).into(),
        score_warp: Some(score_runtime.clone()),
        link_dev: Some(link_runtime.clone()),
        influence_absorber: None,
        time_linear_constraints: None,
        time_wiggle_knots: Some(knots),
        time_wiggle_degree: Some(degree),
        time_wiggle_ncols: wiggle_width,
        intercept_warm_starts: None,
        auto_subsample_phase_counter: Arc::new(AtomicUsize::new(0)),
        auto_subsample_last_rho: Arc::new(Mutex::new(None)),
    };

    let mut beta_time = Array1::zeros(time_width);
    beta_time[0] = 0.10;
    for local in 0..wiggle_width {
        beta_time[1 + local] = 0.006 + 0.002 * local as f64;
    }
    let beta_marginal = array![0.12];
    let beta_logslope = array![0.20];
    let beta_score =
        Array1::from_iter((0..score_runtime.basis_dim()).map(|axis| 0.004 * (1.0 + axis as f64)));
    let beta_link =
        Array1::from_iter((0..link_runtime.basis_dim()).map(|axis| -0.003 + 0.001 * axis as f64));
    let states = vec![
        ParameterBlockState {
            eta: exit_design.dot(&beta_time) + geometry.offset_exit.clone(),
            beta: beta_time,
        },
        ParameterBlockState {
            eta: marginal_design.dot(&beta_marginal),
            beta: beta_marginal,
        },
        ParameterBlockState {
            eta: logslope_design.dot(&beta_logslope),
            beta: beta_logslope,
        },
        ParameterBlockState {
            beta: beta_score,
            eta: Array1::zeros(1),
        },
        ParameterBlockState {
            beta: beta_link,
            eta: Array1::zeros(1),
        },
    ];
    let derivative_blocks = vec![
        Vec::new(),
        vec![crate::custom_family::CustomFamilyBlockPsiDerivative::new(
            None,
            array![[0.17]],
            Array2::zeros((1, 1)),
            None,
            None,
            None,
            None,
        )],
        Vec::new(),
        Vec::new(),
        Vec::new(),
    ];
    let mut values = vec![0.0];
    values.extend(geometry.theta.iter().copied());
    let layout = crate::custom_family::CustomFamilyHyperLayout::new(
        derivative_blocks,
        (0..geometry.theta.len()).collect(),
        Array1::from_vec(values),
    )
    .expect("build FLEX baseline/design hyper layout");
    let specs = states
        .iter()
        .map(|state| dummy_blockspec(state.beta.len()))
        .collect();
    (family, states, specs, layout)
}

fn assert_psi_first_terms_match(
    direct: &ExactNewtonJointPsiTerms,
    workspace: &ExactNewtonJointPsiTerms,
) {
    assert_close(
        direct.objective_psi,
        workspace.objective_psi,
        1e-13,
        "baseline first objective direct/workspace",
    );
    assert!(
        rel_diff_array1_survival(&direct.score_psi, &workspace.score_psi) < 1e-13,
        "baseline first score direct/workspace mismatch",
    );
    let direct_hessian = direct
        .hessian_psi_operator
        .as_ref()
        .expect("direct baseline Hessian operator")
        .to_dense();
    let workspace_hessian = workspace
        .hessian_psi_operator
        .as_ref()
        .expect("workspace baseline Hessian operator")
        .to_dense();
    assert!(
        rel_diff_array2_survival(&direct_hessian, &workspace_hessian) < 1e-13,
        "baseline first Hessian direct/workspace mismatch",
    );
}

fn assert_psi_second_terms_match(
    direct: &ExactNewtonJointPsiSecondOrderTerms,
    workspace: &ExactNewtonJointPsiSecondOrderTerms,
) {
    assert_close(
        direct.objective_psi_psi,
        workspace.objective_psi_psi,
        1e-13,
        "baseline pair objective direct/workspace",
    );
    assert!(
        rel_diff_array1_survival(&direct.score_psi_psi, &workspace.score_psi_psi) < 1e-13,
        "baseline pair score direct/workspace mismatch",
    );
    let direct_hessian = direct
        .hessian_psi_psi_operator
        .as_ref()
        .expect("direct baseline-pair Hessian operator")
        .to_dense();
    let workspace_hessian = workspace
        .hessian_psi_psi_operator
        .as_ref()
        .expect("workspace baseline-pair Hessian operator")
        .to_dense();
    assert!(
        rel_diff_array2_survival(&direct_hessian, &workspace_hessian) < 1e-13,
        "baseline pair Hessian direct/workspace mismatch",
    );
}

#[test]
fn rigid_baseline_dispatch_matches_direct_and_owned_workspace_without_fd() {
    let (family, hyper_layout) = make_rigid_baseline_psi_test_family(12);
    let states = block_psi_test_block_states(&family, 0.15, 0.25);
    let specs = vec![dummy_blockspec(0), dummy_blockspec(1), dummy_blockspec(1)];
    let workspace = family
        .exact_newton_joint_psi_workspace_with_options(
            &states,
            &specs,
            &hyper_layout,
            &BlockwiseFitOptions::default(),
        )
        .expect("construct rigid baseline workspace")
        .expect("rigid baseline workspace is available");

    for axis in 0..hyper_layout.len() {
        let direct = family
            .exact_newton_joint_psi_terms(&states, &specs, &hyper_layout, axis)
            .expect("direct rigid baseline first terms")
            .expect("direct rigid baseline first terms available");
        let owned = workspace
            .first_order_terms(axis)
            .expect("workspace rigid baseline first terms")
            .expect("workspace rigid baseline first terms available");
        assert_psi_first_terms_match(&direct, &owned);
    }

    let axis = 0;
    let other_axis = usize::from(hyper_layout.len() > 1);
    let direct_pair = family
        .exact_newton_joint_psisecond_order_terms(&states, &specs, &hyper_layout, axis, other_axis)
        .expect("direct rigid baseline pair terms")
        .expect("direct rigid baseline pair terms available");
    let owned_pair = workspace
        .second_order_terms(axis, other_axis)
        .expect("workspace rigid baseline pair terms")
        .expect("workspace rigid baseline pair terms available");
    assert_psi_second_terms_match(&direct_pair, &owned_pair);

    let direction = array![0.17, -0.09];
    let direct_drift = family
        .exact_newton_joint_psihessian_directional_derivative(
            &states,
            &specs,
            &hyper_layout,
            axis,
            &direction,
        )
        .expect("direct rigid baseline Hessian drift")
        .expect("direct rigid baseline Hessian drift available");
    let owned_drift = workspace
        .hessian_directional_derivative(axis, &direction)
        .expect("workspace rigid baseline Hessian drift")
        .expect("workspace rigid baseline Hessian drift available");
    let gam_problem::DriftDerivResult::Dense(owned_drift) = owned_drift else {
        panic!("rigid baseline workspace drift must preserve the dense exact carrier");
    };
    assert!(
        rel_diff_array2_survival(&direct_drift, &owned_drift) < 1e-13,
        "baseline Hessian drift direct/workspace mismatch",
    );
}

#[test]
fn flex_timewiggle_baseline_public_workspace_owns_family_and_design_pairs_without_fd() {
    let (family, states, specs, hyper_layout) = make_flex_baseline_psi_test_fixture();
    assert!(family.flex_active());
    assert!(family.flex_timewiggle_active());
    assert!(family.score_warp.is_some());
    assert!(family.link_dev.is_some());
    let slices = block_slices(&family, &states);
    let dimension = slices.total;
    let workspace = family
        .exact_newton_joint_psi_workspace_with_options(
            &states,
            &specs,
            &hyper_layout,
            &BlockwiseFitOptions::default(),
        )
        .expect("construct FLEX baseline workspace")
        .expect("FLEX baseline workspace is available");
    let design_axis = 0;
    let baseline_axis = hyper_layout.design_axis_count();
    let other_baseline_axis = baseline_axis + 1;

    let first = workspace
        .first_order_terms(baseline_axis)
        .expect("FLEX baseline first callback")
        .expect("FLEX baseline first terms are present");
    assert_eq!(first.score_psi.len(), dimension);
    let first_hessian = first
        .hessian_psi_operator
        .as_ref()
        .expect("FLEX baseline first Hessian operator")
        .to_dense();
    assert_eq!(first_hessian.dim(), (dimension, dimension));
    assert!(first.score_psi.iter().all(|value| value.is_finite()));
    assert!(first_hessian.iter().all(|value| value.is_finite()));

    let baseline_pair = workspace
        .second_order_terms(baseline_axis, other_baseline_axis)
        .expect("FLEX baseline-pair callback")
        .expect("FLEX baseline-pair terms are present");
    assert_eq!(baseline_pair.score_psi_psi.len(), dimension);
    let baseline_pair_hessian = baseline_pair
        .hessian_psi_psi_operator
        .as_ref()
        .expect("FLEX baseline-pair Hessian operator")
        .to_dense();
    assert_eq!(baseline_pair_hessian.dim(), (dimension, dimension));
    assert!(
        baseline_pair
            .score_psi_psi
            .iter()
            .all(|value| value.is_finite())
    );

    let beta_direction =
        Array1::from_iter((0..dimension).map(|axis| -0.025 + 0.006 * (axis % 7) as f64));
    let drift = workspace
        .hessian_directional_derivative(baseline_axis, &beta_direction)
        .expect("FLEX baseline Hessian-drift callback")
        .expect("FLEX baseline Hessian drift is present");
    let gam_problem::DriftDerivResult::Dense(drift) = drift else {
        panic!("FLEX baseline workspace drift must preserve the dense exact carrier");
    };
    assert_eq!(drift.dim(), (dimension, dimension));
    assert!(drift.iter().all(|value| value.is_finite()));
    assert!(drift.iter().any(|value| *value != 0.0));

    // This is the large-scale family×design seam: the global design axis is a
    // real DesignPenalty manifest entry, not a coefficient direction.  The
    // workspace must route it through X_psi beta + X_psi Jet3 seeding and
    // return the complete fixed-beta pair.
    let mixed = workspace
        .second_order_terms(baseline_axis, design_axis)
        .expect("FLEX baseline-by-design callback")
        .expect("FLEX baseline-by-design terms are present");
    assert_eq!(mixed.score_psi_psi.len(), dimension);
    assert!(mixed.objective_psi_psi.is_finite());
    assert!(mixed.score_psi_psi.iter().all(|value| value.is_finite()));
    assert!(
        mixed.score_psi_psi.iter().any(|value| value.abs() > 1e-12),
        "active FLEX/timewiggle baseline-by-design score must not collapse to zero"
    );
    let mixed_hessian = mixed
        .hessian_psi_psi_operator
        .as_ref()
        .expect("FLEX baseline-by-design Hessian operator")
        .to_dense();
    assert_eq!(mixed_hessian.dim(), (dimension, dimension));
    assert!(mixed_hessian.iter().all(|value| value.is_finite()));
    assert!(mixed_hessian.iter().any(|value| *value != 0.0));
}

/// Derivative blocks with a single ψ on the marginal block (block_idx=1).
/// `x_psi` has shape (n, 1) so the test family gets a per-row psi map.
fn block_psi_test_marginal_derivative_blocks(
    n: usize,
) -> Vec<Vec<crate::custom_family::CustomFamilyBlockPsiDerivative>> {
    let x_psi = Array2::from_shape_fn((n, 1), |(i, _)| {
        0.4 + 0.3 * (((i * 41 + 13) % n) as f64) / (n as f64)
    });
    vec![
        Vec::new(),
        vec![crate::custom_family::CustomFamilyBlockPsiDerivative::new(
            None,
            x_psi,
            Array2::zeros((1, 1)),
            None,
            None,
            None,
            None,
        )],
        Vec::new(),
    ]
}

/// Derivative blocks with one ψ on marginal (block 1) and one on logslope
/// (block 2), so second-order terms can mix.
fn block_psi_test_dual_derivative_blocks(
    n: usize,
) -> Vec<Vec<crate::custom_family::CustomFamilyBlockPsiDerivative>> {
    let x_psi_m = Array2::from_shape_fn((n, 1), |(i, _)| {
        0.4 + 0.3 * (((i * 41 + 13) % n) as f64) / (n as f64)
    });
    let x_psi_g = Array2::from_shape_fn((n, 1), |(i, _)| {
        0.2 + 0.5 * (((i * 43 + 17) % n) as f64) / (n as f64)
    });
    vec![
        Vec::new(),
        vec![crate::custom_family::CustomFamilyBlockPsiDerivative::new(
            None,
            x_psi_m,
            Array2::zeros((1, 1)),
            None,
            None,
            None,
            None,
        )],
        vec![crate::custom_family::CustomFamilyBlockPsiDerivative::new(
            None,
            x_psi_g,
            Array2::zeros((1, 1)),
            None,
            None,
            None,
            None,
        )],
    ]
}

#[test]
fn survival_psi_terms_inner_subsample_full_equals_unsampled() {
    use crate::outer_subsample::OuterScoreSubsample;
    let n = 200usize;
    let family = make_block_psi_test_family(n);
    let states = block_psi_test_block_states(&family, 0.15, 0.25);
    let derivative_blocks = block_psi_test_marginal_derivative_blocks(n);

    let baseline = family
        .psi_terms_inner(&states, &derivative_blocks, 0, None)
        .expect("baseline psi terms")
        .expect("some");

    let mut opts_full = BlockwiseFitOptions::default();
    opts_full.outer_score_subsample = Some(Arc::new(
        OuterScoreSubsample::from_uniform_inclusion_mask((0..n).collect(), n, 0xDEADBEEF),
    ));
    let with_full = family
        .psi_terms_inner_with_options(&states, &derivative_blocks, 0, None, &opts_full)
        .expect("with full mask")
        .expect("some");

    let obj_rel = ((with_full.objective_psi - baseline.objective_psi)
        / baseline.objective_psi.abs().max(1.0))
    .abs();
    assert!(obj_rel < 1e-12, "objective_psi rel {}", obj_rel);
    let score_rel = rel_diff_array1_survival(&with_full.score_psi, &baseline.score_psi);
    assert!(score_rel < 1e-12, "score_psi rel {}", score_rel);
}

#[test]
fn survival_psi_terms_inner_subsample_half_scales_correctly() {
    use crate::outer_subsample::OuterScoreSubsample;
    let n = 200usize;
    let family = make_block_psi_test_family(n);
    let states = block_psi_test_block_states(&family, 0.15, 0.25);
    let derivative_blocks = block_psi_test_marginal_derivative_blocks(n);

    let even_mask: Vec<usize> = (0..n).filter(|i| i % 2 == 0).collect();
    let m = even_mask.len();

    let mut opts_half = BlockwiseFitOptions::default();
    opts_half.outer_score_subsample = Some(Arc::new(
        OuterScoreSubsample::from_uniform_inclusion_mask(even_mask.clone(), n, 0xCAFE),
    ));
    let scaled = family
        .psi_terms_inner_with_options(&states, &derivative_blocks, 0, None, &opts_half)
        .expect("scaled")
        .expect("some");

    let mut opts_raw = BlockwiseFitOptions::default();
    opts_raw.outer_score_subsample = Some(Arc::new(OuterScoreSubsample::with_uniform_weight(
        even_mask, m, 0, 1.0,
    )));
    let raw = family
        .psi_terms_inner_with_options(&states, &derivative_blocks, 0, None, &opts_raw)
        .expect("raw")
        .expect("some");

    let factor = n as f64 / m as f64;
    let exp_obj = factor * raw.objective_psi;
    let obj_rel = ((scaled.objective_psi - exp_obj) / exp_obj.abs().max(1.0)).abs();
    assert!(obj_rel < 1e-12, "objective_psi rel {}", obj_rel);
    let exp_score = &raw.score_psi * factor;
    let score_rel = rel_diff_array1_survival(&scaled.score_psi, &exp_score);
    assert!(score_rel < 1e-12, "score_psi rel {}", score_rel);
}

#[test]
fn survival_psi_terms_inner_batched_matches_per_axis() {
    // The batched first-order ψ row pass shares the per-row primary
    // gradient/Hessian across axes; this asserts it produces the same
    // ExactNewtonJointPsiTerms (objective, score, Hessian-operator action)
    // as K serial calls to `psi_terms_inner_with_options`.
    let n = 200usize;
    let family = make_block_psi_test_family(n);
    let states = block_psi_test_block_states(&family, 0.15, 0.25);
    let derivative_blocks = block_psi_test_dual_derivative_blocks(n);
    let opts = BlockwiseFitOptions::default();

    let per_axis_0 = family
        .psi_terms_inner_with_options(&states, &derivative_blocks, 0, None, &opts)
        .expect("per-axis 0")
        .expect("some");
    let per_axis_1 = family
        .psi_terms_inner_with_options(&states, &derivative_blocks, 1, None, &opts)
        .expect("per-axis 1")
        .expect("some");

    let batched = family
        .psi_terms_inner_batched_with_options(&states, &derivative_blocks, &[0, 1], None, &opts)
        .expect("batched")
        .expect("batched simple-spatial path returned None unexpectedly");
    assert_eq!(batched.len(), 2, "batched should yield one term per axis");

    let per_axis = [&per_axis_0, &per_axis_1];
    for (i, (lhs, rhs)) in per_axis.iter().zip(batched.iter()).enumerate() {
        let obj_rel =
            ((rhs.objective_psi - lhs.objective_psi) / lhs.objective_psi.abs().max(1.0)).abs();
        assert!(
            obj_rel < 1e-12,
            "axis {i} objective_psi rel {obj_rel} (per-axis={}, batched={})",
            lhs.objective_psi,
            rhs.objective_psi,
        );
        let score_rel = rel_diff_array1_survival(&rhs.score_psi, &lhs.score_psi);
        assert!(score_rel < 1e-12, "axis {i} score_psi rel {score_rel}");

        let op_a = lhs
            .hessian_psi_operator
            .as_ref()
            .expect("per-axis Hessian operator");
        let op_b = rhs
            .hessian_psi_operator
            .as_ref()
            .expect("batched Hessian operator");
        assert_eq!(op_a.dim(), op_b.dim(), "axis {i} operator dim mismatch");
        let dim = op_a.dim();
        let probe = Array1::from_shape_fn(dim, |j| {
            ((j as i64 * 37 + 11).rem_euclid(7)) as f64 * 0.1 - 0.3
        });
        let a = op_a.mul_vec(&probe);
        let b = op_b.mul_vec(&probe);
        let op_rel = rel_diff_array1_survival(&a, &b);
        assert!(op_rel < 1e-12, "axis {i} Hessian-action rel {op_rel}");
    }
}

#[test]
fn survival_psi_terms_inner_batched_subsample_matches_per_axis() {
    // Same equivalence under a half-row Horvitz-Thompson mask, exercising
    // the per-row weight branch of the batched fast path.
    use crate::outer_subsample::OuterScoreSubsample;
    let n = 200usize;
    let family = make_block_psi_test_family(n);
    let states = block_psi_test_block_states(&family, 0.15, 0.25);
    let derivative_blocks = block_psi_test_dual_derivative_blocks(n);

    let even_mask: Vec<usize> = (0..n).filter(|i| i % 2 == 0).collect();
    let mut opts = BlockwiseFitOptions::default();
    opts.outer_score_subsample = Some(Arc::new(OuterScoreSubsample::from_uniform_inclusion_mask(
        even_mask, n, 0xC0FFEE,
    )));

    let per_axis_0 = family
        .psi_terms_inner_with_options(&states, &derivative_blocks, 0, None, &opts)
        .expect("per-axis 0")
        .expect("some");
    let per_axis_1 = family
        .psi_terms_inner_with_options(&states, &derivative_blocks, 1, None, &opts)
        .expect("per-axis 1")
        .expect("some");

    let batched = family
        .psi_terms_inner_batched_with_options(&states, &derivative_blocks, &[0, 1], None, &opts)
        .expect("batched")
        .expect("batched simple-spatial path returned None under subsample");
    assert_eq!(batched.len(), 2);

    let per_axis = [&per_axis_0, &per_axis_1];
    for (i, (lhs, rhs)) in per_axis.iter().zip(batched.iter()).enumerate() {
        let obj_rel =
            ((rhs.objective_psi - lhs.objective_psi) / lhs.objective_psi.abs().max(1.0)).abs();
        assert!(
            obj_rel < 1e-12,
            "axis {i} subsample objective_psi rel {obj_rel}"
        );
        let score_rel = rel_diff_array1_survival(&rhs.score_psi, &lhs.score_psi);
        assert!(
            score_rel < 1e-12,
            "axis {i} subsample score_psi rel {score_rel}"
        );

        let op_a = lhs.hessian_psi_operator.as_ref().unwrap();
        let op_b = rhs.hessian_psi_operator.as_ref().unwrap();
        let dim = op_a.dim();
        let probe = Array1::from_shape_fn(dim, |j| {
            ((j as i64 * 41 + 5).rem_euclid(11)) as f64 * 0.07 - 0.4
        });
        let a = op_a.mul_vec(&probe);
        let b = op_b.mul_vec(&probe);
        let op_rel = rel_diff_array1_survival(&a, &b);
        assert!(
            op_rel < 1e-12,
            "axis {i} subsample Hessian-action rel {op_rel}"
        );
    }
}

#[test]
fn survival_psi_second_order_terms_inner_subsample_full_equals_unsampled() {
    use crate::outer_subsample::OuterScoreSubsample;
    let n = 200usize;
    let family = make_block_psi_test_family(n);
    let states = block_psi_test_block_states(&family, 0.15, 0.25);
    let derivative_blocks = block_psi_test_dual_derivative_blocks(n);

    let baseline = family
        .psi_second_order_terms_inner(&states, &derivative_blocks, 0, 1, None)
        .expect("baseline psi second-order")
        .expect("some");

    let mut opts_full = BlockwiseFitOptions::default();
    opts_full.outer_score_subsample = Some(Arc::new(
        OuterScoreSubsample::from_uniform_inclusion_mask((0..n).collect(), n, 0xDEADBEEF),
    ));
    let with_full = family
        .psi_second_order_terms_inner_with_options(
            &states,
            &derivative_blocks,
            0,
            1,
            None,
            &opts_full,
        )
        .expect("with full")
        .expect("some");

    let obj_rel = ((with_full.objective_psi_psi - baseline.objective_psi_psi)
        / baseline.objective_psi_psi.abs().max(1.0))
    .abs();
    assert!(obj_rel < 1e-12, "objective rel {}", obj_rel);
    let score_rel = rel_diff_array1_survival(&with_full.score_psi_psi, &baseline.score_psi_psi);
    assert!(score_rel < 1e-12, "score rel {}", score_rel);
}

#[test]
fn survival_psi_second_order_terms_inner_subsample_half_scales_correctly() {
    use crate::outer_subsample::OuterScoreSubsample;
    let n = 200usize;
    let family = make_block_psi_test_family(n);
    let states = block_psi_test_block_states(&family, 0.15, 0.25);
    let derivative_blocks = block_psi_test_dual_derivative_blocks(n);

    let even_mask: Vec<usize> = (0..n).filter(|i| i % 2 == 0).collect();
    let m = even_mask.len();

    let mut opts_half = BlockwiseFitOptions::default();
    opts_half.outer_score_subsample = Some(Arc::new(
        OuterScoreSubsample::from_uniform_inclusion_mask(even_mask.clone(), n, 0xCAFE),
    ));
    let scaled = family
        .psi_second_order_terms_inner_with_options(
            &states,
            &derivative_blocks,
            0,
            1,
            None,
            &opts_half,
        )
        .expect("scaled")
        .expect("some");

    let mut opts_raw = BlockwiseFitOptions::default();
    opts_raw.outer_score_subsample = Some(Arc::new(OuterScoreSubsample::with_uniform_weight(
        even_mask, m, 0, 1.0,
    )));
    let raw = family
        .psi_second_order_terms_inner_with_options(
            &states,
            &derivative_blocks,
            0,
            1,
            None,
            &opts_raw,
        )
        .expect("raw")
        .expect("some");

    let factor = n as f64 / m as f64;
    let exp_obj = factor * raw.objective_psi_psi;
    let obj_rel = ((scaled.objective_psi_psi - exp_obj) / exp_obj.abs().max(1.0)).abs();
    assert!(obj_rel < 1e-12, "objective rel {}", obj_rel);
    let exp_score = &raw.score_psi_psi * factor;
    let score_rel = rel_diff_array1_survival(&scaled.score_psi_psi, &exp_score);
    assert!(score_rel < 1e-12, "score rel {}", score_rel);
}

#[test]
fn survival_psi_hessian_directional_derivative_subsample_full_equals_unsampled() {
    use crate::outer_subsample::OuterScoreSubsample;
    let n = 200usize;
    let family = make_block_psi_test_family(n);
    let states = block_psi_test_block_states(&family, 0.15, 0.25);
    let derivative_blocks = block_psi_test_marginal_derivative_blocks(n);
    let slices = block_slices(&family, &states);
    let mut d_beta_flat = Array1::<f64>::zeros(slices.total);
    d_beta_flat[slices.marginal.start] = 0.05;
    d_beta_flat[slices.logslope.start] = -0.04;

    let baseline = family
        .psi_hessian_directional_derivative(&states, &derivative_blocks, 0, &d_beta_flat)
        .expect("baseline psi-Hessian directional derivative")
        .expect("some");

    let mut opts_full = BlockwiseFitOptions::default();
    opts_full.outer_score_subsample = Some(Arc::new(
        OuterScoreSubsample::from_uniform_inclusion_mask((0..n).collect(), n, 0xDEADBEEF),
    ));
    let with_full = family
        .psi_hessian_directional_derivative_with_options(
            &states,
            &derivative_blocks,
            0,
            &d_beta_flat,
            &opts_full,
        )
        .expect("with full mask")
        .expect("some");

    let rel = rel_diff_array2_survival(&with_full, &baseline);
    assert!(rel < 1e-12, "drift rel {}", rel);
}

#[test]
fn survival_psi_hessian_directional_derivative_subsample_half_scales_correctly() {
    use crate::outer_subsample::OuterScoreSubsample;
    let n = 200usize;
    let family = make_block_psi_test_family(n);
    let states = block_psi_test_block_states(&family, 0.15, 0.25);
    let derivative_blocks = block_psi_test_marginal_derivative_blocks(n);
    let slices = block_slices(&family, &states);
    let mut d_beta_flat = Array1::<f64>::zeros(slices.total);
    d_beta_flat[slices.marginal.start] = 0.05;
    d_beta_flat[slices.logslope.start] = -0.04;

    let even_mask: Vec<usize> = (0..n).filter(|i| i % 2 == 0).collect();
    let m = even_mask.len();

    let mut opts_half = BlockwiseFitOptions::default();
    opts_half.outer_score_subsample = Some(Arc::new(
        OuterScoreSubsample::from_uniform_inclusion_mask(even_mask.clone(), n, 0xCAFE),
    ));
    let scaled = family
        .psi_hessian_directional_derivative_with_options(
            &states,
            &derivative_blocks,
            0,
            &d_beta_flat,
            &opts_half,
        )
        .expect("scaled")
        .expect("some");

    let mut opts_raw = BlockwiseFitOptions::default();
    opts_raw.outer_score_subsample = Some(Arc::new(OuterScoreSubsample::with_uniform_weight(
        even_mask, m, 0, 1.0,
    )));
    let raw = family
        .psi_hessian_directional_derivative_with_options(
            &states,
            &derivative_blocks,
            0,
            &d_beta_flat,
            &opts_raw,
        )
        .expect("raw")
        .expect("some");

    let factor = n as f64 / m as f64;
    let exp = &raw * factor;
    let rel = rel_diff_array2_survival(&scaled, &exp);
    assert!(rel < 1e-12, "drift rel {}", rel);
}

#[test]
fn survival_psi_workspace_hessian_directional_derivative_is_operator_and_matches_dense() {
    let n = 40usize;
    let family = make_block_psi_test_family(n);
    let states = block_psi_test_block_states(&family, 0.15, 0.25);
    let derivative_blocks = block_psi_test_marginal_derivative_blocks(n);
    let specs = vec![dummy_blockspec(0), dummy_blockspec(1), dummy_blockspec(1)];
    let slices = block_slices(&family, &states);
    let mut d_beta_flat = Array1::<f64>::zeros(slices.total);
    d_beta_flat[slices.marginal.start] = 0.05;
    d_beta_flat[slices.logslope.start] = -0.04;

    let dense = family
        .psi_hessian_directional_derivative_with_options(
            &states,
            &derivative_blocks,
            0,
            &d_beta_flat,
            &BlockwiseFitOptions::default(),
        )
        .expect("dense drift")
        .expect("dense drift available");
    let hyper_layout = crate::custom_family::CustomFamilyHyperLayout::new(
        derivative_blocks,
        Vec::new(),
        Array1::zeros(1),
    )
    .expect("build typed design-hyper layout");
    let workspace = family
        .exact_newton_joint_psi_workspace_with_options(
            &states,
            &specs,
            &hyper_layout,
            &BlockwiseFitOptions::default(),
        )
        .expect("workspace")
        .expect("workspace available");
    let result = workspace
        .hessian_directional_derivative(0, &d_beta_flat)
        .expect("workspace drift")
        .expect("workspace drift available");

    let gam_problem::DriftDerivResult::Operator(op) = result else {
        panic!("survival psi drift should use operator representation");
    };
    assert_eq!(op.dim(), dense.nrows());
    let operator_dense = op.to_dense();
    let rel = rel_diff_array2_survival(&operator_dense, &dense);
    assert!(rel < 1e-12, "operator/dense drift rel {rel}");
}

#[test]
fn survival_psi_hessian_directional_derivative_operator_subsample_full_equals_unsampled() {
    use crate::outer_subsample::OuterScoreSubsample;
    let n = 200usize;
    let family = make_block_psi_test_family(n);
    let states = block_psi_test_block_states(&family, 0.15, 0.25);
    let derivative_blocks = block_psi_test_marginal_derivative_blocks(n);
    let slices = block_slices(&family, &states);
    let mut d_beta_flat = Array1::<f64>::zeros(slices.total);
    d_beta_flat[slices.marginal.start] = 0.05;
    d_beta_flat[slices.logslope.start] = -0.04;

    let baseline = family
        .psi_hessian_directional_derivative_operator_with_options(
            &states,
            &derivative_blocks,
            0,
            &d_beta_flat,
            &BlockwiseFitOptions::default(),
        )
        .expect("baseline operator")
        .expect("some");
    let baseline_dense = baseline.to_dense();

    let mut opts_full = BlockwiseFitOptions::default();
    opts_full.outer_score_subsample = Some(Arc::new(
        OuterScoreSubsample::from_uniform_inclusion_mask((0..n).collect(), n, 0xDEADBEEF),
    ));
    let with_full = family
        .psi_hessian_directional_derivative_operator_with_options(
            &states,
            &derivative_blocks,
            0,
            &d_beta_flat,
            &opts_full,
        )
        .expect("with full mask")
        .expect("some");
    let with_full_dense = with_full.to_dense();

    let rel = rel_diff_array2_survival(&with_full_dense, &baseline_dense);
    assert!(rel < 1e-12, "operator drift rel {}", rel);
}

#[test]
fn survival_psi_hessian_directional_derivative_operator_subsample_half_scales_correctly() {
    use crate::outer_subsample::OuterScoreSubsample;
    let n = 200usize;
    let family = make_block_psi_test_family(n);
    let states = block_psi_test_block_states(&family, 0.15, 0.25);
    let derivative_blocks = block_psi_test_marginal_derivative_blocks(n);
    let slices = block_slices(&family, &states);
    let mut d_beta_flat = Array1::<f64>::zeros(slices.total);
    d_beta_flat[slices.marginal.start] = 0.05;
    d_beta_flat[slices.logslope.start] = -0.04;

    let even_mask: Vec<usize> = (0..n).filter(|i| i % 2 == 0).collect();
    let m = even_mask.len();

    let mut opts_half = BlockwiseFitOptions::default();
    opts_half.outer_score_subsample = Some(Arc::new(
        OuterScoreSubsample::from_uniform_inclusion_mask(even_mask.clone(), n, 0xCAFE),
    ));
    let scaled = family
        .psi_hessian_directional_derivative_operator_with_options(
            &states,
            &derivative_blocks,
            0,
            &d_beta_flat,
            &opts_half,
        )
        .expect("scaled")
        .expect("some");
    let scaled_dense = scaled.to_dense();

    let mut opts_raw = BlockwiseFitOptions::default();
    opts_raw.outer_score_subsample = Some(Arc::new(OuterScoreSubsample::with_uniform_weight(
        even_mask, m, 0, 1.0,
    )));
    let raw = family
        .psi_hessian_directional_derivative_operator_with_options(
            &states,
            &derivative_blocks,
            0,
            &d_beta_flat,
            &opts_raw,
        )
        .expect("raw")
        .expect("some");
    let raw_dense = raw.to_dense();

    let factor = n as f64 / m as f64;
    let exp = &raw_dense * factor;
    let rel = rel_diff_array2_survival(&scaled_dense, &exp);
    assert!(rel < 1e-12, "operator drift rel {}", rel);
}

// ── Phase 7: joint-Hessian flex-no-wiggle directional-derivative
// operator subsample tests. The flex-no-wiggle helpers are the path
// taken by the joint-Hessian workspace's `directional_derivative_operator`
// when `effective_flex_active(states)` is true and timewiggle is off.
// We exercise the helpers directly through a flex-active fixture so the
// outer subsample threading is verified end-to-end.

fn make_flex_no_wiggle_test_family(n: usize) -> SurvivalMarginalSlopeFamily {
    let score_runtime = test_deviation_runtime();
    let event: Array1<f64> =
        Array1::from_iter((0..n).map(|i| if (i * 31 + 7) % 5 >= 3 { 1.0 } else { 0.0 }));
    let weights: Array1<f64> =
        Array1::from_iter((0..n).map(|i| 0.5 + ((i * 13 + 4) % 5) as f64 * 0.1));
    let z: Array1<f64> = Array1::from_iter(
        (0..n).map(|i| -1.0 + 2.0 * (((i * 17 + 5) % n) as f64 + 0.5) / (n as f64)),
    );
    let offset_entry: Array1<f64> = Array1::from_iter(
        (0..n).map(|i| -0.4 + 0.7 * (((i * 11 + 3) % n) as f64 + 0.5) / (n as f64)),
    );
    let offset_exit: Array1<f64> = Array1::from_iter(
        (0..n).map(|i| 0.1 + 0.6 * (((i * 19 + 7) % n) as f64 + 0.5) / (n as f64)),
    );
    let derivative_offset_exit: Array1<f64> =
        Array1::from_iter((0..n).map(|i| 0.5 + 0.05 * ((i * 23 + 1) % 3) as f64));
    let marginal_design = Array2::from_shape_fn((n, 1), |(i, _)| {
        0.3 + 0.4 * (((i * 29 + 11) % n) as f64) / (n as f64)
    });
    let logslope_design = Array2::from_shape_fn((n, 1), |(i, _)| {
        0.2 + 0.5 * (((i * 37 + 9) % n) as f64) / (n as f64)
    });
    SurvivalMarginalSlopeFamily {
        n,
        event: Arc::new(event),
        weights: Arc::new(weights),
        z: Arc::new(z.insert_axis(Axis(1))),
        score_covariance: unit_score_covariance(),
        gaussian_frailty_sd: None,
        family_hyper: SurvivalMarginalSlopeFamilyHyperState::default(),
        derivative_guard: 1e-6,
        design_entry: DesignMatrix::from(Array2::zeros((n, 0))),
        design_exit: DesignMatrix::from(Array2::zeros((n, 0))),
        design_derivative_exit: DesignMatrix::from(Array2::zeros((n, 0))),
        offset_entry: Arc::new(offset_entry),
        offset_exit: Arc::new(offset_exit),
        derivative_offset_exit: Arc::new(derivative_offset_exit),
        marginal_design: DesignMatrix::from(marginal_design),
        logslope_layout: (DesignMatrix::from(logslope_design)).into(),
        score_warp: Some(score_runtime),
        link_dev: None,
        influence_absorber: None,
        time_linear_constraints: None,
        time_wiggle_knots: None,
        time_wiggle_degree: None,
        time_wiggle_ncols: 0,
        intercept_warm_starts: None,
        auto_subsample_phase_counter: Arc::new(AtomicUsize::new(0)),
        auto_subsample_last_rho: Arc::new(Mutex::new(None)),
    }
}

fn flex_no_wiggle_test_block_states(
    family: &SurvivalMarginalSlopeFamily,
) -> Vec<ParameterBlockState> {
    let n = family.n;
    let m_design = family.marginal_design.to_dense().to_owned();
    let g_design = family
        .logslope_layout
        .coefficient_design()
        .to_dense()
        .to_owned();
    let m_beta = 0.15_f64;
    let g_beta = 0.25_f64;
    let m_eta = m_design.dot(&array![m_beta]);
    let g_eta = g_design.dot(&array![g_beta]);
    let score_dim = family
        .score_warp
        .as_ref()
        .map(|w| w.basis_dim())
        .unwrap_or(0);
    vec![
        ParameterBlockState {
            beta: Array1::zeros(0),
            eta: Array1::zeros(n),
        },
        ParameterBlockState {
            beta: array![m_beta],
            eta: m_eta,
        },
        ParameterBlockState {
            beta: array![g_beta],
            eta: g_eta,
        },
        ParameterBlockState {
            beta: Array1::zeros(score_dim),
            eta: Array1::zeros(n),
        },
    ]
}

#[test]
fn survival_jointhessian_flex_no_wiggle_operator_subsample_full_equals_unsampled() {
    use crate::outer_subsample::OuterScoreSubsample;
    let n = 40usize;
    let family = make_flex_no_wiggle_test_family(n);
    let states = flex_no_wiggle_test_block_states(&family);
    assert!(family.effective_flex_active(&states).unwrap());
    assert!(!family.flex_timewiggle_active());
    let slices = block_slices(&family, &states);
    let mut d_beta_flat = Array1::<f64>::zeros(slices.total);
    d_beta_flat[slices.marginal.start] = 0.05;
    d_beta_flat[slices.logslope.start] = -0.04;

    let baseline = family
        .exact_newton_joint_hessian_directional_derivative_operator_flex_no_wiggle_with_options(
            &states,
            &d_beta_flat,
            &BlockwiseFitOptions::default(),
        )
        .expect("baseline operator");
    let baseline_dense = baseline.to_dense();

    let mut opts_full = BlockwiseFitOptions::default();
    opts_full.outer_score_subsample = Some(Arc::new(
        OuterScoreSubsample::from_uniform_inclusion_mask((0..n).collect(), n, 0xDEADBEEF),
    ));
    let with_full = family
        .exact_newton_joint_hessian_directional_derivative_operator_flex_no_wiggle_with_options(
            &states,
            &d_beta_flat,
            &opts_full,
        )
        .expect("with full mask");
    let with_full_dense = with_full.to_dense();

    let rel = rel_diff_array2_survival(&with_full_dense, &baseline_dense);
    assert!(
        rel < 1e-10,
        "joint Hessian flex-no-wiggle dH operator drift rel {}",
        rel
    );
}

/// #979 flex hot path: the build-once all-axes directional-derivative sweep
/// (`..._flex_no_wiggle_all_axes`) must reproduce calling the single-direction
/// routine once per coordinate axis. The all-axes path hoists the
/// direction-independent per-row geometry (intercept solves, cached partitions,
/// base timepoints) out of the per-axis loop; this asserts the optimization
/// changes only the cost, never the result. The two paths feed the same per-row
/// assemblers but run independent rayon reductions, so equality is to a tight
/// relative tolerance (the same convention the other flex-no-wiggle operator
/// tests use), not bit-for-bit.
#[test]
fn survival_jointhessian_flex_no_wiggle_all_axes_matches_per_axis() {
    let n = 40usize;
    let family = make_flex_no_wiggle_test_family(n);
    let states = flex_no_wiggle_test_block_states(&family);
    assert!(family.effective_flex_active(&states).unwrap());
    assert!(!family.flex_timewiggle_active());
    let slices = block_slices(&family, &states);
    let p = slices.total;
    assert!(p >= 2, "test family must exercise multiple axes, got p={p}");

    let all_axes = family
        .exact_newton_joint_hessian_directional_derivative_flex_no_wiggle_all_axes(&states)
        .expect("all-axes sweep");
    assert_eq!(all_axes.len(), p);

    for axis_idx in 0..p {
        let mut axis = Array1::<f64>::zeros(p);
        axis[axis_idx] = 1.0;
        let per_axis = family
            .exact_newton_joint_hessian_directional_derivative_flex_no_wiggle(&states, &axis)
            .expect("per-axis directional derivative");
        let batched = &all_axes[axis_idx];
        assert_eq!(batched.dim(), per_axis.dim());
        let rel = rel_diff_array2_survival(batched, &per_axis);
        assert!(
            rel < 1e-10,
            "axis {axis_idx}: build-once all-axes path diverged from per-axis sweep, rel {rel}"
        );
    }
}

#[test]
fn survival_jointhessian_flex_no_wiggle_operator_subsample_half_scales_correctly() {
    use crate::outer_subsample::OuterScoreSubsample;
    let n = 40usize;
    let family = make_flex_no_wiggle_test_family(n);
    let states = flex_no_wiggle_test_block_states(&family);
    assert!(family.effective_flex_active(&states).unwrap());
    assert!(!family.flex_timewiggle_active());
    let slices = block_slices(&family, &states);
    let mut d_beta_flat = Array1::<f64>::zeros(slices.total);
    d_beta_flat[slices.marginal.start] = 0.05;
    d_beta_flat[slices.logslope.start] = -0.04;

    let even_mask: Vec<usize> = (0..n).filter(|i| i % 2 == 0).collect();
    let m = even_mask.len();

    let mut opts_half = BlockwiseFitOptions::default();
    opts_half.outer_score_subsample = Some(Arc::new(
        OuterScoreSubsample::from_uniform_inclusion_mask(even_mask.clone(), n, 0xCAFE),
    ));
    let scaled = family
        .exact_newton_joint_hessian_directional_derivative_operator_flex_no_wiggle_with_options(
            &states,
            &d_beta_flat,
            &opts_half,
        )
        .expect("scaled");
    let scaled_dense = scaled.to_dense();

    let mut opts_raw = BlockwiseFitOptions::default();
    opts_raw.outer_score_subsample = Some(Arc::new(OuterScoreSubsample::with_uniform_weight(
        even_mask, m, 0, 1.0,
    )));
    let raw = family
        .exact_newton_joint_hessian_directional_derivative_operator_flex_no_wiggle_with_options(
            &states,
            &d_beta_flat,
            &opts_raw,
        )
        .expect("raw");
    let raw_dense = raw.to_dense();

    let factor = n as f64 / m as f64;
    let exp = &raw_dense * factor;
    let rel = rel_diff_array2_survival(&scaled_dense, &exp);
    assert!(
        rel < 1e-10,
        "joint Hessian flex-no-wiggle dH operator HT rel {}",
        rel
    );
}

#[test]
fn survival_auto_subsample_phase_counter_field_initializes_to_zero() {
    let family = make_closed_form_test_family(8);
    assert_eq!(
        family
            .auto_subsample_phase_counter
            .load(std::sync::atomic::Ordering::SeqCst),
        0,
        "fresh family must start at Phase-1 step 0"
    );
    assert!(
        family
            .auto_subsample_last_rho
            .lock()
            .expect("auto_subsample_last_rho mutex poisoned")
            .is_none(),
        "fresh family must have no recorded last-rho proxy"
    );
}

// ────────────────────────────────────────────────────────────────────
// Independent fourth-contraction finite-difference fixtures.
// ────────────────────────────────────────────────────────────────────

#[derive(Clone, Copy)]
struct FlexContractionFixture {
    label: &'static str,
    event: f64,
    weight: f64,
    z: f64,
    q0: f64,
    q1: f64,
    qd1: f64,
    score_eta: f64,
    h_scale: f64,
    w_scale: f64,
}

const FLEX_CONTRACTION_FIXTURES: &[FlexContractionFixture] = &[
    FlexContractionFixture {
        label: "event_nonzero_warps",
        event: 1.0,
        weight: 0.75,
        z: -0.2,
        q0: -0.4,
        q1: 0.6,
        qd1: 0.85,
        score_eta: 0.32,
        h_scale: 0.05,
        w_scale: 0.04,
    },
    FlexContractionFixture {
        label: "censored_left_tail",
        event: 0.0,
        weight: 1.35,
        z: -1.15,
        q0: -1.35,
        q1: -0.9,
        qd1: 0.42,
        score_eta: -0.55,
        h_scale: -0.035,
        w_scale: 0.025,
    },
    FlexContractionFixture {
        label: "near_boundary_derivative",
        event: 1.0,
        weight: 0.2,
        z: 0.95,
        q0: 0.15,
        q1: 1.05,
        qd1: 0.08,
        score_eta: 0.72,
        h_scale: 0.015,
        w_scale: -0.02,
    },
    FlexContractionFixture {
        label: "zero_warp_edge",
        event: 0.0,
        weight: 0.9,
        z: 0.0,
        q0: -0.05,
        q1: 0.25,
        qd1: 1.2,
        score_eta: 0.0,
        h_scale: 0.0,
        w_scale: 0.0,
    },
];

fn flex_contraction_fixture_family(
    fixture: FlexContractionFixture,
) -> (SurvivalMarginalSlopeFamily, Vec<ParameterBlockState>) {
    let score_runtime = test_deviation_runtime();
    let link_runtime = test_deviation_runtime();
    let family = SurvivalMarginalSlopeFamily {
        n: 1,
        event: Arc::new(array![fixture.event]),
        weights: Arc::new(array![fixture.weight]),
        z: Arc::new(array![fixture.z].insert_axis(Axis(1))),
        score_covariance: unit_score_covariance(),
        gaussian_frailty_sd: None,
        family_hyper: SurvivalMarginalSlopeFamilyHyperState::default(),
        derivative_guard: 1e-6,
        design_entry: DesignMatrix::from(Array2::zeros((1, 1))),
        design_exit: DesignMatrix::from(Array2::zeros((1, 1))),
        design_derivative_exit: DesignMatrix::from(Array2::zeros((1, 1))),
        offset_entry: Arc::new(array![fixture.q0]),
        offset_exit: Arc::new(array![fixture.q1]),
        derivative_offset_exit: Arc::new(array![fixture.qd1]),
        marginal_design: DesignMatrix::from(Array2::zeros((1, 0))),
        logslope_layout: (DesignMatrix::from(Array2::zeros((1, 0)))).into(),
        score_warp: Some(score_runtime.clone()),
        link_dev: Some(link_runtime.clone()),
        influence_absorber: None,
        time_linear_constraints: None,
        time_wiggle_knots: None,
        time_wiggle_degree: None,
        time_wiggle_ncols: 0,
        intercept_warm_starts: None,
        auto_subsample_phase_counter: Arc::new(AtomicUsize::new(0)),
        auto_subsample_last_rho: Arc::new(Mutex::new(None)),
    };
    let h_dim = score_runtime.basis_dim();
    let w_dim = link_runtime.basis_dim();
    let h_beta: Array1<f64> = (0..h_dim)
        .map(|k| fixture.h_scale * ((k as f64 + 1.0).sin()))
        .collect::<Vec<_>>()
        .into();
    let w_beta: Array1<f64> = (0..w_dim)
        .map(|k| fixture.w_scale * ((k as f64 + 1.0).cos()))
        .collect::<Vec<_>>()
        .into();
    let block_states = vec![
        ParameterBlockState {
            beta: Array1::zeros(1),
            eta: Array1::zeros(1),
        },
        ParameterBlockState {
            beta: Array1::zeros(0),
            eta: Array1::zeros(1),
        },
        ParameterBlockState {
            beta: Array1::zeros(0),
            eta: array![fixture.score_eta],
        },
        ParameterBlockState {
            beta: h_beta,
            eta: Array1::zeros(1),
        },
        ParameterBlockState {
            beta: w_beta,
            eta: Array1::zeros(1),
        },
    ];
    (family, block_states)
}

fn flex_contraction_directions(p: usize) -> Vec<(&'static str, Array1<f64>)> {
    let mixed: Array1<f64> = (0..p)
        .map(|k| 0.1 + 0.07 * ((k as f64 + 1.7).sin()))
        .collect::<Vec<_>>()
        .into();
    let alternating: Array1<f64> = (0..p)
        .map(|k| if k % 2 == 0 { 0.16 } else { -0.11 })
        .collect::<Vec<_>>()
        .into();
    let mut qd_axis = Array1::zeros(p);
    if p > 2 {
        qd_axis[2] = 1.0;
    }
    let zero = Array1::zeros(p);
    vec![
        ("mixed", mixed),
        ("alternating", alternating),
        ("qd_axis", qd_axis),
        ("zero", zero),
    ]
}

/// Shift the g/h/w primary axes of a flex `block_states` by `step · dir` (the
/// q0/q1/qd1 offset axes are left fixed). `dir` is a full primary-space vector;
/// only its `g` (block_states[2].eta), `h` (block_states[3].beta), and `w`
/// (block_states[4].beta) components are applied — exactly the axes the
/// directional contraction differentiates through the score/link coefficient
/// chain (the #1454-sensitive cross channels).
fn perturb_flex_ghw_block_states(
    block_states: &[ParameterBlockState],
    primary: &FlexPrimarySlices,
    dir: &Array1<f64>,
    step: f64,
) -> Vec<ParameterBlockState> {
    let mut bs: Vec<ParameterBlockState> = block_states.to_vec();
    // g axis: marginal slope eta (block index 2).
    bs[2].eta[0] += step * dir[primary.g];
    // h axis: score-warp beta (block index 3).
    if let Some(h_range) = primary.h.as_ref() {
        for (local, idx) in h_range.clone().enumerate() {
            bs[3].beta[local] += step * dir[idx];
        }
    }
    // w axis: link-dev beta (block index 4).
    if let Some(w_range) = primary.w.as_ref() {
        for (local, idx) in w_range.clone().enumerate() {
            bs[4].beta[local] += step * dir[idx];
        }
    }
    bs
}

/// #932-2 / #1454: the production fourth contraction
/// `D_u D_v H = row_flex_primary_fourth_contracted_exact` is the single-source
/// jet path (`flex_jet` `flex_timepoint_inputs_generic` at `Jet4`). It is checked
/// here against an INDEPENDENT scalar finite difference of the production THIRD
/// contraction `D_u H = row_flex_primary_third_contracted_exact` along the second
/// direction `v` (g/h/w axes), re-solving the moving-boundary intercept exactly at
/// every perturbed point. The third contraction is
/// itself pinned to the same FD ground truth by
/// `flex_contracted_tower_matches_independent_fd_witness_nonzero_deviation`, so a
/// central difference of it is a faithful fourth-order ground truth.
#[test]
fn flex_production_fourth_contraction_matches_scalar_fd_witness() {
    for &fixture in FLEX_CONTRACTION_FIXTURES {
        let (family, block_states) = flex_contraction_fixture_family(fixture);
        let primary = flex_primary_slices(&family);
        let p = primary.total;
        let dirs = flex_contraction_directions(p);
        let pairs = [(0usize, 0usize), (0, 1), (1, 0), (1, 2), (2, 3), (3, 0)];
        for &(u_idx, v_idx) in &pairs {
            let (u_label, dir_u) = &dirs[u_idx];
            let (v_label, dir_v_full) = &dirs[v_idx];
            // Restrict the FD (second) direction to the g/h/w block-state axes (the
            // q0/q1/qd1 offset axes are not block-state-perturbable here); the
            // production fourth is contracted against the SAME restricted direction so
            // both sides see the identical `v`.
            let mut dir_v = Array1::<f64>::zeros(p);
            dir_v[primary.g] = dir_v_full[primary.g];
            if let Some(h_range) = primary.h.as_ref() {
                for idx in h_range.clone() {
                    dir_v[idx] = dir_v_full[idx];
                }
            }
            if let Some(w_range) = primary.w.as_ref() {
                for idx in w_range.clone() {
                    dir_v[idx] = dir_v_full[idx];
                }
            }
            if dir_v.iter().all(|x| x.abs() == 0.0) {
                continue;
            }

            let production = family
                .row_flex_primary_fourth_contracted_exact(0, &block_states, dir_u, &dir_v)
                .unwrap_or_else(|err| {
                    panic!(
                        "{} / {u_label}->{v_label}: production fourth contraction failed: {err}",
                        fixture.label
                    )
                });

            // Central difference (Richardson) of the production THIRD contraction
            // `D_u H` along the g/h/w direction `dir_v`.
            let third_at = |step: f64| -> ndarray::Array2<f64> {
                let bs = perturb_flex_ghw_block_states(&block_states, &primary, &dir_v, step);
                family
                    .row_flex_primary_third_contracted_exact(0, &bs, dir_u)
                    .unwrap_or_else(|err| {
                        panic!(
                            "{} / {u_label}->{v_label}: perturbed third contraction failed: {err}",
                            fixture.label
                        )
                    })
            };
            let central =
                |h: f64| -> ndarray::Array2<f64> { (&third_at(h) - &third_at(-h)) / (2.0 * h) };
            let h0 = 4e-3;
            let coarse = central(h0);
            let fine = central(h0 * 0.5);
            // Richardson extrapolation (O(h⁴)) of the central difference.
            let fd_fourth = (&fine * 4.0 - &coarse) / 3.0;

            for u in 0..p {
                for v in 0..p {
                    let got = production[[u, v]];
                    let want = fd_fourth[[u, v]];
                    let scale = want.abs().max(1.0);
                    // Richardson convergence guard: if the coarse (h0) and fine
                    // (h0/2) central differences of the third contraction disagree
                    // beyond the assert tolerance, the FD witness has NOT converged —
                    // the third contraction is non-smooth in the block-state
                    // perturbation at this point (e.g. the all-zero `zero_warp_edge`
                    // degenerate fixture: censored event=0, z=0, score_eta=0, every
                    // warp coeff 0, where the re-solved moving-boundary intercept is
                    // non-differentiable). The FD witness is unreliable there, not the
                    // production path (which the #1454 gate FD-validates on smooth
                    // fixtures). Skip only those provably-non-convergent entries; every
                    // entry where the FD converges is still asserted strictly.
                    let fd_unconverged =
                        (fine[[u, v]] - coarse[[u, v]]).abs() > 2e-2 * scale + 1e-5;
                    if fd_unconverged {
                        eprintln!(
                            "#932 flex fourth[{u},{v}] {}/{u_label}->{v_label}: FD witness \
                             unconverged (coarse {:+.4e}, fine {:+.4e}); skipping — production \
                             {got:+.4e}",
                            fixture.label,
                            coarse[[u, v]],
                            fine[[u, v]]
                        );
                        continue;
                    }
                    assert!(
                        (got - want).abs() <= 2e-2 * scale + 1e-5,
                        "{} / {u_label}->{v_label} fourth[{u},{v}]: production {got:+.6e} != \
                         scalar-FD-of-third {want:+.6e}",
                        fixture.label
                    );
                }
            }
        }
    }
}

/// Build a one-row time-block family whose single derivative-design row has
/// the given coefficient and offset, so we can drive
/// `validate_time_qd1_feasible` with a controlled raw `qd1` and constraint
/// row scaling. Mirrors the field layout of the other in-module fixtures.
fn make_time_guard_family(deriv_coeff: f64, deriv_offset: f64) -> SurvivalMarginalSlopeFamily {
    SurvivalMarginalSlopeFamily {
        n: 1,
        event: Arc::new(array![1.0]),
        weights: Arc::new(array![1.0]),
        z: Arc::new(array![0.0].insert_axis(Axis(1))),
        score_covariance: unit_score_covariance(),
        gaussian_frailty_sd: None,
        family_hyper: SurvivalMarginalSlopeFamilyHyperState::default(),
        derivative_guard: 1e-6,
        design_entry: DesignMatrix::from(Array2::zeros((1, 1))),
        design_exit: DesignMatrix::from(Array2::zeros((1, 1))),
        design_derivative_exit: DesignMatrix::from(array![[deriv_coeff]]),
        offset_entry: Arc::new(array![0.0]),
        offset_exit: Arc::new(array![0.0]),
        derivative_offset_exit: Arc::new(array![deriv_offset]),
        marginal_design: DesignMatrix::from(Array2::zeros((1, 1))),
        logslope_layout: (DesignMatrix::from(array![[1.0]])).into(),
        score_warp: None,
        link_dev: None,
        influence_absorber: None,
        time_linear_constraints: None,
        time_wiggle_knots: None,
        time_wiggle_degree: None,
        time_wiggle_ncols: 0,
        intercept_warm_starts: None,
        auto_subsample_phase_counter: Arc::new(AtomicUsize::new(0)),
        auto_subsample_last_rho: Arc::new(Mutex::new(None)),
    }
}

/// Regression for #379: a full-rank PH marginal-slope fit aborted because
/// `validate_time_qd1_feasible` rejected a constrained time-block iterate
/// that the inequality-constrained active-set Newton solver considered
/// feasible. The solver only guarantees primal feasibility to
/// `ACTIVE_SET_PRIMAL_FEASIBILITY_TOL` in the *scaled* constraint-row units
/// produced by `time_derivative_guard_constraints` (row scale
/// `max(||design_row||, |guard - offset|, 1)`). On a design row with a
/// large norm, a scaled slack of ~1e-8 maps to a raw `qd1` shortfall of
/// ~1e-5 below the 1e-6 guard — exactly the overshoot in the issue. The
/// validator must accept that boundary iterate while still rejecting a
/// genuine monotonicity divergence.
#[test]
fn validate_time_qd1_accepts_scaled_boundary_overshoot_rejects_real_violation() {
    let guard = 1e-6_f64;

    // Boundary case: design row norm ~6000, so the solver's 1e-8 scaled
    // primal tolerance permits a raw qd1 shortfall of ~6e-5 below the guard.
    let deriv_coeff = 6000.0_f64;
    let family = make_time_guard_family(deriv_coeff, 0.0);
    // Choose beta so qd1 = guard - 6e-5 (raw overshoot ~6e-5, as observed).
    let target_qd1 = guard - 6e-5;
    let beta = array![target_qd1 / deriv_coeff];
    // Confirm the fixture actually reproduces the issue's regime: the raw
    // qd1 trips the historic 256·eps guard band (so the *old* validator
    // would have rejected this iterate and aborted the whole fit).
    assert!(
        super::survival_derivative_guard_violated(target_qd1, guard),
        "fixture must reproduce the sub-guard raw overshoot the old validator rejected",
    );
    // Scaled violation = 6e-5 / max(6000, |guard|, 1) = 1e-8, below the
    // solver-consistent feasibility band, so the iterate must be accepted.
    family
        .validate_time_qd1_feasible(&beta, "proposed")
        .expect("solver-feasible boundary iterate must be accepted");

    // Genuine violation: unit-scale row, qd1 driven far below the guard.
    let bad_family = make_time_guard_family(1.0, 0.0);
    let bad_beta = array![-0.5]; // qd1 = -0.5, scaled violation ~0.5 >> band
    let err = bad_family
        .validate_time_qd1_feasible(&bad_beta, "proposed")
        .expect_err("a true monotonicity divergence must still hard-error");
    assert!(
        err.contains("violates monotonicity"),
        "expected a monotonicity error, got: {err}",
    );
}

/// Fill every one of the 15 cross-block matrices of a
/// `BlockHessianAccumulator` with distinct, deterministic values keyed by
/// (block-pair, local-row, local-col) so that a transposed or mis-placed
/// block is impossible to miss in the dense/operator parity assertions.
fn fill_block_hessian_accumulator(
    p_t: usize,
    p_m: usize,
    p_g: usize,
    p_h: usize,
    p_w: usize,
) -> BlockHessianAccumulator {
    let mut acc = BlockHessianAccumulator::new(p_t, p_m, p_g, p_h, p_w, 0);
    // Per-pair base offsets keep the diagonal blocks symmetric (required
    // for a valid Hessian) while making the off-diagonal blocks distinct
    // and asymmetric, so `to_dense` must place each transpose correctly.
    let sym = |m: &mut Array2<f64>, base: f64| {
        let (r, c) = m.dim();
        for i in 0..r {
            for j in 0..c {
                m[[i, j]] = base + (i.min(j) as f64) * 0.5 + (i.max(j) as f64) * 0.25;
            }
        }
    };
    sym(&mut acc.h_tt, 1.0);
    sym(&mut acc.h_mm, 2.0);
    sym(&mut acc.h_gg, 3.0);
    sym(&mut acc.h_hh, 4.0);
    sym(&mut acc.h_ww, 5.0);
    let rect = |m: &mut Array2<f64>, base: f64| {
        let (r, c) = m.dim();
        for i in 0..r {
            for j in 0..c {
                m[[i, j]] = base + (i as f64) * 1.0 + (j as f64) * 0.1;
            }
        }
    };
    rect(&mut acc.h_tm, 10.0);
    rect(&mut acc.h_tg, 20.0);
    rect(&mut acc.h_th, 30.0);
    rect(&mut acc.h_tw, 40.0);
    rect(&mut acc.h_mg, 50.0);
    rect(&mut acc.h_mh, 60.0);
    rect(&mut acc.h_mw, 70.0);
    rect(&mut acc.h_gh, 80.0);
    rect(&mut acc.h_gw, 90.0);
    rect(&mut acc.h_hw, 100.0);
    acc
}

fn full_block_slices(p_t: usize, p_m: usize, p_g: usize, p_h: usize, p_w: usize) -> BlockSlices {
    let time = 0..p_t;
    let marginal = time.end..time.end + p_m;
    let logslope = marginal.end..marginal.end + p_g;
    let score_warp = logslope.end..logslope.end + p_h;
    let link_dev = score_warp.end..score_warp.end + p_w;
    let total = link_dev.end;
    BlockSlices {
        time,
        marginal,
        logslope,
        score_warp: Some(score_warp),
        link_dev: Some(link_dev),
        influence: None,
        total,
    }
}

/// Parity guard for issue #428: the dense and operator block-Hessian
/// assembly are now one storage abstraction with a single scatter, so
/// `BlockHessianOperator` (the operator wrapping the accumulator) must
/// agree *exactly* with the dense scatter — to_dense, matvec, and the
/// bilinear form — across the full five-block layout (time, marginal,
/// logslope, score_warp, link_dev) that the directional-derivative path
/// produces. Any future reintroduction of a second, divergent block
/// layout would break this test.
#[test]
fn block_hessian_dense_operator_parity_all_five_blocks() {
    let (p_t, p_m, p_g, p_h, p_w) = (3usize, 2, 2, 3, 2);
    let slices = full_block_slices(p_t, p_m, p_g, p_h, p_w);
    let acc = fill_block_hessian_accumulator(p_t, p_m, p_g, p_h, p_w);

    // The dense reference scatter (single source of truth).
    let dense = acc.to_dense(&slices);
    assert_eq!(dense.dim(), (slices.total, slices.total));

    // The full block Hessian must be symmetric: each off-diagonal block
    // appears with its transpose at the mirrored location.
    for i in 0..slices.total {
        for j in 0..slices.total {
            assert_relative_eq!(dense[[i, j]], dense[[j, i]], max_relative = 1e-12);
        }
    }

    // Operator densification must equal the dense scatter bit-for-bit:
    // both now route through `BlockHessianAccumulator::to_dense`.
    let op = acc.into_operator(slices.clone());
    let op_dense = op.to_dense();
    assert_eq!(op_dense.dim(), dense.dim());
    for i in 0..slices.total {
        for j in 0..slices.total {
            assert_eq!(
                op_dense[[i, j]],
                dense[[i, j]],
                "operator/dense mismatch at ({i}, {j})",
            );
        }
    }

    // Matvec parity: operator.mul_vec(v) == dense.dot(v) for arbitrary v,
    // including the directional-derivative use (the operator is exactly
    // what the operator-variant directional path returns).
    let v: Array1<f64> =
        Array1::from_iter((0..slices.total).map(|k| 0.7 + (k as f64) * 0.31 - 0.05 * k as f64));
    let u: Array1<f64> = Array1::from_iter((0..slices.total).map(|k| -1.3 + (k as f64) * 0.17));
    let mv = op.mul_vec(&v);
    let dense_mv = dense.dot(&v);
    for k in 0..slices.total {
        assert_relative_eq!(mv[k], dense_mv[k], max_relative = 1e-12);
    }
    // The view-based matvec must match the owned one exactly.
    let mv_view = op.mul_vec_view(v.view());
    for k in 0..slices.total {
        assert_eq!(mv_view[k], mv[k]);
    }

    // Bilinear parity: operator.bilinear(v, u) == vᵀ · dense · u.
    let bil = op.bilinear(&v, &u);
    let dense_bil = v.dot(&dense.dot(&u));
    assert_relative_eq!(bil, dense_bil, max_relative = 1e-12);
}

#[test]
fn zz_diag_failure1_flex_vs_rigid_vs_fdhess() {
    use gam_math::jet_tower::program_third_contracted;
    // FAILURE 1 fixture row.
    let event = 1.0_f64;
    let weight = 0.75_f64;
    let zr = -0.2_f64;
    let q0 = -0.4_f64;
    let q1 = 0.6_f64;
    let qd1 = 0.85_f64;
    let gv = 0.32_f64;

    let make = |q0: f64, q1: f64, qd1: f64, g: f64| {
        let score_runtime = test_deviation_runtime();
        let link_runtime = test_deviation_runtime();
        let family = SurvivalMarginalSlopeFamily {
            n: 1,
            event: Arc::new(array![event]),
            weights: Arc::new(array![weight]),
            z: Arc::new(array![zr].insert_axis(Axis(1))),
            score_covariance: unit_score_covariance(),
            gaussian_frailty_sd: None,
            family_hyper: SurvivalMarginalSlopeFamilyHyperState::default(),
            derivative_guard: 1e-6,
            design_entry: DesignMatrix::from(Array2::zeros((1, 1))),
            design_exit: DesignMatrix::from(Array2::zeros((1, 1))),
            design_derivative_exit: DesignMatrix::from(Array2::zeros((1, 1))),
            offset_entry: Arc::new(array![q0]),
            offset_exit: Arc::new(array![q1]),
            derivative_offset_exit: Arc::new(array![qd1]),
            marginal_design: DesignMatrix::from(Array2::zeros((1, 0))),
            logslope_layout: (DesignMatrix::from(Array2::zeros((1, 0)))).into(),
            score_warp: Some(score_runtime.clone()),
            link_dev: Some(link_runtime.clone()),
            influence_absorber: None,
            time_linear_constraints: None,
            time_wiggle_knots: None,
            time_wiggle_degree: None,
            time_wiggle_ncols: 0,
            intercept_warm_starts: None,
            auto_subsample_phase_counter: Arc::new(AtomicUsize::new(0)),
            auto_subsample_last_rho: Arc::new(Mutex::new(None)),
        };
        let sd = score_runtime.basis_dim();
        let ld = link_runtime.basis_dim();
        let bs = vec![
            ParameterBlockState {
                beta: Array1::zeros(1),
                eta: Array1::zeros(1),
            },
            ParameterBlockState {
                beta: Array1::zeros(0),
                eta: Array1::zeros(1),
            },
            ParameterBlockState {
                beta: Array1::zeros(0),
                eta: array![g],
            },
            ParameterBlockState {
                beta: Array1::zeros(sd),
                eta: Array1::zeros(1),
            },
            ParameterBlockState {
                beta: Array1::zeros(ld),
                eta: Array1::zeros(1),
            },
        ];
        (family, bs)
    };

    let (family, bs) = make(q0, q1, qd1, gv);
    let primary = flex_primary_slices(&family);
    let p = primary.total;
    let bidx = [primary.q0, primary.q1, primary.qd1, primary.g];

    // Flex Hessian at a primary point (q0,q1,qd1,g).
    let flex_hess = |q0: f64, q1: f64, qd1: f64, g: f64| -> Array2<f64> {
        let (fam, bs) = make(q0, q1, qd1, g);
        let qg = fam.row_dynamic_q_geometry(0, &bs).unwrap();
        let pr = flex_primary_slices(&fam);
        let (_, _, h) = fam
            .compute_row_flex_primary_gradient_hessian_exact(0, &bs, &qg, &pr)
            .unwrap();
        h
    };

    let dir4 = [0.7f64, -1.3, 0.5, 0.9];
    let dirvec = {
        let mut v = Array1::zeros(p);
        for (k, &s) in bidx.iter().enumerate() {
            v[s] = dir4[k];
        }
        v
    };

    // Production flex third-contracted.
    let flex_third = family
        .row_flex_primary_third_contracted_exact(0, &bs, &dirvec)
        .unwrap();

    // Independent FD of flex Hessian along dir4 (central, Richardson).
    let fd_dir = |h: f64| -> Array2<f64> {
        let hp = flex_hess(
            q0 + h * dir4[0],
            q1 + h * dir4[1],
            qd1 + h * dir4[2],
            gv + h * dir4[3],
        );
        let hm = flex_hess(
            q0 - h * dir4[0],
            q1 - h * dir4[1],
            qd1 - h * dir4[2],
            gv - h * dir4[3],
        );
        (&hp - &hm) / (2.0 * h)
    };
    let fd_coarse = fd_dir(1e-3);
    let fd_fine = fd_dir(5e-4);
    let fd_rich = (&fd_fine * 4.0 - &fd_coarse) / 3.0;

    // Rigid tower.
    let program = SurvivalMarginalSlopeRigidNllProgram {
        primaries: vec![[q0, q1, qd1, gv]],
        z: vec![zr],
        w: vec![weight],
        d: vec![event],
        probit_scale: family.probit_frailty_scale(),
    };
    let rigid = program_third_contracted(&program, 0, &dir4).unwrap();

    const THIRD_CONTRACTED_TOL: f64 = 1e-5;
    for (u, &bu) in bidx.iter().enumerate() {
        for (v, &bv) in bidx.iter().enumerate() {
            let flex = flex_third[[bu, bv]];
            let fd = fd_rich[[bu, bv]];
            let rigid = rigid[u][v];

            assert_close(
                flex,
                fd,
                THIRD_CONTRACTED_TOL,
                &format!("third-contracted flex vs FD at primary block ({u}, {v})"),
            );
            assert_close(
                flex,
                rigid,
                THIRD_CONTRACTED_TOL,
                &format!("third-contracted flex vs rigid at primary block ({u}, {v})"),
            );
            assert_close(
                fd,
                rigid,
                THIRD_CONTRACTED_TOL,
                &format!("third-contracted FD vs rigid at primary block ({u}, {v})"),
            );
        }
    }
}

/// gam#979 build-once equality contract for the rigid survival marginal-slope
/// kernel.
///
/// The inner-Newton Jeffreys/Firth term needs, each cycle, the directional
/// derivative of the joint Hessian for every canonical axis `e_a` (and the
/// outer-REML `H_Φ` drift needs the second-directional analogue). Before this
/// fix the rigid survival kernel implemented NEITHER
/// `directional_derivative_all_axes_dense_override` NOR its second-order
/// sibling, so `row_kernel_directional_derivative_all_axes` fell into the
/// generic per-axis fall-back — `p` independent full-data sweeps, each
/// rebuilding the per-row `Tower4<4>` for every row (`n·p` tower evaluations).
/// That redundant tower rebuild is the survival #979 hot path.
///
/// The new overrides build each row's tower ONCE and contract every axis off
/// that single build. A batched override is only a valid optimisation if it is
/// bit-for-bit what the per-axis sweep returns. This gate builds a
/// representative rigid fixture (non-trivial time / marginal / logslope designs,
/// `n = 300` rows so the `ARROW_ROW_CHUNK = 256` chunked reduction spans more
/// than one tile, mixed event / censored, with and without Gaussian frailty so
/// the probit scale ≠ 1) and asserts, for EVERY coefficient axis, that the
/// build-once override equals the per-axis sweep to machine precision — the
/// exact contract a wrong override would silently violate while corrupting
/// survival derivatives.
#[test]
fn rigid_survival_all_axes_build_once_equals_per_axis_sweep_979() {
    use crate::row_kernel::{
        RowKernel, RowSet, row_kernel_directional_derivative,
        row_kernel_directional_derivative_all_axes, row_kernel_second_directional_derivative,
        row_kernel_second_directional_derivative_all_axes,
    };

    let n = 300usize;
    // Mixed event / censored rows, deterministic and finite-margin.
    let z: Vec<f64> = (0..n).map(|r| ((r as f64) * 0.37).sin() * 1.1).collect();
    let weights: Vec<f64> = (0..n).map(|r| 0.7 + 0.5 * ((r % 5) as f64) / 5.0).collect();
    let event: Vec<f64> = (0..n).map(|r| ((r % 3 == 0) as u8) as f64).collect();

    // Non-trivial marginal (2-col) and logslope (2-col) designs so the
    // all-axes sweep exercises the coupled marginal block (which feeds BOTH
    // the entry and exit primaries through `jacobian_action`) and the logslope
    // block, not just the single time axis.
    let p_m = 2usize;
    let p_g = 2usize;
    let marginal_design = Array2::from_shape_fn((n, p_m), |(r, j)| {
        0.2 + 0.05 * (r as f64).cos() + 0.11 * (j as f64) - 0.013 * (r as f64) / (n as f64)
    });
    let logslope_design = Array2::from_shape_fn((n, p_g), |(r, j)| {
        0.1 + 0.07 * (r as f64).sin() - 0.09 * (j as f64) + 0.004 * (r as f64) / (n as f64)
    });
    let beta_marginal = Array1::from_vec(vec![0.18, -0.12]);
    let beta_logslope = Array1::from_vec(vec![-0.2, 0.13]);
    // Realized linear predictors carried on the block states (the marginal /
    // logslope eta channels the rigid primaries read), exactly as a real fit
    // installs them: eta = design · beta.
    let marginal_eta = marginal_design.dot(&beta_marginal);
    let logslope_eta = logslope_design.dot(&beta_logslope);

    for frailty in [None, Some(0.55_f64)] {
        let mut family = oracle_rigid_family(n, &z, &weights, &event, frailty);
        family.marginal_design = DesignMatrix::from(marginal_design.clone());
        family
            .logslope_layout
            .replace_coefficient_design(DesignMatrix::from(logslope_design.clone()));

        let beta_time = array![0.65];
        let block_states = vec![
            ParameterBlockState {
                beta: beta_time.clone(),
                eta: Array1::zeros(n),
            },
            ParameterBlockState {
                beta: beta_marginal.clone(),
                eta: marginal_eta.clone(),
            },
            ParameterBlockState {
                beta: beta_logslope.clone(),
                eta: logslope_eta.clone(),
            },
        ];

        let kernel = SurvivalMarginalSlopeRowKernel::new(family, block_states);
        let p = RowKernel::n_coefficients(&kernel);
        assert_eq!(
            p,
            1 + p_m + p_g,
            "fixture coefficient width should be time(1)+marginal({p_m})+logslope({p_g})"
        );

        // ---- FIRST directional derivative: override vs per-axis sweep --------
        let batched = row_kernel_directional_derivative_all_axes(&kernel, &RowSet::All)
            .expect("build-once all-axes first directional derivative");
        assert_eq!(
            batched.len(),
            p,
            "the batched first-directional sweep returns one p×p matrix per axis"
        );
        for a in 0..p {
            let mut e_a = vec![0.0_f64; p];
            e_a[a] = 1.0;
            let per_axis = row_kernel_directional_derivative(&kernel, &RowSet::All, &e_a)
                .expect("per-axis first directional derivative");
            assert_eq!(batched[a].dim(), (p, p));
            assert_eq!(per_axis.dim(), (p, p));
            let mut max_gap = 0.0_f64;
            for r in 0..p {
                for c in 0..p {
                    max_gap = max_gap.max((batched[a][[r, c]] - per_axis[[r, c]]).abs());
                }
            }
            assert!(
                max_gap < 1e-9,
                "frailty {frailty:?} axis {a}: #979 build-once FIRST directional override \
                 diverged from the per-axis sweep by {max_gap:e}; the optimisation changed \
                 the math, not just the schedule"
            );
        }

        // ---- SECOND directional derivative: override vs per-axis sweep -------
        // Fix a non-trivial direction `u` touching every block.
        let d_u = vec![0.5, 0.3, -0.7, 0.9, -0.4];
        let batched2 =
            row_kernel_second_directional_derivative_all_axes(&kernel, &RowSet::All, &d_u)
                .expect("build-once all-axes second directional derivative");
        assert_eq!(
            batched2.len(),
            p,
            "the batched second-directional sweep returns one p×p matrix per axis"
        );
        for a in 0..p {
            let mut e_a = vec![0.0_f64; p];
            e_a[a] = 1.0;
            let per_axis =
                row_kernel_second_directional_derivative(&kernel, &RowSet::All, &d_u, &e_a)
                    .expect("per-axis second directional derivative");
            let mut max_gap = 0.0_f64;
            for r in 0..p {
                for c in 0..p {
                    max_gap = max_gap.max((batched2[a][[r, c]] - per_axis[[r, c]]).abs());
                }
            }
            assert!(
                max_gap < 1e-9,
                "frailty {frailty:?} axis {a}: #979 build-once SECOND directional override \
                 diverged from the per-axis sweep by {max_gap:e}"
            );
        }
    }
}

/// gam#979 Jeffreys wide-p contracted-trace-Hessian FD verification (survival
/// twin of the BMS gate `bernoulli_jeffreys_contracted_trace_hessian_matches_fd_of_trace`).
///
/// Builds the same kind of non-trivial rigid fixture as the build-once gate
/// above (`oracle_rigid_family` + 2-column marginal/logslope designs, so the
/// hook's `time` block genuinely shares 3 different design rows with the
/// `q0/q1/qd1` primaries and `marginal_design` genuinely couples `q0` and
/// `q1`), picks a fixed deterministic symmetric 5×5 trace weight `W`, and
/// checks `family.joint_jeffreys_information_contracted_trace_hessian_with_specs`
/// against central second differences of `tr(W · H(β))` (using the existing
/// `exact_newton_joint_hessian` to get `H` at perturbed β) over several
/// directions spanning all three blocks.
#[test]
fn survival_jeffreys_contracted_trace_hessian_matches_fd_of_trace() {
    let n = 120usize;
    let z: Vec<f64> = (0..n).map(|r| ((r as f64) * 0.29).sin() * 0.9).collect();
    let weights: Vec<f64> = (0..n).map(|r| 0.6 + 0.4 * ((r % 5) as f64) / 5.0).collect();
    let event: Vec<f64> = (0..n).map(|r| ((r % 3 == 0) as u8) as f64).collect();

    let p_m = 2usize;
    let p_g = 2usize;
    let marginal_design = Array2::from_shape_fn((n, p_m), |(r, j)| {
        0.2 + 0.05 * (r as f64).cos() + 0.11 * (j as f64) - 0.013 * (r as f64) / (n as f64)
    });
    let logslope_design = Array2::from_shape_fn((n, p_g), |(r, j)| {
        0.1 + 0.07 * (r as f64).sin() - 0.09 * (j as f64) + 0.004 * (r as f64) / (n as f64)
    });

    let mut family = oracle_rigid_family(n, &z, &weights, &event, None);
    family.marginal_design = DesignMatrix::from(marginal_design.clone());
    family
        .logslope_layout
        .replace_coefficient_design(DesignMatrix::from(logslope_design.clone()));

    let total = 1 + p_m + p_g;
    let specs = vec![
        dummy_blockspec(1),
        dummy_blockspec(p_m),
        dummy_blockspec(p_g),
    ];

    // beta_flat = [time(1), marginal(2), logslope(2)].
    let states_at = |beta_flat: &Array1<f64>| -> Vec<ParameterBlockState> {
        let beta_time = beta_flat.slice(ndarray::s![0..1]).to_owned();
        let beta_marginal = beta_flat.slice(ndarray::s![1..1 + p_m]).to_owned();
        let beta_logslope = beta_flat.slice(ndarray::s![1 + p_m..total]).to_owned();
        let marginal_eta = marginal_design.dot(&beta_marginal);
        let logslope_eta = logslope_design.dot(&beta_logslope);
        vec![
            ParameterBlockState {
                beta: beta_time,
                // Unused by the closed-form likelihood (recomputed from
                // beta_time directly via the 3 time designs); zeros satisfy
                // the CustomFamily interface shape contract only.
                eta: Array1::zeros(n),
            },
            ParameterBlockState {
                beta: beta_marginal,
                eta: marginal_eta,
            },
            ParameterBlockState {
                beta: beta_logslope,
                eta: logslope_eta,
            },
        ]
    };

    let beta0 = array![0.6, 0.18, -0.12, -0.2, 0.13];

    // Fixed asymmetric-then-symmetrized 5x5 trace weight `W` (deterministic
    // pseudo-noise pattern, no RNG dependency).
    let mut w_raw = Array2::<f64>::zeros((total, total));
    for i in 0..total {
        for j in 0..total {
            w_raw[[i, j]] = ((i * 7 + j * 11 + 2) % 13) as f64 * 0.1 - 0.6;
        }
    }
    let w_raw_t = w_raw.t().to_owned();
    let w = (&w_raw + &w_raw_t).mapv(|v| v * 0.5);

    let states0 = states_at(&beta0);
    let analytic = family
        .joint_jeffreys_information_contracted_trace_hessian_with_specs(&states0, &specs, &w)
        .expect("contracted trace hessian call")
        .expect("rigid path must supply the contracted completion");
    assert_eq!(analytic.dim(), (total, total));

    let trace_of_hessian_at = |beta_flat: &Array1<f64>| -> f64 {
        let states = states_at(beta_flat);
        let h = family
            .exact_newton_joint_hessian(&states)
            .expect("exact_newton_joint_hessian")
            .expect("exact_newton_joint_hessian some");
        // tr(W H) = Σ_ij W_ij H_ij for symmetric W, H (H_ji = H_ij).
        (&w * &h).sum()
    };

    let directions = [
        array![1.0, 0.0, 0.0, 0.0, 0.0],
        array![0.0, 1.0, 0.0, 0.0, 0.0],
        array![0.0, 0.0, 1.0, 0.0, 0.0],
        array![0.0, 0.0, 0.0, 1.0, 0.0],
        array![0.0, 0.0, 0.0, 0.0, 1.0],
        array![0.5, -0.4, 0.3, 0.6, -0.2],
        array![-0.3, 0.5, -0.6, 0.2, 0.4],
    ];

    // ── Truncation-free analytic cross-check: the AUTHORITATIVE assembly gate ──
    // The finite-difference loop further below is discretization-limited (see
    // its comment): `tr(W·H)` is already second-order in β, so its second
    // difference probes the FOURTH derivative and the Richardson residual floors
    // at `O(h⁴·f⁽⁸⁾)`, which the neglog-Φ Mills-ratio tails push to ~3e-4
    // relative on the high-magnitude time direction — a pure truncation artifact
    // that no achievable `h` removes. So the FD can only catch an O(1) formula
    // blunder; it cannot pin the assembly.
    //
    // This block pins the hook's FULL assembly (`primary_trace_weight`
    // W-projection + full-`t4` contraction + `add_pullback_hessian`) to machine
    // precision with NO finite-difference truncation, using the identity
    //     uᵀ · (∇²_β tr(W·H)) · u  ==  tr(W · ∂²H/∂β_u²).
    // The right-hand side is built by
    // `exact_newton_joint_hessiansecond_directional_derivative` — the rank-1
    // `fourth_contracted` second-directional path that the outer-REML Jeffreys
    // drift already relies on. It shares neither `primary_trace_weight` nor the
    // direct full-`t4` read with the hook: it assembles the whole `∂²H` matrix
    // first and contracts `W` in COEFFICIENT space, whereas the hook projects
    // `W` into PRIMARY space per row before contracting `t4`. Their agreement
    // therefore validates that the primary-space trace projection and the
    // Jacobian pullback commute with the coefficient-space trace. Both sides are
    // exact (no `h`), so the tolerance is tight (1e-9), and this — not the FD —
    // is what guarantees the #979 completion is correct.
    for (idx, dir) in directions.iter().enumerate() {
        let huu = family
            .exact_newton_joint_hessiansecond_directional_derivative(&states0, dir, dir)
            .expect("second directional derivative call")
            .expect("rigid path must supply the second directional derivative");
        let trace_w_huu = (&w * &huu).sum();
        let analytic_quad = dir.dot(&analytic.dot(dir));
        let rel = (trace_w_huu - analytic_quad).abs() / trace_w_huu.abs().max(1.0);
        assert!(
            rel < 1e-9,
            "direction {idx}: contracted trace-Hessian assembly vs independent \
             tr(W·∂²H/∂β_u²) mismatch: analytic={analytic_quad:.12e} \
             independent={trace_w_huu:.12e} rel={rel:.3e}"
        );
    }

    // `tr(W·H(β))` is a second derivative of the row NLL, so its central
    // second difference probes the FOURTH β-derivative and carries an
    // `(h²/12)·(sixth-derivative-of-NLL)` truncation term. For the probit /
    // log-φ composite that sixth derivative summed over rows is large enough
    // that a single-`h` difference at `h=1e-3` sits right at the 1e-5 relative
    // tolerance on the `g`-spanning directions. Richardson-extrapolate two
    // central differences (`h`, `h/2`) to cancel the O(h²) term and validate
    // the analytic to O(h⁴). This is a strictly stronger check than the raw
    // difference: it cannot mask a genuine O(1) analytic error (both `D(h)`
    // and `D(h/2)` would then be wrong by ≈the same O(1) amount and the
    // combination stays wrong) — it only removes the discretization artifact.
    let eps = 1e-3;
    let center = trace_of_hessian_at(&beta0);
    for (idx, dir) in directions.iter().enumerate() {
        let central_second = |h: f64| -> f64 {
            let step: Array1<f64> = dir.mapv(|v| v * h);
            let plus = trace_of_hessian_at(&(&beta0 + &step));
            let minus = trace_of_hessian_at(&(&beta0 - &step));
            (plus - 2.0 * center + minus) / (h * h)
        };
        let d_h = central_second(eps);
        let d_h2 = central_second(eps * 0.5);
        let fd_second = (4.0 * d_h2 - d_h) / 3.0;
        let analytic_quad = dir.dot(&analytic.dot(dir));
        let rel = (fd_second - analytic_quad).abs() / fd_second.abs().max(1.0);
        // Discretization-limited cross-check, NOT the authoritative correctness
        // gate. The trace `tr(W·H)` is already second-order in β, so this second
        // difference probes the FOURTH derivative; the residual after Richardson
        // is `O(h⁴·f⁽⁸⁾)`, and survival's neglog-Φ Mills-ratio tails carry an
        // `f⁽⁸⁾` large enough (≳1e11 on the fixture's high-magnitude directions,
        // e.g. direction 0 at analytic≈1.06e3) that even the extrapolated FD
        // floors near ~3e-4 relative — a pure truncation artifact, not an
        // analytic error. The EXACT correctness of the contracted-trace tower is
        // pinned to machine precision by `survival_sparse_tower4_full_t4_matches_
        // dense_oracle_979` (dense-oracle agreement to 1e-9); this FD gate only
        // guards against an O(1) formula blunder, for which 1e-3 is ample.
        assert!(
            rel < 1e-3,
            "direction {idx}: survival contracted trace-Hessian FD mismatch: \
             analytic={analytic_quad:.10e} fd={fd_second:.10e} rel={rel:.3e}"
        );
    }
}

/// gam#979 perf datapoint + large-`p` correctness cross-check: the O(n·p²)
/// contracted-trace hook vs the `p(p+1)/2` pairwise second-directional
/// completion it replaces, at the issue's repro scale (`matern(...,
/// centers≈20)` ⇒ `p_marginal = p_logslope = 20`, total `p = 41`).
///
/// This is the mechanism behind the #979 rescope: the exact Firth/Jeffreys
/// second-order completion is gated on
/// `joint_jeffreys_information_contracted_trace_hessian_available()`. Before the
/// hook that gate was `false` and the completion NEVER ran (the inner Newton
/// under-modelled curvature near a Firth-active mode ⇒ the near-separation
/// crawl / 2400 s large_scale timeouts). The completion could not simply be
/// switched on with the generic pairwise fallback because that costs
/// `p(p+1)/2` full-data row-streamed second-directional Hessian passes
/// (`O(n·p⁴)`) — "hundreds of passes" at production p, far too slow to run every
/// endgame cycle. The hook produces the identical completion in ONE `O(n·p²)`
/// family pass, cheap enough to run every cycle.
///
/// The wall-clock speedup is REPORTED (the `[979 perf]` line) but not asserted:
/// shared-node CI timing is non-deterministic, so the gate is the deterministic
/// CORRECTNESS cross-check — the cheap `O(n·p²)` hook must reproduce the
/// expensive `p(p+1)/2` pairwise assembly over the FULL `p×p` matrix
/// (truncation-free, via `uᵀ∇²tr(W·H)v = tr(W·∂²H/∂β_u∂β_v)`), a strictly
/// larger correctness surface than the 7-direction gate above. `n` is kept
/// modest so the `p(p+1)/2`-pass reference stays CI-fast (the `#[ignore]`
/// wall-clock-benchmark form is banned by the workspace hygiene scanner).
#[test]
fn survival_jeffreys_contracted_trace_hook_beats_pairwise_979() {
    let n = 800usize;
    let z: Vec<f64> = (0..n).map(|r| ((r as f64) * 0.29).sin() * 0.9).collect();
    let weights: Vec<f64> = (0..n).map(|r| 0.6 + 0.4 * ((r % 5) as f64) / 5.0).collect();
    let event: Vec<f64> = (0..n).map(|r| ((r % 3 == 0) as u8) as f64).collect();

    let p_m = 20usize;
    let p_g = 20usize;
    let marginal_design = Array2::from_shape_fn((n, p_m), |(r, j)| {
        0.2 + 0.05 * ((r + j) as f64).cos() + 0.011 * (j as f64) - 0.003 * (r as f64) / (n as f64)
    });
    let logslope_design = Array2::from_shape_fn((n, p_g), |(r, j)| {
        0.1 + 0.07 * ((r + 2 * j) as f64).sin() - 0.009 * (j as f64)
            + 0.004 * (r as f64) / (n as f64)
    });

    let mut family = oracle_rigid_family(n, &z, &weights, &event, None);
    family.marginal_design = DesignMatrix::from(marginal_design.clone());
    family
        .logslope_layout
        .replace_coefficient_design(DesignMatrix::from(logslope_design.clone()));

    let total = 1 + p_m + p_g;
    let specs = vec![
        dummy_blockspec(1),
        dummy_blockspec(p_m),
        dummy_blockspec(p_g),
    ];

    let beta0 = Array1::from_shape_fn(total, |i| 0.1 + 0.03 * ((i as f64) * 1.7).sin());
    let beta_time = beta0.slice(ndarray::s![0..1]).to_owned();
    let beta_marginal = beta0.slice(ndarray::s![1..1 + p_m]).to_owned();
    let beta_logslope = beta0.slice(ndarray::s![1 + p_m..total]).to_owned();
    let states0 = vec![
        ParameterBlockState {
            beta: beta_time,
            eta: Array1::zeros(n),
        },
        ParameterBlockState {
            beta: beta_marginal.clone(),
            eta: marginal_design.dot(&beta_marginal),
        },
        ParameterBlockState {
            beta: beta_logslope.clone(),
            eta: logslope_design.dot(&beta_logslope),
        },
    ];

    let mut w_raw = Array2::<f64>::zeros((total, total));
    for i in 0..total {
        for j in 0..total {
            w_raw[[i, j]] = ((i * 7 + j * 11 + 2) % 13) as f64 * 0.1 - 0.6;
        }
    }
    let w = (&w_raw + &w_raw.t()).mapv(|v| v * 0.5);

    // ── Hook: ONE O(n·p²) family pass ────────────────────────────────────
    let t_hook = std::time::Instant::now();
    let hook = family
        .joint_jeffreys_information_contracted_trace_hessian_with_specs(&states0, &specs, &w)
        .expect("contracted trace hessian call")
        .expect("rigid path must supply the contracted completion");
    let hook_secs = t_hook.elapsed().as_secs_f64();
    assert_eq!(hook.dim(), (total, total));

    // ── Pairwise: p(p+1)/2 full second-directional passes (what the hook
    // replaces), reconstructing the SAME ∇²_β tr(W·H) via the truncation-free
    // identity uᵀ∇²tr(W·H)v = tr(W·∂²H/∂β_u∂β_v). ─────────────────────────
    let unit = |a: usize| -> Array1<f64> {
        let mut e = Array1::<f64>::zeros(total);
        e[a] = 1.0;
        e
    };
    let t_pair = std::time::Instant::now();
    let mut pairwise = Array2::<f64>::zeros((total, total));
    for a in 0..total {
        let ea = unit(a);
        for b in a..total {
            let eb = unit(b);
            let huv = family
                .exact_newton_joint_hessiansecond_directional_derivative(&states0, &ea, &eb)
                .expect("second directional derivative call")
                .expect("rigid path must supply the second directional derivative");
            let val = (&w * &huv).sum();
            pairwise[[a, b]] = val;
            pairwise[[b, a]] = val;
        }
    }
    let pair_secs = t_pair.elapsed().as_secs_f64();

    // Correctness: the cheap hook equals the expensive pairwise assembly over
    // the full p×p matrix, truncation-free.
    let mut max_rel = 0.0_f64;
    for a in 0..total {
        for b in 0..total {
            let rel = (hook[[a, b]] - pairwise[[a, b]]).abs() / pairwise[[a, b]].abs().max(1.0);
            max_rel = max_rel.max(rel);
        }
    }
    let n_pairwise_passes = total * (total + 1) / 2;
    eprintln!(
        "[979 perf] n={n} p={total} (p_m={p_m} p_g={p_g}) | hook={hook_secs:.4}s (1 pass) | \
         pairwise={pair_secs:.4}s ({n_pairwise_passes} passes) | speedup={:.1}× | \
         hook-vs-pairwise max_rel={max_rel:.3e}",
        pair_secs / hook_secs.max(1e-12)
    );
    assert!(
        max_rel < 1e-9,
        "hook completion disagrees with pairwise assembly at large p: max_rel={max_rel:.3e}"
    );
    // The hook must be at least as fast as the pairwise fallback it replaces.
    // The theoretical ratio is ~p(p+1)/2; we only assert the hook is not SLOWER
    // (a loose sanity bound robust to shared-node timing noise), and report the
    // measured speedup above for the perf datapoint.
    assert!(
        hook_secs <= pair_secs,
        "hook slower than the pairwise fallback it replaces: hook={hook_secs:.4}s pairwise={pair_secs:.4}s"
    );
}

/// #932 release speed gate for the rigid contracted third/fourth towers: the
/// specialized directional seedings that production consumes
/// (`OneSeed<4>` behind `row_primary_third_contracted_tower`, `TwoSeed<4>`
/// behind `row_primary_fourth_contracted_tower`) must beat the generic dense
/// `Tower4<4>` program evaluation they specialize (full 256-entry `t4` build
/// plus dense contraction). No hand-derived third/fourth tower ever existed
/// for this family, so — exactly as in the multinomial and cause-specific
/// release cells — the honest fail-closed baseline is the generic
/// AD tower, and the emitted `hand_over_production` token carries
/// `generic_tower_ns / production_ns` for the MSI release harness to fail
/// closed on any cell `<= 1`.
#[test]
fn release_measure_rigid_contracted_towers_vs_generic_tower_932() {
    use super::row_kernel::{RigidRowInputs, rigid_row_nll};
    use gam_math::jet_scalar::{OneSeed, TwoSeed};
    use gam_math::jet_tower::program_full_tower;
    use std::time::Instant;

    // One ordinary interior row per event branch (censored / event are
    // distinct live derivative stacks).
    let cases: [[f64; 8]; 2] = [
        [-0.7, 0.4, 0.8, -0.3, 0.6, 1.0, 0.0, 0.75],
        [0.2, -0.5, 1.4, 0.9, -1.1, 0.8, 1.0, 1.0],
    ];
    let dir_u = [0.7_f64, -1.3, 0.4, 0.6];
    let dir_v = [-0.4_f64, 0.6, 1.1, -0.2];

    // Feedback-coupled timing barrier (no `std::hint::black_box`): each
    // iteration nudges the log-slope primary by a negligible multiple of the
    // running checksum, so the pure tower call can be neither hoisted nor
    // dropped while the measured regime stays bit-adjacent to the fixture.
    fn best_ns<F: FnMut(f64) -> f64>(iterations: usize, base_g: f64, mut evaluate: F) -> f64 {
        let mut best = f64::INFINITY;
        for _ in 0..5 {
            let mut checksum = 0.0_f64;
            let started = Instant::now();
            for _ in 0..iterations {
                checksum += evaluate(base_g + checksum * 1e-18);
            }
            assert!(
                checksum.is_finite(),
                "contracted-tower release-measure checksum must stay finite"
            );
            best = best.min(started.elapsed().as_secs_f64());
        }
        best * 1e9 / iterations as f64
    }

    let iterations = 200_000usize;
    for &[q0, q1, qd1, g, z, w, d, probit_scale] in &cases {
        let inputs = RigidRowInputs {
            row: 0,
            wi: w,
            di: d,
            z_sum: z,
            covariance_ones: 1.0,
            probit_scale,
            qd1_lower: 1.0e-8,
        };
        let mut program = SurvivalMarginalSlopeRigidNllProgram {
            primaries: vec![[q0, q1, qd1, g]],
            z: vec![z],
            w: vec![w],
            d: vec![d],
            probit_scale,
        };

        // Parity pin on the exact benchmarked inputs: the specialized
        // contractions must equal the dense tower's contractions.
        let dense = program_full_tower(&program, 0).expect("dense tower");
        let third_vars: [OneSeed<4>; 4] =
            std::array::from_fn(|a| OneSeed::seed_direction([q0, q1, qd1, g][a], a, dir_u[a]));
        let third = rigid_row_nll(&third_vars, &inputs)
            .expect("specialized third")
            .contracted_third();
        let fourth_vars: [TwoSeed<4>; 4] =
            std::array::from_fn(|a| TwoSeed::seed([q0, q1, qd1, g][a], a, dir_u[a], dir_v[a]));
        let fourth = rigid_row_nll(&fourth_vars, &inputs)
            .expect("specialized fourth")
            .contracted_fourth();
        let dense_third = dense.third_contracted(&dir_u);
        let dense_fourth = dense.fourth_contracted(&dir_u, &dir_v);
        for a in 0..4 {
            for b in 0..4 {
                let band = 1e-11 * third[a][b].abs().max(dense_third[a][b].abs()).max(1.0);
                assert!(
                    (third[a][b] - dense_third[a][b]).abs() <= band,
                    "event={d:.0} third[{a}][{b}]: specialized {:+.15e} vs dense {:+.15e}",
                    third[a][b],
                    dense_third[a][b],
                );
                let band = 1e-11 * fourth[a][b].abs().max(dense_fourth[a][b].abs()).max(1.0);
                assert!(
                    (fourth[a][b] - dense_fourth[a][b]).abs() <= band,
                    "event={d:.0} fourth[{a}][{b}]: specialized {:+.15e} vs dense {:+.15e}",
                    fourth[a][b],
                    dense_fourth[a][b],
                );
            }
        }

        let third_production_ns = best_ns(iterations, g, |perturbed_g| {
            let vars: [OneSeed<4>; 4] = std::array::from_fn(|a| {
                OneSeed::seed_direction([q0, q1, qd1, perturbed_g][a], a, dir_u[a])
            });
            let t = rigid_row_nll(&vars, &inputs)
                .expect("specialized third")
                .contracted_third();
            t[0][0] + t[3][3]
        });
        let third_generic_ns = best_ns(iterations, g, |perturbed_g| {
            program.primaries[0][3] = perturbed_g;
            let t = program_full_tower(&program, 0)
                .expect("dense tower")
                .third_contracted(&dir_u);
            t[0][0] + t[3][3]
        });
        eprintln!(
            "RIGID-CONTRACTED-932 order=3 event={d:.0} production={third_production_ns:.2} ns/row \
             generic_tower={third_generic_ns:.2} ns/row hand_over_production={:.6}",
            third_generic_ns / third_production_ns,
        );

        let fourth_production_ns = best_ns(iterations, g, |perturbed_g| {
            let vars: [TwoSeed<4>; 4] = std::array::from_fn(|a| {
                TwoSeed::seed([q0, q1, qd1, perturbed_g][a], a, dir_u[a], dir_v[a])
            });
            let t = rigid_row_nll(&vars, &inputs)
                .expect("specialized fourth")
                .contracted_fourth();
            t[0][0] + t[3][3]
        });
        let fourth_generic_ns = best_ns(iterations, g, |perturbed_g| {
            program.primaries[0][3] = perturbed_g;
            let t = program_full_tower(&program, 0)
                .expect("dense tower")
                .fourth_contracted(&dir_u, &dir_v);
            t[0][0] + t[3][3]
        });
        eprintln!(
            "RIGID-CONTRACTED-932 order=4 event={d:.0} production={fourth_production_ns:.2} ns/row \
             generic_tower={fourth_generic_ns:.2} ns/row hand_over_production={:.6}",
            fourth_generic_ns / fourth_production_ns,
        );
    }
}

/// gam#979 isolation gate: does `SurvivalMarginalSlopeRowKernel`'s
/// static-sparsity `SparseTower4<RIGID_LINEAR_MASK>` build the SAME full
/// `t4` (all 256 entries, not just the 3 `fourth_contracted(u,v)` direction
/// pairs `rigid_row_kernel_agrees_with_jet_tower_program_all_channels`
/// checks) as the dense `Tower4<4>` oracle, on the EXACT fixture/row the
/// contracted-trace-Hessian FD gate above exercises?
///
/// My `contracted_trace_hessian` hook is the FIRST production consumer to
/// read `t4[a][b][c][d]` directly for every `(a,b)` pair (a genuine
/// non-rank-1 bilinear contraction against a full 4×4 weight matrix,
/// `Σ_ab w_row[a,b]·t4[a][b][c][d]`) rather than through a single
/// `fourth_contracted(u, v)` rank-1 direction pair — every prior consumer
/// only ever exercised the latter. This test empirically checks whether that
/// previously-untested full-tensor read path agrees with the dense oracle,
/// independent of whether the contracted-trace-Hessian FD gate above passes
/// or fails.
#[test]
fn survival_sparse_tower4_full_t4_matches_dense_oracle_979() {
    use super::row_kernel::{
        RIGID_LINEAR_MASK, SparseTower4, rigid_row_inputs, rigid_row_kernel_primaries,
        rigid_row_nll,
    };
    use gam_math::jet_scalar::JetScalar;
    use gam_math::jet_tower::program_full_tower;

    let n = 120usize;
    let z: Vec<f64> = (0..n).map(|r| ((r as f64) * 0.29).sin() * 0.9).collect();
    let weights: Vec<f64> = (0..n).map(|r| 0.6 + 0.4 * ((r % 5) as f64) / 5.0).collect();
    let event: Vec<f64> = (0..n).map(|r| ((r % 3 == 0) as u8) as f64).collect();

    let p_m = 2usize;
    let p_g = 2usize;
    let marginal_design = Array2::from_shape_fn((n, p_m), |(r, j)| {
        0.2 + 0.05 * (r as f64).cos() + 0.11 * (j as f64) - 0.013 * (r as f64) / (n as f64)
    });
    let logslope_design = Array2::from_shape_fn((n, p_g), |(r, j)| {
        0.1 + 0.07 * (r as f64).sin() - 0.09 * (j as f64) + 0.004 * (r as f64) / (n as f64)
    });

    let mut family = oracle_rigid_family(n, &z, &weights, &event, None);
    family.marginal_design = DesignMatrix::from(marginal_design.clone());
    family
        .logslope_layout
        .replace_coefficient_design(DesignMatrix::from(logslope_design.clone()));

    let beta_time = array![0.6];
    let beta_marginal = array![0.18, -0.12];
    let beta_logslope = array![-0.2, 0.13];
    let marginal_eta = marginal_design.dot(&beta_marginal);
    let logslope_eta = logslope_design.dot(&beta_logslope);
    let block_states = vec![
        ParameterBlockState {
            beta: beta_time.clone(),
            eta: Array1::zeros(n),
        },
        ParameterBlockState {
            beta: beta_marginal.clone(),
            eta: marginal_eta.clone(),
        },
        ParameterBlockState {
            beta: beta_logslope.clone(),
            eta: logslope_eta.clone(),
        },
    ];
    let probit_scale = family.probit_frailty_scale();

    let mut primaries = Vec::with_capacity(n);
    for row in 0..n {
        let q0 = family.design_entry.dot_row(row, &beta_time)
            + family.offset_entry[row]
            + marginal_eta[row];
        let q1 = family.design_exit.dot_row(row, &beta_time)
            + family.offset_exit[row]
            + marginal_eta[row];
        let qd1 = family.design_derivative_exit.dot_row(row, &beta_time)
            + family.derivative_offset_exit[row];
        let g = logslope_eta[row];
        primaries.push([q0, q1, qd1, g]);
    }
    let program = SurvivalMarginalSlopeRigidNllProgram {
        primaries,
        z: z.clone(),
        w: weights.clone(),
        d: event.clone(),
        probit_scale,
    };

    let mut max_abs_gap = 0.0_f64;
    let mut max_rel_gap = 0.0_f64;
    for row in 0..n {
        let dense = program_full_tower(&program, row).expect("dense tower4 oracle");

        let p = rigid_row_kernel_primaries(&family, &block_states, row).expect("primaries");
        let inputs = rigid_row_inputs(
            &family,
            &block_states,
            row,
            "sparse-vs-dense t4 isolation test",
        )
        .expect("rigid row inputs");
        let vars: [SparseTower4<RIGID_LINEAR_MASK>; 4] =
            std::array::from_fn(|a| SparseTower4::variable(p[a], a));
        let sparse = rigid_row_nll(&vars, &inputs).expect("sparse tower4");

        for a in 0..4 {
            for b in 0..4 {
                for c in 0..4 {
                    for d in 0..4 {
                        let dv = dense.t4[a][b][c][d];
                        let sv = sparse.t4[a][b][c][d];
                        let abs_gap = (dv - sv).abs();
                        let rel_gap = abs_gap / dv.abs().max(1.0);
                        if abs_gap > max_abs_gap {
                            max_abs_gap = abs_gap;
                        }
                        if rel_gap > max_rel_gap {
                            max_rel_gap = rel_gap;
                        }
                        assert!(
                            rel_gap < 1e-9,
                            "row {row} t4[{a}][{b}][{c}][{d}]: sparse={sv:.12e} dense={dv:.12e} \
                             abs_gap={abs_gap:e} rel_gap={rel_gap:e}"
                        );
                    }
                }
            }
        }
    }
    println!(
        "[979 isolation] full t4 max_abs_gap={max_abs_gap:e} max_rel_gap={max_rel_gap:e} over {n} rows"
    );
}

// ── #2352 ports: logslope block-jacobian production contract ────────────────
//
// Moved from `tests/survival/misc/frailty_scale_audit_plumbing.rs` and
// `tests/survival/survival/survival_marginal_slope_jacobian_hyperbolic_correction.rs`
// when `LogslopeBlockJacobian` construction went crate-internal (layout +
// covariance record; probit scale read from the linearization state). The
// old root-tree contract "Err when family_scalars is None at nonzero β" was
// deliberately replaced by origin-linearization (q ≡ 0) for the pre-fit
// structural audit, so the ports assert the CURRENT contract.

fn logslope_port_design(n: usize, p: usize, seed: u64) -> Array2<f64> {
    let mut out = Array2::<f64>::zeros((n, p));
    let mut state = seed;
    for i in 0..n {
        for j in 0..p {
            state = state
                .wrapping_mul(6364136223846793005)
                .wrapping_add(1442695040888963407);
            out[[i, j]] = ((state >> 33) as f64) / (u32::MAX as f64) - 0.5;
        }
    }
    out
}

fn logslope_port_z(n: usize, seed: u64) -> Vec<f64> {
    let mut state = seed ^ 0xdeadbeef;
    (0..n)
        .map(|_| {
            state = state
                .wrapping_mul(6364136223846793005)
                .wrapping_add(1442695040888963407);
            0.3 + ((state >> 33) as f64) / (u32::MAX as f64) * 1.4
        })
        .collect()
}

fn logslope_port_jacobian(
    design: &Array2<f64>,
    z: &[f64],
) -> super::block_jacobians::LogslopeBlockJacobian {
    let n = design.nrows();
    let layout = LogslopeLayout::shared(DesignMatrix::from(design.clone()), Array1::zeros(n));
    let z_mat = Array2::from_shape_fn((n, 1), |(i, _)| z[i]);
    let covariance = MarginalSlopeCovariance::diagonal(array![1.0]).expect("unit covariance");
    super::block_jacobians::LogslopeBlockJacobian::new(layout, Arc::new(z_mat), covariance)
        .expect("logslope jacobian construction")
}

/// At β = 0 the logslope Jacobian's η0/η1 channels are `s·z_i·G[i,j]` with `s`
/// read from `state.probit_frailty_scale` (NOT baked in at construction), and
/// the ad1 channel is zero; two states with different `s` scale exactly.
#[test]
fn logslope_jacobian_reads_probit_scale_from_state_at_beta_zero() {
    use crate::custom_family::{BlockEffectiveJacobian, FamilyLinearizationState};
    let n = 40;
    let p = 5;
    let design = logslope_port_design(n, p, 42);
    let z = logslope_port_z(n, 42);
    let s_f = 0.75_f64;
    let cb = logslope_port_jacobian(&design, &z);

    let beta_zero = vec![0.0f64; p];
    let jac_at = |s: f64| {
        let state = FamilyLinearizationState {
            beta: &beta_zero,
            family_scalars: None,
            channel_hessian: None,
            probit_frailty_scale: s,
        };
        cb.effective_jacobian_at(&state)
            .expect("beta=0 logslope jacobian")
    };
    let jac_sf = jac_at(s_f);
    let jac_1 = jac_at(1.0);

    assert_eq!(jac_sf.nrows(), 3 * n, "jacobian must have 3*n rows");
    assert_eq!(jac_sf.ncols(), p, "jacobian must have p cols");
    for i in 0..n {
        for j in 0..p {
            let expected = s_f * z[i] * design[[i, j]];
            let got = jac_sf[[i, j]];
            let err = (got - expected).abs();
            assert!(
                err / expected.abs().max(1e-14) < 1e-10 || err < 1e-12,
                "eta0[{i},{j}]: got {got:.6e} expected {expected:.6e}"
            );
            let ad1 = jac_sf[[2 * n + i, j]];
            assert!(ad1.abs() < 1e-12, "ad1[{i},{j}] must be 0 at beta=0: {ad1:.3e}");
            if got.abs() > 1e-14 {
                let ratio = jac_1[[i, j]] / got;
                assert!(
                    (ratio - 1.0 / s_f).abs() < 1e-10,
                    "state probit scale must set the jacobian scale: ratio {ratio:.12} != {:.12}",
                    1.0 / s_f
                );
            }
        }
    }
}

/// With populated `SurvivalMarginalSlopeFamilyScalars` at moderate β, the
/// production logslope Jacobian carries the full hyperbolic correction
/// `(q·dc/dg + s·z)·G` — FD-verified against the frozen-q η stack. Without
/// scalars the same point linearizes q at the origin (pre-fit audit
/// convention) and returns the pure `s·z_i·G` rows.
#[test]
fn logslope_jacobian_hyperbolic_correction_matches_fd_with_scalars() {
    use crate::custom_family::{BlockEffectiveJacobian, FamilyLinearizationState};
    use std::any::Any;
    let n = 40;
    let p = 5;
    let design = logslope_port_design(n, p, 31415);
    let z = logslope_port_z(n, 31415);
    let s_f = 0.8_f64;
    let cb = logslope_port_jacobian(&design, &z);
    let covariance = MarginalSlopeCovariance::diagonal(array![1.0]).expect("unit covariance");

    // Moderate deterministic beta so g_i != 0 on every row scale.
    let beta: Vec<f64> = (0..p)
        .map(|j| 0.35 * ((j as f64 + 1.0) * 0.7).sin())
        .collect();
    let g: Vec<f64> = (0..n)
        .map(|i| {
            design
                .row(i)
                .iter()
                .zip(beta.iter())
                .map(|(&x, &b)| x * b)
                .sum()
        })
        .collect();
    assert!(g.iter().any(|&v| v.abs() > 0.05), "fixture must reach nonzero g");

    // Frozen primary scalars (pilot q values held fixed for the FD functional).
    let q0: Vec<f64> = (0..n).map(|i| -0.5 + 0.3 * z[i]).collect();
    let q1: Vec<f64> = (0..n).map(|i| 0.2 + 0.4 * z[i]).collect();
    let qd1: Vec<f64> = (0..n).map(|i| 0.5 + 0.1 * (z[i] * z[i]).min(2.0)).collect();
    let slopes = Array2::from_shape_fn((n, 1), |(i, _)| g[i]);
    let scalars: Arc<dyn Any + Send + Sync> = Arc::new(
        SurvivalMarginalSlopeFamilyScalars::new(
            q0.clone(),
            q1.clone(),
            qd1.clone(),
            slopes,
            None,
            s_f,
            &covariance,
        )
        .expect("family scalars"),
    );
    let state = FamilyLinearizationState {
        beta: &beta,
        family_scalars: Some(scalars),
        channel_hessian: None,
        probit_frailty_scale: s_f,
    };
    let jac = cb
        .effective_jacobian_at(&state)
        .expect("logslope jacobian with scalars");
    assert_eq!(jac.nrows(), 3 * n);
    assert_eq!(jac.ncols(), p);

    // Central-difference reference of the frozen-q eta stack
    //   eta0 = q0·c(β) + s·g(β)·z, eta1 likewise, ad1 = qd1·c(β),
    //   c(β) = sqrt(1 + s²·g(β)²).
    let eta_stack = |b: &[f64]| -> Vec<f64> {
        let mut out = vec![0.0; 3 * n];
        for i in 0..n {
            let gi: f64 = design
                .row(i)
                .iter()
                .zip(b.iter())
                .map(|(&x, &bb)| x * bb)
                .sum();
            let c = (1.0 + s_f * s_f * gi * gi).sqrt();
            out[i] = q0[i] * c + s_f * gi * z[i];
            out[n + i] = q1[i] * c + s_f * gi * z[i];
            out[2 * n + i] = qd1[i] * c;
        }
        out
    };
    let h = 1e-6;
    let mut max_rel = 0.0_f64;
    for col in 0..p {
        let mut bp = beta.clone();
        let mut bm = beta.clone();
        bp[col] += h;
        bm[col] -= h;
        let ep = eta_stack(&bp);
        let em = eta_stack(&bm);
        for row in 0..3 * n {
            let fd = (ep[row] - em[row]) / (2.0 * h);
            let an = jac[[row, col]];
            let denom = fd.abs().max(1e-8);
            max_rel = max_rel.max((an - fd).abs() / denom);
        }
    }
    assert!(
        max_rel < 1e-5,
        "logslope jacobian must carry the hyperbolic correction: max rel err vs FD = {max_rel:.3e}"
    );

    // Origin-linearized fallback: scalars=None at the SAME nonzero beta gives
    // the pure s·z·G rows (q ≡ 0) with a zero ad1 channel.
    let state_none = FamilyLinearizationState {
        beta: &beta,
        family_scalars: None,
        channel_hessian: None,
        probit_frailty_scale: s_f,
    };
    let jac_none = cb
        .effective_jacobian_at(&state_none)
        .expect("origin-linearized logslope jacobian");
    for i in 0..n {
        for j in 0..p {
            let expected = s_f * z[i] * design[[i, j]];
            let got = jac_none[[i, j]];
            assert!(
                (got - expected).abs() < 1e-10 * (1.0 + expected.abs()),
                "origin-linearized eta0[{i},{j}]: got {got:.6e} expected {expected:.6e}"
            );
            assert!(
                jac_none[[2 * n + i, j]].abs() < 1e-12,
                "origin-linearized ad1 channel must be zero"
            );
        }
    }
}