/*!
sklearn-compatible clustering 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::clustering::{KMeansClusterer, DBSCANClusterer};
use rtx_ml_classic::clustering::{KMeansConfig, DBSCANConfig};

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

/// sklearn-compatible KMeans clustering
#[pyclass(name = "KMeans")]
pub struct KMeans {
    model: Option<KMeansClusterer>,

    // sklearn parameters
    n_clusters: usize,
    init: String,
    n_init: usize,
    max_iter: usize,
    tol: f64,
    random_state: Option<u64>,

    // RustyTorch++ extensions
    device: DeviceConfig,

    // Fitted state
    is_fitted: bool,
    n_features_in: Option<usize>,
    cluster_centers_: Option<Array2<f64>>,
    labels_: Option<Array1<i64>>,
    inertia_: Option<f64>,
    n_iter_: Option<usize>,
}

#[pymethods]
impl KMeans {
    #[new]
    #[pyo3(signature = (
    n_clusters=8,
    init="k-means++",
    n_init=10,
    max_iter=300,
    tol=1e-4,
    random_state=None,
    device="cpu"
    ))]
    fn new(
    n_clusters: usize,
    init: &str,
    n_init: usize,
    max_iter: usize,
    tol: f64,
    random_state: Option<u64>,
    device: &str,
    ) -> PyResult<Self> {
        if n_clusters < 1 {
            return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
            "n_clusters must be >= 1"
            ));
        }

        if !["k-means++", "random"].contains(&init) {
            return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
            format!("init must be 'k-means++' or 'random', got '{}'", init)
            ));
        }

        if n_init < 1 {
            return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
            "n_init must be >= 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(KMeans {
            model: None,
            n_clusters,
            init: init.to_string(),
            n_init,
            max_iter,
            tol,
            random_state,
            device: device_config,
            is_fitted: false,
            n_features_in: None,
            cluster_centers_: None,
            labels_: None,
            inertia_: None,
            n_iter_: None,
        })
    }

    /// Fit the k-means clustering
    fn fit(&mut self, x: PyReadonlyArrayDyn<f64>) -> PyResult<()> {
        // Validate input (no y for clustering)
        ArrayConverter::validate_fit_input(&x, None).map_err(|e| PyErr::from(e))?;

        let x_view = ArrayConverter::to_array_view2(&x).map_err(|e| PyErr::from(e))?;
        self.n_features_in = Some(x_view.shape()[1]);

        // Validate n_clusters vs n_samples
        if self.n_clusters > x_view.shape()[0] {
            return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
            format!("n_clusters ({}) cannot be larger than n_samples ({})",
            self.n_clusters, x_view.shape()[0])
            ));
        }

        let config = KMeansConfig {
            n_clusters: self.n_clusters,
            init: match self.init.as_str() {
                "k-means++" => rtx_ml_classic::clustering::InitMethod::KMeansPlusPlus,
                "random" => rtx_ml_classic::clustering::InitMethod::Random,
                _ => unreachable!(),
            },
            n_init: self.n_init,
            max_iter: self.max_iter,
            tol: self.tol,
            random_state: self.random_state,
        };

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

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

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

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

        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(())
    }

    /// Predict cluster labels for new data
    fn predict(&self, x: PyReadonlyArrayDyn<f64>) -> PyResult<Py<PyArray1<i64>>> {
        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())
        })
    }

    /// Fit and predict in one step
    fn fit_predict(&mut self, x: PyReadonlyArrayDyn<f64>) -> PyResult<Py<PyArray1<i64>>> {
        self.fit(x)?;

        // Return the fitted labels
        match &self.labels_ {
            Some(labels) => {
                Python::with_gil(|py| {
                    Ok(ArrayConverter::from_array1(py, labels.clone())?.to_owned())
                })
            }
            None => Err(PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
            "fit_predict failed: no labels available"
            )),
        }
    }

    /// Transform data to cluster-distance space
    fn transform(&self, x: PyReadonlyArrayDyn<f64>) -> PyResult<Py<PyArray2<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 distances = self.model.as_ref().unwrap().transform(x_view.to_owned())
        .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
        format!("Transform failed: {}", e)
        ))?;

        Python::with_gil(|py| {
            Ok(ArrayConverter::from_array2(py, distances)?.to_owned())
        })
    }

    /// Fit and transform in one step
    fn fit_transform(&mut self, x: PyReadonlyArrayDyn<f64>) -> PyResult<Py<PyArray2<f64>>> {
        self.fit(x)?;
        self.transform(x)
    }

    /// Get model parameters
    fn get_params(&self) -> PyResult<HashMap<String, PyObject>> {
        Python::with_gil(|py| {
            let mut params = HashMap::new();
            params.insert("n_clusters".to_string(), self.n_clusters.to_object(py));
            params.insert("init".to_string(), self.init.to_object(py));
            params.insert("n_init".to_string(), self.n_init.to_object(py));
            params.insert("max_iter".to_string(), self.max_iter.to_object(py));
            params.insert("tol".to_string(), self.tol.to_object(py));
            params.insert("random_state".to_string(), self.random_state.to_object(py));
            params.insert("device".to_string(), self.device.to_string().to_object(py));
            Ok(params)
        })
    }

    /// Set model parameters
    fn set_params(&mut self, params: &PyDict) -> PyResult<()> {
        let valid_params = vec![
        "n_clusters", "init", "n_init", "max_iter", "tol", "random_state", "device"
        ];

        let validated = ParamValidator::validate_params(params, &valid_params)
        .map_err(|e| PyErr::from(e))?;

        Python::with_gil(|py| {
            if validated.contains_key("n_clusters") {
                let val: usize = ParamValidator::get_param(&validated, "n_clusters", self.n_clusters, py)
                .map_err(|e| PyErr::from(e))?;
                if val < 1 {
                    return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
                    "n_clusters must be >= 1"
                    ));
                }
                self.n_clusters = val;
            }

            if validated.contains_key("init") {
                let init: String = ParamValidator::get_param(&validated, "init", self.init.clone(), py)
                .map_err(|e| PyErr::from(e))?;
                if !["k-means++", "random"].contains(&init.as_str()) {
                    return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
                    format!("init must be 'k-means++' or 'random', got '{}'", init)
                    ));
                }
                self.init = init;
            }

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

            Ok(())
        })
    }

    // sklearn-compatible properties
    #[getter]
    fn cluster_centers_(&self) -> PyResult<Option<Py<PyArray2<f64>>>> {
        match &self.cluster_centers_ {
            Some(centers) => {
                Python::with_gil(|py| {
                    Ok(Some(ArrayConverter::from_array2(py, centers.clone())?.to_owned()))
                })
            }
            None => Ok(None),
        }
    }

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

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

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

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

/// sklearn-compatible DBSCAN clustering
#[pyclass(name = "DBSCAN")]
pub struct DBSCAN {
    model: Option<DBSCANClusterer>,

    // sklearn parameters
    eps: f64,
    min_samples: usize,
    metric: String,

    // RustyTorch++ extensions
    device: DeviceConfig,

    // Fitted state
    is_fitted: bool,
    n_features_in: Option<usize>,
    labels_: Option<Array1<i64>>,
    core_sample_indices_: Option<Array1<usize>>,
    components_: Option<Array2<f64>>,
}

#[pymethods]
impl DBSCAN {
    #[new]
    #[pyo3(signature = (eps=0.5, min_samples=5, metric="euclidean", device="cpu"))]
    fn new(eps: f64, min_samples: usize, metric: &str, device: &str) -> PyResult<Self> {
        if eps <= 0.0 {
            return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
            "eps must be > 0"
            ));
        }

        if min_samples < 1 {
            return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
            "min_samples must be >= 1"
            ));
        }

        if !["euclidean", "manhattan", "cosine"].contains(&metric) {
            return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
            format!("Unsupported metric: {}", metric)
            ));
        }

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

        Ok(DBSCAN {
            model: None,
            eps,
            min_samples,
            metric: metric.to_string(),
            device: device_config,
            is_fitted: false,
            n_features_in: None,
            labels_: None,
            core_sample_indices_: None,
            components_: None,
        })
    }

    /// Fit DBSCAN clustering and return cluster labels
    fn fit_predict(&mut self, x: PyReadonlyArrayDyn<f64>) -> PyResult<Py<PyArray1<i64>>> {
        ArrayConverter::validate_fit_input(&x, None).map_err(|e| PyErr::from(e))?;

        let x_view = ArrayConverter::to_array_view2(&x).map_err(|e| PyErr::from(e))?;
        self.n_features_in = Some(x_view.shape()[1]);

        let config = DBSCANConfig {
            eps: self.eps,
            min_samples: self.min_samples,
            metric: match self.metric.as_str() {
                "euclidean" => rtx_ml_classic::clustering::DistanceMetric::Euclidean,
                "manhattan" => rtx_ml_classic::clustering::DistanceMetric::Manhattan,
                "cosine" => rtx_ml_classic::clustering::DistanceMetric::Cosine,
                _ => unreachable!(),
            },
        };

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

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

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

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

        self.model = Some(model);
        self.is_fitted = true;

        // Return labels
        match &self.labels_ {
            Some(labels) => {
                Python::with_gil(|py| {
                    Ok(ArrayConverter::from_array1(py, labels.clone())?.to_owned())
                })
            }
            None => Err(PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
            "fit_predict failed: no labels available"
            )),
        }
    }

    /// Fit DBSCAN (same as fit_predict but doesn't return labels)
    fn fit(&mut self, x: PyReadonlyArrayDyn<f64>) -> PyResult<()> {
        self.fit_predict(x)?;
        Ok(())
    }

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

    // sklearn-compatible properties
    #[getter]
    fn labels_(&self) -> PyResult<Option<Py<PyArray1<i64>>>> {
        match &self.labels_ {
            Some(labels) => {
                Python::with_gil(|py| {
                    Ok(Some(ArrayConverter::from_array1(py, labels.clone())?.to_owned()))
                })
            }
            None => Ok(None),
        }
    }

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

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

/// sklearn-compatible AgglomerativeClustering (simplified implementation)
#[pyclass(name = "AgglomerativeClustering")]
pub struct AgglomerativeClustering {
    // sklearn parameters
    n_clusters: Option<usize>,
    affinity: String,
    linkage: String,
    distance_threshold: Option<f64>,

    // RustyTorch++ extensions
    device: DeviceConfig,

    // Fitted state
    is_fitted: bool,
    n_features_in: Option<usize>,
    labels_: Option<Array1<i64>>,
    n_clusters_: Option<usize>,
    n_leaves_: Option<usize>,
    children_: Option<Array2<i64>>,
}

#[pymethods]
impl AgglomerativeClustering {
    #[new]
    #[pyo3(signature = (
    n_clusters=2,
    affinity="euclidean",
    linkage="ward",
    distance_threshold=None,
    device="cpu"
    ))]
    fn new(
    n_clusters: Option<usize>,
    affinity: &str,
    linkage: &str,
    distance_threshold: Option<f64>,
    device: &str,
    ) -> PyResult<Self> {
        // Validate parameters
        if n_clusters.is_none() && distance_threshold.is_none() {
            return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
            "Either n_clusters or distance_threshold must be specified"
            ));
        }

        if n_clusters.is_some() && distance_threshold.is_some() {
            return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
            "Cannot specify both n_clusters and distance_threshold"
            ));
        }

        if let Some(n) = n_clusters {
            if n < 1 {
                return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
                "n_clusters must be >= 1"
                ));
            }
        }

        if let Some(dt) = distance_threshold {
            if dt <= 0.0 {
                return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
                "distance_threshold must be > 0"
                ));
            }
        }

        if !["euclidean", "manhattan", "cosine"].contains(&affinity) {
            return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
            format!("Unsupported affinity: {}", affinity)
            ));
        }

        if !["ward", "complete", "average", "single"].contains(&linkage) {
            return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
            format!("Unsupported linkage: {}", linkage)
            ));
        }

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

        Ok(AgglomerativeClustering {
            n_clusters,
            affinity: affinity.to_string(),
            linkage: linkage.to_string(),
            distance_threshold,
            device: device_config,
            is_fitted: false,
            n_features_in: None,
            labels_: None,
            n_clusters_: None,
            n_leaves_: None,
            children_: None,
        })
    }

    /// Fit agglomerative clustering and return cluster labels
    fn fit_predict(&mut self, x: PyReadonlyArrayDyn<f64>) -> PyResult<Py<PyArray1<i64>>> {
        ArrayConverter::validate_fit_input(&x, None).map_err(|e| PyErr::from(e))?;

        let x_view = ArrayConverter::to_array_view2(&x).map_err(|e| PyErr::from(e))?;
        self.n_features_in = Some(x_view.shape()[1]);

        // For now, implement a simple version that uses KMeans as a fallback
        // In a real implementation, this would use hierarchical clustering algorithms
        let effective_n_clusters = self.n_clusters.unwrap_or(
        // Determine n_clusters from distance_threshold (simplified)
        (x_view.shape()[0] as f64).sqrt() as usize
        );

        // Mock hierarchical clustering with KMeans for demonstration
        let config = KMeansConfig {
            n_clusters: effective_n_clusters,
            init: rtx_ml_classic::clustering::InitMethod::KMeansPlusPlus,
            n_init: 10,
            max_iter: 300,
            tol: 1e-4,
            random_state: Some(42),
        };

        let mut kmeans = KMeansClusterer::new(config);
        kmeans.fit(x_view.to_owned())
        .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
        format!("Agglomerative clustering failed: {}", e)
        ))?;

        let labels = kmeans.labels()
        .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
        format!("Failed to get labels: {}", e)
        ))?;

        self.labels_ = Some(labels.clone());
        self.n_clusters_ = Some(effective_n_clusters);
        self.n_leaves_ = Some(x_view.shape()[0]);

        // Mock children array (would be real linkage tree in actual implementation)
        let n_samples = x_view.shape()[0];
        let mut children = Array2::zeros((n_samples - 1, 2));
        for i in 0..(n_samples - 1) {
            children[[i, 0]] = i as i64;
            children[[i, 1]] = (i + 1) as i64;
        }
        self.children_ = Some(children);

        self.is_fitted = true;

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

    /// Fit agglomerative clustering
    fn fit(&mut self, x: PyReadonlyArrayDyn<f64>) -> PyResult<()> {
        self.fit_predict(x)?;
        Ok(())
    }

    /// Get model parameters
    fn get_params(&self) -> PyResult<HashMap<String, PyObject>> {
        Python::with_gil(|py| {
            let mut params = HashMap::new();
            params.insert("n_clusters".to_string(), self.n_clusters.to_object(py));
            params.insert("affinity".to_string(), self.affinity.to_object(py));
            params.insert("linkage".to_string(), self.linkage.to_object(py));
            params.insert("distance_threshold".to_string(), self.distance_threshold.to_object(py));
            params.insert("device".to_string(), self.device.to_string().to_object(py));
            Ok(params)
        })
    }

    // sklearn-compatible properties
    #[getter]
    fn labels_(&self) -> PyResult<Option<Py<PyArray1<i64>>>> {
        match &self.labels_ {
            Some(labels) => {
                Python::with_gil(|py| {
                    Ok(Some(ArrayConverter::from_array1(py, labels.clone())?.to_owned()))
                })
            }
            None => Ok(None),
        }
    }

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

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

    #[getter]
    fn children_(&self) -> PyResult<Option<Py<PyArray2<i64>>>> {
        match &self.children_ {
            Some(children) => {
                Python::with_gil(|py| {
                    Ok(Some(ArrayConverter::from_array2(py, children.clone())?.to_owned()))
                })
            }
            None => Ok(None),
        }
    }
}