//! K-Nearest Neighbors implementations
//!
//! KNN is a lazy learning algorithm that makes predictions based on the
//! k nearest neighbors in the feature space.

use crate::error::Result;
use rtx_tensor::Tensor;
use std::collections::HashMap;

/// Distance metrics for KNN
#[derive(Debug, Clone, Copy)]
pub enum Distance {
    Euclidean,
    Manhattan,
    Minkowski { p: f32 },
}

/// K-Nearest Neighbors Classifier
#[derive(Debug, Clone)]
pub struct KNeighborsClassifier {
    /// Number of neighbors
    n_neighbors: usize,
    /// Distance metric
    metric: Distance,
    /// Training data
    x_train: Option<Tensor>,
    /// Training labels
    y_train: Option<Tensor>,
}

impl KNeighborsClassifier {
    /// Create new KNN classifier
    pub fn new(n_neighbors: usize) -> Self {
        Self {
            n_neighbors,
            metric: Distance::Euclidean,
            x_train: None,
            y_train: None,
        }
    }

    /// Set distance metric
    pub fn metric(mut self, metric: Distance) -> Self {
        self.metric = metric;
        self
    }

    /// Fit the classifier (store training data)
    pub fn fit(&mut self, x: &Tensor, y: &Tensor) -> Result<&mut Self> {
        // Validate input
        if x.shape().ndim() != 2 {
            return Err(crate::error::MLError::invalid_input(
                "X must be 2-dimensional".to_string(),
            ));
        }

        if y.shape().ndim() != 1 {
            return Err(crate::error::MLError::invalid_input(
                "y must be 1-dimensional".to_string(),
            ));
        }

        let n_samples = x.shape().dims()[0];
        if y.shape().dims()[0] != n_samples {
            return Err(crate::error::MLError::invalid_input(
                "X and y must have same number of samples".to_string(),
            ));
        }

        if n_samples < self.n_neighbors {
            return Err(crate::error::MLError::invalid_input(format!(
                "Number of samples ({}) must be >= n_neighbors ({})",
                n_samples, self.n_neighbors
            )));
        }

        // Store training data (KNN is lazy learning)
        self.x_train = Some(x.clone());
        self.y_train = Some(y.clone());

        Ok(self)
    }

    /// Predict class labels for test data
    pub fn predict(&self, x: &Tensor) -> Result<Tensor> {
        let x_train = self.x_train.as_ref().ok_or_else(|| {
            crate::error::MLError::not_fitted("Model has not been fitted yet".to_string())
        })?;

        let y_train = self.y_train.as_ref().unwrap();

        // Validate input
        if x.shape().ndim() != 2 {
            return Err(crate::error::MLError::invalid_input(
                "X must be 2-dimensional".to_string(),
            ));
        }

        let n_test_samples = x.shape().dims()[0];
        let n_features = x.shape().dims()[1];

        if x_train.shape().dims()[1] != n_features {
            return Err(crate::error::MLError::invalid_input(format!(
                "X has {} features but model was fitted with {} features",
                n_features,
                x_train.shape().dims()[1]
            )));
        }

        // Get data as CPU tensors
        let x_test_data = x.to_cpu()?;
        let x_train_data = x_train.to_cpu()?;
        let y_train_data = y_train.to_cpu()?;

        let n_train_samples = x_train.shape().dims()[0];
        let mut predictions = vec![0.0f32; n_test_samples];

        // For each test sample
        for i in 0..n_test_samples {
            // Compute distances to all training samples
            let mut distances: Vec<(f32, usize)> = Vec::with_capacity(n_train_samples);

            for j in 0..n_train_samples {
                let distance = self.compute_distance(
                    &x_test_data[i * n_features..(i + 1) * n_features],
                    &x_train_data[j * n_features..(j + 1) * n_features],
                );
                distances.push((distance, j));
            }

            // Sort by distance and take k nearest
            distances.sort_by(|a, b| a.0.total_cmp(&b.0));
            let k_nearest: Vec<usize> = distances
                .iter()
                .take(self.n_neighbors)
                .map(|(_, idx)| *idx)
                .collect();

            // Majority vote among k nearest neighbors
            let mut class_votes: HashMap<i32, usize> = HashMap::new();
            for &neighbor_idx in &k_nearest {
                let class = y_train_data[neighbor_idx] as i32;
                *class_votes.entry(class).or_insert(0) += 1;
            }

            // Find class with most votes
            let predicted_class = class_votes
                .iter()
                .max_by_key(|(_, count)| **count)
                .map_or(0, |(&class, _)| class) as f32;

            predictions[i] = predicted_class;
        }

        Ok(Tensor::from_data(
            predictions,
            vec![n_test_samples],
            x.device(),
        )?)
    }

    /// Compute distance between two points
    fn compute_distance(&self, point1: &[f32], point2: &[f32]) -> f32 {
        match self.metric {
            Distance::Euclidean => point1
                .iter()
                .zip(point2.iter())
                .map(|(&x1, &x2)| (x1 - x2).powi(2))
                .sum::<f32>()
                .sqrt(),
            Distance::Manhattan => point1
                .iter()
                .zip(point2.iter())
                .map(|(&x1, &x2)| (x1 - x2).abs())
                .sum::<f32>(),
            Distance::Minkowski { p } => point1
                .iter()
                .zip(point2.iter())
                .map(|(&x1, &x2)| (x1 - x2).abs().powf(p))
                .sum::<f32>()
                .powf(1.0 / p),
        }
    }

    /// Get parameters
    pub fn params(&self) -> (usize, Distance) {
        (self.n_neighbors, self.metric)
    }

    /// Check if model is fitted
    pub fn is_fitted(&self) -> bool {
        self.x_train.is_some() && self.y_train.is_some()
    }
}

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

/// K-Nearest Neighbors Regressor
#[derive(Debug, Clone)]
pub struct KNeighborsRegressor {
    /// Number of neighbors
    n_neighbors: usize,
    /// Distance metric
    metric: Distance,
    /// Weights for neighbors
    weights: Weights,
    /// Training data
    x_train: Option<Tensor>,
    /// Training targets
    y_train: Option<Tensor>,
}

/// Weighting schemes for KNN regression
#[derive(Debug, Clone, Copy)]
pub enum Weights {
    Uniform,  // All neighbors have equal weight
    Distance, // Weight by inverse distance
}

impl KNeighborsRegressor {
    /// Create new KNN regressor
    pub fn new(n_neighbors: usize) -> Self {
        Self {
            n_neighbors,
            metric: Distance::Euclidean,
            weights: Weights::Uniform,
            x_train: None,
            y_train: None,
        }
    }

    /// Set distance metric
    pub fn metric(mut self, metric: Distance) -> Self {
        self.metric = metric;
        self
    }

    /// Set weighting scheme
    pub fn weights(mut self, weights: Weights) -> Self {
        self.weights = weights;
        self
    }

    /// Fit the regressor (store training data)
    pub fn fit(&mut self, x: &Tensor, y: &Tensor) -> Result<&mut Self> {
        // Validate input
        if x.shape().ndim() != 2 {
            return Err(crate::error::MLError::invalid_input(
                "X must be 2-dimensional".to_string(),
            ));
        }

        if y.shape().ndim() != 1 {
            return Err(crate::error::MLError::invalid_input(
                "y must be 1-dimensional".to_string(),
            ));
        }

        let n_samples = x.shape().dims()[0];
        if y.shape().dims()[0] != n_samples {
            return Err(crate::error::MLError::invalid_input(
                "X and y must have same number of samples".to_string(),
            ));
        }

        if n_samples < self.n_neighbors {
            return Err(crate::error::MLError::invalid_input(format!(
                "Number of samples ({}) must be >= n_neighbors ({})",
                n_samples, self.n_neighbors
            )));
        }

        // Store training data (KNN is lazy learning)
        self.x_train = Some(x.clone());
        self.y_train = Some(y.clone());

        Ok(self)
    }

    /// Predict continuous values for test data
    pub fn predict(&self, x: &Tensor) -> Result<Tensor> {
        let x_train = self.x_train.as_ref().ok_or_else(|| {
            crate::error::MLError::not_fitted("Model has not been fitted yet".to_string())
        })?;

        let y_train = self.y_train.as_ref().unwrap();

        // Validate input
        if x.shape().ndim() != 2 {
            return Err(crate::error::MLError::invalid_input(
                "X must be 2-dimensional".to_string(),
            ));
        }

        let n_test_samples = x.shape().dims()[0];
        let n_features = x.shape().dims()[1];

        if x_train.shape().dims()[1] != n_features {
            return Err(crate::error::MLError::invalid_input(format!(
                "X has {} features but model was fitted with {} features",
                n_features,
                x_train.shape().dims()[1]
            )));
        }

        // Get data as CPU tensors
        let x_test_data = x.to_cpu()?;
        let x_train_data = x_train.to_cpu()?;
        let y_train_data = y_train.to_cpu()?;

        let n_train_samples = x_train.shape().dims()[0];
        let mut predictions = vec![0.0f32; n_test_samples];

        // For each test sample
        for i in 0..n_test_samples {
            // Compute distances to all training samples
            let mut distances: Vec<(f32, usize)> = Vec::with_capacity(n_train_samples);

            for j in 0..n_train_samples {
                let distance = self.compute_distance(
                    &x_test_data[i * n_features..(i + 1) * n_features],
                    &x_train_data[j * n_features..(j + 1) * n_features],
                );
                distances.push((distance, j));
            }

            // Sort by distance and take k nearest
            distances.sort_by(|a, b| a.0.total_cmp(&b.0));
            let k_nearest: Vec<(f32, usize)> =
                distances.into_iter().take(self.n_neighbors).collect();

            // Compute weighted average
            let prediction = match self.weights {
                Weights::Uniform => {
                    // Simple average
                    let sum: f32 = k_nearest.iter().map(|(_, idx)| y_train_data[*idx]).sum();
                    sum / self.n_neighbors as f32
                }
                Weights::Distance => {
                    // Weighted by inverse distance
                    let mut weighted_sum = 0.0f32;
                    let mut weight_sum = 0.0f32;

                    for (distance, idx) in k_nearest {
                        let weight = if distance > 0.0 { 1.0 / distance } else { 1e6 }; // Handle exact matches
                        weighted_sum += weight * y_train_data[idx];
                        weight_sum += weight;
                    }

                    if weight_sum > 0.0 {
                        weighted_sum / weight_sum
                    } else {
                        0.0
                    }
                }
            };

            predictions[i] = prediction;
        }

        Ok(Tensor::from_data(
            predictions,
            vec![n_test_samples],
            x.device(),
        )?)
    }

    /// Compute distance between two points
    fn compute_distance(&self, point1: &[f32], point2: &[f32]) -> f32 {
        match self.metric {
            Distance::Euclidean => point1
                .iter()
                .zip(point2.iter())
                .map(|(&x1, &x2)| (x1 - x2).powi(2))
                .sum::<f32>()
                .sqrt(),
            Distance::Manhattan => point1
                .iter()
                .zip(point2.iter())
                .map(|(&x1, &x2)| (x1 - x2).abs())
                .sum::<f32>(),
            Distance::Minkowski { p } => point1
                .iter()
                .zip(point2.iter())
                .map(|(&x1, &x2)| (x1 - x2).abs().powf(p))
                .sum::<f32>()
                .powf(1.0 / p),
        }
    }

    /// Get parameters
    pub fn params(&self) -> (usize, Distance, Weights) {
        (self.n_neighbors, self.metric, self.weights)
    }

    /// Check if model is fitted
    pub fn is_fitted(&self) -> bool {
        self.x_train.is_some() && self.y_train.is_some()
    }
}

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