/*!
sklearn-compatible regressor wrappers
*/

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

use rtx_ml_classic::linear::{LinearRegressor, ElasticNetRegressor, RidgeRegressor, LassoRegressor};
use rtx_ml_classic::linear::{LinearConfig, ElasticNetConfig, RidgeConfig, LassoConfig};

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

/// sklearn-compatible LinearRegression
#[pyclass(name = "LinearRegression")]
pub struct LinearRegression {
    model: Option<LinearRegressor>,
    fit_intercept: bool,
    normalize: 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, normalize=false, device="cpu"))]
    fn new(fit_intercept: bool, normalize: bool, device: &str) -> PyResult<Self> {
        let device_config = DeviceConfig::new(device).map_err(|e| PyErr::from(e))?;

        Ok(LinearRegression {
            model: None,
            fit_intercept,
            normalize,
            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]);

        let config = LinearConfig {
            fit_intercept: self.fit_intercept,
            normalize: self.normalize,
        };

        let mut model = LinearRegressor::new(config);
        model.fit(x_view.to_owned(), y_view.to_owned())
        .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
        format!("Training failed: {}", e)
        ))?;

        // Extract coefficients
        self.coef_ = Some(model.coefficients()
        .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
        format!("Failed to get coefficients: {}", e)
        ))?);

        self.intercept_ = if self.fit_intercept {
            Some(model.intercept()
            .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
            format!("Failed to get intercept: {}", e)
            ))?)
            } else {
                Some(0.0)
            };

            self.model = Some(model);
            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))?;

            ArrayConverter::validate_predict_input(&x, self.n_features_in)
            .map_err(|e| PyErr::from(e))?;

            let x_view = ArrayConverter::to_array_view2(&x).map_err(|e| PyErr::from(e))?;
            let predictions = self.model.as_ref().unwrap().predict(x_view.to_owned())
            .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
            format!("Prediction failed: {}", e)
            ))?;

            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_array_view1(&pred_array)
                .map_err(|e| PyErr::from(e))?;
                let y_view = ArrayConverter::to_array_view1(&y)
                .map_err(|e| PyErr::from(e))?;

                // Calculate R² score
                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();

                let r2 = 1.0 - ss_res / ss_tot;
                Ok(r2)
            })
        }

        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("normalize".to_string(), self.normalize.to_object(py));
                params.insert("device".to_string(), self.device.to_string().to_object(py));
                Ok(params)
            })
        }

        fn set_params(&mut self, params: &PyDict) -> PyResult<()> {
            let valid_params = vec!["fit_intercept", "normalize", "device"];
            let validated = ParamValidator::validate_params(params, &valid_params)
            .map_err(|e| PyErr::from(e))?;

            Python::with_gil(|py| {
                if validated.contains_key("fit_intercept") {
                    self.fit_intercept = ParamValidator::get_param(&validated, "fit_intercept", self.fit_intercept, py)
                    .map_err(|e| PyErr::from(e))?;
                }

                if validated.contains_key("normalize") {
                    self.normalize = ParamValidator::get_param(&validated, "normalize", self.normalize, py)
                    .map_err(|e| PyErr::from(e))?;
                }

                if validated.contains_key("device") {
                    let device_str: String = ParamValidator::get_param(&validated, "device", self.device.to_string(), py)
                    .map_err(|e| PyErr::from(e))?;
                    self.device = DeviceConfig::new(&device_str)
                    .map_err(|e| PyErr::from(e))?;
                }

                // Reset fitted state if parameters changed
                if !validated.is_empty() {
                    self.is_fitted = false;
                    self.model = None;
                    self.coef_ = None;
                    self.intercept_ = None;
                }

                Ok(())
            })
        }

        #[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_
        }

        #[getter]
        fn n_features_in_(&self) -> Option<usize> {
            self.n_features_in
        }
    }

    /// sklearn-compatible Ridge regression
    #[pyclass(name = "Ridge")]
    pub struct Ridge {
        model: Option<RidgeRegressor>,
        alpha: f64,
        fit_intercept: bool,
        normalize: bool,
        max_iter: Option<usize>,
        device: DeviceConfig,

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

    #[pymethods]
    impl Ridge {
        #[new]
        #[pyo3(signature = (
        alpha=1.0,
        fit_intercept=true,
        normalize=false,
        max_iter=None,
        device="cpu"
        ))]
        fn new(
        alpha: f64,
        fit_intercept: bool,
        normalize: bool,
        max_iter: Option<usize>,
        device: &str,
        ) -> PyResult<Self> {
            if alpha < 0.0 {
                return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
                "alpha must be >= 0"
                ));
            }

            let device_config = DeviceConfig::new(device).map_err(|e| PyErr::from(e))?;

            Ok(Ridge {
                model: None,
                alpha,
                fit_intercept,
                normalize,
                max_iter,
                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]);

            let config = RidgeConfig {
                alpha: self.alpha,
                fit_intercept: self.fit_intercept,
                normalize: self.normalize,
                max_iter: self.max_iter,
            };

            let mut model = RidgeRegressor::new(config);
            model.fit(x_view.to_owned(), y_view.to_owned())
            .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
            format!("Training failed: {}", e)
            ))?;

            self.coef_ = Some(model.coefficients()
            .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
            format!("Failed to get coefficients: {}", e)
            ))?);

            self.intercept_ = if self.fit_intercept {
                Some(model.intercept()
                .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
                format!("Failed to get intercept: {}", e)
                ))?)
                } else {
                    Some(0.0)
                };

                self.model = Some(model);
                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 predictions = self.model.as_ref().unwrap().predict(x_view.to_owned())
                .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
                format!("Prediction failed: {}", e)
                ))?;

                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_array_view1(&pred_array)
                    .map_err(|e| PyErr::from(e))?;
                    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)
                })
            }

            #[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_
            }
        }

        /// sklearn-compatible Lasso regression
        #[pyclass(name = "Lasso")]
        pub struct Lasso {
            model: Option<LassoRegressor>,
            alpha: f64,
            fit_intercept: bool,
            normalize: bool,
            max_iter: usize,
            tol: f64,
            device: DeviceConfig,

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

        #[pymethods]
        impl Lasso {
            #[new]
            #[pyo3(signature = (
            alpha=1.0,
            fit_intercept=true,
            normalize=false,
            max_iter=1000,
            tol=1e-4,
            device="cpu"
            ))]
            fn new(
            alpha: f64,
            fit_intercept: bool,
            normalize: bool,
            max_iter: usize,
            tol: f64,
            device: &str,
            ) -> PyResult<Self> {
                if alpha < 0.0 {
                    return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>("alpha must be >= 0"));
                }
                if max_iter < 1 {
                    return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>("max_iter must be >= 1"));
                }
                if tol <= 0.0 {
                    return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>("tol must be > 0"));
                }

                let device_config = DeviceConfig::new(device).map_err(|e| PyErr::from(e))?;

                Ok(Lasso {
                    model: None,
                    alpha,
                    fit_intercept,
                    normalize,
                    max_iter,
                    tol,
                    device: device_config,
                    is_fitted: false,
                    n_features_in: None,
                    coef_: None,
                    intercept_: None,
                    n_iter_: 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]);

                let config = LassoConfig {
                    alpha: self.alpha,
                    fit_intercept: self.fit_intercept,
                    normalize: self.normalize,
                    max_iter: self.max_iter,
                    tol: self.tol,
                };

                let mut model = LassoRegressor::new(config);
                model.fit(x_view.to_owned(), y_view.to_owned())
                .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
                format!("Training failed: {}", e)
                ))?;

                self.coef_ = Some(model.coefficients()
                .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
                format!("Failed to get coefficients: {}", e)
                ))?);

                self.intercept_ = if self.fit_intercept {
                    Some(model.intercept()
                    .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
                    format!("Failed to get intercept: {}", e)
                    ))?)
                    } else {
                        Some(0.0)
                    };

                    self.n_iter_ = Some(model.n_iterations()
                    .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
                    format!("Failed to get iteration count: {}", e)
                    ))?);

                    self.model = Some(model);
                    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 predictions = self.model.as_ref().unwrap().predict(x_view.to_owned())
                    .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
                    format!("Prediction failed: {}", e)
                    ))?;

                    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_array_view1(&pred_array)
                        .map_err(|e| PyErr::from(e))?;
                        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)
                    })
                }

                #[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_
                }

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

            /// sklearn-compatible ElasticNet regression
            #[pyclass(name = "ElasticNet")]
            pub struct ElasticNet {
                model: Option<ElasticNetRegressor>,
                alpha: f64,
                l1_ratio: f64,
                fit_intercept: bool,
                normalize: bool,
                max_iter: usize,
                tol: f64,
                device: DeviceConfig,

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

            #[pymethods]
            impl ElasticNet {
                #[new]
                #[pyo3(signature = (
                alpha=1.0,
                l1_ratio=0.5,
                fit_intercept=true,
                normalize=false,
                max_iter=1000,
                tol=1e-4,
                device="cpu"
                ))]
                fn new(
                alpha: f64,
                l1_ratio: f64,
                fit_intercept: bool,
                normalize: bool,
                max_iter: usize,
                tol: f64,
                device: &str,
                ) -> PyResult<Self> {
                    if alpha < 0.0 {
                        return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>("alpha must be >= 0"));
                    }
                    if l1_ratio < 0.0 || l1_ratio > 1.0 {
                        return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>("l1_ratio must be in [0, 1]"));
                    }
                    if max_iter < 1 {
                        return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>("max_iter must be >= 1"));
                    }
                    if tol <= 0.0 {
                        return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>("tol must be > 0"));
                    }

                    let device_config = DeviceConfig::new(device).map_err(|e| PyErr::from(e))?;

                    Ok(ElasticNet {
                        model: None,
                        alpha,
                        l1_ratio,
                        fit_intercept,
                        normalize,
                        max_iter,
                        tol,
                        device: device_config,
                        is_fitted: false,
                        n_features_in: None,
                        coef_: None,
                        intercept_: None,
                        n_iter_: 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]);

                    let config = ElasticNetConfig {
                        alpha: self.alpha,
                        l1_ratio: self.l1_ratio,
                        fit_intercept: self.fit_intercept,
                        normalize: self.normalize,
                        max_iter: self.max_iter,
                        tol: self.tol,
                    };

                    let mut model = ElasticNetRegressor::new(config);
                    model.fit(x_view.to_owned(), y_view.to_owned())
                    .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
                    format!("Training failed: {}", e)
                    ))?;

                    self.coef_ = Some(model.coefficients()
                    .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
                    format!("Failed to get coefficients: {}", e)
                    ))?);

                    self.intercept_ = if self.fit_intercept {
                        Some(model.intercept()
                        .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
                        format!("Failed to get intercept: {}", e)
                        ))?)
                        } else {
                            Some(0.0)
                        };

                        self.n_iter_ = Some(model.n_iterations()
                        .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
                        format!("Failed to get iteration count: {}", e)
                        ))?);

                        self.model = Some(model);
                        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 predictions = self.model.as_ref().unwrap().predict(x_view.to_owned())
                        .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
                        format!("Prediction failed: {}", e)
                        ))?;

                        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_array_view1(&pred_array)
                            .map_err(|e| PyErr::from(e))?;
                            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)
                        })
                    }

                    #[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_
                    }

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