#![allow(dead_code)]
#![allow(clippy::redundant_closure)]
use crate::core::survpenal::{self, MatrixBuffers, PenaltyParams, PenaltyResult};
use crate::regression::survregc1::SurvivalDist;
use crate::utilities::matrix::cholesky_solve_with_fallback;
use ndarray::{Array1, Array2, ArrayView1, ArrayView2};
use ndarray_linalg::Inverse;
#[derive(Debug)]
pub struct SurvivalResult {
    pub coefficients: Vec<f64>,
    pub iterations: usize,
    pub h_matrix: Array2<f64>,
    pub h_inv: Array2<f64>,
    pub h_diag: Vec<f64>,
    pub log_likelihood: f64,
    pub score: Vec<f64>,
    pub penalty: f64,
    pub convergence_flag: i32,
}
#[allow(clippy::too_many_arguments)]
pub fn survreg(
    max_iter: usize,
    nvar: usize,
    y: &Array2<f64>,
    covariates: &Array2<f64>,
    weights: &Array1<f64>,
    offsets: &Array1<f64>,
    mut beta: Vec<f64>,
    nstrat: usize,
    strata: &[usize],
    eps: f64,
    tol_chol: f64,
    distribution: Distribution,
    ptype: PenaltyType,
    pdiag: bool,
    nfrail: usize,
    fgrp: &[usize],
) -> Result<SurvivalResult, Box<dyn std::error::Error>> {
    let n = y.nrows();
    let ny = y.ncols();
    let nvar2 = nvar + nstrat;
    let nvar3 = nvar2 + nfrail;
    let mut hmat = Array2::zeros((nvar3, nvar2));
    let mut jj = Array2::zeros((nvar3, nvar2));
    let mut hdiag = Array1::zeros(nvar3);
    let mut jdiag = Array1::zeros(nfrail);
    let mut u = Array1::zeros(nvar3);
    let mut newbeta = beta.clone();
    let mut flag = 0;
    let time1_vec: Vec<f64> = y.column(0).iter().cloned().collect();
    let status_vec: Vec<f64> = if ny == 2 {
        y.column(1).iter().cloned().collect()
    } else {
        y.column(2).iter().cloned().collect()
    };
    let time2_vec: Option<Vec<f64>> = if ny == 3 {
        Some(y.column(1).iter().cloned().collect())
    } else {
        None
    };
    let time1_arr = Array1::from_vec(time1_vec);
    let status_arr = Array1::from_vec(status_vec);
    let time2_arr = time2_vec.map(|v| Array1::from_vec(v));
    let time1_view = time1_arr.view();
    let status_view = status_arr.view();
    let time2_view = time2_arr.as_ref().map(|v| v.view());
    let mut loglik = calculate_likelihood(
        n,
        nvar,
        nstrat,
        &beta,
        &distribution,
        strata,
        &offsets.view(),
        &time1_view,
        time2_view.as_ref(),
        &status_view,
        &weights.view(),
        &covariates.view(),
        &mut hmat,
        &mut jj,
        &mut u,
        &mut hdiag,
        &mut jdiag,
        nfrail,
        fgrp,
    )?;
    let mut penalty_val = apply_penalties(
        &mut hmat, &mut jj, &mut hdiag, &mut jdiag, &mut u, &mut beta, nvar, nfrail, ptype, pdiag,
    )?;
    loglik += penalty_val;
    let mut iter = 0;
    while iter < max_iter {
        let delta = cholesky_solve_with_fallback(&hmat, &jj, &u, tol_chol)?;
        newbeta
            .iter_mut()
            .zip(beta.iter().zip(delta.iter()))
            .for_each(|(nb, (b, d))| *nb = b + d);
        let newlik = calculate_likelihood(
            n,
            nvar,
            nstrat,
            &newbeta,
            &distribution,
            strata,
            &offsets.view(),
            &time1_view,
            time2_view.as_ref(),
            &status_view,
            &weights.view(),
            &covariates.view(),
            &mut hmat,
            &mut jj,
            &mut u,
            &mut hdiag,
            &mut jdiag,
            nfrail,
            fgrp,
        )?;
        let new_penalty = apply_penalties(
            &mut hmat,
            &mut jj,
            &mut hdiag,
            &mut jdiag,
            &mut u,
            &mut newbeta,
            nvar,
            nfrail,
            ptype,
            pdiag,
        )?;
        let newlik = newlik + new_penalty;
        if (1.0 - (loglik / newlik)).abs() <= eps {
            loglik = newlik;
            penalty_val = new_penalty;
            beta.copy_from_slice(&newbeta);
            flag = 0;
            break;
        }
        if newlik < loglik {
            let alpha = golden_section_search(
                &beta,
                &newbeta,
                loglik,
                n,
                nvar,
                nstrat,
                strata,
                &offsets.view(),
                &time1_view,
                time2_view.as_ref(),
                &status_view,
                &weights.view(),
                &covariates.view(),
                &distribution,
                nfrail,
                fgrp,
            )?;
            newbeta
                .iter_mut()
                .zip(beta.iter())
                .for_each(|(nb, b)| *nb = b + alpha * (*nb - b));
        }
        beta.copy_from_slice(&newbeta);
        loglik = newlik;
        penalty_val = new_penalty;
        iter += 1;
    }
    if iter >= max_iter {
        flag = 1;
    }
    let h_inv = calculate_inverse(&hmat, nvar3, nfrail, &hdiag, tol_chol)?;
    Ok(SurvivalResult {
        coefficients: beta,
        iterations: iter,
        h_matrix: hmat,
        h_inv,
        h_diag: hdiag.to_vec(),
        log_likelihood: loglik,
        score: u.to_vec(),
        penalty: penalty_val,
        convergence_flag: flag,
    })
}
#[allow(clippy::too_many_arguments)]
fn calculate_likelihood(
    n: usize,
    nvar: usize,
    nstrat: usize,
    beta: &[f64],
    distribution: &Distribution,
    strata: &[usize],
    offsets: &ArrayView1<f64>,
    time1: &ArrayView1<f64>,
    time2: Option<&ArrayView1<f64>>,
    status: &ArrayView1<f64>,
    weights: &ArrayView1<f64>,
    covariates: &ArrayView2<f64>,
    hmat: &mut Array2<f64>,
    jj: &mut Array2<f64>,
    u: &mut Array1<f64>,
    hdiag: &mut Array1<f64>,
    _jdiag: &mut Array1<f64>,
    nfrail: usize,
    fgrp: &[usize],
) -> Result<f64, Box<dyn std::error::Error>> {
    use crate::regression::survregc1::survregc1;
    let dist = match distribution {
        Distribution::ExtremeValue => SurvivalDist::ExtremeValue,
        Distribution::Logistic => SurvivalDist::Logistic,
        Distribution::Gaussian => SurvivalDist::Gaussian,
        Distribution::Custom(_) => {
            return Err("Custom distributions not yet supported in calculate_likelihood".into());
        }
    };
    let strat_vec: Vec<i32> = strata.iter().map(|&s| (s + 1) as i32).collect();
    let strat_arr = Array1::from_vec(strat_vec);
    let status_vec: Vec<i32> = status.iter().map(|&s| s as i32).collect();
    let status_arr = Array1::from_vec(status_vec);
    let beta_arr = Array1::from_vec(beta.to_vec());
    let beta_view = beta_arr.view();
    let frail_vec: Vec<i32> = if nfrail > 0 && !fgrp.is_empty() {
        fgrp.iter().map(|&g| (g + 1) as i32).collect()
    } else {
        vec![0; n]
    };
    let frail_arr = Array1::from_vec(frail_vec);
    let frail_view = frail_arr.view();
    let result = survregc1(
        n,
        nvar,
        nstrat,
        false,
        &beta_view,
        dist,
        &strat_arr.view(),
        offsets,
        time1,
        time2,
        &status_arr.view(),
        weights,
        covariates,
        nfrail,
        &frail_view,
    )?;
    let nvar2 = nvar + nstrat;
    let nvar3 = nvar2 + nfrail;
    for i in 0..nvar3.min(result.u.len()) {
        u[i] = result.u[i];
    }
    for i in 0..nvar3.min(hmat.nrows()) {
        for j in 0..nvar2.min(hmat.ncols()) {
            if i < result.imat.ncols() && j < result.imat.nrows() {
                hmat[[i, j]] = -result.imat[[j, i]];
            }
        }
    }
    for i in 0..nvar3.min(jj.nrows()) {
        for j in 0..nvar2.min(jj.ncols()) {
            if i < result.jj.ncols() && j < result.jj.nrows() {
                jj[[i, j]] = result.jj[[j, i]];
            }
        }
    }
    for i in 0..nvar3.min(hdiag.len()) {
        if i < result.imat.nrows() && i < result.imat.ncols() {
            hdiag[i] = -result.imat[[i, i]];
        }
    }
    if nfrail > 0 {
        for i in 0..nfrail.min(result.fdiag.len()) {
            hdiag[i] = -result.fdiag[i];
        }
    }
    Ok(result.loglik)
}
#[allow(clippy::too_many_arguments)]
fn apply_penalties(
    hmat: &mut Array2<f64>,
    jj: &mut Array2<f64>,
    hdiag: &mut Array1<f64>,
    jdiag: &mut Array1<f64>,
    u: &mut Array1<f64>,
    beta: &mut [f64],
    nvar: usize,
    nfrail: usize,
    ptype: PenaltyType,
    pdiag: bool,
) -> Result<f64, Box<dyn std::error::Error>> {
    let ptype_int = match ptype {
        PenaltyType::None => 0,
        PenaltyType::Sparse => 1,
        PenaltyType::Dense => 2,
        PenaltyType::Both => 3,
    };
    if ptype_int == 0 {
        return Ok(0.0);
    }
    let pdiag_int = if pdiag { 1 } else { 0 };
    let whichcase = 0;
    let hmat_slice = hmat.as_slice_mut().ok_or("Failed to get hmat slice")?;
    let jj_slice = jj.as_slice_mut().ok_or("Failed to get jj slice")?;
    let hdiag_slice = hdiag.as_slice_mut().ok_or("Failed to get hdiag slice")?;
    let jdiag_slice = jdiag.as_slice_mut().ok_or("Failed to get jdiag slice")?;
    let u_slice = u.as_slice_mut().ok_or("Failed to get u slice")?;
    const LAMBDA: f64 = 0.1;
    let sparse_penalty = |coef: &[f64]| -> PenaltyResult {
        let n = coef.len();
        let mut first_deriv = vec![0.0; n];
        let mut second_deriv = vec![0.0; n];
        let mut loglik_penalty = 0.0;
        for i in 0..n {
            loglik_penalty += LAMBDA * 0.5 * coef[i].powi(2);
            first_deriv[i] = LAMBDA * coef[i];
            second_deriv[i] = LAMBDA;
        }
        PenaltyResult {
            new_coef: coef.to_vec(),
            first_deriv,
            second_deriv,
            loglik_penalty,
            flags: vec![0; n],
        }
    };
    let dense_penalty = |coef: &[f64]| -> PenaltyResult {
        let n = coef.len();
        let mut first_deriv = vec![0.0; n];
        let mut second_deriv = vec![0.0; n];
        let mut loglik_penalty = 0.0;
        for i in 0..n {
            loglik_penalty += LAMBDA * 0.5 * coef[i].powi(2);
            first_deriv[i] = LAMBDA * coef[i];
            second_deriv[i] = LAMBDA;
        }
        PenaltyResult {
            new_coef: coef.to_vec(),
            first_deriv,
            second_deriv,
            loglik_penalty,
            flags: vec![0; n],
        }
    };
    let params = PenaltyParams {
        whichcase,
        nfrail,
        nvar,
        ptype: ptype_int,
        pdiag: pdiag_int,
    };
    let matrices = MatrixBuffers {
        hmat: hmat_slice,
        JJ: jj_slice,
        hdiag: hdiag_slice,
        jdiag: jdiag_slice,
        u: u_slice,
        beta,
    };
    let mut penalty = 0.0;
    survpenal::survpenal(
        params,
        matrices,
        &mut penalty,
        sparse_penalty,
        dense_penalty,
    );
    Ok(penalty)
}
#[allow(clippy::too_many_arguments)]
fn golden_section_search(
    beta: &[f64],
    newbeta: &[f64],
    _current_loglik: f64,
    n: usize,
    nvar: usize,
    nstrat: usize,
    strata: &[usize],
    offsets: &ArrayView1<f64>,
    time1: &ArrayView1<f64>,
    time2: Option<&ArrayView1<f64>>,
    status: &ArrayView1<f64>,
    weights: &ArrayView1<f64>,
    covariates: &ArrayView2<f64>,
    distribution: &Distribution,
    nfrail: usize,
    fgrp: &[usize],
) -> Result<f64, Box<dyn std::error::Error>> {
    const GOLDEN_RATIO: f64 = 0.6180339887498949;
    const TOL: f64 = 1e-6;
    const MAX_ITER: usize = 50;
    let mut a = 0.0;
    let mut b = 1.0;
    let mut c = b - GOLDEN_RATIO * (b - a);
    let mut d = a + GOLDEN_RATIO * (b - a);
    let mut fc = evaluate_likelihood_at_alpha(
        beta,
        newbeta,
        c,
        n,
        nvar,
        nstrat,
        strata,
        offsets,
        time1,
        time2,
        status,
        weights,
        covariates,
        distribution,
        nfrail,
        fgrp,
    )?;
    let mut fd = evaluate_likelihood_at_alpha(
        beta,
        newbeta,
        d,
        n,
        nvar,
        nstrat,
        strata,
        offsets,
        time1,
        time2,
        status,
        weights,
        covariates,
        distribution,
        nfrail,
        fgrp,
    )?;
    let mut iter = 0;
    while (b - a).abs() > TOL && iter < MAX_ITER {
        if fc > fd {
            b = d;
            d = c;
            fd = fc;
            c = b - GOLDEN_RATIO * (b - a);
            fc = evaluate_likelihood_at_alpha(
                beta,
                newbeta,
                c,
                n,
                nvar,
                nstrat,
                strata,
                offsets,
                time1,
                time2,
                status,
                weights,
                covariates,
                distribution,
                nfrail,
                fgrp,
            )?;
        } else {
            a = c;
            c = d;
            fc = fd;
            d = a + GOLDEN_RATIO * (b - a);
            fd = evaluate_likelihood_at_alpha(
                beta,
                newbeta,
                d,
                n,
                nvar,
                nstrat,
                strata,
                offsets,
                time1,
                time2,
                status,
                weights,
                covariates,
                distribution,
                nfrail,
                fgrp,
            )?;
        }
        iter += 1;
    }
    Ok((a + b) / 2.0)
}
#[allow(clippy::too_many_arguments)]
fn evaluate_likelihood_at_alpha(
    beta: &[f64],
    newbeta: &[f64],
    alpha: f64,
    n: usize,
    nvar: usize,
    nstrat: usize,
    strata: &[usize],
    offsets: &ArrayView1<f64>,
    time1: &ArrayView1<f64>,
    time2: Option<&ArrayView1<f64>>,
    status: &ArrayView1<f64>,
    weights: &ArrayView1<f64>,
    covariates: &ArrayView2<f64>,
    distribution: &Distribution,
    nfrail: usize,
    fgrp: &[usize],
) -> Result<f64, Box<dyn std::error::Error>> {
    let beta_alpha: Vec<f64> = beta
        .iter()
        .zip(newbeta.iter())
        .map(|(b, nb)| b + alpha * (nb - b))
        .collect();
    let nvar2 = nvar + nstrat;
    let nvar3 = nvar2 + nfrail;
    let mut hmat_temp = Array2::zeros((nvar3, nvar2));
    let mut jj_temp = Array2::zeros((nvar3, nvar2));
    let mut u_temp = Array1::zeros(nvar3);
    let mut hdiag_temp = Array1::zeros(nvar3);
    let mut jdiag_temp = Array1::zeros(nfrail);
    calculate_likelihood(
        n,
        nvar,
        nstrat,
        &beta_alpha,
        distribution,
        strata,
        offsets,
        time1,
        time2,
        status,
        weights,
        covariates,
        &mut hmat_temp,
        &mut jj_temp,
        &mut u_temp,
        &mut hdiag_temp,
        &mut jdiag_temp,
        nfrail,
        fgrp,
    )
}
fn calculate_inverse(
    hmat: &Array2<f64>,
    _nvar3: usize,
    _nfrail: usize,
    _hdiag: &Array1<f64>,
    _tol_chol: f64,
) -> Result<Array2<f64>, Box<dyn std::error::Error>> {
    Ok(hmat
        .inv()
        .map_err(|_| "Matrix inversion failed".to_string())?)
}
pub enum Distribution {
    ExtremeValue,
    Logistic,
    Gaussian,
    Custom(Box<dyn Fn(f64) -> f64 + Send + Sync>),
}
#[derive(Debug, Clone, Copy)]
pub enum PenaltyType {
    None,
    Sparse,
    Dense,
    Both,
}
