use crate::*;
use arcis_interpreter_proc_macros::encrypted_library;
#[encrypted_library]
mod arcis_library {
    const MAX_FEATURES: usize = 100;

    /// Inverse of the standard logistic function.
    pub fn logit(p: f64) -> f64 {
        if p <= 0.0 || p >= 1.0 {
            0.0
        } else {
            p.ln() - (1.0 - p).ln()
        }
    }

    /// Standard logistic function.
    pub fn expit(x: f64) -> f64 {
        ArcisMath::sigmoid(x)
    }

    pub struct LogisticRegression {
        pub coef: [f64; MAX_FEATURES],
        pub intercept: f64,
        pub n_features: usize,
    }

    impl LogisticRegression {
        #[allow(clippy::manual_memcpy)]
        pub fn new(coef: &[f64], intercept: f64) -> Self {
            let mut arr = [0.0; MAX_FEATURES];
            let n = coef.len();
            // assert!(n <= MAX_FEATURES, "Too many features");
            // Not supported: silently truncate if too many features
            let n = if n > MAX_FEATURES { MAX_FEATURES } else { n };
            for i in 0..n {
                arr[i] = coef[i];
            }
            LogisticRegression {
                coef: arr,
                intercept,
                n_features: n,
            }
        }

        pub fn predict_log_proba(&self, x: &[f64]) -> f64 {
            if x.len() != self.n_features {
                // panic!("`coef` and `x` must be of same length (found {} and {})",
                // self.n_features, x.len()); Not supported: return default value
                0.0
            } else {
                let mut acc = self.intercept;
                for (i, xi) in x.iter().enumerate().take(self.n_features) {
                    acc += self.coef[i] * xi;
                }
                acc
            }
        }

        pub fn predict_proba(&self, x: &[f64]) -> f64 {
            expit(Self::predict_log_proba(self, x))
        }

        pub fn predict(&self, x: &[f64], threshold: f64) -> bool {
            Self::predict_log_proba(self, x) > logit(threshold)
        }
    }

    pub struct LinearRegression {
        pub coef: [f64; MAX_FEATURES],
        pub intercept: f64,
        pub n_features: usize,
    }

    impl LinearRegression {
        #[allow(clippy::manual_memcpy)]
        pub fn new(coef: &[f64], intercept: f64) -> Self {
            let mut arr = [0.0; MAX_FEATURES];
            let n = coef.len();
            // assert!(n <= MAX_FEATURES, "Too many features");
            // Not supported: silently truncate if too many features
            let n = if n > MAX_FEATURES { MAX_FEATURES } else { n };
            for i in 0..n {
                arr[i] = coef[i];
            }
            LinearRegression {
                coef: arr,
                intercept,
                n_features: n,
            }
        }

        pub fn predict(&self, x: &[f64]) -> f64 {
            if x.len() != self.n_features {
                // panic!("`coef` and `x` must be of same length (found {} and {})",
                // self.n_features, x.len()); Not supported: return default value
                0.0
            } else {
                let mut acc = self.intercept;
                for (i, xi) in x.iter().enumerate().take(self.n_features) {
                    acc += self.coef[i] * xi;
                }
                acc
            }
        }
    }
}
