#![allow(dead_code)]
use pyo3::prelude::*;
use thiserror::Error;

#[derive(Error, Debug)]
pub enum SurvivalModelError {
    #[error(
        "Insufficient data: expected {expected_people} observations with {expected_vars} covariates, got {actual_people} observations with {actual_vars} covariates"
    )]
    InsufficientData {
        expected_people: usize,
        expected_vars: usize,
        actual_people: usize,
        actual_vars: usize,
    },
}

struct SurvivalData {
    start_time: f64,
    stop_time: f64,
    status: i32,
}

struct Covariate {
    values: Vec<f64>,
}

#[pyclass]
struct SurvivalModel {
    maxiter: usize,
    nused: usize,
    nvar: usize,
    yy: Vec<SurvivalData>,
    covar: Vec<Covariate>,
    strata: Vec<usize>,
    sort: Vec<(usize, usize)>,
    offset: Vec<f64>,
    weights: Vec<f64>,
    eps: f64,
    tolerch: f64,
    method: u8,
    ptype: u8,
    nfrail: usize,
    frail: Vec<f64>,
    fbeta: Vec<f64>,
    pdiag: u8,

    means: Vec<f64>,
    beta: Vec<f64>,
    u: Vec<f64>,
    imat: Vec<Vec<f64>>,
    loglik: f64,
    flag: i32,
    fdiag: Vec<f64>,
    jmat: Vec<Vec<f64>>,
    expect: Vec<f64>,
}

impl SurvivalModel {
    #[allow(clippy::too_many_arguments)]
    fn new(
        maxiter: usize,
        nused: usize,
        nvar: usize,
        yy: Vec<SurvivalData>,
        covar: Vec<Covariate>,
        strata: Vec<usize>,
        sort: Vec<(usize, usize)>,
        offset: Vec<f64>,
        weights: Vec<f64>,
        eps: f64,
        tolerch: f64,
        method: u8,
        ptype: u8,
        nfrail: usize,
        frail: Vec<f64>,
        fbeta: Vec<f64>,
        pdiag: u8,
    ) -> Self {
        Self {
            maxiter,
            nused,
            nvar,
            yy,
            covar,
            strata,
            sort,
            offset,
            weights,
            eps,
            tolerch,
            method,
            ptype,
            nfrail,
            frail,
            fbeta,
            pdiag,

            means: vec![0.0; nvar],
            beta: vec![0.0; nvar],
            u: vec![0.0; nvar],
            imat: vec![vec![0.0; nvar]; nvar],
            loglik: 0.0,
            flag: 0,
            fdiag: vec![0.0; nfrail + nvar],
            jmat: vec![vec![0.0; nvar]; nvar],
            expect: vec![0.0; nused],
        }
    }
    pub fn gfit5a(&mut self) -> Result<(), SurvivalModelError> {
        if self.yy.len() != self.nused || self.covar.len() != self.nvar {
            return Err(SurvivalModelError::InsufficientData {
                expected_people: self.nused,
                expected_vars: self.nvar,
                actual_people: self.yy.len(),
                actual_vars: self.covar.len(),
            });
        }

        self.means = vec![0.0; self.nvar];
        self.beta = vec![0.0; self.nvar];

        for covariate in &self.covar {
            for (i, &val) in covariate.values.iter().enumerate() {
                self.means[i] += val / self.nused as f64;
            }
        }

        self.u = vec![0.0; self.nvar];
        self.imat = vec![vec![0.0; self.nvar]; self.nvar];

        for i in 0..self.nvar {
            self.u[i] = 0.0;
            for j in 0..self.nvar {
                self.imat[i][j] = if i == j { 1.0 } else { 0.0 };
            }
        }

        Ok(())
    }
    pub fn agfit5b(&mut self) {
        if self.maxiter == 0 || self.nused == 0 || self.nvar == 0 {
            return;
        }

        let mut iter = 0;
        while iter < self.maxiter {
            self.update_beta();
            self.calculate_score_vector();
            self.update_infromation_matrix();
            iter += 1;
            if self.converged() {
                break;
            }
        }
    }
    pub fn update_beta(&mut self) {
        let mut denom = 0.0;
        for i in 0..self.nvar {
            denom += self.imat[i][i];
        }
        for i in 0..self.nvar {
            self.beta[i] += self.u[i] / denom;
        }
    }
    pub fn calculate_score_vector(&mut self) {
        for i in 0..self.nvar {
            self.u[i] = 0.0;
            for j in 0..self.nvar {
                self.u[i] += self.imat[i][j] * self.beta[j];
            }
        }
    }
    pub fn update_infromation_matrix(&mut self) {
        for i in 0..self.nvar {
            for j in 0..self.nvar {
                self.imat[i][j] = 0.0;
                for k in 0..self.nvar {
                    self.imat[i][j] += self.imat[i][k] * self.imat[k][j];
                }
            }
        }
    }

    pub fn converged(&self) -> bool {
        let mut max = 0.0;
        for i in 0..self.nvar {
            let abs = self.u[i].abs();
            if abs > max {
                max = abs;
            }
        }
        max < self.eps
    }
}

#[pymodule]
#[pyo3(name = "pySurvivalModel")]
fn py_survival_model(_py: Python, m: Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<SurvivalModel>()?;
    Ok(())
}
