#![allow(clippy::redundant_closure)]
use crate::regression::survregc1::{SurvivalDist, survregc1};
use crate::utilities::matrix::cholesky_solve;
use ndarray::{Array1, Array2, ArrayView1};
use pyo3::prelude::*;
#[derive(Debug, Clone)]
#[pyclass]
pub struct SurvivalFit {
    #[pyo3(get)]
    pub coefficients: Vec<f64>,
    #[pyo3(get)]
    pub iterations: usize,
    #[pyo3(get)]
    pub variance_matrix: Vec<Vec<f64>>,
    #[pyo3(get)]
    pub log_likelihood: f64,
    #[pyo3(get)]
    pub convergence_flag: i32,
    #[pyo3(get)]
    pub score_vector: Vec<f64>,
}
#[allow(clippy::too_many_arguments)]
fn calculate_likelihood(
    n: usize,
    nvar: usize,
    nstrat: usize,
    beta: &[f64],
    distribution: &DistributionType,
    strata: &[usize],
    offsets: &Array1<f64>,
    time1: &ArrayView1<f64>,
    time2: Option<&ArrayView1<f64>>,
    status: &ArrayView1<f64>,
    weights: &Array1<f64>,
    covariates: &Array2<f64>,
    imat: &mut Array2<f64>,
    jj: &mut Array2<f64>,
    u: &mut Array1<f64>,
) -> Result<f64, Box<dyn std::error::Error>> {
    let dist = match distribution {
        DistributionType::ExtremeValue => SurvivalDist::ExtremeValue,
        DistributionType::Logistic => SurvivalDist::Logistic,
        DistributionType::Gaussian => SurvivalDist::Gaussian,
        DistributionType::Weibull => SurvivalDist::Weibull,
        DistributionType::LogNormal => SurvivalDist::LogNormal,
    };
    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 frail_arr = Array1::from_vec(vec![0i32; n]);
    let nvar2 = nvar + nstrat;
    let result = survregc1(
        n,
        nvar,
        nstrat,
        false,
        &beta_arr.view(),
        dist,
        &strat_arr.view(),
        &offsets.view(),
        time1,
        time2,
        &status_arr.view(),
        &weights.view(),
        &covariates.view(),
        0,
        &frail_arr.view(),
    )?;
    for i in 0..nvar2.min(u.len()) {
        if i < result.u.len() {
            u[i] = result.u[i];
        }
    }
    for i in 0..nvar2.min(imat.nrows()) {
        for j in 0..nvar2.min(imat.ncols()) {
            if i < result.imat.nrows() && j < result.imat.ncols() {
                imat[[i, j]] = -result.imat[[i, j]];
            }
        }
    }
    for i in 0..nvar2.min(jj.nrows()) {
        for j in 0..nvar2.min(jj.ncols()) {
            if i < result.jj.nrows() && j < result.jj.ncols() {
                jj[[i, j]] = result.jj[[i, j]];
            }
        }
    }
    Ok(result.loglik)
}
fn check_convergence(old: f64, new: f64, eps: f64) -> bool {
    (1.0 - new / old).abs() <= eps || (old - new).abs() <= eps
}
fn adjust_strata(newbeta: &mut [f64], beta: &[f64], nvar: usize, nstrat: usize) {
    for i in 0..nstrat {
        let idx = nvar + i;
        if beta[idx] - newbeta[idx] > 1.1 {
            newbeta[idx] = beta[idx] - 1.1;
        }
    }
}
fn calculate_variance_matrix(
    imat: Array2<f64>,
    _nvar2: usize,
    _tol_chol: f64,
) -> Result<Array2<f64>, Box<dyn std::error::Error>> {
    use ndarray_linalg::Inverse;
    if imat.nrows() == 0 || imat.ncols() == 0 {
        return Ok(imat);
    }
    let max_val = imat.iter().map(|&x| x.abs()).fold(0.0f64, f64::max);
    if max_val < 1e-10 {
        return Ok(imat);
    }
    match imat.inv() {
        Ok(inv) => Ok(inv),
        Err(_) => Ok(imat),
    }
}
#[derive(Debug, Clone, Copy)]
#[pyclass]
pub enum DistributionType {
    #[pyo3(name = "extreme_value")]
    ExtremeValue,
    #[pyo3(name = "logistic")]
    Logistic,
    #[pyo3(name = "gaussian")]
    Gaussian,
    #[pyo3(name = "weibull")]
    Weibull,
    #[pyo3(name = "lognormal")]
    LogNormal,
}
#[pyfunction]
#[allow(clippy::too_many_arguments)]
pub fn survreg(
    time: Vec<f64>,
    status: Vec<f64>,
    covariates: Vec<Vec<f64>>,
    weights: Option<Vec<f64>>,
    offsets: Option<Vec<f64>>,
    initial_beta: Option<Vec<f64>>,
    strata: Option<Vec<usize>>,
    distribution: Option<&str>,
    max_iter: Option<usize>,
    eps: Option<f64>,
    tol_chol: Option<f64>,
) -> PyResult<SurvivalFit> {
    let n = time.len();
    if status.len() != n {
        return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
            "time and status must have the same length",
        ));
    }
    let nvar = if !covariates.is_empty() {
        covariates[0].len()
    } else {
        0
    };
    if !covariates.is_empty() && covariates.len() != n {
        return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
            "covariates must have the same number of rows as time",
        ));
    }
    let weights = weights.unwrap_or_else(|| vec![1.0; n]);
    let offsets = offsets.unwrap_or_else(|| vec![0.0; n]);
    let strata = strata.unwrap_or_else(|| vec![0; n]);
    let max_iter = max_iter.unwrap_or(20);
    let eps = eps.unwrap_or(1e-5);
    let tol_chol = tol_chol.unwrap_or(1e-9);
    let dist_type = match distribution {
        Some("logistic") | Some("Logistic") => DistributionType::Logistic,
        Some("gaussian") | Some("Gaussian") | Some("normal") | Some("Normal") => {
            DistributionType::Gaussian
        }
        Some("weibull") | Some("Weibull") => DistributionType::Weibull,
        Some("lognormal") | Some("LogNormal") | Some("lognorm") | Some("LogNorm") => {
            DistributionType::LogNormal
        }
        _ => DistributionType::ExtremeValue,
    };
    let nstrat = if strata.is_empty() {
        1
    } else {
        strata.iter().max().copied().unwrap_or(0) + 1
    };
    let initial_beta = initial_beta.unwrap_or_else(|| vec![0.0; nvar + nstrat]);
    let y = {
        let mut y_data = Vec::new();
        for i in 0..n {
            y_data.push(vec![time[i], status[i]]);
        }
        Array2::from_shape_vec((n, 2), y_data.into_iter().flatten().collect())
            .map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(format!("{}", e)))?
    };
    let cov_array = if nvar > 0 {
        let flat: Vec<f64> = covariates.into_iter().flatten().collect();
        let temp = Array2::from_shape_vec((n, nvar), flat)
            .map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(format!("{}", e)))?;
        temp.t().to_owned()
    } else {
        Array2::zeros((0, n))
    };
    let weights_arr = Array1::from_vec(weights);
    let offsets_arr = Array1::from_vec(offsets);
    let result = compute_survreg(
        max_iter,
        nvar,
        &y,
        &cov_array,
        &weights_arr,
        &offsets_arr,
        initial_beta,
        nstrat,
        &strata,
        eps,
        tol_chol,
        dist_type,
    )
    .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(format!("{}", e)))?;
    let variance_matrix = result
        .variance_matrix
        .outer_iter()
        .map(|row| row.iter().cloned().collect())
        .collect();
    Ok(SurvivalFit {
        coefficients: result.coefficients,
        iterations: result.iterations,
        variance_matrix,
        log_likelihood: result.log_likelihood,
        convergence_flag: result.convergence_flag,
        score_vector: result.score_vector,
    })
}
#[allow(clippy::too_many_arguments)]
fn compute_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: DistributionType,
) -> Result<SurvivalFitComputed, Box<dyn std::error::Error>> {
    let n = y.nrows();
    let ny = y.ncols();
    let nvar2 = nvar + nstrat;
    let mut imat = Array2::zeros((nvar2, nvar2));
    let mut jj = Array2::zeros((nvar2, nvar2));
    let mut u = Array1::zeros(nvar2);
    let mut newbeta = beta.clone();
    let mut usave = Array1::zeros(nvar2);
    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 = time1_arr.view();
    let status = status_arr.view();
    let time2_view: Option<ArrayView1<f64>> = time2_arr.as_ref().map(|v| v.view());
    let mut loglik = calculate_likelihood(
        n,
        nvar,
        nstrat,
        &beta,
        &distribution,
        strata,
        offsets,
        &time1,
        time2_view.as_ref(),
        &status,
        weights,
        covariates,
        &mut imat,
        &mut jj,
        &mut u,
    )?;
    usave.assign(&u);
    let mut iter = 0;
    let mut halving = 0;
    while iter < max_iter {
        let chol_result = cholesky_solve(&imat, &u, tol_chol);
        let delta = match chol_result {
            Ok(d) => d,
            Err(_) => cholesky_solve(&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,
            &time1,
            time2_view.as_ref(),
            &status,
            weights,
            covariates,
            &mut imat,
            &mut jj,
            &mut u,
        )?;
        if check_convergence(loglik, newlik, eps) && halving == 0 {
            loglik = newlik;
            beta = newbeta.clone();
            iter += 1;
            break;
        }
        if newlik.is_nan() || newlik < loglik {
            halving += 1;
            newbeta
                .iter_mut()
                .zip(&beta)
                .for_each(|(nb, b)| *nb = (*nb + 2.0 * b) / 3.0);
            if halving == 1 {
                adjust_strata(&mut newbeta, &beta, nvar, nstrat);
            }
        } else {
            halving = 0;
            loglik = newlik;
            beta = newbeta.clone();
        }
        iter += 1;
    }
    let converged = iter < max_iter;
    let convergence_flag = if converged { 0 } else { -1 };
    let variance = calculate_variance_matrix(imat, nvar2, tol_chol)?;
    Ok(SurvivalFitComputed {
        coefficients: beta,
        iterations: iter,
        variance_matrix: variance,
        log_likelihood: loglik,
        convergence_flag,
        score_vector: usave.to_vec(),
    })
}
pub(crate) struct SurvivalFitComputed {
    coefficients: Vec<f64>,
    iterations: usize,
    variance_matrix: Array2<f64>,
    log_likelihood: f64,
    convergence_flag: i32,
    score_vector: Vec<f64>,
}
#[pymodule]
#[pyo3(name = "survreg")]
fn survreg_module(_py: Python, m: Bound<'_, PyModule>) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(survreg, &m)?)?;
    m.add_class::<SurvivalFit>()?;
    m.add_class::<DistributionType>()?;
    Ok(())
}
