use crate::op::{ComputeContext, GradientContext, Op, OpError};
use crate::tensor::Tensor;
use crate::tensor_ops;
use crate::Float;
use scirs2_core::ndarray::{Array1, Array2, ArrayView2, Ix2};

/// Frobenius norm operation with improved gradient computation
pub struct FrobeniusNormOp;

impl<F: Float> Op<F> for FrobeniusNormOp {
    fn name(&self) -> &'static str {
        "FrobeniusNorm"
    }

    fn compute(&self, ctx: &mut ComputeContext<F>) -> Result<(), OpError> {
        let input = ctx.input(0);
        let input_array = input.view();

        let mut sum_squared = F::zero();
        for &elem in input_array.iter() {
            sum_squared += elem * elem;
        }

        let norm = sum_squared.sqrt();
        ctx.append_output(scirs2_core::ndarray::arr0(norm).into_dyn());
        Ok(())
    }

    fn grad(&self, ctx: &mut GradientContext<F>) {
        let grad_output = ctx.output_grad();
        let input = ctx.input(0);
        let output = ctx.output();
        let g = ctx.graph();
        
        // Use tensor operations to maintain gradient flow
        let epsilon = tensor_ops::scalar(F::epsilon() * F::from(10.0).expect("Failed to convert constant to float"), g);
        let safe_norm = tensor_ops::maximum(&output, &epsilon);
        
        // Compute gradient: (input / norm) * grad_output
        let grad_input = tensor_ops::mul(
            &tensor_ops::div(&input, &safe_norm),
            &grad_output
        );
        
        ctx.append_input_grad(0, Some(grad_input));
    }
}

/// Spectral norm operation with proper gradient computation through SVD
pub struct SpectralNormOp;

impl<F: Float + scirs2_core::ndarray::ScalarOperand> Op<F> for SpectralNormOp {
    fn name(&self) -> &'static str {
        "SpectralNorm"
    }

    fn compute(&self, ctx: &mut ComputeContext<F>) -> Result<(), OpError> {
        let input = ctx.input(0);
        let shape = input.shape();

        if shape.len() != 2 {
            return Err(OpError::IncompatibleShape(
                "Spectral norm requires 2D matrix".into(),
            ));
        }

        // Convert input to 2D matrix
        let matrix = input
            .view()
            .into_dimensionality::<Ix2>()
            .map_err(|_| OpError::IncompatibleShape("Failed to convert to 2D array".into()))?;

        // Use power iteration to find the largest singular value
        let (_, sigma_max) = power_iteration_spectral(&matrix, 50, F::from(1e-8).expect("Failed to convert constant to float"));

        ctx.append_output(scirs2_core::ndarray::arr0(sigma_max).into_dyn());
        Ok(())
    }

    fn grad(&self, ctx: &mut GradientContext<F>) {
        let grad_output = ctx.output_grad();
        let input = ctx.input(0);
        let g = ctx.graph();

        // Evaluate the input to work with concrete values for SVD computation
        let input_array = match input.eval(g) {
            Ok(arr) => arr,
            Err(_) => {
                ctx.append_input_grad(0, None);
                return;
            }
        };

        let grad_output_array = match grad_output.eval(g) {
            Ok(arr) => arr,
            Err(_) => {
                ctx.append_input_grad(0, None);
                return;
            }
        };

        let grad_scalar = grad_output_array[[]];

        if let Ok(matrix) = input_array.view().into_dimensionality::<Ix2>() {
            let grad_matrix = compute_spectral_norm_gradient(&matrix, grad_scalar);
            let grad_tensor = tensor_ops::convert_to_tensor(grad_matrix.into_dyn(), g);
            ctx.append_input_grad(0, Some(grad_tensor));
            return;
        }

        // Fallback
        ctx.append_input_grad(0, None);
    }
}

/// Nuclear norm operation with proper gradient computation through SVD
pub struct NuclearNormOp;

impl<F: Float + scirs2_core::ndarray::ScalarOperand> Op<F> for NuclearNormOp {
    fn name(&self) -> &'static str {
        "NuclearNorm"
    }

    fn compute(&self, ctx: &mut ComputeContext<F>) -> Result<(), OpError> {
        let input = ctx.input(0);
        let shape = input.shape();

        if shape.len() != 2 {
            return Err(OpError::IncompatibleShape(
                "Nuclear norm requires 2D matrix".into(),
            ));
        }

        // Convert input to 2D matrix
        let matrix = input
            .view()
            .into_dimensionality::<Ix2>()
            .map_err(|_| OpError::IncompatibleShape("Failed to convert to 2D array".into()))?;

        // Check for diagonal matrix special case
        if is_diagonal_matrix(&matrix) {
            let nuclear_norm = compute_diagonal_nuclear_norm(&matrix);
            ctx.append_output(scirs2_core::ndarray::arr0(nuclear_norm).into_dyn());
            return Ok(());
        }

        // For general matrices, compute nuclear norm as sum of singular values
        let nuclear_norm = compute_nuclear_norm_improved(&matrix);

        ctx.append_output(scirs2_core::ndarray::arr0(nuclear_norm).into_dyn());
        Ok(())
    }

    fn grad(&self, ctx: &mut GradientContext<F>) {
        let grad_output = ctx.output_grad();
        let input = ctx.input(0);
        let g = ctx.graph();

        // Evaluate inputs to work with concrete values
        let input_array = match input.eval(g) {
            Ok(arr) => arr,
            Err(_) => {
                ctx.append_input_grad(0, None);
                return;
            }
        };

        let grad_output_array = match grad_output.eval(g) {
            Ok(arr) => arr,
            Err(_) => {
                ctx.append_input_grad(0, None);
                return;
            }
        };

        let grad_scalar = grad_output_array[[]];

        if let Ok(matrix) = input_array.view().into_dimensionality::<Ix2>() {
            let grad_matrix = compute_nuclear_norm_gradient_improved(&matrix, grad_scalar);
            let grad_tensor = tensor_ops::convert_to_tensor(grad_matrix.into_dyn(), g);
            ctx.append_input_grad(0, Some(grad_tensor));
            return;
        }

        // Fallback
        ctx.append_input_grad(0, None);
    }
}

// Helper functions

/// Check if matrix is diagonal
#[allow(dead_code)]
fn is_diagonal_matrix<F: Float>(matrix: &ArrayView2<F>) -> bool {
    let (m, n) = matrix.dim();
    for i in 0..m {
        for j in 0..n {
            if i != j && matrix[[i, j]].abs() > F::epsilon() {
                return false;
            }
        }
    }
    true
}

/// Compute nuclear norm for diagonal matrix
#[allow(dead_code)]
fn compute_diagonal_nuclear_norm<F: Float>(matrix: &ArrayView2<F>) -> F {
    let (m, n) = matrix.dim();
    let mut sum = F::zero();
    let min_dim = m.min(n);

    for i in 0..min_dim {
        sum += matrix[[i, i]].abs();
    }

    sum
}

/// Compute sign gradient for diagonal matrix
#[allow(dead_code)]
fn compute_diagonal_sign_gradient<F: Float>(matrix: &ArrayView2<F>) -> Array2<F> {
    let (m, n) = matrix.dim();
    let mut grad_matrix = Array2::zeros((m, n));
    let min_dim = m.min(n);

    for i in 0..min_dim {
        let diag_val = matrix[[i, i]];
        grad_matrix[[i, i]] = if diag_val > F::zero() {
            F::one()
        } else if diag_val < F::zero() {
            -F::one()
        } else {
            F::zero()
        };
    }

    grad_matrix
}

/// Power iteration for spectral norm
#[allow(dead_code)]
fn power_iteration_spectral<F: Float + scirs2_core::ndarray::ScalarOperand>(
    matrix: &ArrayView2<F>,
    max_iter: usize,
    tol: F,
) -> (Array1<F>, F) {
    let (m_n) = matrix.dim();

    // Initialize with normalized vector
    let mut u = Array1::<F>::zeros(m);
    u[0] = F::one();

    // Add some perturbation to avoid getting stuck
    for i in 1..m {
        u[i] = F::from(0.01).expect("Failed to convert constant to float") * F::from(i as f64).expect("Failed to convert to float");
    }

    // Normalize
    let norm = u.iter().fold(F::zero(), |acc, &x| acc + x * x).sqrt();
    if norm > F::epsilon() {
        u.mapv_inplace(|x| x / norm);
    }

    let mut prev_sigma = F::zero();

    for _iter in 0..max_iter {
        // A * u
        let au = matrix.dot(&u);

        // A^T * (A * u)
        let atau = matrix.t().dot(&au);

        // Compute norm (approximate eigenvalue of A^T * A)
        let sigma = atau.iter().fold(F::zero(), |acc, &x| acc + x * x).sqrt();

        // Check convergence
        if (sigma - prev_sigma).abs() < tol {
            // Final computation of actual singular value
            let au_final = matrix.dot(&u);
            let sigma_final = au_final
                .iter()
                .fold(F::zero(), |acc, &x| acc + x * x)
                .sqrt();
            return (u, sigma_final);
        }

        prev_sigma = sigma;

        // Normalize for next iteration
        let norm = atau.iter().fold(F::zero(), |acc, &x| acc + x * x).sqrt();
        if norm > F::epsilon() {
            u = atau.mapv(|x| x / norm);
        }
    }

    // Final estimate
    let au = matrix.dot(&u);
    let sigma = au.iter().fold(F::zero(), |acc, &x| acc + x * x).sqrt();
    (u, sigma)
}

/// Compute gradient for spectral norm using proper SVD gradient computation
#[allow(dead_code)]
fn compute_spectral_norm_gradient<F: Float + scirs2_core::ndarray::ScalarOperand>(
    matrix: &ArrayView2<F>,
    grad_scalar: F,
) -> Array2<F> {
    let (m, n) = matrix.dim();

    // Special handling for diagonal matrices
    if is_diagonal_matrix(matrix) {
        let mut grad_matrix = Array2::zeros((m, n));
        let min_dim = m.min(n);

        // Find the largest diagonal element
        let mut max_idx = 0;
        let mut max_val = F::zero();
        for i in 0..min_dim {
            let abs_val = matrix[[i, i]].abs();
            if abs_val > max_val {
                max_val = abs_val;
                max_idx = i;
            }
        }

        // Gradient is 1 at the position of the largest singular value
        grad_matrix[[max_idx, max_idx]] = grad_scalar;

        return grad_matrix;
    }

    // For general matrices, recompute the singular vectors
    let (u, sigma) = power_iteration_spectral(matrix, 50, F::from(1e-8).expect("Failed to convert constant to float"));

    // Compute v = A^T * u / sigma
    let v = if sigma > F::epsilon() {
        matrix.t().dot(&u) / sigma
    } else {
        Array1::zeros(n)
    };

    // Create outer product u * v^T
    let mut grad_matrix = Array2::zeros((m, n));
    for i in 0..m {
        for j in 0..n {
            grad_matrix[[i, j]] = u[i] * v[j] * grad_scalar;
        }
    }

    grad_matrix
}

/// Improved nuclear norm computation using better SVD approximation
#[allow(dead_code)]
fn compute_nuclear_norm_improved<F: Float + scirs2_core::ndarray::ScalarOperand>(
    matrix: &ArrayView2<F>,
) -> F {
    let (m, n) = matrix.dim();
    let min_dim = m.min(n);

    // For small matrices, use a simple approximation
    if min_dim <= 2 {
        // Sum of absolute values of diagonal elements as approximation
        let mut nuclear_norm = F::zero();
        for i in 0..min_dim {
            nuclear_norm += matrix[[i, i]].abs();
        }
        return nuclear_norm;
    }

    // For larger matrices, use power iteration to estimate singular values
    let mut working_matrix = matrix.to_owned();
    let mut nuclear_norm = F::zero();
    let max_rank = (min_dim.min(5)) as usize; // Limit iterations for performance

    for _ in 0..max_rank {
        let (u, sigma) = power_iteration_spectral(&working_matrix.view(), 20, F::from(1e-6).expect("Failed to convert constant to float"));

        if sigma < F::epsilon() * F::from(10.0).expect("Failed to convert constant to float") {
            break;
        }

        nuclear_norm += sigma;

        // Simple deflation: subtract a rank-1 approximation
        let at = working_matrix.t();
        let v = at.dot(&u) / sigma;

        // Deflate: A = A - sigma * u * v^T
        for i in 0..m {
            for j in 0..n {
                working_matrix[[i, j]] -= sigma * u[i] * v[j];
            }
        }
    }

    nuclear_norm
}

/// Improved nuclear norm gradient computation
#[allow(dead_code)]
fn compute_nuclear_norm_gradient_improved<F: Float + scirs2_core::ndarray::ScalarOperand>(
    matrix: &ArrayView2<F>,
    grad_scalar: F,
) -> Array2<F> {
    let (m, n) = matrix.dim();

    // Handle diagonal matrix case
    if is_diagonal_matrix(matrix) {
        let mut grad_matrix = Array2::zeros((m, n));
        let min_dim = m.min(n);

        // Gradient is sign of diagonal elements
        for i in 0..min_dim {
            let diag_val = matrix[[i, i]];
            grad_matrix[[i, i]] = if diag_val > F::zero() {
                grad_scalar
            } else if diag_val < F::zero() {
                -grad_scalar
            } else {
                F::zero()
            };
        }

        return grad_matrix;
    }

    // For general matrices, use approximate SVD-based gradient
    // This is a simplified version that accumulates gradients from multiple singular vectors
    let mut grad_matrix = Array2::zeros((m, n));
    let mut working_matrix = matrix.to_owned();
    let min_dim = m.min(n);
    let max_rank = (min_dim.min(3)) as usize; // Limit for performance

    for _ in 0..max_rank {
        let (u, sigma) = power_iteration_spectral(&working_matrix.view(), 10, F::from(1e-6).expect("Failed to convert constant to float"));

        if sigma < F::epsilon() * F::from(10.0).expect("Failed to convert constant to float") {
            break;
        }

        // Compute v = A^T * u / sigma
        let v = if sigma > F::epsilon() {
            working_matrix.t().dot(&u) / sigma
        } else {
            Array1::zeros(n)
        };

        // Add contribution from this singular vector pair
        for i in 0..m {
            for j in 0..n {
                grad_matrix[[i, j]] += u[i] * v[j] * grad_scalar;
            }
        }

        // Simple deflation for next iteration
        for i in 0..m {
            for j in 0..n {
                working_matrix[[i, j]] -= sigma * u[i] * v[j];
            }
        }
    }

    grad_matrix
}

// Public API functions

#[allow(dead_code)]
pub fn frobenius_norm<'g, F: Float>(matrix: &Tensor<'g, F>) -> Tensor<'g, F> {
    let g = matrix.graph();
    Tensor::builder(g)
        .append_input(_matrix, false)
        .build(FrobeniusNormOp)
}

#[allow(dead_code)]
pub fn spectral_norm<'g, F: Float + scirs2_core::ndarray::ScalarOperand>(
    matrix: &Tensor<'g, F>,
) -> Tensor<'g, F> {
    let g = matrix.graph();
    Tensor::builder(g)
        .append_input(matrix, false)
        .build(SpectralNormOp)
}

#[allow(dead_code)]
pub fn nuclear_norm<'g, F: Float + scirs2_core::ndarray::ScalarOperand>(
    matrix: &Tensor<'g, F>,
) -> Tensor<'g, F> {
    let g = matrix.graph();
    Tensor::builder(g)
        .append_input(matrix, false)
        .build(NuclearNormOp)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::graph::Graph;
    use crate::tensor_ops;
    use scirs2_core::ndarray::array;

    #[test]
    fn test_frobenius_norm_gradient() {
        let g = Graph::<f64>::new();
        
        // Create test matrix
        let matrix = tensor_ops::convert_to_tensor(array![[3.0, 4.0], [0.0, 0.0]].into_dyn(), &g);
        
        // Compute Frobenius norm
        let norm = frobenius_norm(&matrix);
        
        // Expected norm: sqrt(3^2 + 4^2) = 5.0
        let result = norm.eval(&g).expect("Operation failed");
        assert!((result[[]] - 5.0).abs() < 1e-6);
        
        // Test gradient
        let grad_norm = tensor_ops::scalar(1.0, &g);
        let grads = g.grad(&norm, &[matrix.clone()]);
        
        // Expected gradient: [3/5, 4/5; 0, 0]
        let grad_result = grads[0].as_ref().expect("Operation failed").eval(&g).expect("Operation failed");
        let expected = array![[0.6, 0.8], [0.0, 0.0]];
        
        for i in 0..2 {
            for j in 0..2 {
                assert!((grad_result[[i, j]] - expected[[i, j]]).abs() < 1e-6);
            }
        }
    }

    #[test]
    fn test_diagonal_nuclear_norm_gradient() {
        let g = Graph::<f64>::new();
        
        // Create diagonal matrix
        let matrix = tensor_ops::convert_to_tensor(array![[2.0, 0.0], [0.0, -3.0]].into_dyn(), &g);
        
        // Compute nuclear norm
        let norm = nuclear_norm(&matrix);
        
        // Expected norm: |2| + |-3| = 5.0
        let result = norm.eval(&g).expect("Operation failed");
        assert!((result[[]] - 5.0).abs() < 1e-6);
        
        // Test gradient
        let grads = g.grad(&norm, &[matrix.clone()]);
        
        // Expected gradient: sign(diag) = [1, 0; 0, -1]
        let grad_result = grads[0].as_ref().expect("Operation failed").eval(&g).expect("Operation failed");
        let expected = array![[1.0, 0.0], [0.0, -1.0]];
        
        for i in 0..2 {
            for j in 0..2 {
                assert!((grad_result[[i, j]] - expected[[i, j]]).abs() < 1e-6);
            }
        }
    }
}
