use ndarray::{Array2, ArrayView2};
use std::collections::HashMap;

/// Function type for custom distance metrics.
/// Takes a slice of unique values and returns a distance matrix.
pub type DistanceFunc = fn(&[f64]) -> Array2<f64>;

/// Distance metric type for Krippendorff's alpha.
#[derive(Clone)]
pub enum Distance {
    Nominal,
    Ordinal,
    Interval,
    Ratio,
    CustomFunc(DistanceFunc),
    CustomMatrix(Array2<f64>),
}

impl Distance {
    /// Parse distance from string.
    ///
    /// # Errors
    ///
    /// Returns an error if `s` does not match a supported distance type.
    pub fn parse(s: &str) -> Result<Self, String> {
        match s {
            "nominal" => Ok(Self::Nominal),
            "ordinal" => Ok(Self::Ordinal),
            "interval" => Ok(Self::Interval),
            "ratio" => Ok(Self::Ratio),
            _ => Err(format!("Unknown distance '{s}'")),
        }
    }
}

/// Factorize values in the data matrix, mapping unique values to integer codes.
///
/// # Arguments
///
/// * `data` - 2D array where NaN represents missing values
///
/// # Returns
///
/// A tuple of (`coded_matrix`, `unique_values`) where missing values are coded as `-1`.
fn factorize_values(data: &ArrayView2<f64>) -> (Array2<i64>, Vec<f64>) {
    let mut value_to_code: HashMap<u64, usize> = HashMap::new();
    let mut unique_values: Vec<f64> = Vec::new();
    let mut coded = Array2::from_elem(data.dim(), -1i64);

    for &val in data.iter() {
        if !val.is_nan() {
            let bits = val.to_bits();
            value_to_code.entry(bits).or_insert_with(|| {
                let code = unique_values.len();
                unique_values.push(val);
                code
            });
        }
    }

    unique_values.sort_by(|a, b| a.partial_cmp(b).unwrap());

    let mut remap: HashMap<u64, i64> = HashMap::new();
    for (new_idx, &val) in unique_values.iter().enumerate() {
        remap.insert(
            val.to_bits(),
            i64::try_from(new_idx).expect("index should fit into i64"),
        );
    }

    for ((i, j), &val) in data.indexed_iter() {
        if !val.is_nan() {
            coded[(i, j)] = remap[&val.to_bits()];
        }
    }

    (coded, unique_values)
}

/// Build the coincidence matrix from coded unit values.
///
/// # Arguments
///
/// * `matrix_indices` - Coded matrix with -1 for missing values
/// * `n_unique` - Number of unique non-missing values
///
/// # Returns
///
/// The coincidence matrix.
fn coincidence_matrix(matrix_indices: &ArrayView2<i64>, n_unique: usize) -> Array2<f64> {
    let mut coincidence = Array2::<f64>::zeros((n_unique, n_unique));
    let min_raters = 2;
    let mut counts = vec![0.0; n_unique];

    for unit_row in matrix_indices.rows() {
        counts.fill(0.0);
        let mut m_u = 0usize;

        for &val in unit_row {
            if val >= 0 {
                if let Ok(index) = usize::try_from(val) {
                    counts[index] += 1.0;
                }
                m_u += 1;
            }
        }

        if m_u < min_raters {
            continue;
        }

        let weight = 1.0 / (m_u - 1) as f64;

        for i in 0..n_unique {
            let count_i = counts[i];
            if count_i == 0.0 {
                continue;
            }

            coincidence[[i, i]] += weight * count_i * (count_i - 1.0);

            for j in (i + 1)..n_unique {
                let count_j = counts[j];
                if count_j == 0.0 {
                    continue;
                }

                let contribution = weight * count_i * count_j;
                coincidence[[i, j]] += contribution;
                coincidence[[j, i]] += contribution;
            }
        }
    }

    coincidence
}

/// Compute nominal distance matrix.
fn nominal_distance(n_unique: usize) -> Array2<f64> {
    let mut delta = Array2::<f64>::ones((n_unique, n_unique));
    for i in 0..n_unique {
        delta[[i, i]] = 0.0;
    }
    delta
}

/// Compute ordinal distance matrix based on cumulative frequencies.
fn ordinal_distance(coincidence: &ArrayView2<f64>) -> Array2<f64> {
    let n_c: Vec<f64> = (0..coincidence.nrows())
        .map(|i| coincidence.row(i).sum())
        .collect();

    let mut cum_freqs = Vec::with_capacity(n_c.len());
    let mut cumsum = 0.0;
    for &freq in &n_c {
        cumsum += freq;
        cum_freqs.push(cumsum);
    }

    let midpoints: Vec<f64> = cum_freqs
        .iter()
        .zip(&n_c)
        .map(|(&cum, &freq)| cum - freq / 2.0)
        .collect();

    let n = midpoints.len();
    let mut delta = Array2::<f64>::zeros((n, n));
    for i in 0..n {
        for j in 0..n {
            delta[[i, j]] = (midpoints[i] - midpoints[j]).powi(2);
        }
    }

    delta
}

/// Compute interval distance matrix.
fn interval_distance(unique_values: &[f64]) -> Array2<f64> {
    let n = unique_values.len();
    let mut delta = Array2::<f64>::zeros((n, n));
    for i in 0..n {
        for j in 0..n {
            delta[[i, j]] = (unique_values[i] - unique_values[j]).powi(2);
        }
    }
    delta
}

/// Compute ratio distance matrix.
fn ratio_distance(unique_values: &[f64]) -> Array2<f64> {
    let n = unique_values.len();
    let mut delta = Array2::<f64>::zeros((n, n));
    for i in 0..n {
        for j in 0..n {
            let sum = unique_values[i] + unique_values[j];
            let diff = unique_values[i] - unique_values[j];
            if sum != 0.0 {
                delta[[i, j]] = (diff / sum).powi(2);
            }
        }
    }
    delta
}

/// Compute the expected disagreement matrix.
fn compute_expected_matrix(coincidence: &ArrayView2<f64>) -> Array2<f64> {
    let n_c: Vec<f64> = (0..coincidence.nrows())
        .map(|i| coincidence.row(i).sum())
        .collect();
    let n_total: f64 = n_c.iter().sum();

    if n_total > 1.0 {
        let n = n_c.len();
        let mut expected = Array2::<f64>::zeros((n, n));
        for i in 0..n {
            for j in 0..n {
                expected[[i, j]] = if i == j {
                    n_c[i] * (n_c[i] - 1.0)
                } else {
                    n_c[i] * n_c[j]
                };
            }
        }
        expected / (n_total - 1.0)
    } else {
        Array2::<f64>::zeros(coincidence.dim())
    }
}

/// Compute Krippendorff's alpha from factorized data.
///
/// # Arguments
///
/// * `matrix_indices` - Coded matrix with -1 for missing values
/// * `unique_values` - Unique values corresponding to the codes
/// * `distance` - Distance metric to use
///
/// # Returns
///
/// A tuple of (`alpha`, `observed_disagreement`, `expected_disagreement`).
///
/// # Errors
///
/// Returns an error when the provided distance value cannot be processed.
pub fn alpha_from_factorized(
    matrix_indices: &ArrayView2<i64>,
    unique_values: &[f64],
    distance: Distance,
) -> Result<(f64, f64, f64), String> {
    let n_unique = unique_values.len();
    let coincidence = coincidence_matrix(matrix_indices, n_unique);
    let expected = compute_expected_matrix(&coincidence.view());

    let delta = match distance {
        Distance::Nominal => nominal_distance(n_unique),
        Distance::Ordinal => ordinal_distance(&coincidence.view()),
        Distance::Interval => interval_distance(unique_values),
        Distance::Ratio => ratio_distance(unique_values),
        Distance::CustomFunc(func) => func(unique_values),
        Distance::CustomMatrix(matrix) => matrix,
    };

    let observed_disagreement: f64 = (&coincidence * &delta).sum();
    let expected_disagreement: f64 = (&expected * &delta).sum();

    let alpha_value = if expected_disagreement == 0.0 {
        0.0
    } else {
        1.0 - observed_disagreement / expected_disagreement
    };

    Ok((alpha_value, observed_disagreement, expected_disagreement))
}

/// Compute Krippendorff's alpha.
///
/// # Arguments
///
/// * `data` - 2D array with units as rows and raters as columns; NaN for missing values
/// * `distance` - Distance metric to use
///
/// # Returns
///
/// A tuple of (`alpha`, `observed_disagreement`, `expected_disagreement`).
///
/// # Errors
///
/// Returns an error if no units have at least 2 ratings or if the distance is unknown.
pub fn alpha(data: &ArrayView2<f64>, distance: Distance) -> Result<(f64, f64, f64), String> {
    let mut has_valid_unit = false;
    for row in data.rows() {
        let valid_count = row.iter().filter(|&&v| !v.is_nan()).count();
        if valid_count >= 2 {
            has_valid_unit = true;
            break;
        }
    }

    if !has_valid_unit {
        return Err("No units have at least 2 ratings.".to_string());
    }

    let (matrix_indices, unique_values) = factorize_values(data);

    alpha_from_factorized(&matrix_indices.view(), &unique_values, distance)
}

/// Example custom distance function: squared difference.
/// This is equivalent to the interval distance metric.
#[must_use]
pub fn custom_squared_diff(unique_values: &[f64]) -> Array2<f64> {
    let n = unique_values.len();
    let mut delta = Array2::<f64>::zeros((n, n));
    for i in 0..n {
        for j in 0..n {
            let diff = unique_values[i] - unique_values[j];
            delta[[i, j]] = diff * diff;
        }
    }
    delta
}

/// Example custom distance function: absolute difference.
#[must_use]
pub fn custom_abs_diff(unique_values: &[f64]) -> Array2<f64> {
    let n = unique_values.len();
    let mut delta = Array2::<f64>::zeros((n, n));
    for i in 0..n {
        for j in 0..n {
            let diff = (unique_values[i] - unique_values[j]).abs();
            delta[[i, j]] = diff;
        }
    }
    delta
}

#[cfg(test)]
mod tests {
    use super::*;
    use approx::assert_abs_diff_eq as assert_approx_eq;
    use ndarray::array;

    #[test]
    fn test_factorize_values() {
        let data = array![[1.0, 2.0], [2.0, 3.0], [f64::NAN, 1.0]];
        let (coded, unique) = factorize_values(&data.view());

        assert_eq!(unique, vec![1.0, 2.0, 3.0]);
        assert_eq!(coded[[0, 0]], 0);
        assert_eq!(coded[[0, 1]], 1);
        assert_eq!(coded[[1, 0]], 1);
        assert_eq!(coded[[1, 1]], 2);
        assert_eq!(coded[[2, 0]], -1);
        assert_eq!(coded[[2, 1]], 0);
    }

    #[test]
    fn test_nominal_distance() {
        let delta = nominal_distance(3);
        assert_approx_eq!(delta[[0, 0]], 0.0, epsilon = f64::EPSILON);
        assert_approx_eq!(delta[[0, 1]], 1.0, epsilon = f64::EPSILON);
        assert_approx_eq!(delta[[1, 2]], 1.0, epsilon = f64::EPSILON);
    }

    #[test]
    fn test_custom_distance_squared_diff() {
        let data = array![[1.0, 2.0], [2.0, 3.0], [3.0, 4.0]];

        let result_interval = alpha(&data.view(), Distance::Interval).unwrap();
        let result_custom = alpha(&data.view(), Distance::CustomFunc(custom_squared_diff)).unwrap();

        assert_approx_eq!(result_interval.0, result_custom.0, epsilon = 1e-10);
        assert_approx_eq!(result_interval.1, result_custom.1, epsilon = 1e-10);
        assert_approx_eq!(result_interval.2, result_custom.2, epsilon = 1e-10);
    }

    #[test]
    fn test_custom_distance_abs_diff() {
        let data = array![[1.0, 2.0], [2.0, 3.0], [3.0, 4.0]];

        let result = alpha(&data.view(), Distance::CustomFunc(custom_abs_diff)).unwrap();

        assert!(result.0.is_finite());
        assert!(result.1 >= 0.0);
        assert!(result.2 >= 0.0);
    }

    #[test]
    fn test_custom_distance_with_missing_values() {
        let data = array![[1.0, f64::NAN], [2.0, 3.0], [3.0, 4.0]];

        let result = alpha(&data.view(), Distance::CustomFunc(custom_squared_diff)).unwrap();

        assert!(result.0.is_finite());
        assert!(result.1 >= 0.0);
        assert!(result.2 >= 0.0);
    }
}
