gam-models 0.3.127

//! Flex primary third/fourth contracted exact tensors.
//!
//! Builds the third- and fourth-order directional contractions of the
//! primary-space NLL Hessian on top of the exact timepoint evaluations, and
//! routes the general (flex vs rigid) entry points to either the exact flex
//! path or the rigid fallback.

use super::*;

impl SurvivalMarginalSlopeFamily {
    /// Exact third-order directional contraction for the flexible survival
    /// path.  Returns D_dir H[u,v] where H is the primary-space NLL Hessian.
    pub(crate) fn row_flex_primary_third_contracted_exact(
        &self,
        row: usize,
        block_states: &[ParameterBlockState],
        dir: &Array1<f64>,
    ) -> Result<Array2<f64>, String> {
        self.ensure_scalar_flex_exact_score_geometry("row_flex_primary_third_contracted_exact")?;
        let primary = flex_primary_slices(self);
        let p = primary.total;
        if dir.len() != p {
            return Err(SurvivalMarginalSlopeError::IncompatibleDimensions {
                reason: format!(
                    "survival third contracted: dir length {} != primary dimension {p}",
                    dir.len()
                ),
            }
            .into());
        }
        if dir.iter().all(|v| v.abs() == 0.0) {
            return Ok(Array2::<f64>::zeros((p, p)));
        }

        let base = self.build_row_flex_third_base_with_states(row, block_states, &primary)?;
        self.row_flex_third_contract_from_base(&base, dir)
    }

    /// Build the direction-independent per-row geometry that
    /// [`Self::row_flex_third_contract_from_base`] reuses across every
    /// coefficient axis of a Jeffreys all-axes sweep.
    ///
    /// The intercept solves, cached partitions, and exact base timepoints
    /// depend only on the row (its `q`-geometry and the current `β`), not on
    /// the contraction direction `dir`. Hoisting them out of the per-axis loop
    /// turns the all-axes flex third contraction from a `p`-fold rebuild into a
    /// build-once + `p` cheap directional contractions — the #979 flex hot path.
    pub(crate) fn build_row_flex_third_base_with_states(
        &self,
        row: usize,
        block_states: &[ParameterBlockState],
        primary: &FlexPrimarySlices,
    ) -> Result<FlexThirdRowBase, String> {
        let q_geom = self.row_dynamic_q_geometry(row, block_states)?;
        let q0 = q_geom.q0;
        let q1 = q_geom.q1;
        let qd1 = q_geom.qd1;
        let g = block_states[2].eta[row];
        let beta_h = self.flex_score_beta(block_states)?;
        let beta_w = self.flex_link_beta(block_states)?;
        let o_infl = self.influence_index_offset(row, block_states)?;

        if survival_derivative_guard_violated(qd1, self.derivative_guard) {
            return Err(SurvivalMarginalSlopeError::MonotonicityViolation {
                reason: format!(
                    "survival third contracted monotonicity violated at row {row}: qd1={qd1:.3e}"
                ),
            }
            .into());
        }

        // The intercept solve's density check `d0`/`d1` is recomputed inside the jet
        // base builder (`evaluate_survival_denom_d`); only the intercepts are needed.
        let (a0, _) = self.solve_row_survival_intercept_with_slot(
            q0,
            g,
            beta_h,
            beta_w,
            Some((row, SurvivalInterceptSlotKind::Entry)),
        )?;
        let (a1, _) = self.solve_row_survival_intercept_with_slot(
            q1,
            g,
            beta_h,
            beta_w,
            Some((row, SurvivalInterceptSlotKind::Exit)),
        )?;

        let entry_cached = self.build_cached_partition(primary, a0, g, beta_h, beta_w)?;
        let exit_cached = self.build_cached_partition(primary, a1, g, beta_h, beta_w)?;

        // #932-2 increment 3: the contracted BASE timepoint (value/grad/Hessian seed of
        // the Jet3/Jet4 contraction) now comes from the single-source jet builder, not
        // the hand `compute_survival_timepoint_exact_from_cached`.
        let entry = self.compute_survival_timepoint_exact_jet_from_cached(
            row,
            primary,
            q0,
            primary.q0,
            a0,
            g,
            beta_h,
            beta_w,
            o_infl,
            &entry_cached,
        )?;
        let exit = self.compute_survival_timepoint_exact_jet_from_cached(
            row,
            primary,
            q1,
            primary.q1,
            a1,
            g,
            beta_h,
            beta_w,
            o_infl,
            &exit_cached,
        )?;

        if !exit.chi.is_finite() || exit.chi <= 0.0 {
            return Err(SurvivalMarginalSlopeError::NumericalFailure {
                reason: format!(
                    "survival third contracted row {row}: non-positive chi1={:.3e}",
                    exit.chi,
                ),
            }
            .into());
        }

        Ok(FlexThirdRowBase {
            row,
            p: primary.total,
            qd1,
            q0,
            q1,
            q0_index: primary.q0,
            q1_index: primary.q1,
            a0,
            a1,
            g,
            beta_h: beta_h.cloned(),
            beta_w: beta_w.cloned(),
            entry_cached,
            exit_cached,
            entry_base: block10_pack_base(&entry),
            exit_base: block10_pack_base(&exit),
        })
    }

    /// Contract the third-order tensor of a row against a single direction,
    /// reusing the direction-independent [`FlexThirdRowBase`]. Only the
    /// directional timepoint extensions and the Block 10 third-contraction
    /// assembly are recomputed per axis. Bit-identical to the inline path that
    /// [`Self::row_flex_primary_third_contracted_exact`] previously ran.
    pub(crate) fn row_flex_third_contract_from_base(
        &self,
        base: &FlexThirdRowBase,
        dir: &Array1<f64>,
    ) -> Result<Array2<f64>, String> {
        let p = base.p;
        if dir.iter().all(|v| v.abs() == 0.0) {
            return Ok(Array2::<f64>::zeros((p, p)));
        }
        let primary = flex_primary_slices(self);
        let beta_h = base.beta_h.as_ref();
        let beta_w = base.beta_w.as_ref();

        // #932-2 PRODUCTION cutover (increment 2): the directional timepoint
        // extension (Block-10 directional pack) comes from the single-source
        // `flex_timepoint_inputs_generic` jet builder at `Jet3`, replacing the hand
        // `compute_survival_timepoint_directional_exact_from_cached` + `block10_pack_dir`.
        let entry_ext = self.compute_survival_timepoint_directional_jet_from_cached(
            base.row,
            &primary,
            base.q0,
            base.q0_index,
            base.a0,
            base.g,
            beta_h,
            beta_w,
            &base.entry_cached,
            dir,
        )?;
        let exit_ext = self.compute_survival_timepoint_directional_jet_from_cached(
            base.row,
            &primary,
            base.q1,
            base.q1_index,
            base.a1,
            base.g,
            beta_h,
            beta_w,
            &base.exit_cached,
            dir,
        )?;

        // #932 single-source: the contracted third `Σ_c ℓ_{abc} dir_c =
        // (D_dir H)[a][b]` is the ε-Hessian channel of the ONE generic flex
        // row-NLL expression (`flex_row_nll`) instantiated at the one-seed jet
        // `Jet3`, seeded from the base + directional timepoint packs. Replaces
        // the bespoke probit-chain / quotient-rule assembly that the Block-10
        // `cpu_oracle_third_contraction` hand-coded.
        self.flex_row_nll_third_contracted(
            base.row,
            &primary,
            base.q1,
            base.qd1,
            dir.as_slice()
                .ok_or_else(|| "third contraction: dir must be contiguous".to_string())?,
            super::flex_jet::FlexThirdPacks {
                entry_base: &base.entry_base,
                exit_base: &base.exit_base,
                entry_ext: &entry_ext,
                exit_ext: &exit_ext,
            },
        )
    }

    /// Fourth-order directional contraction for the flexible survival path.
    ///
    /// The mixed second-directional timepoint transport is carried exactly
    /// through the implicit intercept solve, the observed-point eta/chi jets,
    /// and the cellwise density-normalization integrand.
    pub(crate) fn row_flex_primary_fourth_contracted_exact(
        &self,
        row: usize,
        block_states: &[ParameterBlockState],
        dir_u: &Array1<f64>,
        dir_v: &Array1<f64>,
    ) -> Result<Array2<f64>, String> {
        self.ensure_scalar_flex_exact_score_geometry("row_flex_primary_fourth_contracted_exact")?;
        let primary = flex_primary_slices(self);
        let p = primary.total;
        if dir_u.len() != p || dir_v.len() != p {
            return Err(SurvivalMarginalSlopeError::IncompatibleDimensions {
                reason: format!(
                    "survival fourth contracted: dir lengths ({},{}) != {p}",
                    dir_u.len(),
                    dir_v.len(),
                ),
            }
            .into());
        }
        if dir_u.iter().all(|v| v.abs() == 0.0) || dir_v.iter().all(|v| v.abs() == 0.0) {
            return Ok(Array2::<f64>::zeros((p, p)));
        }

        let q_geom = self.row_dynamic_q_geometry(row, block_states)?;
        let q0 = q_geom.q0;
        let q1 = q_geom.q1;
        let qd1 = q_geom.qd1;
        let g = block_states[2].eta[row];
        let beta_h = self.flex_score_beta(block_states)?;
        let beta_w = self.flex_link_beta(block_states)?;
        let o_infl = self.influence_index_offset(row, block_states)?;

        if survival_derivative_guard_violated(qd1, self.derivative_guard) {
            return Err(SurvivalMarginalSlopeError::MonotonicityViolation {
                reason: format!(
                    "survival fourth contracted monotonicity violated at row {row}: qd1={qd1:.3e}"
                ),
            }
            .into());
        }

        // Only the solved intercepts are needed; the jet base builder recomputes the
        // density check internally.
        let (a0, _) = self.solve_row_survival_intercept_with_slot(
            q0,
            g,
            beta_h,
            beta_w,
            Some((row, SurvivalInterceptSlotKind::Entry)),
        )?;
        let (a1, _) = self.solve_row_survival_intercept_with_slot(
            q1,
            g,
            beta_h,
            beta_w,
            Some((row, SurvivalInterceptSlotKind::Exit)),
        )?;

        let entry_cached = self.build_cached_partition(&primary, a0, g, beta_h, beta_w)?;
        let exit_cached = self.build_cached_partition(&primary, a1, g, beta_h, beta_w)?;

        // #932-2 increment 3: contracted-fourth BASE timepoint via the single-source
        // jet builder (replaces the hand `compute_survival_timepoint_exact_from_cached`).
        let entry_base = self.compute_survival_timepoint_exact_jet_from_cached(
            row,
            &primary,
            q0,
            primary.q0,
            a0,
            g,
            beta_h,
            beta_w,
            o_infl,
            &entry_cached,
        )?;
        let exit_base = self.compute_survival_timepoint_exact_jet_from_cached(
            row,
            &primary,
            q1,
            primary.q1,
            a1,
            g,
            beta_h,
            beta_w,
            o_infl,
            &exit_cached,
        )?;

        if !exit_base.chi.is_finite() || exit_base.chi <= 0.0 {
            return Err(SurvivalMarginalSlopeError::NumericalFailure {
                reason: format!(
                    "survival fourth contracted row {row}: non-positive chi1={:.3e}",
                    exit_base.chi,
                ),
            }
            .into());
        }

        // #932-2 PRODUCTION cutover (increment 2): both the directional (u, v) and
        // the mixed second-directional (bi) timepoint extensions come from the
        // single-source `flex_timepoint_inputs_generic` jet builder at `Jet3`/`Jet4`,
        // returning the Block-10 packs directly, replacing the hand
        // `compute_survival_timepoint_{directional,bidirectional}_exact_from_cached`.
        let entry_ext_u = self.compute_survival_timepoint_directional_jet_from_cached(
            row,
            &primary,
            q0,
            primary.q0,
            a0,
            g,
            beta_h,
            beta_w,
            &entry_cached,
            dir_u,
        )?;
        let entry_ext_v = self.compute_survival_timepoint_directional_jet_from_cached(
            row,
            &primary,
            q0,
            primary.q0,
            a0,
            g,
            beta_h,
            beta_w,
            &entry_cached,
            dir_v,
        )?;
        let exit_ext_u = self.compute_survival_timepoint_directional_jet_from_cached(
            row,
            &primary,
            q1,
            primary.q1,
            a1,
            g,
            beta_h,
            beta_w,
            &exit_cached,
            dir_u,
        )?;
        let exit_ext_v = self.compute_survival_timepoint_directional_jet_from_cached(
            row,
            &primary,
            q1,
            primary.q1,
            a1,
            g,
            beta_h,
            beta_w,
            &exit_cached,
            dir_v,
        )?;

        // Bidirectional extensions D_{d1} D_{d2} (η_uv, χ_uv, D_uv).
        let entry_bi = self.compute_survival_timepoint_bidirectional_jet_from_cached(
            row,
            &primary,
            q0,
            primary.q0,
            a0,
            g,
            beta_h,
            beta_w,
            &entry_cached,
            dir_u,
            dir_v,
        )?;
        let exit_bi = self.compute_survival_timepoint_bidirectional_jet_from_cached(
            row,
            &primary,
            q1,
            primary.q1,
            a1,
            g,
            beta_h,
            beta_w,
            &exit_cached,
            dir_u,
            dir_v,
        )?;

        // #932 single-source: the contracted fourth `Σ_cd ℓ_{abcd} u_c v_d` is
        // the εδ-Hessian channel of the ONE generic flex row-NLL expression
        // (`flex_row_nll`) instantiated at the two-seed jet `Jet4`, seeded from
        // the base + both directional + bidirectional timepoint packs. Replaces
        // the bespoke per-(u,v) probit-chain / quotient-rule assembly that the
        // Block-10 `cpu_oracle_fourth_contraction` hand-coded.
        self.flex_row_nll_fourth_contracted(
            row,
            &primary,
            q1,
            qd1,
            dir_u
                .as_slice()
                .ok_or_else(|| "fourth contraction: dir_u must be contiguous".to_string())?,
            dir_v
                .as_slice()
                .ok_or_else(|| "fourth contraction: dir_v must be contiguous".to_string())?,
            super::flex_jet::FlexFourthPacks {
                entry_base: &block10_pack_base(&entry_base),
                exit_base: &block10_pack_base(&exit_base),
                entry_ext_u: &entry_ext_u,
                exit_ext_u: &exit_ext_u,
                entry_ext_v: &entry_ext_v,
                exit_ext_v: &exit_ext_v,
                entry_bi: &entry_bi,
                exit_bi: &exit_bi,
            },
        )
    }

    pub(crate) fn row_primary_third_contracted_general(
        &self,
        row: usize,
        block_states: &[ParameterBlockState],
        dir: &Array1<f64>,
    ) -> Result<Array2<f64>, String> {
        if self.effective_flex_active(block_states)? {
            self.row_flex_primary_third_contracted_exact(row, block_states, dir)
        } else {
            self.row_primary_third_contracted(row, block_states, dir.view())
        }
    }

    pub(crate) fn row_primary_fourth_contracted(
        &self,
        row: usize,
        block_states: &[ParameterBlockState],
        dir_u: ArrayView1<'_, f64>,
        dir_v: ArrayView1<'_, f64>,
    ) -> Result<Array2<f64>, String> {
        // Batched path delegating to the shared k=6 jet helper.
        let r = self.row_primary_fourth_contracted_tower(row, block_states, dir_u, dir_v)?;
        let mut out = Array2::<f64>::zeros((N_PRIMARY, N_PRIMARY));
        for a in 0..N_PRIMARY {
            for b in 0..N_PRIMARY {
                out[[a, b]] = r[a][b];
            }
        }
        Ok(out)
    }

    pub(crate) fn row_primary_fourth_contracted_general(
        &self,
        row: usize,
        block_states: &[ParameterBlockState],
        dir_u: &Array1<f64>,
        dir_v: &Array1<f64>,
    ) -> Result<Array2<f64>, String> {
        if self.effective_flex_active(block_states)? {
            self.row_flex_primary_fourth_contracted_exact(row, block_states, dir_u, dir_v)
        } else {
            self.row_primary_fourth_contracted(row, block_states, dir_u.view(), dir_v.view())
        }
    }

    // ── Pullback through design matrices ──────────────────────────────
}