gam-models 0.3.151

//! Fit-setup plumbing: the per-surface blockspec builders, guard-constraint
//! construction, hyperparameter joint-setup and log-lambda seeding, penalty
//! transformations, per-z surface combination, and the pooled survival
//! baseline solve.

use super::*;

// ── Building block specs ──────────────────────────────────────────────

pub(crate) fn build_time_blockspec(
    time_block: &TimeBlockInput,
    design_exit: &DesignMatrix,
    rho: Array1<f64>,
    beta_hint: Option<Array1<f64>>,
) -> ParameterBlockSpec {
    // Diagnostic (fires once per fit): the identifiability audit has reported the
    // `time_surface` design as intra-block rank-deficient on degenerate fixtures.
    // Dump the per-column span (max−min over exit rows) and mean so a rank
    // collapse is traceable to constant/collinear columns vs a genuine full-rank
    // design. Cheap (n×p pass, once), never on a hot path.
    if let Ok(dense) = design_exit.try_to_dense_arc("build_time_blockspec::diag") {
        let d = dense.as_ref();
        let (n_rows, p_cols) = d.dim();
        let mut spans: Vec<f64> = Vec::with_capacity(p_cols);
        let mut norms: Vec<f64> = Vec::with_capacity(p_cols);
        let mut degenerate_cols = 0usize;
        for j in 0..p_cols {
            let col = d.column(j);
            let lo = col.iter().copied().fold(f64::INFINITY, f64::min);
            let hi = col.iter().copied().fold(f64::NEG_INFINITY, f64::max);
            let span = hi - lo;
            if span.abs() <= 1e-10 {
                degenerate_cols += 1;
            }
            spans.push(span);
            norms.push(col.dot(&col).sqrt());
        }
        // Collinearity probe: pick the highest-norm column as the reference and
        // report |cos angle| of every other column against it (after mean-
        // centering, so a shared constant offset doesn't inflate the cosine).
        // Near-1 across the board ⇒ all columns collinear ⇒ the rank-1 collapse
        // is genuine near-collinearity, not just constant columns.
        let ref_col = norms
            .iter()
            .enumerate()
            .max_by(|a, b| a.1.total_cmp(b.1))
            .map(|(j, _)| j)
            .unwrap_or(0);
        let centered = |j: usize| -> Array1<f64> {
            let col = d.column(j).to_owned();
            let mean = col.sum() / (n_rows.max(1) as f64);
            col.mapv(|v| v - mean)
        };
        let r = centered(ref_col);
        let rn = r.dot(&r).sqrt().max(1e-300);
        let cosines: Vec<String> = (0..p_cols)
            .map(|j| {
                let c = centered(j);
                let cn = c.dot(&c).sqrt();
                if cn <= 1e-300 {
                    "const".to_string()
                } else {
                    format!("{:.4}", (c.dot(&r) / (cn * rn)).abs())
                }
            })
            .collect();
        log::info!(
            "[marginal-slope/time_surface-diag] design_exit {n_rows}x{p_cols}; near-constant cols={degenerate_cols}/{p_cols}; ref_col={ref_col}; per-col span={:?}; |cos vs ref (mean-centered)|={:?}",
            spans.iter().map(|s| format!("{s:.3e}")).collect::<Vec<_>>(),
            cosines,
        );
    }
    // Share the three dense design matrices with the multi-output Jacobian
    // via `Arc` — `try_to_dense_arc` is zero-copy for materialized designs,
    // so the callback retains no duplicate `n × p` storage. Falls back to no
    // callback if densification fails.
    let jac_cb: Option<Arc<dyn crate::custom_family::BlockEffectiveJacobian>> = (|| {
        let d_entry = time_block
            .design_entry
            .try_to_dense_arc("build_time_blockspec::entry")
            .ok()?;
        let d_exit = design_exit
            .try_to_dense_arc("build_time_blockspec::exit")
            .ok()?;
        let d_deriv = time_block
            .design_derivative_exit
            .try_to_dense_arc("build_time_blockspec::deriv")
            .ok()?;
        if d_entry.dim() != d_exit.dim() || d_entry.dim() != d_deriv.dim() {
            return None;
        }
        Some(Arc::new(TimeBlockJacobian::new(d_entry, d_exit, d_deriv))
            as Arc<dyn crate::custom_family::BlockEffectiveJacobian>)
    })();

    ParameterBlockSpec {
        name: "time_surface".to_string(),
        design: design_exit.clone(),
        offset: Array1::zeros(design_exit.nrows()),
        penalties: time_block
            .penalties
            .iter()
            .cloned()
            .map(PenaltyMatrix::Dense)
            .collect(),
        nullspace_dims: time_block.nullspace_dims.clone(),
        initial_log_lambdas: rho,
        initial_beta: beta_hint,
        gauge_priority: 200,
        jacobian_callback: jac_cb,
        stacked_design: None,
        stacked_offset: None,
    }
}

pub(crate) fn build_logslope_blockspec(
    design: &TermCollectionDesign,
    layout: &LogslopeLayout,
    baseline: f64,
    offset: &Array1<f64>,
    rho: Array1<f64>,
    beta_hint: Option<Array1<f64>>,
    z: Arc<Array2<f64>>,
    covariance: MarginalSlopeCovariance,
) -> Result<ParameterBlockSpec, String> {
    let jac_cb: Option<Arc<dyn crate::custom_family::BlockEffectiveJacobian>> = Some(Arc::new(
        LogslopeBlockJacobian::new(layout.clone(), z, covariance)?,
    ));

    Ok(ParameterBlockSpec {
        name: "logslope_surface".to_string(),
        design: layout.coefficient_design().clone(),
        offset: offset + baseline,
        penalties: design.penalties_as_penalty_matrix(),
        nullspace_dims: design.nullspace_dims.clone(),
        initial_log_lambdas: rho,
        initial_beta: beta_hint,
        gauge_priority: 120,
        jacobian_callback: jac_cb,
        stacked_design: None,
        stacked_offset: None,
    })
}

pub(crate) fn build_marginal_blockspec(
    design: &TermCollectionDesign,
    offset: &Array1<f64>,
    rho: Array1<f64>,
    beta_hint: Option<Array1<f64>>,
) -> ParameterBlockSpec {
    let jac_cb: Option<Arc<dyn crate::custom_family::BlockEffectiveJacobian>> = design
        .design
        .try_to_dense_arc("build_marginal_blockspec")
        .ok()
        .map(|d| {
            Arc::new(MarginalBlockJacobian::new(d))
                as Arc<dyn crate::custom_family::BlockEffectiveJacobian>
        });

    ParameterBlockSpec {
        name: "marginal_surface".to_string(),
        design: design.design.clone(),
        offset: offset.clone(),
        penalties: design.penalties_as_penalty_matrix(),
        nullspace_dims: design.nullspace_dims.clone(),
        initial_log_lambdas: rho,
        initial_beta: beta_hint,
        gauge_priority: 150,
        jacobian_callback: jac_cb,
        stacked_design: None,
        stacked_offset: None,
    }
}

pub(crate) fn inner_fit(
    family: &SurvivalMarginalSlopeFamily,
    blocks: &[ParameterBlockSpec],
    options: &BlockwiseFitOptions,
) -> Result<UnifiedFitResult, String> {
    fit_custom_family(family, blocks, options).map_err(|e| e.to_string())
}

pub(crate) fn inner_fit_from_certified_outer(
    family: &SurvivalMarginalSlopeFamily,
    blocks: &[ParameterBlockSpec],
    options: &BlockwiseFitOptions,
    mode: CustomFamilyJointHyperModeSelection,
    theta: &Array1<f64>,
    outer: &gam_solve::rho_optimizer::CertifiedOuterResult,
) -> Result<UnifiedFitResult, String> {
    let options = crate::outer_subsample::exact_outer_options_for_row_set(
        options,
        &crate::row_kernel::RowSet::All,
    );
    fit_custom_family_fixed_log_lambdas_from_mode_selection(
        family, blocks, &options, mode, theta, outer,
    )
    .map_err(|error| error.to_string())
}

/// Marginal-slope guard policy: the guard is required to be strictly positive
/// (`q'(t) ≥ guard > 0`), because the row-wise representation here is the *only*
/// place the monotonicity barrier lives — a zero guard would silently collapse
/// it. Coefficient-free row feasibility uses the family's epsilon-scaled slack
/// (`survival_derivative_guard_tolerance`).
pub(crate) const MARGINAL_SLOPE_GUARD_POLICY: GuardConstraintPolicy = GuardConstraintPolicy {
    guard_policy: GuardPolicy::Positive,
    feasibility: FeasibilityTolerance::EpsilonScaled,
};

pub(crate) fn time_derivative_guard_constraints(
    design_derivative_exit: &DesignMatrix,
    derivative_offset_exit: &Array1<f64>,
    derivative_guard: f64,
) -> Result<Option<LinearInequalityConstraints>, String> {
    build_time_derivative_guard_constraints(
        design_derivative_exit,
        derivative_offset_exit,
        derivative_guard,
        MARGINAL_SLOPE_GUARD_POLICY,
    )
    .map_err(map_guard_constraint_failure)
}

/// Render a shared guard-constraint failure into the marginal-slope error
/// vocabulary, preserving the family's historical wording.
pub(crate) fn map_guard_constraint_failure(failure: GuardConstraintFailure) -> String {
    match failure {
        GuardConstraintFailure::RowOffsetMismatch { rows, offsets } => {
            SurvivalMarginalSlopeError::IncompatibleDimensions {
                reason: format!(
                    "survival marginal-slope derivative guard constraints require matching rows/offsets: rows={rows}, offsets={offsets}"
                ),
            }
            .into()
        }
        GuardConstraintFailure::GuardOutOfRange { guard, range } => {
            SurvivalMarginalSlopeError::InvalidInput {
                reason: format!(
                    "survival marginal-slope derivative guard must be finite and {range}, got {guard}"
                ),
            }
            .into()
        }
        GuardConstraintFailure::NonFiniteOffset { row, offset } => {
            SurvivalMarginalSlopeError::MonotonicityViolation {
                reason: format!(
                    "survival marginal-slope derivative guard constraints require finite derivative offsets; found offset[{row}]={offset}"
                ),
            }
            .into()
        }
        GuardConstraintFailure::NonFiniteDesign { row, col } => {
            SurvivalMarginalSlopeError::MonotonicityViolation {
                reason: format!(
                    "survival marginal-slope derivative guard constraints require finite derivative design entries; found row {row}, column {col}"
                ),
            }
            .into()
        }
        GuardConstraintFailure::InfeasibleRow {
            row,
            offset,
            guard,
            no_time_coefficients,
        } => {
            let reason = if no_time_coefficients {
                format!(
                    "survival marginal-slope derivative guard is infeasible at row {row}: offset={offset:.3e} < guard={guard:.3e} with no time coefficients"
                )
            } else {
                format!(
                    "survival marginal-slope derivative guard is infeasible at row {row}: zero derivative design row with offset={offset:.3e} < guard={guard:.3e}"
                )
            };
            SurvivalMarginalSlopeError::MonotonicityViolation { reason }.into()
        }
    }
}

pub(crate) fn append_timewiggle_tail_nonnegative_constraints(
    base: Option<LinearInequalityConstraints>,
    p_total: usize,
    time_wiggle_ncols: usize,
) -> Result<Option<LinearInequalityConstraints>, String> {
    let p_wiggle = time_wiggle_ncols.min(p_total);
    if p_wiggle == 0 {
        return Ok(base);
    }
    if let Some(base_constraints) = base.as_ref() {
        if base_constraints.a.ncols() != p_total {
            return Err(SurvivalMarginalSlopeError::IncompatibleDimensions {
                reason: format!(
                    "survival marginal-slope time constraint width mismatch: constraints={}, time block={p_total}",
                    base_constraints.a.ncols()
                ),
            }
            .into());
        }
        if base_constraints.a.nrows() != base_constraints.b.len() {
            return Err(SurvivalMarginalSlopeError::IncompatibleDimensions {
                reason: format!(
                    "survival marginal-slope time constraint row mismatch: A rows={}, b len={}",
                    base_constraints.a.nrows(),
                    base_constraints.b.len()
                ),
            }
            .into());
        }
    }

    let base_rows = base.as_ref().map_or(0, |constraints| constraints.a.nrows());
    let rows = base_rows + p_wiggle;
    let mut a = Array2::<f64>::zeros((rows, p_total));
    let mut b = Array1::<f64>::zeros(rows);

    if let Some(base_constraints) = base {
        a.slice_mut(s![..base_rows, ..]).assign(&base_constraints.a);
        b.slice_mut(s![..base_rows]).assign(&base_constraints.b);
    }

    let tail_start = p_total - p_wiggle;
    for (row_offset, col) in (tail_start..p_total).enumerate() {
        a[[base_rows + row_offset, col]] = 1.0;
    }
    Ok(Some(LinearInequalityConstraints { a, b }))
}

pub(crate) fn mean_abs(values: impl IntoIterator<Item = f64>) -> f64 {
    let mut sum = 0.0;
    let mut count = 0usize;
    for v in values {
        sum += v.abs();
        count += 1;
    }
    if count == 0 { 0.0 } else { sum / count as f64 }
}

pub(crate) fn block_log_lambda_seeds<'a, I>(design: &DesignMatrix, penalty_locals: I) -> Vec<f64>
where
    I: IntoIterator<Item = &'a Array2<f64>>,
{
    let unit_weights = Array1::<f64>::ones(design.nrows());
    let likelihood_scale = match design.diag_gram(&unit_weights) {
        Ok(d) => mean_abs(d.iter().copied()).max(1.0e-8),
        Err(_) => 1.0,
    };
    penalty_locals
        .into_iter()
        .map(|s| {
            let penalty_scale = mean_abs(s.diag().iter().copied()).max(1.0e-8);
            (likelihood_scale / penalty_scale).ln().clamp(-12.0, 12.0)
        })
        .collect()
}

pub(crate) fn joint_setup(
    data: ArrayView2<'_, f64>,
    time_penalties: usize,
    marginalspec: &TermCollectionSpec,
    marginal_penalties: usize,
    logslopespec: &TermCollectionSpec,
    logslope_penalties: usize,
    core_rho0_seed: &[f64],
    extra_rho0: &[f64],
    baseline_initial_theta: &[f64],
    baseline_lower_theta: &[f64],
    baseline_upper_theta: &[f64],
    learned_log_sigma_coordinate: Option<(f64, f64, f64)>,
    kappa_options: &SpatialLengthScaleOptimizationOptions,
) -> Result<ExactJointHyperSetup, gam_terms::basis::BasisError> {
    let (marginal_terms, logslope_terms) = if kappa_options.enabled {
        (
            spatial_length_scale_term_indices(marginalspec),
            spatial_length_scale_term_indices(logslopespec),
        )
    } else {
        (Vec::new(), Vec::new())
    };
    let core_len = time_penalties + marginal_penalties + logslope_penalties;
    let rho_dim = core_len + extra_rho0.len();
    let mut rho0vec = Array1::<f64>::zeros(rho_dim);
    assert_eq!(
        core_rho0_seed.len(),
        core_len,
        "core_rho0_seed length must equal time+marginal+logslope penalty count"
    );
    for (idx, value) in core_rho0_seed.iter().copied().enumerate().take(core_len) {
        rho0vec[idx] = value;
    }
    if !extra_rho0.is_empty() {
        let start = core_len;
        for (idx, value) in extra_rho0.iter().copied().enumerate() {
            rho0vec[start + idx] = value;
        }
    }
    let rho_lower = Array1::<f64>::from_elem(rho_dim, -12.0);
    let rho_upper = Array1::<f64>::from_elem(rho_dim, 12.0);
    // Time block has no spatial length scales (pure B-spline on time)
    let empty_kappa = SpatialLogKappaCoords::new_with_dims(Array1::zeros(0), vec![]);
    let marginal_kappa = SpatialLogKappaCoords::from_length_scales_aniso(
        marginalspec,
        &marginal_terms,
        kappa_options,
    )
    .reseed_from_data(data, marginalspec, &marginal_terms, kappa_options)?;
    let logslope_kappa = SpatialLogKappaCoords::from_length_scales_aniso(
        logslopespec,
        &logslope_terms,
        kappa_options,
    )
    .reseed_from_data(data, logslopespec, &logslope_terms, kappa_options)?;
    let mut values = empty_kappa.as_array().to_vec();
    values.extend(marginal_kappa.as_array().iter());
    values.extend(logslope_kappa.as_array().iter());
    let marginal_dims = marginal_kappa.dims_per_term().to_vec();
    let logslope_dims = logslope_kappa.dims_per_term().to_vec();
    let mut dims = empty_kappa.dims_per_term().to_vec();
    dims.extend(marginal_dims.iter().copied());
    dims.extend(logslope_dims.iter().copied());
    let log_kappa0 =
        SpatialLogKappaCoords::new_with_dims(Array1::from_vec(values.clone()), dims.clone());
    // Bounds: concatenate [empty | marginal data-aware | logslope data-aware]
    let marginal_lower = SpatialLogKappaCoords::lower_bounds_aniso_from_data(
        data,
        marginalspec,
        &marginal_terms,
        &marginal_dims,
        kappa_options,
    )?;
    let logslope_lower = SpatialLogKappaCoords::lower_bounds_aniso_from_data(
        data,
        logslopespec,
        &logslope_terms,
        &logslope_dims,
        kappa_options,
    )?;
    let mut lower_vals = Vec::with_capacity(dims.iter().sum());
    lower_vals.extend(marginal_lower.as_array().iter());
    lower_vals.extend(logslope_lower.as_array().iter());
    let log_kappa_lower =
        SpatialLogKappaCoords::new_with_dims(Array1::from_vec(lower_vals), dims.clone());
    let marginal_upper = SpatialLogKappaCoords::upper_bounds_aniso_from_data(
        data,
        marginalspec,
        &marginal_terms,
        &marginal_dims,
        kappa_options,
    )?;
    let logslope_upper = SpatialLogKappaCoords::upper_bounds_aniso_from_data(
        data,
        logslopespec,
        &logslope_terms,
        &logslope_dims,
        kappa_options,
    )?;
    let mut upper_vals = Vec::with_capacity(dims.iter().sum());
    upper_vals.extend(marginal_upper.as_array().iter());
    upper_vals.extend(logslope_upper.as_array().iter());
    let log_kappa_upper = SpatialLogKappaCoords::new_with_dims(Array1::from_vec(upper_vals), dims);
    // Project seed onto bounds; spec.length_scale is a hint, not a constraint.
    let log_kappa0 = log_kappa0.clamp_to_bounds(&log_kappa_lower, &log_kappa_upper);
    let setup = ExactJointHyperSetup::new(
        rho0vec,
        rho_lower,
        rho_upper,
        log_kappa0,
        log_kappa_lower,
        log_kappa_upper,
    );
    assert_eq!(
        baseline_initial_theta.len(),
        baseline_lower_theta.len(),
        "baseline lower bound length mismatch"
    );
    assert_eq!(
        baseline_initial_theta.len(),
        baseline_upper_theta.len(),
        "baseline upper bound length mismatch"
    );
    let mut auxiliary0 = baseline_initial_theta.to_vec();
    let mut lower = baseline_lower_theta.to_vec();
    let mut upper = baseline_upper_theta.to_vec();
    if let Some((log_sigma, log_sigma_lower, log_sigma_upper)) = learned_log_sigma_coordinate {
        auxiliary0.push(log_sigma);
        lower.push(log_sigma_lower);
        upper.push(log_sigma_upper);
    }
    if !auxiliary0.is_empty() {
        Ok(setup.with_auxiliary(
            Array1::from_vec(auxiliary0),
            Array1::from_vec(lower),
            Array1::from_vec(upper),
        ))
    } else {
        Ok(setup)
    }
}

pub(crate) fn install_time_nullspace_shrinkage_penalty(
    time_block: &mut TimeBlockInput,
) -> Result<bool, String> {
    let p = time_block.design_exit.ncols();
    if p == 0 || time_block.penalties.is_empty() {
        return Ok(false);
    }
    if time_block.nullspace_dims.len() != time_block.penalties.len() {
        return Err(format!(
            "survival-marginal-slope time_block nullspace_dims length {} does not match penalties {}",
            time_block.nullspace_dims.len(),
            time_block.penalties.len(),
        ));
    }

    let mut aggregate = Array2::<f64>::zeros((p, p));
    for (idx, penalty) in time_block.penalties.iter().enumerate() {
        if penalty.nrows() != p || penalty.ncols() != p {
            return Err(format!(
                "survival-marginal-slope time_block penalty {idx} must be {p}x{p}, got {}x{}",
                penalty.nrows(),
                penalty.ncols(),
            ));
        }
        let scale = penalty
            .iter()
            .try_fold(0.0_f64, |acc, &value| {
                value.is_finite().then_some(acc.max(value.abs()))
            })
            .ok_or_else(|| {
                format!(
                    "survival-marginal-slope time_block penalty {idx} contains non-finite values"
                )
            })?;
        if scale > 0.0 {
            ndarray::Zip::from(&mut aggregate)
                .and(penalty)
                .for_each(|agg, &value| *agg += value / scale);
        }
    }

    if time_block.design_entry.ncols() != p {
        return Err(format!(
            "survival-marginal-slope time_block entry design has {} columns, expected {p}",
            time_block.design_entry.ncols(),
        ));
    }
    // Use the average endpoint function measure, not Euclidean coefficient
    // length. Entry and exit evaluations are the two value channels through
    // which this time function enters the survival likelihood; averaging their
    // Grams makes the metric invariant to whole-sample replication while
    // treating both endpoints symmetrically. Under any coefficient chart
    // change M, `G -> M' G M`, so the resulting null ridge transforms by the
    // same congruence instead of changing the represented penalty.
    let entry_mass = time_block.design_entry.nrows();
    let exit_mass = time_block.design_exit.nrows();
    let total_mass = entry_mass.saturating_add(exit_mass);
    if total_mass == 0 {
        return Err(
            "survival-marginal-slope time_block cannot define a function metric from zero endpoint rows"
                .to_string(),
        );
    }
    let entry_gram = time_block
        .design_entry
        .diag_xtw_x(&Array1::ones(entry_mass))
        .map_err(|err| format!("survival-marginal-slope time_block entry function Gram: {err}"))?;
    let exit_gram = time_block
        .design_exit
        .diag_xtw_x(&Array1::ones(exit_mass))
        .map_err(|err| format!("survival-marginal-slope time_block exit function Gram: {err}"))?;
    let function_gram = (entry_gram + exit_gram).mapv(|value| value / total_mass as f64);

    let Some(shrinkage) =
        gam_terms::basis::function_space_nullspace_shrinkage(&aggregate, &function_gram).map_err(
            |err| format!("survival-marginal-slope time_block nullspace shrinkage: {err}"),
        )?
    else {
        return Ok(false);
    };
    if shrinkage.nrows() != p || shrinkage.ncols() != p {
        return Err(format!(
            "survival-marginal-slope time_block nullspace shrinkage penalty must be {p}x{p}, got {}x{}",
            shrinkage.nrows(),
            shrinkage.ncols(),
        ));
    }
    time_block.penalties.push(shrinkage);
    time_block.nullspace_dims.push(0);
    log::info!(
        "[survival-marginal-slope] added time_block nullspace shrinkage penalty (p={p}, penalties={})",
        time_block.penalties.len(),
    );
    Ok(true)
}

/// Enable the Marra & Wood double penalty — a null-space shrinkage ridge with
/// its own REML-selected smoothing parameter — on every smooth in `spec` whose
/// basis carries an unpenalized polynomial null space, so that "trend"
/// direction is identified rather than left flat.
///
/// gam#979: the survival marginal-slope marginal/logslope smooth surfaces are
/// built with `double_penalty=false` (the default), so the general basis
/// builder leaves their penalty null space (the polynomial trend the
/// wiggliness penalty does not touch) unpenalized. At the REML optimum that
/// trend direction is then a large-signal / tiny-curvature mode, and the inner
/// joint-Newton deadlocks on it: the spectral trust-region step DROPS it as a
/// near-null gauge direction (gam#1082) while the convergence certificate
/// REQUIRES it be resolved (gam#1449) — so the solve can neither progress nor
/// certify and grinds to the cycle cap on every outer ρ-evaluation (the
/// measured n=3000 ~1976 s hang). The survival TIME block already avoids this
/// via [`install_time_nullspace_shrinkage_penalty`]; this brings the
/// marginal/logslope surfaces to the same footing through the ordinary basis
/// builder (which keeps the layered penalty representation self-consistent
/// across every probe/frozen/kappa rebuild). Duchon bases are skipped because
/// their function-norm penalty already spans the polynomial null space (and
/// they carry no `double_penalty` field); factor/measure/tensor/PCA bases have
/// no simple trend null space to shrink here and are left untouched.
pub(crate) fn enable_surface_identifiability_double_penalty(spec: &mut TermCollectionSpec) {
    for term in spec.smooth_terms.iter_mut() {
        enable_double_penalty_on_basis(&mut term.basis);
    }
}

fn enable_double_penalty_on_basis(basis: &mut gam_terms::smooth::SmoothBasisSpec) {
    use gam_terms::smooth::SmoothBasisSpec as B;
    match basis {
        B::ByVariable { inner, .. } | B::FactorSumToZero { inner, .. } => {
            enable_double_penalty_on_basis(inner)
        }
        B::BySmooth { smooth, .. } => enable_double_penalty_on_basis(smooth),
        B::BSpline1D { spec, .. } => spec.double_penalty = true,
        B::ThinPlate { spec, .. } => spec.double_penalty = true,
        B::Matern { spec, .. } => spec.double_penalty = true,
        B::Sphere { spec, .. } => spec.double_penalty = true,
        B::ConstantCurvature { spec, .. } => spec.double_penalty = true,
        B::Duchon { .. }
        | B::MeasureJet { .. }
        | B::Pca { .. }
        | B::TensorBSpline { .. }
        | B::FactorSmooth { .. } => {}
    }
}

pub(crate) fn concatenate_term_specs(specs: &[TermCollectionSpec]) -> TermCollectionSpec {
    let mut out = TermCollectionSpec {
        linear_terms: Vec::new(),
        random_effect_terms: Vec::new(),
        smooth_terms: Vec::new(),
    };
    for spec in specs {
        out.linear_terms.extend(spec.linear_terms.clone());
        out.random_effect_terms
            .extend(spec.random_effect_terms.clone());
        out.smooth_terms.extend(spec.smooth_terms.clone());
    }
    out
}

pub(crate) fn shift_penalty(mut penalty: BlockwisePenalty, offset: usize) -> BlockwisePenalty {
    penalty.col_range = (penalty.col_range.start + offset)..(penalty.col_range.end + offset);
    penalty
}

pub(crate) fn combine_logslope_surface_designs(
    mut designs: Vec<TermCollectionDesign>,
    specs: &[TermCollectionSpec],
) -> Result<(TermCollectionDesign, TermCollectionSpec, LogslopeTopology), String> {
    if designs.is_empty() {
        return Err(
            "survival marginal-slope requires at least one logslope surface design".to_string(),
        );
    }
    if designs.len() == 1 {
        let design = designs.remove(0);
        let spec = specs
            .first()
            .cloned()
            .ok_or_else(|| "missing logslope surface spec".to_string())?;
        return Ok((design, spec, LogslopeTopology::shared()));
    }
    if designs.iter().any(|design| {
        design.linear_constraints.is_some() || design.coefficient_lower_bounds.is_some()
    }) {
        return Err(
            "per-z logslope surface concatenation does not support coefficient bounds or linear constraints"
                .to_string(),
        );
    }
    if designs
        .iter()
        .any(|design| design.affine_offset.iter().any(|value| *value != 0.0))
    {
        return Err(
            "per-score logslope surfaces do not support non-zero smooth anchors: each affine \
             surface lift must be coupled to its own latent-score coordinate, but the shared \
             logslope offset channel is scalar"
                .to_string(),
        );
    }

    let mut ranges = Vec::with_capacity(designs.len());
    let mut offset = 0usize;
    let mut blocks = Vec::with_capacity(designs.len());
    let mut penalties = Vec::new();
    let mut nullspace_dims = Vec::new();
    let mut penaltyinfo = Vec::new();
    let mut dropped_penaltyinfo = Vec::new();
    let mut linear_ranges = Vec::new();
    let mut random_effect_ranges = Vec::new();
    let mut random_effect_levels = Vec::new();
    let mut combined = designs[0].clone();
    combined.smooth.term_designs.clear();
    combined.smooth.penalties.clear();
    combined.smooth.nullspace_dims.clear();
    combined.smooth.penaltyinfo.clear();
    combined.smooth.dropped_penaltyinfo.clear();
    combined.smooth.terms.clear();
    combined.smooth.coefficient_lower_bounds = None;
    combined.smooth.linear_constraints = None;

    for (surface_idx, design) in designs.into_iter().enumerate() {
        let width = design.design.ncols();
        ranges.push(offset..offset + width);
        blocks.push(design.design.clone());
        for (local_penalty_idx, penalty) in design.penalties.iter().cloned().enumerate() {
            let global_index = penalties.len();
            penalties.push(shift_penalty(penalty, offset));
            if let Some(info) = design.penaltyinfo.get(local_penalty_idx) {
                let mut info = info.clone();
                info.global_index = global_index;
                if let Some(termname) = info.termname.as_mut() {
                    *termname = format!("logslope[z{surface_idx}]::{termname}");
                }
                penaltyinfo.push(info);
            }
        }
        nullspace_dims.extend(design.nullspace_dims.iter().copied());
        dropped_penaltyinfo.extend(design.dropped_penaltyinfo.iter().cloned());
        linear_ranges.extend(design.linear_ranges.iter().cloned().map(|(name, range)| {
            (
                format!("logslope[z{surface_idx}]::{name}"),
                (range.start + offset)..(range.end + offset),
            )
        }));
        random_effect_ranges.extend(design.random_effect_ranges.iter().cloned().map(
            |(name, range)| {
                (
                    format!("logslope[z{surface_idx}]::{name}"),
                    (range.start + offset)..(range.end + offset),
                )
            },
        ));
        random_effect_levels.extend(design.random_effect_levels.iter().cloned());
        offset += width;
    }
    combined.design = DesignMatrix::hstack(blocks)
        .map_err(|e| format!("survival marginal-slope logslope hstack: {e}"))?;
    combined.penalties = penalties;
    combined.nullspace_dims = nullspace_dims;
    combined.penaltyinfo = penaltyinfo;
    combined.dropped_penaltyinfo = dropped_penaltyinfo;
    combined.coefficient_lower_bounds = None;
    combined.linear_constraints = None;
    combined.intercept_range = 0..0;
    combined.linear_ranges = linear_ranges;
    combined.random_effect_ranges = random_effect_ranges;
    combined.random_effect_levels = random_effect_levels;
    let topology = LogslopeTopology::per_score(ranges, combined.design.ncols())?;
    Ok((combined, concatenate_term_specs(specs), topology))
}

/// Compute a baseline slope from the actual survival marginal-slope likelihood,
/// using the baseline offsets alone as a time-only pilot q(t).
///
/// This is a safeguarded 1D Newton solve on the true row objective. It does not
/// use a coarse fixed grid scan.
pub(crate) fn pooled_survival_baseline(
    event: &Array1<f64>,
    weights: &Array1<f64>,
    z: &Array1<f64>,
    q0: &Array1<f64>,
    q1: &Array1<f64>,
    qd1: &Array1<f64>,
    probit_scale: f64,
) -> f64 {
    let n = event.len();
    if n == 0 {
        return 0.0;
    }
    let objective_grad_hess = |slope: f64| -> Option<(f64, f64, f64)> {
        use rayon::iter::{IntoParallelIterator, ParallelIterator};
        let triples: Option<Vec<(f64, f64, f64)>> = (0..n)
            .into_par_iter()
            .map(|i| {
                let (row_obj, row_grad, row_hess) = row_primary_closed_form(
                    q0[i],
                    q1[i],
                    qd1[i],
                    slope,
                    z[i],
                    weights[i],
                    event[i],
                    0.0,
                    probit_scale,
                )
                .ok()?;
                Some((row_obj, row_grad[3], row_hess[3][3]))
            })
            .collect();
        let triples = triples?;
        Some(
            triples
                .into_iter()
                .fold((0.0_f64, 0.0_f64, 0.0_f64), |(o, g, h), (oi, gi, hi)| {
                    (o + oi, g + gi, h + hi)
                }),
        )
    };

    let Some(state0) = objective_grad_hess(0.0) else {
        return 0.0;
    };
    if !state0.0.is_finite() {
        return 0.0;
    }
    if state0.1.abs() < 1e-8 {
        return 0.0;
    }

    let mut best_slope = 0.0;
    let mut best = state0;

    let mut bracket_lo = if state0.1 <= 0.0 {
        Some((0.0, state0))
    } else {
        None
    };
    let mut bracket_hi = if state0.1 >= 0.0 {
        Some((0.0, state0))
    } else {
        None
    };
    let mut step = 0.5f64;
    for _ in 0..48 {
        for &candidate in &[-step, step] {
            if let Some(state) = objective_grad_hess(candidate) {
                if state.0 < best.0 {
                    best_slope = candidate;
                    best = state;
                }
                if state.1 <= 0.0 {
                    bracket_lo = Some((candidate, state));
                }
                if state.1 >= 0.0 {
                    bracket_hi = Some((candidate, state));
                }
                if let (Some((lo, lo_state)), Some((hi, hi_state))) = (bracket_lo, bracket_hi)
                    && lo < hi
                    && lo_state.1 <= 0.0
                    && hi_state.1 >= 0.0
                {
                    let mut slope = best_slope.clamp(lo, hi);
                    let mut state = if (slope - lo).abs() < f64::EPSILON {
                        lo_state
                    } else if (slope - hi).abs() < f64::EPSILON {
                        hi_state
                    } else {
                        match objective_grad_hess(slope) {
                            Some(s) => s,
                            None => {
                                slope = 0.5 * (lo + hi);
                                objective_grad_hess(slope).unwrap_or(best)
                            }
                        }
                    };

                    let mut bracket_lo = (lo, lo_state);
                    let mut bracket_hi = (hi, hi_state);
                    for _ in 0..60 {
                        if state.1.abs() < 1e-8 || (bracket_hi.0 - bracket_lo.0).abs() < 1e-8 {
                            break;
                        }
                        let mut candidate = 0.5 * (bracket_lo.0 + bracket_hi.0);
                        if state.2.is_finite() && state.2 > 0.0 {
                            let newton = slope - state.1 / state.2;
                            if newton > bracket_lo.0 && newton < bracket_hi.0 {
                                candidate = newton;
                            }
                        }
                        let Some(candidate_state) = objective_grad_hess(candidate) else {
                            candidate = 0.5 * (bracket_lo.0 + bracket_hi.0);
                            let Some(mid_state) = objective_grad_hess(candidate) else {
                                break;
                            };
                            if mid_state.0 < best.0 {
                                best_slope = candidate;
                                best = mid_state;
                            }
                            if mid_state.1 <= 0.0 {
                                bracket_lo = (candidate, mid_state);
                            } else {
                                bracket_hi = (candidate, mid_state);
                            }
                            slope = candidate;
                            state = mid_state;
                            continue;
                        };
                        if candidate_state.0 < best.0 {
                            best_slope = candidate;
                            best = candidate_state;
                        }
                        if candidate_state.1 <= 0.0 {
                            bracket_lo = (candidate, candidate_state);
                        } else {
                            bracket_hi = (candidate, candidate_state);
                        }
                        slope = candidate;
                        state = candidate_state;
                    }
                    return if best.0.is_finite() { best_slope } else { 0.0 };
                }
            }
        }
        step *= 2.0;
    }
    if best.0.is_finite() { best_slope } else { 0.0 }
}

// ── Public fitting function ───────────────────────────────────────────

/// Whether the optional score-warp / link-deviation flex blocks participate in
/// a family build. The rigid warm-start pilot must construct its family and
/// blocks with `OffForRigidPilot` so the cold-start coefficient solve cannot
/// silently activate the survival flex exact-Joint-Newton path.
#[derive(Copy, Clone, PartialEq, Eq, Debug)]
pub(crate) enum FlexActivation {
    OffForRigidPilot,
    On,
}