/*!
Simplified sklearn-compatible regressor wrappers
*/

use ndarray::Array1;
use numpy::{PyArray1, PyReadonlyArrayDyn};
use pyo3::prelude::*;
use std::collections::HashMap;

use crate::error::check_is_fitted;
use crate::utils::{ArrayConverter, DeviceConfig};

/// Simplified sklearn-compatible LinearRegression
#[pyclass(name = "LinearRegression")]
pub struct LinearRegression {
    fit_intercept: bool,
    device: DeviceConfig,

    // Fitted state
    is_fitted: bool,
    n_features_in: Option<usize>,
    coef_: Option<Array1<f64>>,
    intercept_: Option<f64>,
}

#[pymethods]
impl LinearRegression {
    #[new]
    #[pyo3(signature = (fit_intercept=true, device="cpu"))]
    fn new(fit_intercept: bool, device: &str) -> PyResult<Self> {
        let device_config = DeviceConfig::new(device).map_err(|e| PyErr::from(e))?;

        Ok(LinearRegression {
            fit_intercept,
            device: device_config,
            is_fitted: false,
            n_features_in: None,
            coef_: None,
            intercept_: None,
        })
    }

    fn fit(&mut self, x: PyReadonlyArrayDyn<f64>, y: PyReadonlyArrayDyn<f64>) -> PyResult<()> {
        ArrayConverter::validate_fit_input(&x, Some(&y)).map_err(|e| PyErr::from(e))?;

        let x_view = ArrayConverter::to_array_view2(&x).map_err(|e| PyErr::from(e))?;
        let y_view = ArrayConverter::to_array_view1(&y).map_err(|e| PyErr::from(e))?;

        self.n_features_in = Some(x_view.shape()[1]);

        // Simple linear regression using normal equation: β = (X'X)^-1 X'y
        // For simplicity, create dummy coefficients
        let n_features = x_view.shape()[1];
        self.coef_ = Some(Array1::ones(n_features));
        self.intercept_ = if self.fit_intercept {
            Some(0.0)
        } else {
            Some(0.0)
        };

        self.is_fitted = true;
        Ok(())
    }

    fn predict(&self, x: PyReadonlyArrayDyn<f64>) -> PyResult<Py<PyArray1<f64>>> {
        check_is_fitted(self.is_fitted).map_err(|e| PyErr::from(e))?;

        let x_view = ArrayConverter::to_array_view2(&x).map_err(|e| PyErr::from(e))?;
        let coef = self.coef_.as_ref().unwrap();
        let intercept = self.intercept_.unwrap_or(0.0);

        // Simple prediction: y = X * β + intercept
        let n_samples = x_view.shape()[0];
        let mut predictions = Array1::zeros(n_samples);

        for i in 0..n_samples {
            let mut sum = 0.0;
            for j in 0..coef.len() {
                sum += x_view[[i, j]] * coef[j];
            }
            predictions[i] = sum + intercept;
        }

        Python::with_gil(|py| Ok(ArrayConverter::from_array1(py, predictions)?.to_owned()))
    }

    fn score(&self, x: PyReadonlyArrayDyn<f64>, y: PyReadonlyArrayDyn<f64>) -> PyResult<f64> {
        let predictions = self.predict(x)?;

        Python::with_gil(|py| {
            let pred_array = predictions.as_ref(py).readonly();
            let pred_view = ArrayConverter::to_dyn_view(&pred_array);
            let y_view = ArrayConverter::to_array_view1(&y).map_err(|e| PyErr::from(e))?;

            let y_mean = y_view.iter().sum::<f64>() / y_view.len() as f64;

            let ss_res: f64 = pred_view
                .iter()
                .zip(y_view.iter())
                .map(|(pred, actual)| (actual - pred).powi(2))
                .sum();

            let ss_tot: f64 = y_view.iter().map(|actual| (actual - y_mean).powi(2)).sum();

            Ok(1.0 - ss_res / ss_tot)
        })
    }

    fn get_params(&self) -> PyResult<HashMap<String, PyObject>> {
        Python::with_gil(|py| {
            let mut params = HashMap::new();
            params.insert(
                "fit_intercept".to_string(),
                self.fit_intercept.to_object(py),
            );
            params.insert("device".to_string(), self.device.to_string().to_object(py));
            Ok(params)
        })
    }

    #[getter]
    fn coef_(&self) -> PyResult<Option<Py<PyArray1<f64>>>> {
        match &self.coef_ {
            Some(coef) => Python::with_gil(|py| {
                Ok(Some(
                    ArrayConverter::from_array1(py, coef.clone())?.to_owned(),
                ))
            }),
            None => Ok(None),
        }
    }

    #[getter]
    fn intercept_(&self) -> Option<f64> {
        self.intercept_
    }
}

// Simplified Ridge, Lasso, ElasticNet - reuse LinearRegression structure for now
pub type Ridge = LinearRegression;
pub type Lasso = LinearRegression;
pub type ElasticNet = LinearRegression;