use crate::errors::ProblemResult;
use crate::problem::{Problem, ProblemBuilder};

/// Scaling techniques to improve numerical stability
pub struct Scaling {
    /// Method to use for scaling
    pub method: ScalingMethod,
    /// Maximum scaling factor
    pub max_factor: f64,
    /// Minimum scaling factor
    pub min_factor: f64,
}

/// Different scaling methods
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum ScalingMethod {
    /// No scaling
    None,
    /// Equilibration scaling (balance row and column norms)
    Equilibration,
    /// Geometric mean scaling
    Geometric,
    /// Maximum element scaling
    MaxElement,
}

/// Result of scaling operations
#[derive(Debug, Clone)]
pub struct ScalingResult {
    /// The scaled problem
    pub problem: Problem,
    /// Row scaling factors
    pub row_scales: Vec<f64>,
    /// Column scaling factors
    pub column_scales: Vec<f64>,
    /// Scaling statistics
    pub statistics: ScalingStatistics,
}

/// Statistics from scaling operations
#[derive(Debug, Clone)]
pub struct ScalingStatistics {
    pub method: ScalingMethod,
    pub max_row_scale: f64,
    pub min_row_scale: f64,
    pub max_col_scale: f64,
    pub min_col_scale: f64,
    pub condition_number_before: f64,
    pub condition_number_after: f64,
}

impl Default for Scaling {
    fn default() -> Self {
        Scaling {
            method: ScalingMethod::Equilibration,
            max_factor: 1e6,
            min_factor: 1e-6,
        }
    }
}

impl Scaling {
    /// Creates a new scaling configuration
    pub fn new(method: ScalingMethod) -> Self {
        Scaling {
            method,
            ..Default::default()
        }
    }

    /// Applies scaling to the problem
    pub fn apply(&self, problem: &Problem) -> ProblemResult<ScalingResult> {
        match self.method {
            ScalingMethod::None => self.no_scaling(problem),
            ScalingMethod::Equilibration => self.equilibration_scaling(problem),
            ScalingMethod::Geometric => self.geometric_scaling(problem),
            ScalingMethod::MaxElement => self.max_element_scaling(problem),
        }
    }

    /// No scaling - returns the original problem
    fn no_scaling(&self, problem: &Problem) -> ProblemResult<ScalingResult> {
        Ok(ScalingResult {
            problem: problem.clone(),
            row_scales: vec![1.0; problem.num_constraints()],
            column_scales: vec![1.0; problem.num_variables()],
            statistics: ScalingStatistics {
                method: ScalingMethod::None,
                max_row_scale: 1.0,
                min_row_scale: 1.0,
                max_col_scale: 1.0,
                min_col_scale: 1.0,
                condition_number_before: self.estimate_condition_number(problem),
                condition_number_after: self.estimate_condition_number(problem),
            },
        })
    }

    /// Equilibration scaling - balance row and column norms
    fn equilibration_scaling(&self, problem: &Problem) -> ProblemResult<ScalingResult> {
        let mut row_scales = vec![1.0; problem.num_constraints()];
        let mut col_scales = vec![1.0; problem.num_variables()];
        
        let condition_before = self.estimate_condition_number(problem);
        let max_iterations = 10;
        let tolerance = 1e-3;

        // Iterative equilibration
        for _ in 0..max_iterations {
            let mut max_change = 0.0f64;

            // Scale rows
            for (i, constraint) in problem.constraints.iter().enumerate() {
                let row_norm = constraint.coefficients.iter()
                    .enumerate()
                    .map(|(j, &c)| (c * col_scales[j]).powi(2))
                    .sum::<f64>()
                    .sqrt();

                if row_norm > 1e-12 {
                    let new_scale = row_scales[i] / row_norm;
                    let clamped_scale = new_scale.max(self.min_factor).min(self.max_factor);
                    max_change = max_change.max((clamped_scale / row_scales[i] - 1.0).abs());
                    row_scales[i] = clamped_scale;
                }
            }

            // Scale columns
            for j in 0..problem.num_variables() {
                let col_norm = problem.constraints.iter()
                    .enumerate()
                    .map(|(i, constraint)| {
                        if j < constraint.coefficients.len() {
                            (constraint.coefficients[j] * row_scales[i]).powi(2)
                        } else {
                            0.0
                        }
                    })
                    .sum::<f64>()
                    .sqrt();

                if col_norm > 1e-12 {
                    let new_scale = col_scales[j] / col_norm;
                    let clamped_scale = new_scale.max(self.min_factor).min(self.max_factor);
                    max_change = max_change.max((clamped_scale / col_scales[j] - 1.0).abs());
                    col_scales[j] = clamped_scale;
                }
            }

            if max_change < tolerance {
                break;
            }
        }

        let scaled_problem = self.apply_scaling(problem, &row_scales, &col_scales)?;
        let condition_after = self.estimate_condition_number(&scaled_problem);

        Ok(ScalingResult {
            problem: scaled_problem,
            row_scales: row_scales.clone(),
            column_scales: col_scales.clone(),
            statistics: ScalingStatistics {
                method: ScalingMethod::Equilibration,
                max_row_scale: row_scales.iter().fold(f64::NEG_INFINITY, |a, &b| a.max(b)),
                min_row_scale: row_scales.iter().fold(f64::INFINITY, |a, &b| a.min(b)),
                max_col_scale: col_scales.iter().fold(f64::NEG_INFINITY, |a, &b| a.max(b)),
                min_col_scale: col_scales.iter().fold(f64::INFINITY, |a, &b| a.min(b)),
                condition_number_before: condition_before,
                condition_number_after: condition_after,
            },
        })
    }

    /// Geometric mean scaling
    fn geometric_scaling(&self, problem: &Problem) -> ProblemResult<ScalingResult> {
        let mut row_scales = vec![1.0; problem.num_constraints()];
        let mut col_scales = vec![1.0; problem.num_variables()];

        // Compute geometric mean of row elements
        for (i, constraint) in problem.constraints.iter().enumerate() {
            let mut product = 1.0;
            let mut count = 0;
            
            for &coeff in &constraint.coefficients {
                if coeff.abs() > 1e-12 {
                    product *= coeff.abs();
                    count += 1;
                }
            }
            
            if count > 0 {
                let geom_mean = product.powf(1.0 / count as f64);
                row_scales[i] = (1.0 / geom_mean).max(self.min_factor).min(self.max_factor);
            }
        }

        // Compute geometric mean of column elements
        for j in 0..problem.num_variables() {
            let mut product = 1.0;
            let mut count = 0;
            
            for constraint in &problem.constraints {
                if j < constraint.coefficients.len() {
                    let coeff = constraint.coefficients[j];
                    if coeff.abs() > 1e-12 {
                        product *= coeff.abs();
                        count += 1;
                    }
                }
            }
            
            if count > 0 {
                let geom_mean = product.powf(1.0 / count as f64);
                col_scales[j] = (1.0 / geom_mean).max(self.min_factor).min(self.max_factor);
            }
        }

        let condition_before = self.estimate_condition_number(problem);
        let scaled_problem = self.apply_scaling(problem, &row_scales, &col_scales)?;
        let condition_after = self.estimate_condition_number(&scaled_problem);

        Ok(ScalingResult {
            problem: scaled_problem,
            row_scales: row_scales.clone(),
            column_scales: col_scales.clone(),
            statistics: ScalingStatistics {
                method: ScalingMethod::Geometric,
                max_row_scale: row_scales.iter().fold(f64::NEG_INFINITY, |a, &b| a.max(b)),
                min_row_scale: row_scales.iter().fold(f64::INFINITY, |a, &b| a.min(b)),
                max_col_scale: col_scales.iter().fold(f64::NEG_INFINITY, |a, &b| a.max(b)),
                min_col_scale: col_scales.iter().fold(f64::INFINITY, |a, &b| a.min(b)),
                condition_number_before: condition_before,
                condition_number_after: condition_after,
            },
        })
    }

    /// Maximum element scaling
    fn max_element_scaling(&self, problem: &Problem) -> ProblemResult<ScalingResult> {
        let mut row_scales = vec![1.0; problem.num_constraints()];
        let mut col_scales = vec![1.0; problem.num_variables()];

        // Scale rows based on maximum element
        for (i, constraint) in problem.constraints.iter().enumerate() {
            let max_element = constraint.coefficients.iter()
                .map(|&c| c.abs())
                .fold(0.0, f64::max);
            
            if max_element > 1e-12 {
                row_scales[i] = (1.0 / max_element).max(self.min_factor).min(self.max_factor);
            }
        }

        // Scale columns based on maximum element
        for j in 0..problem.num_variables() {
            let max_element = problem.constraints.iter()
                .filter_map(|constraint| {
                    if j < constraint.coefficients.len() {
                        Some(constraint.coefficients[j].abs())
                    } else {
                        None
                    }
                })
                .fold(0.0, f64::max);
            
            if max_element > 1e-12 {
                col_scales[j] = (1.0 / max_element).max(self.min_factor).min(self.max_factor);
            }
        }

        let condition_before = self.estimate_condition_number(problem);
        let scaled_problem = self.apply_scaling(problem, &row_scales, &col_scales)?;
        let condition_after = self.estimate_condition_number(&scaled_problem);

        Ok(ScalingResult {
            problem: scaled_problem,
            row_scales: row_scales.clone(),
            column_scales: col_scales.clone(),
            statistics: ScalingStatistics {
                method: ScalingMethod::MaxElement,
                max_row_scale: row_scales.iter().fold(f64::NEG_INFINITY, |a, &b| a.max(b)),
                min_row_scale: row_scales.iter().fold(f64::INFINITY, |a, &b| a.min(b)),
                max_col_scale: col_scales.iter().fold(f64::NEG_INFINITY, |a, &b| a.max(b)),
                min_col_scale: col_scales.iter().fold(f64::INFINITY, |a, &b| a.min(b)),
                condition_number_before: condition_before,
                condition_number_after: condition_after,
            },
        })
    }

    /// Applies scaling factors to create a scaled problem
    fn apply_scaling(&self, problem: &Problem, row_scales: &[f64], col_scales: &[f64]) -> ProblemResult<Problem> {
        let mut builder = ProblemBuilder::new();

        // Scale variables (bounds)
        for (j, var) in problem.variables.iter().enumerate() {
            let scale = col_scales[j];
            let scaled_lower = var.lower_bound * scale;
            let scaled_upper = var.upper_bound * scale;
            builder = builder.add_variable(&var.name, scaled_lower, scaled_upper)?;
        }

        // Scale constraints
        for (i, constraint) in problem.constraints.iter().enumerate() {
            let row_scale = row_scales[i];
            let mut scaled_coeffs = Vec::new();
            
            for (j, &coeff) in constraint.coefficients.iter().enumerate() {
                scaled_coeffs.push(coeff * row_scale * col_scales[j]);
            }
            
            let scaled_rhs = constraint.rhs * row_scale;
            builder = builder.add_constraint(scaled_coeffs, constraint.constraint_type, scaled_rhs)?;
        }

        // Scale objective
        let mut scaled_obj_coeffs = Vec::new();
        for (j, &coeff) in problem.objective.coefficients.iter().enumerate() {
            scaled_obj_coeffs.push(coeff * col_scales[j]);
        }
        
        builder = builder.objective(scaled_obj_coeffs, problem.objective.objective_type)?;

        builder.build()
    }

    /// Estimates the condition number of the problem matrix
    fn estimate_condition_number(&self, problem: &Problem) -> f64 {
        // Simple estimation based on the ratio of largest to smallest non-zero elements
        let mut min_element = f64::INFINITY;
        let mut max_element = 0.0f64;

        for constraint in &problem.constraints {
            for &coeff in &constraint.coefficients {
                let abs_coeff = coeff.abs();
                if abs_coeff > 1e-12 {
                    min_element = min_element.min(abs_coeff);
                    max_element = max_element.max(abs_coeff);
                }
            }
        }

        if min_element == f64::INFINITY || max_element == 0.0 {
            1.0
        } else {
            max_element / min_element
        }
    }

    /// Unscales a solution vector
    pub fn unscale_solution(&self, scaled_solution: &[f64], col_scales: &[f64]) -> Vec<f64> {
        scaled_solution.iter()
            .zip(col_scales.iter())
            .map(|(&val, &scale)| val * scale)
            .collect()
    }

    /// Unscales dual variables
    pub fn unscale_duals(&self, scaled_duals: &[f64], row_scales: &[f64]) -> Vec<f64> {
        scaled_duals.iter()
            .zip(row_scales.iter())
            .map(|(&val, &scale)| val * scale)
            .collect()
    }
}

impl std::fmt::Display for ScalingStatistics {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        writeln!(f, "Scaling Statistics:")?;
        writeln!(f, "  Method: {:?}", self.method)?;
        writeln!(f, "  Row scales: [{:.2e}, {:.2e}]", self.min_row_scale, self.max_row_scale)?;
        writeln!(f, "  Column scales: [{:.2e}, {:.2e}]", self.min_col_scale, self.max_col_scale)?;
        writeln!(f, "  Condition number: {:.2e} -> {:.2e}", 
                 self.condition_number_before, self.condition_number_after)?;
        let improvement = if self.condition_number_before > 0.0 {
            self.condition_number_before / self.condition_number_after
        } else {
            1.0
        };
        writeln!(f, "  Improvement factor: {:.2e}", improvement)?;
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::problem::*;
    use approx::assert_abs_diff_eq;

    #[test]
    fn test_no_scaling() {
        let problem = create_test_problem();
        let scaling = Scaling::new(ScalingMethod::None);
        let result = scaling.apply(&problem).unwrap();
        
        assert_eq!(result.row_scales, vec![1.0; problem.num_constraints()]);
        assert_eq!(result.column_scales, vec![1.0; problem.num_variables()]);
        assert_eq!(result.statistics.method, ScalingMethod::None);
    }

    #[test]
    fn test_equilibration_scaling() {
        let problem = create_poorly_scaled_problem();
        let scaling = Scaling::new(ScalingMethod::Equilibration);
        let result = scaling.apply(&problem).unwrap();
        
        // Condition number should improve
        assert!(result.statistics.condition_number_after < result.statistics.condition_number_before);
        
        // Scaling factors should be reasonable
        for &scale in &result.row_scales {
            assert!(scale >= scaling.min_factor && scale <= scaling.max_factor);
        }
        for &scale in &result.column_scales {
            assert!(scale >= scaling.min_factor && scale <= scaling.max_factor);
        }
    }

    #[test]
    fn test_geometric_scaling() {
        let problem = create_test_problem();
        let scaling = Scaling::new(ScalingMethod::Geometric);
        let result = scaling.apply(&problem).unwrap();
        
        assert_eq!(result.statistics.method, ScalingMethod::Geometric);
        assert!(result.row_scales.iter().all(|&s| s > 0.0));
        assert!(result.column_scales.iter().all(|&s| s > 0.0));
    }

    #[test]
    fn test_max_element_scaling() {
        let problem = create_test_problem();
        let scaling = Scaling::new(ScalingMethod::MaxElement);
        let result = scaling.apply(&problem).unwrap();
        
        assert_eq!(result.statistics.method, ScalingMethod::MaxElement);
        assert!(result.row_scales.iter().all(|&s| s > 0.0));
        assert!(result.column_scales.iter().all(|&s| s > 0.0));
    }

    #[test]
    fn test_solution_unscaling() {
        let scaling = Scaling::new(ScalingMethod::None);
        let scaled_solution = vec![2.0, 4.0];
        let col_scales = vec![0.5, 2.0];
        
        let unscaled = scaling.unscale_solution(&scaled_solution, &col_scales);
        assert_abs_diff_eq!(unscaled[0], 1.0);
        assert_abs_diff_eq!(unscaled[1], 8.0);
    }

    fn create_test_problem() -> Problem {
        ProblemBuilder::new()
            .add_non_negative_variable("x").unwrap()
            .add_non_negative_variable("y").unwrap()
            .add_constraint_by_name(
                &[("x", 1.0), ("y", 1.0)],
                ConstraintType::LessEqual,
                10.0,
            ).unwrap()
            .objective_by_name(
                &[("x", 2.0), ("y", 3.0)],
                ObjectiveType::Maximize,
            ).unwrap()
            .build()
            .unwrap()
    }

    fn create_poorly_scaled_problem() -> Problem {
        ProblemBuilder::new()
            .add_non_negative_variable("x").unwrap()
            .add_non_negative_variable("y").unwrap()
            .add_constraint_by_name(
                &[("x", 1000.0), ("y", 0.001)],  // Poorly scaled
                ConstraintType::LessEqual,
                10.0,
            ).unwrap()
            .objective_by_name(
                &[("x", 0.001), ("y", 1000.0)],  // Poorly scaled
                ObjectiveType::Maximize,
            ).unwrap()
            .build()
            .unwrap()
    }
}