//! Logistic Regression implementation with GPU acceleration

use crate::error::Result;
use rtx_tensor::Tensor;

/// Logistic Regression for binary classification
#[derive(Debug, Clone)]
pub struct LogisticRegression {
    /// Regularization strength (C = 1/alpha)
    c: f32,
    /// Maximum number of iterations
    max_iter: usize,
    /// Learning rate for gradient descent
    learning_rate: f32,
    /// Tolerance for convergence
    tolerance: f32,
    /// Model weights (coefficients)
    weights: Option<Tensor>,
    /// Model intercept
    intercept: Option<Tensor>,
    /// Whether to fit intercept
    fit_intercept: bool,
}

impl LogisticRegression {
    /// Create a new Logistic regression model
    pub fn new() -> Self {
        Self {
            c: 1.0,
            max_iter: 1000,
            learning_rate: 0.01,
            tolerance: 1e-6,
            weights: None,
            intercept: None,
            fit_intercept: true,
        }
    }

    /// Set regularization strength (C = 1/alpha)
    pub fn c(mut self, c: f32) -> Self {
        self.c = c;
        self
    }

    /// Set maximum iterations
    pub fn max_iter(mut self, max_iter: usize) -> Self {
        self.max_iter = max_iter;
        self
    }

    /// Set learning rate
    pub fn learning_rate(mut self, lr: f32) -> Self {
        self.learning_rate = lr;
        self
    }

    /// Set whether to fit intercept
    pub fn fit_intercept(mut self, fit_intercept: bool) -> Self {
        self.fit_intercept = fit_intercept;
        self
    }

    /// Sigmoid activation function
    fn sigmoid(&self, z: &Tensor) -> Result<Tensor> {
        // sigmoid(z) = 1 / (1 + exp(-z))
        let neg_z = z.neg()?;
        let exp_neg_z = neg_z.exp()?;
        let one_plus_exp = exp_neg_z.add_scalar(1.0)?;
        one_plus_exp.reciprocal().map_err(std::convert::Into::into)
    }

    /// Fit the Logistic regression model using gradient descent
    pub fn fit(&mut self, x: &Tensor, y: &Tensor) -> Result<&mut Self> {
        let (n_samples, n_features) = (x.shape().dims()[0], x.shape().dims()[1]);

        // Add intercept column if needed
        let x_design = if self.fit_intercept {
            let ones = Tensor::ones(vec![n_samples, 1], x.device())?;
            Tensor::cat(&[x.clone(), ones], 1)?
        } else {
            x.clone()
        };

        let n_params = x_design.shape().dims()[1];

        // Initialize weights to small random values
        let mut weights = Tensor::randn(&[n_params, 1], x.device())?.mul_scalar(0.01)?;

        let alpha = 1.0 / self.c; // Convert C to alpha
        let lr = self.learning_rate;
        let n_samples_f = n_samples as f32;

        // Gradient descent training
        for _iter in 0..self.max_iter {
            // Forward pass
            let z = x_design.matmul(&weights)?;
            let predictions = self.sigmoid(&z)?;

            // Calculate loss (cross-entropy + L2 regularization)
            let y_reshaped = y.view([n_samples, 1])?;
            let diff = predictions.subtract(&y_reshaped)?;

            // Gradient calculation
            let gradient = x_design.transpose(0, 1)?.matmul(&diff)?;
            let gradient = gradient.div_scalar(n_samples_f)?;

            // Add L2 regularization to gradient (except bias term if fitting intercept)
            if alpha > 0.0 {
                let reg_weights = if self.fit_intercept {
                    // Don't regularize bias term
                    let weight_part = weights.slice(0, 0, n_features)?;
                    let reg_part = weight_part.mul_scalar(alpha)?;
                    let zero_bias = Tensor::zeros(vec![1, 1], x.device())?;
                    Tensor::cat(&[reg_part, zero_bias], 0)?
                } else {
                    weights.mul_scalar(alpha)?
                };
                let _gradient = gradient.add(&reg_weights)?;
            }

            // Update weights
            let update = gradient.mul_scalar(lr)?;
            weights = weights.subtract(&update)?;

            // Check convergence (simplified - using sum of absolute gradients)
            let grad_sum = gradient.abs()?.sum(None)?;
            if grad_sum.to_scalar::<f32>()? < self.tolerance {
                break;
            }
        }

        if self.fit_intercept {
            // Split weights and intercept
            let weights_slice = weights.slice(0, 0, n_features)?;
            let intercept_slice = weights.slice(0, n_features, n_features + 1)?;

            self.weights = Some(weights_slice);
            self.intercept = Some(intercept_slice);
        } else {
            self.weights = Some(weights);
            self.intercept = None;
        }

        Ok(self)
    }

    /// Predict class probabilities
    pub fn predict_proba(&self, x: &Tensor) -> Result<Tensor> {
        let weights = self
            .weights
            .as_ref()
            .ok_or_else(|| crate::error::MLError::ModelNotFitted)?;

        let mut z = x.matmul(weights)?;

        if let Some(intercept) = &self.intercept {
            z = z.add(intercept)?;
        }

        self.sigmoid(&z)
    }

    /// Make binary predictions (0 or 1)
    pub fn predict(&self, x: &Tensor) -> Result<Tensor> {
        let probabilities = self.predict_proba(x)?;
        // Convert probabilities to binary predictions (threshold = 0.5)
        probabilities
            .ge_scalar(0.5)
            .map_err(std::convert::Into::into)
    }

    /// Get model coefficients
    pub fn coef(&self) -> Option<&Tensor> {
        self.weights.as_ref()
    }

    /// Get model intercept
    pub fn intercept(&self) -> Option<&Tensor> {
        self.intercept.as_ref()
    }
}

impl Default for LogisticRegression {
    fn default() -> Self {
        Self::new()
    }
}