tensor-calc 0.1.3

use serde::{Deserialize, Serialize};
use crate::symbolic::SymbolicExpr;
use crate::TensorError;

pub type MetricTensor = Vec<Vec<SymbolicExpr>>;
pub type ChristoffelSymbols = Vec<Vec<Vec<SymbolicExpr>>>;
pub type RiemannTensor = Vec<Vec<Vec<Vec<SymbolicExpr>>>>;

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TensorComponent {
    pub indices: Vec<usize>,
    pub expression: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChristoffelResult {
    pub symbols: Vec<TensorComponent>,
    pub dimension: usize,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RiemannResult {
    pub components: Vec<TensorComponent>,
    pub dimension: usize,
}

pub fn parse_metric_tensor(metric_strings: Vec<Vec<String>>, _coords: &[String]) -> Result<MetricTensor, TensorError> {
    let n = metric_strings.len();
    
    // Validate square matrix
    for row in &metric_strings {
        if row.len() != n {
            return Err(TensorError::InvalidMetric("Metric tensor must be square".to_string()));
        }
    }
    
    let mut metric = Vec::with_capacity(n);
    
    for row in metric_strings {
        let mut parsed_row = Vec::with_capacity(n);
        for expr_str in row {
            let expr = SymbolicExpr::parse(&expr_str)
                .map_err(|e| TensorError::InvalidMetric(format!("Failed to parse expression '{}': {}", expr_str, e)))?;
            parsed_row.push(expr);
        }
        metric.push(parsed_row);
    }
    
    Ok(metric)
}

pub fn calculate_christoffel_symbols(metric: &MetricTensor, coords: &[String]) -> Result<ChristoffelResult, TensorError> {
    let n = metric.len();
    let mut symbols = Vec::new();
    
    // Calculate metric inverse (simplified - in real implementation would use proper matrix inversion)
    let metric_inv = calculate_metric_inverse(metric)?;
    
    // Calculate Christoffel symbols: Γ^μ_αβ = (1/2) * g^μν * (∂g_νβ/∂x^α + ∂g_να/∂x^β - ∂g_αβ/∂x^ν)
    for mu in 0..n {
        for alpha in 0..n {
            for beta in 0..n {
                let mut christoffel_expr = SymbolicExpr::Zero;
                
                for nu in 0..n {
                    // ∂g_νβ/∂x^α
                    let d_g_nu_beta_d_alpha = metric[nu][beta].derivative(&coords[alpha]);
                    
                    // ∂g_να/∂x^β  
                    let d_g_nu_alpha_d_beta = metric[nu][alpha].derivative(&coords[beta]);
                    
                    // ∂g_αβ/∂x^ν
                    let d_g_alpha_beta_d_nu = metric[alpha][beta].derivative(&coords[nu]);
                    
                    // (∂g_νβ/∂x^α + ∂g_να/∂x^β - ∂g_αβ/∂x^ν)
                    let sum = SymbolicExpr::Subtract(
                        Box::new(SymbolicExpr::Add(
                            Box::new(d_g_nu_beta_d_alpha),
                            Box::new(d_g_nu_alpha_d_beta),
                        )),
                        Box::new(d_g_alpha_beta_d_nu),
                    );
                    
                    // g^μν * (...)
                    let term = SymbolicExpr::Multiply(
                        Box::new(metric_inv[mu][nu].clone()),
                        Box::new(sum),
                    );
                    
                    christoffel_expr = SymbolicExpr::Add(
                        Box::new(christoffel_expr),
                        Box::new(term),
                    );
                }
                
                // Multiply by 1/2
                christoffel_expr = SymbolicExpr::Multiply(
                    Box::new(SymbolicExpr::Constant(0.5)),
                    Box::new(christoffel_expr),
                );
                
                let simplified = christoffel_expr.simplify();
                
                // Only include non-zero components
                if !simplified.is_zero() {
                    symbols.push(TensorComponent {
                        indices: vec![mu, alpha, beta],
                        expression: simplified.to_string(),
                    });
                }
            }
        }
    }
    
    Ok(ChristoffelResult {
        symbols,
        dimension: n,
    })
}

pub fn calculate_riemann_tensor(metric: &MetricTensor, coords: &[String]) -> Result<RiemannResult, TensorError> {
    let n = metric.len();
    let mut components = Vec::new();
    
    // First calculate Christoffel symbols
    let christoffel_result = calculate_christoffel_symbols(metric, coords)?;
    let christoffel = symbols_to_tensor(&christoffel_result, n);
    
    // Calculate Riemann tensor: R^ρ_σμν = ∂Γ^ρ_σν/∂x^μ - ∂Γ^ρ_σμ/∂x^ν + Γ^ρ_λμ*Γ^λ_σν - Γ^ρ_λν*Γ^λ_σμ
    for rho in 0..n {
        for sigma in 0..n {
            for mu in 0..n {
                for nu in 0..n {
                    let mut riemann_expr = SymbolicExpr::Zero;
                    
                    // ∂Γ^ρ_σν/∂x^μ
                    let d_christoffel_rho_sigma_nu_d_mu = christoffel[rho][sigma][nu].derivative(&coords[mu]);
                    
                    // ∂Γ^ρ_σμ/∂x^ν
                    let d_christoffel_rho_sigma_mu_d_nu = christoffel[rho][sigma][mu].derivative(&coords[nu]);
                    
                    // ∂Γ^ρ_σν/∂x^μ - ∂Γ^ρ_σμ/∂x^ν
                    riemann_expr = SymbolicExpr::Add(
                        Box::new(riemann_expr),
                        Box::new(SymbolicExpr::Subtract(
                            Box::new(d_christoffel_rho_sigma_nu_d_mu),
                            Box::new(d_christoffel_rho_sigma_mu_d_nu),
                        )),
                    );
                    
                    // Add quadratic terms: Γ^ρ_λμ*Γ^λ_σν - Γ^ρ_λν*Γ^λ_σμ
                    for lambda in 0..n {
                        let term1 = SymbolicExpr::Multiply(
                            Box::new(christoffel[rho][lambda][mu].clone()),
                            Box::new(christoffel[lambda][sigma][nu].clone()),
                        );
                        
                        let term2 = SymbolicExpr::Multiply(
                            Box::new(christoffel[rho][lambda][nu].clone()),
                            Box::new(christoffel[lambda][sigma][mu].clone()),
                        );
                        
                        riemann_expr = SymbolicExpr::Add(
                            Box::new(riemann_expr),
                            Box::new(SymbolicExpr::Subtract(
                                Box::new(term1),
                                Box::new(term2),
                            )),
                        );
                    }
                    
                    let simplified = riemann_expr.simplify();
                    
                    // Only include non-zero components
                    if !simplified.is_zero() {
                        components.push(TensorComponent {
                            indices: vec![rho, sigma, mu, nu],
                            expression: simplified.to_string(),
                        });
                    }
                }
            }
        }
    }
    
    Ok(RiemannResult {
        components,
        dimension: n,
    })
}

pub fn calculate_ricci_tensor(metric: &MetricTensor, coords: &[String]) -> Result<RiemannResult, TensorError> {
    let n = metric.len();
    let riemann_result = calculate_riemann_tensor(metric, coords)?;
    let riemann = riemann_result_to_tensor(&riemann_result, n);
    
    let mut components = Vec::new();
    
    // Ricci tensor: R_μν = R^ρ_μρν (contraction over first and third indices)
    for mu in 0..n {
        for nu in 0..n {
            let mut ricci_expr = SymbolicExpr::Zero;
            
            for rho in 0..n {
                ricci_expr = SymbolicExpr::Add(
                    Box::new(ricci_expr),
                    Box::new(riemann[rho][mu][rho][nu].clone()),
                );
            }
            
            let simplified = ricci_expr.simplify();
            
            if !simplified.is_zero() {
                components.push(TensorComponent {
                    indices: vec![mu, nu],
                    expression: simplified.to_string(),
                });
            }
        }
    }
    
    Ok(RiemannResult {
        components,
        dimension: n,
    })
}

pub fn calculate_ricci_scalar(metric: &MetricTensor, coords: &[String]) -> Result<TensorComponent, TensorError> {
    let n = metric.len();
    let ricci_result = calculate_ricci_tensor(metric, coords)?;
    let ricci = ricci_result_to_matrix(&ricci_result, n);
    let metric_inv = calculate_metric_inverse(metric)?;
    
    let mut scalar_expr = SymbolicExpr::Zero;
    
    // Ricci scalar: R = g^μν * R_μν
    for mu in 0..n {
        for nu in 0..n {
            let term = SymbolicExpr::Multiply(
                Box::new(metric_inv[mu][nu].clone()),
                Box::new(ricci[mu][nu].clone()),
            );
            
            scalar_expr = SymbolicExpr::Add(
                Box::new(scalar_expr),
                Box::new(term),
            );
        }
    }
    
    let simplified = scalar_expr.simplify();
    
    Ok(TensorComponent {
        indices: vec![],
        expression: simplified.to_string(),
    })
}

pub fn calculate_einstein_tensor(metric: &MetricTensor, coords: &[String]) -> Result<RiemannResult, TensorError> {
    let n = metric.len();
    let ricci_result = calculate_ricci_tensor(metric, coords)?;
    let ricci = ricci_result_to_matrix(&ricci_result, n);
    let ricci_scalar = calculate_ricci_scalar(metric, coords)?;
    let ricci_scalar_expr = SymbolicExpr::parse(&ricci_scalar.expression)
        .map_err(|e| TensorError::ComputationError(format!("Failed to parse Ricci scalar: {}", e)))?;
    
    let mut components = Vec::new();
    
    // Einstein tensor: G_μν = R_μν - (1/2) * g_μν * R
    for mu in 0..n {
        for nu in 0..n {
            let half_metric_scalar = SymbolicExpr::Multiply(
                Box::new(SymbolicExpr::Constant(0.5)),
                Box::new(SymbolicExpr::Multiply(
                    Box::new(metric[mu][nu].clone()),
                    Box::new(ricci_scalar_expr.clone()),
                )),
            );
            
            let einstein_expr = SymbolicExpr::Subtract(
                Box::new(ricci[mu][nu].clone()),
                Box::new(half_metric_scalar),
            );
            
            let simplified = einstein_expr.simplify();
            
            if !simplified.is_zero() {
                components.push(TensorComponent {
                    indices: vec![mu, nu],
                    expression: simplified.to_string(),
                });
            }
        }
    }
    
    Ok(RiemannResult {
        components,
        dimension: n,
    })
}

// Helper functions

fn calculate_metric_inverse(metric: &MetricTensor) -> Result<MetricTensor, TensorError> {
    let n = metric.len();
    
    // For now, implement a simple 2x2 inverse
    // In a complete implementation, we'd use proper symbolic matrix inversion
    if n == 2 {
        let det = SymbolicExpr::Subtract(
            Box::new(SymbolicExpr::Multiply(
                Box::new(metric[0][0].clone()),
                Box::new(metric[1][1].clone()),
            )),
            Box::new(SymbolicExpr::Multiply(
                Box::new(metric[0][1].clone()),
                Box::new(metric[1][0].clone()),
            )),
        );
        
        let inv_det = SymbolicExpr::Divide(
            Box::new(SymbolicExpr::One),
            Box::new(det),
        );
        
        Ok(vec![
            vec![
                SymbolicExpr::Multiply(Box::new(inv_det.clone()), Box::new(metric[1][1].clone())),
                SymbolicExpr::Multiply(
                    Box::new(inv_det.clone()),
                    Box::new(SymbolicExpr::Subtract(
                        Box::new(SymbolicExpr::Zero),
                        Box::new(metric[0][1].clone()),
                    )),
                ),
            ],
            vec![
                SymbolicExpr::Multiply(
                    Box::new(inv_det.clone()),
                    Box::new(SymbolicExpr::Subtract(
                        Box::new(SymbolicExpr::Zero),
                        Box::new(metric[1][0].clone()),
                    )),
                ),
                SymbolicExpr::Multiply(Box::new(inv_det), Box::new(metric[0][0].clone())),
            ],
        ])
    } else {
        // For higher dimensions, return identity for now
        // A complete implementation would use symbolic matrix inversion
        let mut identity = vec![vec![SymbolicExpr::Zero; n]; n];
        for i in 0..n {
            identity[i][i] = SymbolicExpr::One;
        }
        Ok(identity)
    }
}

fn symbols_to_tensor(christoffel_result: &ChristoffelResult, n: usize) -> ChristoffelSymbols {
    let mut tensor = vec![vec![vec![SymbolicExpr::Zero; n]; n]; n];
    
    for component in &christoffel_result.symbols {
        if component.indices.len() == 3 {
            let i = component.indices[0];
            let j = component.indices[1];
            let k = component.indices[2];
            
            if let Ok(expr) = SymbolicExpr::parse(&component.expression) {
                tensor[i][j][k] = expr;
            }
        }
    }
    
    tensor
}

fn riemann_result_to_tensor(riemann_result: &RiemannResult, n: usize) -> RiemannTensor {
    let mut tensor = vec![vec![vec![vec![SymbolicExpr::Zero; n]; n]; n]; n];
    
    for component in &riemann_result.components {
        if component.indices.len() == 4 {
            let i = component.indices[0];
            let j = component.indices[1];
            let k = component.indices[2];
            let l = component.indices[3];
            
            if let Ok(expr) = SymbolicExpr::parse(&component.expression) {
                tensor[i][j][k][l] = expr;
            }
        }
    }
    
    tensor
}

fn ricci_result_to_matrix(ricci_result: &RiemannResult, n: usize) -> MetricTensor {
    let mut matrix = vec![vec![SymbolicExpr::Zero; n]; n];
    
    for component in &ricci_result.components {
        if component.indices.len() == 2 {
            let i = component.indices[0];
            let j = component.indices[1];
            
            if let Ok(expr) = SymbolicExpr::parse(&component.expression) {
                matrix[i][j] = expr;
            }
        }
    }
    
    matrix
}