//! Stain vector estimation for histopathology images
//!
//! Provides standard stain vectors and estimation methods for
//! common staining protocols like H&E and IHC-DAB.

use crate::config::EstimationMethod;
use crate::optical_density::{extract_foreground_pixels, get_tissue_mask};
use anyhow::{Result, bail};
use ndarray::{Array1, Array2, Array3, Axis};
use slidescope_shared::StainType;

/// Standard H&E stain vectors (Ruifrok & Johnston, 2001)
///
/// Rows are RGB components, columns are stains:
/// - Column 0: Hematoxylin (blue-purple)
/// - Column 1: Eosin (pink-red)
/// - Column 2: Residual/background
pub fn he_stain_matrix() -> Array2<f32> {
    let mut matrix = Array2::from_shape_vec(
        (3, 3),
        vec![
            0.650, 0.072, 0.268, // R: Hematoxylin, Eosin, Residual
            0.704, 0.990, 0.570, // G
            0.286, 0.105, 0.776, // B
        ],
    )
    .unwrap();

    // Normalize each column
    normalize_columns(&mut matrix);
    matrix
}

/// Standard IHC-DAB stain vectors
///
/// - Column 0: Hematoxylin (counterstain)
/// - Column 1: DAB (brown chromogen)
/// - Column 2: Residual
pub fn ihc_dab_stain_matrix() -> Array2<f32> {
    let mut matrix = Array2::from_shape_vec(
        (3, 3),
        vec![
            0.650, 0.268, 0.0, // R
            0.704, 0.570, 0.0, // G
            0.286, 0.776, 0.0, // B
        ],
    )
    .unwrap();

    normalize_columns(&mut matrix);
    matrix
}

/// Normalize each column of the matrix to unit length
fn normalize_columns(matrix: &mut Array2<f32>) {
    for mut col in matrix.columns_mut() {
        let norm: f32 = col.iter().map(|x| x * x).sum::<f32>().sqrt();
        if norm > 0.0 {
            col.mapv_inplace(|x| x / norm);
        }
    }
}

/// Get standard stain matrix for a stain type
pub fn get_standard_matrix(stain_type: StainType) -> Array2<f32> {
    match stain_type {
        StainType::HE => he_stain_matrix(),
        StainType::IhcDab => ihc_dab_stain_matrix(),
        StainType::Custom => Array2::eye(3), // Identity matrix for custom
    }
}

/// Result of stain vector estimation
#[derive(Debug, Clone)]
pub struct StainEstimationResult {
    /// Estimated stain matrix (3 x n_components)
    pub stain_matrix: Array2<f32>,
    /// Quality score (0.0 - 1.0)
    pub quality_score: f32,
    /// Detected stain type
    pub stain_type: StainType,
    /// Number of samples used for estimation
    pub n_samples: usize,
}

/// Stain vector estimator
pub struct StainEstimator {
    method: EstimationMethod,
    background_threshold: f32,
    sample_size: Option<usize>,
}

impl StainEstimator {
    /// Create a new stain estimator
    pub fn new(
        method: EstimationMethod,
        background_threshold: f32,
        sample_size: Option<usize>,
    ) -> Self {
        Self {
            method,
            background_threshold,
            sample_size,
        }
    }

    /// Estimate stain vectors from an optical density image
    ///
    /// # Arguments
    /// * `od_image` - Optical density image (3 channels)
    /// * `n_components` - Number of stain components to extract
    pub fn estimate(
        &self,
        od_image: &Array3<f32>,
        n_components: usize,
    ) -> Result<StainEstimationResult> {
        if !(2..=3).contains(&n_components) {
            bail!("n_components must be 2 or 3");
        }

        match self.method {
            EstimationMethod::Ruifrok => self.estimate_ruifrok(n_components),
            EstimationMethod::Macenko => self.estimate_macenko(od_image, n_components),
            EstimationMethod::Robust => self.estimate_robust(od_image, n_components),
        }
    }

    /// Use pre-defined Ruifrok vectors
    fn estimate_ruifrok(&self, n_components: usize) -> Result<StainEstimationResult> {
        let full_matrix = he_stain_matrix();

        // Extract only the requested number of components
        let stain_matrix = full_matrix
            .slice(ndarray::s![.., 0..n_components])
            .to_owned();

        Ok(StainEstimationResult {
            stain_matrix,
            quality_score: 0.95, // High confidence for standard vectors
            stain_type: StainType::HE,
            n_samples: 0,
        })
    }

    /// Macenko's method using PCA
    fn estimate_macenko(
        &self,
        od_image: &Array3<f32>,
        n_components: usize,
    ) -> Result<StainEstimationResult> {
        // Get tissue mask
        let mask = get_tissue_mask(od_image, self.background_threshold);

        // Extract foreground pixels
        let pixels = extract_foreground_pixels(od_image, &mask);
        let n_samples = pixels.nrows();

        if n_samples < 100 {
            // Fall back to standard vectors if not enough tissue
            return self.estimate_ruifrok(n_components);
        }

        // Optionally subsample for speed
        let samples = if let Some(max_samples) = self.sample_size {
            if n_samples > max_samples {
                subsample_rows(&pixels, max_samples)
            } else {
                pixels
            }
        } else {
            pixels
        };

        // Compute covariance matrix
        let cov = compute_covariance(&samples)?;

        // Compute eigenvectors (simple power iteration for top 2)
        let eigenvectors = power_iteration_top_k(&cov, n_components)?;

        // Build stain matrix
        let mut stain_matrix = Array2::zeros((3, n_components));
        for (i, eigvec) in eigenvectors.iter().enumerate() {
            stain_matrix.column_mut(i).assign(eigvec);
        }

        // Ensure positive orientation
        for mut col in stain_matrix.columns_mut() {
            if col.sum() < 0.0 {
                col.mapv_inplace(|x| -x);
            }
        }

        normalize_columns(&mut stain_matrix);

        // Compute quality score based on eigenvalue separation
        let quality_score = 0.8; // Simplified

        Ok(StainEstimationResult {
            stain_matrix,
            quality_score,
            stain_type: StainType::Custom,
            n_samples: samples.nrows(),
        })
    }

    /// Robust estimation combining methods
    fn estimate_robust(
        &self,
        od_image: &Array3<f32>,
        n_components: usize,
    ) -> Result<StainEstimationResult> {
        // Try Macenko first
        let macenko_result = self.estimate_macenko(od_image, n_components)?;

        // Get standard Ruifrok vectors
        let ruifrok_result = self.estimate_ruifrok(n_components)?;

        // Compare with standard and choose best
        let similarity =
            matrix_similarity(&macenko_result.stain_matrix, &ruifrok_result.stain_matrix);

        if similarity > 0.7 && macenko_result.quality_score > 0.6 {
            // Macenko result is consistent with standard
            Ok(macenko_result)
        } else {
            // Fall back to standard vectors
            Ok(ruifrok_result)
        }
    }
}

/// Subsample rows from a matrix
fn subsample_rows(matrix: &Array2<f32>, n_samples: usize) -> Array2<f32> {
    use rand::seq::SliceRandom;
    let mut rng = rand::thread_rng();

    let n_rows = matrix.nrows();
    let mut indices: Vec<usize> = (0..n_rows).collect();
    indices.shuffle(&mut rng);
    indices.truncate(n_samples);

    let mut result = Array2::zeros((n_samples, matrix.ncols()));
    for (i, &idx) in indices.iter().enumerate() {
        result.row_mut(i).assign(&matrix.row(idx));
    }
    result
}

/// Compute covariance matrix
fn compute_covariance(data: &Array2<f32>) -> Result<Array2<f32>> {
    let n_samples = data.nrows() as f32;
    if n_samples < 2.0 {
        bail!("Need at least 2 samples for covariance");
    }

    let mean = data.mean_axis(Axis(0)).unwrap();
    let centered = data - &mean;
    let cov = centered.t().dot(&centered) / (n_samples - 1.0);

    Ok(cov)
}

/// Power iteration to find top k eigenvectors
fn power_iteration_top_k(matrix: &Array2<f32>, k: usize) -> Result<Vec<Array1<f32>>> {
    let n = matrix.nrows();
    let mut eigenvectors = Vec::with_capacity(k);
    let mut deflated = matrix.clone();

    for _ in 0..k {
        // Random initial vector
        let mut v = Array1::from_iter((0..n).map(|i| (i as f32 + 1.0).sin()));

        // Power iteration
        for _ in 0..100 {
            let new_v = deflated.dot(&v);
            let norm: f32 = new_v.iter().map(|x| x * x).sum::<f32>().sqrt();
            if norm < 1e-10 {
                break;
            }
            v = new_v / norm;
        }

        eigenvectors.push(v.clone());

        // Deflate matrix
        let eigenvalue = v.dot(&deflated.dot(&v));
        let outer = outer_product(&v, &v);
        deflated = deflated - eigenvalue * outer;
    }

    Ok(eigenvectors)
}

/// Compute outer product of two vectors
fn outer_product(a: &Array1<f32>, b: &Array1<f32>) -> Array2<f32> {
    let n = a.len();
    let m = b.len();
    let mut result = Array2::zeros((n, m));
    for i in 0..n {
        for j in 0..m {
            result[[i, j]] = a[i] * b[j];
        }
    }
    result
}

/// Compute similarity between two matrices (cosine of angle)
fn matrix_similarity(a: &Array2<f32>, b: &Array2<f32>) -> f32 {
    if a.shape() != b.shape() {
        return 0.0;
    }

    let dot: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
    let norm_a: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
    let norm_b: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();

    if norm_a < 1e-10 || norm_b < 1e-10 {
        return 0.0;
    }

    (dot / (norm_a * norm_b)).abs()
}

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

    #[test]
    fn test_he_matrix_normalized() {
        let matrix = he_stain_matrix();

        // Each column should have unit norm
        for col in matrix.columns() {
            let norm: f32 = col.iter().map(|x| x * x).sum::<f32>().sqrt();
            assert!((norm - 1.0).abs() < 1e-5);
        }
    }

    #[test]
    fn test_ruifrok_estimation() {
        let estimator = StainEstimator::new(EstimationMethod::Ruifrok, 0.8, None);
        let result = estimator.estimate_ruifrok(2).unwrap();

        assert_eq!(result.stain_matrix.shape(), &[3, 2]);
        assert!(result.quality_score > 0.9);
    }

    #[test]
    fn test_covariance() {
        let data = Array2::from_shape_vec(
            (4, 3),
            vec![1.0, 2.0, 3.0, 2.0, 3.0, 4.0, 3.0, 4.0, 5.0, 4.0, 5.0, 6.0],
        )
        .unwrap();

        let cov = compute_covariance(&data).unwrap();
        assert_eq!(cov.shape(), &[3, 3]);

        // Covariance matrix should be symmetric
        for i in 0..3 {
            for j in 0..3 {
                assert!((cov[[i, j]] - cov[[j, i]]).abs() < 1e-5);
            }
        }
    }
}