gam 0.2.3

use std::collections::HashMap;

use ndarray::{Array1, Array2};
use rayon::iter::{IntoParallelIterator, ParallelIterator};

use crate::basis::{
    BasisOptions, Dense, KnotSource, create_basis, create_ispline_derivative_dense,
};
use crate::estimate::{BlockRole, PredictInput};
use crate::families::bernoulli_marginal_slope::LatentMeasureKind;
use crate::families::scale_design::{build_scale_deviation_operator, scale_transform_from_payload};
use crate::families::survival_predict::SurvivalPredictError;
use crate::families::survival_predict::{
    fit_result_from_saved_model_for_prediction, resolve_termspec_for_prediction,
};
use crate::families::transformation_normal::{
    TRANSFORMATION_MONOTONICITY_EPS, TRANSFORMATION_NORMAL_H_ABS_MAX,
    transformation_normal_pit_score,
};
use crate::inference::model::{
    FittedModel, FittedModelError, PredictModelClass, append_deployment_extension_columns,
};
use crate::linalg::utils::inf_norm;
use crate::matrix::DesignMatrix;
use crate::smooth::build_term_collection_design;
use crate::term_builder::resolve_role_col;

/// Typed errors emitted while assembling a [`PredictInput`] from a saved model.
///
/// Each variant carries a pre-formatted `reason` string so `Display` is
/// byte-equivalent to the original `format!(...)` outputs the module used
/// before the typed-error migration. The category split lets callers
/// pattern-match on the failure kind without dragging the string apart.
#[derive(Debug, Clone)]
pub enum PredictInputError {
    /// Request-level input did not satisfy the predict contract: bad offset
    /// lengths, non-finite covariates, unsupported predict options for the
    /// saved model class, or unparseable model metadata at the boundary.
    InvalidInput { reason: String },
    /// Rebuilt prediction designs disagree with saved coefficient blocks or
    /// transform matrices (model/design column counts, basis shapes,
    /// reshape failures).
    DimensionMismatch { reason: String },
    /// The saved model is missing payload metadata required to drive the
    /// prediction (response knots, transform, degree, calibration block,
    /// unified fit, z column, etc.).
    MissingMetadata { reason: String },
    /// Survival-specific prediction assembly failed below this layer; the
    /// source error keeps its own semantic variant instead of being flattened
    /// into a generic predict-input bucket.
    SurvivalPrediction {
        context: &'static str,
        source: SurvivalPredictError,
    },
    /// Saved-model payload validation failed below this layer; the source
    /// error keeps its model-layer category and payload context.
    ModelPayload {
        context: &'static str,
        source: FittedModelError,
    },
}

impl std::fmt::Display for PredictInputError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            PredictInputError::InvalidInput { reason }
            | PredictInputError::DimensionMismatch { reason }
            | PredictInputError::MissingMetadata { reason } => f.write_str(reason),
            PredictInputError::SurvivalPrediction { context, source } => {
                write!(f, "{context}: {source}")
            }
            PredictInputError::ModelPayload { context, source } => {
                write!(f, "{context}: {source}")
            }
        }
    }
}

impl std::error::Error for PredictInputError {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        match self {
            PredictInputError::SurvivalPrediction { source, .. } => Some(source),
            PredictInputError::ModelPayload { source, .. } => Some(source),
            PredictInputError::InvalidInput { .. }
            | PredictInputError::DimensionMismatch { .. }
            | PredictInputError::MissingMetadata { .. } => None,
        }
    }
}

impl From<PredictInputError> for String {
    fn from(err: PredictInputError) -> String {
        err.to_string()
    }
}

impl From<String> for PredictInputError {
    /// Inbound conversion from the many `Result<_, String>` helpers this
    /// module still calls into (basis builders, term-collection assembly,
    /// fit deserializers). The text is preserved verbatim; we only pick a
    /// category so external messages flow through `?` without per-callsite
    /// `.map_err`.
    fn from(reason: String) -> PredictInputError {
        PredictInputError::InvalidInput { reason }
    }
}

impl From<SurvivalPredictError> for PredictInputError {
    /// Survival-prediction helpers (`resolve_termspec_for_prediction`,
    /// `fit_result_from_saved_model_for_prediction`) emit their own typed
    /// errors; keep that typed source so `?` preserves the layer that failed.
    fn from(err: SurvivalPredictError) -> PredictInputError {
        PredictInputError::SurvivalPrediction {
            context: "predict-input survival assembly",
            source: err,
        }
    }
}

impl From<FittedModelError> for PredictInputError {
    /// `FittedModel` payload helpers (deployment extension assembly,
    /// calibration validation) surface model-layer errors that remain
    /// chained here instead of being recategorized as request input.
    fn from(err: FittedModelError) -> PredictInputError {
        PredictInputError::ModelPayload {
            context: "predict-input model payload",
            source: err,
        }
    }
}

fn build_marginal_slope_local_auxiliary_matrix(
    model: &FittedModel,
    data: ndarray::ArrayView2<'_, f64>,
    col_map: &HashMap<String, usize>,
) -> Result<Option<Array2<f64>>, PredictInputError> {
    let Some(LatentMeasureKind::LocalEmpirical {
        feature_cols,
        input_scales,
        ..
    }) = model.latent_measure.as_ref()
    else {
        return Ok(None);
    };
    let n = data.nrows();
    let d = feature_cols.len();
    let mut out = Array2::<f64>::zeros((n, d));
    let training_headers = model.training_headers.as_ref();
    for (local_col, &fit_col) in feature_cols.iter().enumerate() {
        let prediction_col = training_headers
            .and_then(|headers| headers.get(fit_col))
            .and_then(|name| col_map.get(name))
            .copied()
            .unwrap_or(fit_col);
        if prediction_col >= data.ncols() {
            return Err(PredictInputError::DimensionMismatch {
                reason: format!(
                    "local empirical marginal-slope prediction feature column {fit_col} is out of bounds for {} columns",
                    data.ncols()
                ),
            });
        }
        out.column_mut(local_col)
            .assign(&data.column(prediction_col));
    }
    if let Some(scales) = input_scales.as_ref() {
        if scales.len() != d {
            return Err(PredictInputError::DimensionMismatch {
                reason: format!(
                    "local empirical marginal-slope prediction input scale dimension mismatch: scales={}, features={d}",
                    scales.len()
                ),
            });
        }
        for (col, &scale) in scales.iter().enumerate() {
            if !(scale.is_finite() && scale > 0.0) {
                return Err(PredictInputError::InvalidInput {
                    reason: format!(
                        "local empirical marginal-slope prediction input scale {col} must be finite and positive, got {scale}"
                    ),
                });
            }
            out.column_mut(col).mapv_inplace(|value| value / scale);
        }
    }
    if out.iter().any(|value| !value.is_finite()) {
        return Err(PredictInputError::InvalidInput {
            reason: "local empirical marginal-slope prediction conditioning values must be finite"
                .to_string(),
        });
    }
    Ok(Some(out))
}

fn build_predict_input_for_model_inner(
    model: &FittedModel,
    data: ndarray::ArrayView2<'_, f64>,
    col_map: &HashMap<String, usize>,
    training_headers: Option<&Vec<String>>,
    offset: &Array1<f64>,
    offset_noise: &Array1<f64>,
    noise_offset_supplied: bool,
) -> Result<PredictInput, PredictInputError> {
    let spec = resolve_termspec_for_prediction(
        &model.resolved_termspec,
        training_headers,
        col_map,
        "resolved_termspec",
    )?;
    let clipped = model.axis_clip_to_training_ranges(data, col_map);
    let design_input = clipped.as_ref().map_or(data, |arr| arr.view());
    let design = build_term_collection_design(design_input, &spec).map_err(|e| {
        PredictInputError::InvalidInput {
            reason: format!("failed to build prediction design: {e}"),
        }
    })?;
    let n = data.nrows();
    if offset.len() != n || offset_noise.len() != n {
        return Err(PredictInputError::DimensionMismatch {
            reason: format!(
                "prediction offset length mismatch: rows={n}, offset={}, noise_offset={}",
                offset.len(),
                offset_noise.len()
            ),
        });
    }

    match model.predict_model_class() {
        PredictModelClass::Standard => {
            if noise_offset_supplied {
                return Err(PredictInputError::InvalidInput {
                    reason: "--noise-offset-column is not supported for standard prediction"
                        .to_string(),
                });
            }
            let fit_saved = fit_result_from_saved_model_for_prediction(model)?;
            let beta = if model.has_link_wiggle() {
                fit_saved
                    .block_by_role(BlockRole::Mean)
                    .ok_or_else(|| PredictInputError::MissingMetadata {
                        reason: "standard link-wiggle model is missing Mean coefficient block"
                            .to_string(),
                    })?
                    .beta
                    .clone()
            } else {
                fit_saved.beta.clone()
            };
            let mean_design = if model.deployment_extensions.is_empty() {
                design.design.clone()
            } else {
                DesignMatrix::from(append_deployment_extension_columns(
                    model.payload(),
                    design_input,
                    col_map,
                    training_headers,
                    design.design.to_dense(),
                )?)
            };
            if beta.len() != mean_design.ncols() {
                return Err(PredictInputError::DimensionMismatch {
                    reason: format!(
                        "model/design mismatch: model beta has {} coefficients but new-data design has {} columns",
                        beta.len(),
                        mean_design.ncols()
                    ),
                });
            }
            Ok(PredictInput {
                design: mean_design,
                offset: offset.clone(),
                design_noise: None,
                offset_noise: None,
                auxiliary_scalar: None,
                auxiliary_matrix: None,
            })
        }
        PredictModelClass::GaussianLocationScale | PredictModelClass::BinomialLocationScale => {
            let spec_noise = resolve_termspec_for_prediction(
                &model.resolved_termspec_noise,
                training_headers,
                col_map,
                "resolved_termspec_noise",
            )?;
            let design_noise_raw = build_term_collection_design(design_input, &spec_noise)
                .map_err(|e| PredictInputError::InvalidInput {
                    reason: format!("failed to build noise prediction design: {e}"),
                })?;

            let noise_transform = scale_transform_from_payload(
                &model.noise_projection,
                &model.noise_center,
                &model.noise_scale,
                model.noise_non_intercept_start,
                model.noise_projection_ridge_alpha,
            )?;
            let prepared_noise_design = if let Some(transform) = noise_transform.as_ref() {
                build_scale_deviation_operator(
                    design.design.clone(),
                    design_noise_raw.design.clone(),
                    transform,
                )?
            } else {
                design_noise_raw.design.clone()
            };

            Ok(PredictInput {
                design: design.design.clone(),
                offset: offset.clone(),
                design_noise: Some(prepared_noise_design),
                offset_noise: Some(offset_noise.clone()),
                auxiliary_scalar: None,
                auxiliary_matrix: None,
            })
        }
        PredictModelClass::BernoulliMarginalSlope => {
            let z_name =
                model
                    .z_column
                    .as_ref()
                    .ok_or_else(|| PredictInputError::MissingMetadata {
                        reason: "marginal-slope model is missing z_column".to_string(),
                    })?;
            let z_col = resolve_role_col(col_map, z_name, "z")?;
            let z = data.column(z_col).to_owned();
            let spec_logslope = resolve_termspec_for_prediction(
                &model.resolved_termspec_logslope.as_ref().cloned(),
                training_headers,
                col_map,
                "resolved_termspec_logslope",
            )?;
            let design_logslope = build_term_collection_design(design_input, &spec_logslope)
                .map_err(|e| PredictInputError::InvalidInput {
                    reason: format!("failed to build logslope prediction design: {e}"),
                })?;
            let auxiliary_matrix =
                build_marginal_slope_local_auxiliary_matrix(model, design_input, col_map)?;
            Ok(PredictInput {
                design: design.design.clone(),
                offset: offset.clone(),
                design_noise: Some(design_logslope.design.clone()),
                offset_noise: Some(offset_noise.clone()),
                auxiliary_scalar: Some(z),
                auxiliary_matrix,
            })
        }
        PredictModelClass::Survival => Err(PredictInputError::InvalidInput {
            reason: "build_predict_input_for_model should not be called for survival models"
                .to_string(),
        }),
        PredictModelClass::TransformationNormal => {
            if noise_offset_supplied {
                return Err(PredictInputError::InvalidInput {
                    reason:
                        "--noise-offset-column is not supported for transformation-normal prediction"
                            .to_string(),
                });
            }
            let payload = model.payload();
            let response_knots =
                payload
                    .transformation_response_knots
                    .as_ref()
                    .ok_or_else(|| PredictInputError::MissingMetadata {
                        reason: "saved transformation-normal model missing response_knots"
                            .to_string(),
                    })?;
            let response_transform_vecs = payload
                .transformation_response_transform
                .as_ref()
                .ok_or_else(|| PredictInputError::MissingMetadata {
                    reason: "saved transformation-normal model missing response_transform"
                        .to_string(),
                })?;
            let response_degree = payload.transformation_response_degree.ok_or_else(|| {
                PredictInputError::MissingMetadata {
                    reason: "saved transformation-normal model missing response_degree".to_string(),
                }
            })?;
            let response_median = payload.transformation_response_median.ok_or_else(|| {
                PredictInputError::MissingMetadata {
                    reason: "saved transformation-normal model missing response_median".to_string(),
                }
            })?;

            let t_rows = response_transform_vecs.len();
            let t_cols = if t_rows > 0 {
                response_transform_vecs[0].len()
            } else {
                0
            };
            let mut resp_transform = ndarray::Array2::<f64>::zeros((t_rows, t_cols));
            for (i, row) in response_transform_vecs.iter().enumerate() {
                for (j, &v) in row.iter().enumerate() {
                    resp_transform[[i, j]] = v;
                }
            }
            let resp_knots = ndarray::Array1::from_vec(response_knots.clone());

            let response_col_name = payload
                .formula
                .split('~')
                .next()
                .map(str::trim)
                .ok_or_else(|| PredictInputError::InvalidInput {
                    reason: "cannot parse response column from formula".to_string(),
                })?;
            let response_col_idx = resolve_role_col(col_map, response_col_name, "response")?;
            let response_new = data.column(response_col_idx).to_owned();
            for value in response_new.iter().copied() {
                if !value.is_finite() {
                    return Err(PredictInputError::InvalidInput {
                        reason: format!(
                            "transformation-normal response value in prediction data is not finite: {value}"
                        ),
                    });
                }
            }

            let (raw_val_basis, _) = create_basis::<Dense>(
                response_new.view(),
                KnotSource::Provided(resp_knots.view()),
                response_degree,
                BasisOptions::i_spline(),
            )
            .map_err(|e| PredictInputError::InvalidInput {
                reason: e.to_string(),
            })?;
            let raw_val = raw_val_basis.as_ref().clone();
            if raw_val.ncols() != resp_transform.nrows() {
                return Err(PredictInputError::DimensionMismatch {
                    reason: format!(
                        "saved transformation-normal response transform shape mismatch: raw I-spline cols={} transform rows={}",
                        raw_val.ncols(),
                        resp_transform.nrows()
                    ),
                });
            }
            let shape_val = raw_val.dot(&resp_transform);
            let p_shape = resp_transform.ncols();
            let p_resp = 1 + p_shape;
            let mut resp_val = ndarray::Array2::<f64>::zeros((n, p_resp));
            resp_val.column_mut(0).fill(1.0);
            resp_val.slice_mut(ndarray::s![.., 1..]).assign(&shape_val);

            let raw_deriv = create_ispline_derivative_dense(
                response_new.view(),
                &resp_knots,
                response_degree,
                1,
            )
            .map_err(|e| PredictInputError::InvalidInput {
                reason: e.to_string(),
            })?;
            if raw_deriv.ncols() != resp_transform.nrows() {
                return Err(PredictInputError::DimensionMismatch {
                    reason: format!(
                        "saved transformation-normal derivative transform shape mismatch: raw M-spline cols={} transform rows={}",
                        raw_deriv.ncols(),
                        resp_transform.nrows()
                    ),
                });
            }
            let shape_deriv = raw_deriv.dot(&resp_transform);
            let mut resp_deriv = ndarray::Array2::<f64>::zeros((n, p_resp));
            resp_deriv
                .slice_mut(ndarray::s![.., 1..])
                .assign(&shape_deriv);

            let fit_saved = model
                .unified()
                .ok_or_else(|| PredictInputError::MissingMetadata {
                    reason: "saved transformation-normal model missing unified fit".to_string(),
                })?;
            let beta = &fit_saved.blocks[0].beta;
            let p_cov = design.design.ncols();
            if beta.len() != p_resp * p_cov {
                return Err(PredictInputError::DimensionMismatch {
                    reason: format!(
                        "beta length {} != p_resp({}) * p_cov({})",
                        beta.len(),
                        p_resp,
                        p_cov
                    ),
                });
            }
            let beta_mat = beta
                .view()
                .into_shape_with_order((p_resp, p_cov))
                .map_err(|e| PredictInputError::DimensionMismatch {
                    reason: format!("beta reshape failed: {e}"),
                })?;
            let cov_mat =
                design
                    .design
                    .try_row_chunk(0..n)
                    .map_err(|e| PredictInputError::InvalidInput {
                        reason: e.to_string(),
                    })?;
            let calibration = payload
                .transformation_score_calibration
                .as_ref()
                .ok_or_else(|| PredictInputError::MissingMetadata {
                    reason: "saved transformation-normal model missing score calibration"
                        .to_string(),
                })?;
            calibration.validate("saved transformation-normal score calibration")?;

            if resp_knots.is_empty() {
                return Err(PredictInputError::MissingMetadata {
                    reason: "saved transformation-normal response knots are empty".to_string(),
                });
            }
            let mut response_lower_basis = vec![0.0; p_resp];
            let mut response_upper_basis = vec![0.0; p_resp];
            response_lower_basis[0] = 1.0;
            response_upper_basis[0] = 1.0;
            for col in 0..p_shape {
                response_upper_basis[col + 1] = resp_transform.column(col).sum();
            }
            let response_lower_floor_offset =
                TRANSFORMATION_MONOTONICITY_EPS * (resp_knots[0] - response_median);
            let response_upper_floor_offset = TRANSFORMATION_MONOTONICITY_EPS
                * (resp_knots[resp_knots.len() - 1] - response_median);

            // Under SCOP-CTN with I-spline shape components,
            // `h'(y, x) = ε + Σ_{r≥1} M_r(y) · γ_r(x)²`. Both M_r and γ_r²
            // are non-negative for every (β, x, y), and ε is the fixed
            // derivative floor serialized through the model definition.
            let monotonicity_eps = TRANSFORMATION_MONOTONICITY_EPS;
            let beta_mat_ref = &beta_mat;
            let cov_mat_ref = &cov_mat;
            let resp_deriv_ref = &resp_deriv;
            let min_h_prime: f64 = (0..n)
                .into_par_iter()
                .map(|i| {
                    let cov_row = cov_mat_ref.row(i);
                    let resp_row = resp_deriv_ref.row(i);
                    let mut hp = resp_row[0] * beta_mat_ref.row(0).dot(&cov_row);
                    for r in 1..p_resp {
                        let gamma = beta_mat_ref.row(r).dot(&cov_row);
                        hp += resp_row[r] * gamma * gamma;
                    }
                    hp + monotonicity_eps
                })
                .reduce(|| f64::INFINITY, f64::min);
            if min_h_prime < monotonicity_eps {
                return Err(PredictInputError::InvalidInput {
                    reason: format!(
                        "prediction failed: transformation-normal h'(y, x) numerical floor \
                         violated. Minimum evaluated h'(y, x) is {min_h_prime:.3e}, threshold \
                         {monotonicity_eps:.0e}. Under SCOP h' = ε + Σ M_r γ_r² holds \
                         structurally, so this indicates floating-point cancellation below \
                         the fixed derivative floor."
                    ),
                });
            }

            // h_i and finite-support endpoints share the same γ_r(x_i). The
            // prediction score is the fitted PIT, not a post-h location/scale
            // normalization.
            let pit_vec: Vec<Result<f64, String>> = (0..n)
                .into_par_iter()
                .map(|i| {
                    let resp_row = resp_val.row(i);
                    let cov_row = cov_mat.row(i);
                    let gamma0 = beta_mat.row(0).dot(&cov_row);
                    let mut val = resp_row[0] * gamma0;
                    let mut lower = response_lower_basis[0] * gamma0;
                    let mut upper = response_upper_basis[0] * gamma0;
                    let mut max_abs_gamma = gamma0.abs();
                    for r in 1..p_resp {
                        let gamma = beta_mat.row(r).dot(&cov_row);
                        max_abs_gamma = max_abs_gamma.max(gamma.abs());
                        val += resp_row[r] * gamma * gamma;
                        lower += response_lower_basis[r] * gamma * gamma;
                        upper += response_upper_basis[r] * gamma * gamma;
                    }
                    let h = val
                        + offset[i]
                        + monotonicity_eps * (response_new[i] - response_median);
                    let h_lower = lower + offset[i] + response_lower_floor_offset;
                    let h_upper = upper + offset[i] + response_upper_floor_offset;
                    if !h.is_finite() || !h_lower.is_finite() || !h_upper.is_finite() {
                        let max_abs_cov = inf_norm(cov_row.iter().copied());
                        return Err(format!(
                            "prediction failed: transformation-normal finite-support scores at row {i} are not finite: h={h:.6e}, lower={h_lower:.6e}, upper={h_upper:.6e}; max_abs_covariate_basis={max_abs_cov:.6e}, max_abs_gamma={max_abs_gamma:.6e}"
                        ));
                    }
                    transformation_normal_pit_score(h, h_lower, h_upper, calibration.clip_eps)
                        .map_err(|err| format!("prediction failed at row {i}: {err}"))
                })
                .collect();
            let calibrated = ndarray::Array1::<f64>::from_vec(
                pit_vec.into_iter().collect::<Result<Vec<_>, _>>()?,
            );
            if calibrated
                .iter()
                .any(|value| !value.is_finite() || value.abs() > TRANSFORMATION_NORMAL_H_ABS_MAX)
            {
                return Err(PredictInputError::InvalidInput {
                    reason:
                        "prediction failed: transformation-normal PIT produced non-finite or out-of-range z values"
                            .to_string(),
                });
            }
            Ok(PredictInput {
                design: DesignMatrix::from(ndarray::Array2::from_shape_fn((n, 1), |_| 1.0)),
                offset: calibrated,
                design_noise: None,
                offset_noise: None,
                auxiliary_scalar: None,
                auxiliary_matrix: None,
            })
        }
    }
}

/// Build a `PredictInput` for model types backed directly by `PredictableModel`.
///
/// Survival prediction has its own design assembly because it needs entry/exit
/// time geometry before it can call the same predictor/output machinery.
pub fn build_predict_input_for_model(
    model: &FittedModel,
    data: ndarray::ArrayView2<'_, f64>,
    col_map: &HashMap<String, usize>,
    training_headers: Option<&Vec<String>>,
    offset: &Array1<f64>,
    offset_noise: &Array1<f64>,
    noise_offset_supplied: bool,
) -> Result<PredictInput, String> {
    build_predict_input_for_model_inner(
        model,
        data,
        col_map,
        training_headers,
        offset,
        offset_noise,
        noise_offset_supplied,
    )
    .map_err(Into::into)
}