//! Affinity Propagation clustering implementation
//!
//! Affinity Propagation is a clustering algorithm that identifies exemplars (cluster centers)
//! among the data points and forms clusters of data points around these exemplars.
//! It's particularly useful when the number of clusters is not known in advance,
//! and works well for non-flat geometries.

use scirs2_core::ndarray::{Array1, Array2, ArrayView1, ArrayView2};
use scirs2_core::numeric::{Float, FromPrimitive};
use std::cmp::Ordering;
use std::fmt::Debug;

use crate::error::{ClusteringError, Result};

/// Options for Affinity Propagation clustering
#[derive(Debug, Clone)]
pub struct AffinityPropagationOptions<F: Float> {
    /// Damping factor (between 0.5 and 1.0) to avoid numerical oscillations
    pub damping: F,

    /// Maximum number of iterations to run
    pub max_iter: usize,

    /// Number of iterations with no change in the number of estimated clusters
    /// that stops the convergence
    pub convergence_iter: usize,

    /// Preference value used to decide the number of exemplars
    /// If None, the median of non-diagonal elements in the similarity matrix is used
    pub preference: Option<F>,

    /// Whether the similarity matrix is precomputed
    pub affinity: String,

    /// Maximum allowed similarity between two exemplars to be considered distinct
    pub max_affinity_iterations: usize,
}

impl<F: Float + FromPrimitive> Default for AffinityPropagationOptions<F> {
    fn default() -> Self {
        Self {
            damping: F::from(0.5).expect("Failed to convert constant to float"),
            max_iter: 200,
            convergence_iter: 15,
            preference: None,
            affinity: "euclidean".to_string(),
            max_affinity_iterations: 100,
        }
    }
}

/// Compute pairwise similarity matrix based on negative squared Euclidean distance
///
/// # Arguments
///
/// * `data` - Input data
///
/// # Returns
///
/// * A similarity matrix where similarity(i, j) = -||x_i - x_j||^2
#[allow(dead_code)]
fn compute_similarity<F>(data: ArrayView2<F>) -> Result<Array2<F>>
where
    F: Float + FromPrimitive + Debug + PartialOrd,
{
    let n_samples = data.shape()[0];
    let n_features = data.shape()[1];

    let mut similarity = Array2::zeros((n_samples, n_samples));

    for i in 0..n_samples {
        for j in i..n_samples {
            if i == j {
                // Diagonal elements will be set to preference later
                similarity[[i, i]] = F::zero();
            } else {
                // Compute negative squared Euclidean distance
                let mut dist_sq = F::zero();
                for k in 0..n_features {
                    let diff = data[[i, k]] - data[[j, k]];
                    dist_sq = dist_sq + diff * diff;
                }

                let sim = -dist_sq; // Negative squared distance

                similarity[[i, j]] = sim;
                similarity[[j, i]] = sim; // Symmetric
            }
        }
    }

    Ok(similarity)
}

/// Compute preference values for the similarity matrix
///
/// # Arguments
///
/// * `similarity` - Similarity matrix
/// * `preference` - Optional preference value
///
/// # Returns
///
/// * Updated similarity matrix with diagonal elements set to preference
#[allow(dead_code)]
fn compute_preference<F>(mut similarity: Array2<F>, preference: Option<F>) -> Result<Array2<F>>
where
    F: Float + FromPrimitive + Debug + PartialOrd,
{
    let n_samples = similarity.shape()[0];

    // If preference is provided, use that
    if let Some(pref) = preference {
        for i in 0..n_samples {
            similarity[[i, i]] = pref;
        }
        return Ok(similarity);
    }

    // Otherwise, use the median of non-diagonal similarities
    let mut non_diag_similarities = Vec::new();
    for i in 0..n_samples {
        for j in 0..n_samples {
            if i != j {
                non_diag_similarities.push(similarity[[i, j]]);
            }
        }
    }

    if non_diag_similarities.is_empty() {
        return Err(ClusteringError::ComputationError(
            "Could not compute preferences, no non-diagonal similarities found".to_string(),
        ));
    }

    // Sort the similarities
    non_diag_similarities.sort_by(|a, b| a.partial_cmp(b).unwrap_or(Ordering::Equal));

    // Get the median
    let n = non_diag_similarities.len();
    let median = if n % 2 == 0 {
        // Even number of elements
        let mid1 = non_diag_similarities[n / 2 - 1];
        let mid2 = non_diag_similarities[n / 2];
        (mid1 + mid2) / F::from(2.0).expect("Failed to convert constant to float")
    } else {
        // Odd number of elements
        non_diag_similarities[n / 2]
    };

    // Set diagonal elements to median
    for i in 0..n_samples {
        similarity[[i, i]] = median;
    }

    Ok(similarity)
}

/// Run the Affinity Propagation algorithm
///
/// # Arguments
///
/// * `similarity` - Similarity matrix
/// * `options` - Algorithm parameters
///
/// # Returns
///
/// * Tuple of (cluster_centers_indices, labels)
#[allow(dead_code)]
fn run_affinity_propagation<F>(
    similarity: Array2<F>,
    options: &AffinityPropagationOptions<F>,
) -> Result<(Vec<usize>, Array1<i32>)>
where
    F: Float + FromPrimitive + Debug + PartialOrd,
{
    let n_samples = similarity.shape()[0];

    // Verify damping is in the correct range
    if options.damping < F::from(0.5).expect("Failed to convert constant to float")
        || options.damping > F::one()
    {
        return Err(ClusteringError::InvalidInput(
            "Damping factor must be between 0.5 and 1.0".to_string(),
        ));
    }

    // Initialize messages
    let mut responsibility = Array2::zeros((n_samples, n_samples));
    let mut availability = Array2::zeros((n_samples, n_samples));

    // Initialize temporary copy of messages
    let mut old_responsibility = responsibility.clone();
    let mut old_availability = availability.clone();

    // Initialize convergence monitoring
    let mut convergence_count = 0;
    let mut last_labels: Option<Array1<i32>> = None;

    // Main loop
    for _iter in 0..options.max_iter {
        // Update responsibility matrix
        // r(i, k) = s(i, k) - max_{k' != k} { a(i, k') + s(i, k') }
        old_responsibility.assign(&responsibility);

        for i in 0..n_samples {
            for k in 0..n_samples {
                // Find the maximum value of a(i, k') + s(i, k') for k' != k
                let mut max_val = F::neg_infinity();
                for k_prime in 0..n_samples {
                    if k_prime != k {
                        let val = availability[[i, k_prime]] + similarity[[i, k_prime]];
                        if val > max_val {
                            max_val = val;
                        }
                    }
                }

                // Calculate r(i, k)
                responsibility[[i, k]] = similarity[[i, k]] - max_val;
            }
        }

        // Apply damping to responsibility matrix
        for i in 0..n_samples {
            for k in 0..n_samples {
                responsibility[[i, k]] = options.damping * old_responsibility[[i, k]]
                    + (F::one() - options.damping) * responsibility[[i, k]];
            }
        }

        // Update availability matrix
        // a(i, k) = min { 0, r(k, k) + sum_{i' != i, k} max { 0, r(i', k) } } if i != k
        // a(k, k) = sum_{i' != k} max { 0, r(i', k) } if i == k
        old_availability.assign(&availability);

        for i in 0..n_samples {
            for k in 0..n_samples {
                if i != k {
                    // Case i != k
                    let mut sum = F::zero();
                    for i_prime in 0..n_samples {
                        if i_prime != i && i_prime != k {
                            sum = sum + F::max(F::zero(), responsibility[[i_prime, k]]);
                        }
                    }

                    availability[[i, k]] = F::min(F::zero(), responsibility[[k, k]] + sum);
                } else {
                    // Case i == k (self-availability)
                    let mut sum = F::zero();
                    for i_prime in 0..n_samples {
                        if i_prime != k {
                            sum = sum + F::max(F::zero(), responsibility[[i_prime, k]]);
                        }
                    }

                    availability[[k, k]] = sum;
                }
            }
        }

        // Apply damping to availability matrix
        for i in 0..n_samples {
            for k in 0..n_samples {
                availability[[i, k]] = options.damping * old_availability[[i, k]]
                    + (F::one() - options.damping) * availability[[i, k]];
            }
        }

        // Check for convergence
        let (cluster_centers_indices, labels) = extract_clusters(&responsibility, &availability)?;

        if let Some(old_labels) = &last_labels {
            if compare_labels(old_labels.view(), labels.view()) {
                convergence_count += 1;
            } else {
                convergence_count = 0;
            }
        }

        // Check if algorithm has converged
        if convergence_count >= options.convergence_iter {
            return Ok((cluster_centers_indices, labels));
        }

        last_labels = Some(labels);
    }

    // Return last result if max_iter is reached without convergence
    let (cluster_centers_indices, labels) = extract_clusters(&responsibility, &availability)?;
    Ok((cluster_centers_indices, labels))
}

/// Extract clusters from responsibility and availability matrices
///
/// # Arguments
///
/// * `responsibility` - Responsibility matrix
/// * `availability` - Availability matrix
///
/// # Returns
///
/// * Tuple of (cluster_centers_indices, labels)
#[allow(dead_code)]
fn extract_clusters<F>(
    responsibility: &Array2<F>,
    availability: &Array2<F>,
) -> Result<(Vec<usize>, Array1<i32>)>
where
    F: Float + FromPrimitive + Debug + PartialOrd,
{
    let n_samples = responsibility.shape()[0];

    // Compute criterion matrix (sum of responsibility and availability)
    let mut criterion = Array2::zeros((n_samples, n_samples));
    for i in 0..n_samples {
        for k in 0..n_samples {
            criterion[[i, k]] = responsibility[[i, k]] + availability[[i, k]];
        }
    }

    // Find exemplars (cluster centers)
    let mut cluster_centers_indices = Vec::new();
    for k in 0..n_samples {
        if criterion[[k, k]] > F::zero() {
            cluster_centers_indices.push(k);
        }
    }

    // If no exemplars were identified, select point with highest self-criterion
    if cluster_centers_indices.is_empty() {
        let mut max_criterion = F::neg_infinity();
        let mut max_idx = 0;

        for k in 0..n_samples {
            if criterion[[k, k]] > max_criterion {
                max_criterion = criterion[[k, k]];
                max_idx = k;
            }
        }

        cluster_centers_indices.push(max_idx);
    }

    // Assign labels based on similarity to exemplars
    let mut labels = Array1::from_vec(vec![-1; n_samples]);

    for i in 0..n_samples {
        let mut max_similarity = F::neg_infinity();
        let mut best_cluster = -1;

        for (cluster_idx, &exemplar) in cluster_centers_indices.iter().enumerate() {
            if criterion[[i, exemplar]] > max_similarity {
                max_similarity = criterion[[i, exemplar]];
                best_cluster = cluster_idx as i32;
            }
        }

        labels[i] = best_cluster;
    }

    Ok((cluster_centers_indices, labels))
}

/// Compare two label assignments to check if they are the same
///
/// # Arguments
///
/// * `labels1` - First label assignment
/// * `labels2` - Second label assignment
///
/// # Returns
///
/// * True if the label assignments are the same, false otherwise
#[allow(dead_code)]
fn compare_labels(labels1: ArrayView1<i32>, labels2: ArrayView1<i32>) -> bool {
    if labels1.len() != labels2.len() {
        return false;
    }

    for i in 0..labels1.len() {
        if labels1[i] != labels2[i] {
            return false;
        }
    }

    true
}

/// Affinity Propagation clustering algorithm
///
/// Affinity Propagation finds exemplars (cluster centers) among the data points
/// and forms clusters of data points around these exemplars. The algorithm determines
/// the number of clusters based on the input preference value.
///
/// # Arguments
///
/// * `data` - Input data (n_samples × n_features) or precomputed similarity matrix
/// * `precomputed` - Whether data is a precomputed similarity matrix
/// * `options` - Optional parameters
///
/// # Returns
///
/// * Tuple of (cluster_centers_indices, labels) where:
///   - cluster_centers_indices: Indices of cluster centers
///   - labels: Array of shape (n_samples,) with cluster assignments
///
/// # Examples
///
/// ```
/// use scirs2_core::ndarray::{ArrayView1, Array2, ArrayView2};
/// use scirs2_cluster::affinity::{affinity_propagation, AffinityPropagationOptions};
///
/// // Example data with two clusters
/// let data = Array2::from_shape_vec((6, 2), vec![
///     1.0, 2.0,
///     1.2, 1.8,
///     0.8, 1.9,
///     4.0, 5.0,
///     4.2, 4.8,
///     3.9, 5.1,
/// ]).expect("Operation failed");
///
/// // Run affinity propagation
/// let options = AffinityPropagationOptions {
///     damping: 0.9,
///     ..Default::default()
/// };
///
/// let result = affinity_propagation(data.view(), false, Some(options));
/// if let Ok((centers, labels)) = result {
///     println!("Cluster centers: {:?}", centers);
///     println!("Cluster assignments: {:?}", labels);
/// }
/// ```
#[allow(dead_code)]
pub fn affinity_propagation<F>(
    data: ArrayView2<F>,
    precomputed: bool,
    options: Option<AffinityPropagationOptions<F>>,
) -> Result<(Vec<usize>, Array1<i32>)>
where
    F: Float + FromPrimitive + Debug + PartialOrd,
{
    let opts = options.unwrap_or_default();

    // Validate input
    let n_samples = data.shape()[0];
    if n_samples == 0 {
        return Err(ClusteringError::InvalidInput("Empty input data".into()));
    }

    if precomputed {
        // Check if data is a square matrix (required for precomputed similarity)
        if data.shape()[0] != data.shape()[1] {
            return Err(ClusteringError::InvalidInput(
                "Precomputed similarity matrix must be square".into(),
            ));
        }

        // Use the provided similarity matrix
        let similarity = compute_preference(data.to_owned(), opts.preference)?;
        run_affinity_propagation(similarity, &opts)
    } else {
        // Compute similarity matrix from data
        let similarity = compute_similarity(data)?;
        let similarity = compute_preference(similarity, opts.preference)?;
        run_affinity_propagation(similarity, &opts)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use scirs2_core::ndarray::Array2;

    #[test]
    fn test_affinity_propagation_basic() {
        // Create a dataset with 2 well-separated clusters
        let data = Array2::from_shape_vec(
            (6, 2),
            vec![1.0, 2.0, 1.2, 1.8, 0.8, 1.9, 5.0, 6.0, 5.2, 5.8, 4.8, 6.1],
        )
        .expect("Operation failed");

        // Run affinity propagation with tuned parameters
        let options = AffinityPropagationOptions {
            damping: 0.5,            // Lower damping for faster convergence
            preference: Some(-50.0), // Much lower preference to encourage fewer clusters
            max_iter: 200,
            convergence_iter: 15,
            ..Default::default()
        };

        let result = affinity_propagation(data.view(), false, Some(options));
        assert!(result.is_ok());

        let (centers, labels) = result.expect("Operation failed");

        // We should have at least 1 cluster
        assert!(!centers.is_empty());
        assert!(centers.len() <= 6);

        // Check dimensions
        assert_eq!(labels.len(), 6);

        // Check that all points are assigned to clusters
        for &label in labels.iter() {
            assert!(label >= 0);
            assert!((label as usize) < centers.len());
        }
    }

    #[test]
    fn test_affinity_propagation_precomputed() {
        // Create a precomputed similarity matrix for 4 points
        // Higher values indicate more similarity
        let similarity = Array2::from_shape_vec(
            (4, 4),
            vec![
                0.0, -1.0, -3.0, -5.0, -1.0, 0.0, -2.0, -4.0, -3.0, -2.0, 0.0, -6.0, -5.0, -4.0,
                -6.0, 0.0,
            ],
        )
        .expect("Operation failed");

        // Run affinity propagation with precomputed similarity
        let options = AffinityPropagationOptions {
            damping: 0.9,
            preference: Some(-5.0), // Adjusted preference
            ..Default::default()
        };

        let result = affinity_propagation(similarity.view(), true, Some(options));
        assert!(result.is_ok());

        let (centers, labels) = result.expect("Operation failed");

        // Check dimensions
        assert_eq!(labels.len(), 4);
        assert!(!centers.is_empty());

        // Check that at least some points are in the same cluster
        assert!(labels.iter().any(|&l| l == labels[0]));
    }
}
