use crate::{LPProblem, Variable, Constraint, Objective, ConstraintType, ObjectiveType};

/// Represents a Simplex Tableau for a Linear Programming problem.
#[derive(Debug, Clone)]
pub struct SimplexTableau {
    pub matrix: Vec<Vec<f64>>,
    pub num_vars: usize, // Number of original decision variables
    pub num_constraints: usize, // Number of constraints (rows in the tableau, excluding objective)
    pub basic_variables: Vec<usize>, // Indices of basic variables
}

impl SimplexTableau {
    /// Creates a SimplexTableau from an LPProblem in standard form.
    /// Assumes the problem is already in standard form (all constraints are Eq, objective is Minimize).
    pub fn from_lp_problem(problem: &LPProblem, initial_basic_vars_indices: Vec<usize>) -> Self {
        let num_constraints = problem.constraints.len();
        let total_variables = problem.variables.len(); // Includes original, slack, surplus, artificial

        let num_cols = total_variables + 1; // +1 for RHS column
        let num_rows = num_constraints + 1; // +1 for objective row

        let mut matrix = vec![vec![0.0; num_cols]; num_rows];
        
        // Fill objective row (row 0)
        // Coefficients are already negated if original was Maximize
        for (j, &coeff) in problem.objective.coefficients.iter().enumerate() {
            matrix[0][j] = coeff;
        }

        // Fill constraint rows (row 1 to num_constraints)
        for (i, constraint) in problem.constraints.iter().enumerate() {
            for (j, &coeff) in constraint.coefficients.iter().enumerate() {
                matrix[i + 1][j] = coeff;
            }
            matrix[i + 1][num_cols - 1] = constraint.rhs; // RHS column
        }

        SimplexTableau {
            matrix,
            num_vars: total_variables, // Total variables in the tableau
            num_constraints,
            basic_variables: initial_basic_vars_indices,
        }
    }

    /// Runs Phase 1 of the simplex method to find an initial basic feasible solution.
    /// Returns true if a feasible solution is found, false otherwise.
    pub fn run_phase1(&mut self, problem_variables: &[Variable]) -> bool {
        // 1. Create a temporary auxiliary objective row (minimize sum of artificial variables)
        let mut aux_objective_coeffs = vec![0.0; self.num_vars];
        let mut initial_aux_objective_value = 0.0;

        let mut artificial_vars_present = false;

        for (var_idx, var) in problem_variables.iter().enumerate() {
            if var.name.starts_with('a') {
                aux_objective_coeffs[var_idx] = 1.0;
                artificial_vars_present = true;
            }
        }

        if !artificial_vars_present {
            // No artificial variables, Phase 1 is not needed. Problem is already in basic feasible form.
            return true;
        }

        // Adjust the auxiliary objective for basic artificial variables
        // For each constraint row (from 1 to num_constraints), if the basic variable for that row is artificial,
        // subtract that constraint row from the auxiliary objective row to make the coefficient 0 for basic artificial variables.
        for i in 0..self.num_constraints {
            let basic_var_idx = self.basic_variables[i];
            if problem_variables[basic_var_idx].name.starts_with('a') {
                for j in 0..self.num_vars {
                    aux_objective_coeffs[j] -= self.matrix[i + 1][j];
                }
                initial_aux_objective_value -= self.matrix[i + 1][self.num_vars]; // Subtract RHS from objective
            }
        }

        // Replace the current objective row (matrix[0]) with the auxiliary objective
        for j in 0..self.num_vars {
            self.matrix[0][j] = aux_objective_coeffs[j];
        }
        self.matrix[0][self.num_vars] = initial_aux_objective_value; // Set the RHS of the objective row

        // Now, run simplex iterations on this auxiliary problem (Phase 1)
        loop {
            if self.is_optimal_phase1() {
                break;
            }

            let pivot_col = match self.find_pivot_column() {
                Some(col) => col,
                None => return false, // Should not happen in a feasible Phase 1 problem
            };

            let pivot_row = match self.find_pivot_row(pivot_col) {
                Some(row) => row,
                None => return false, // Auxiliary problem is unbounded, implies original problem is infeasible
            };

            self.pivot(pivot_row, pivot_col);
        }

        // After Phase 1, check the auxiliary objective value. If it's zero, a BFS for the original problem is found.
        let epsilon = 1e-9;
        if (self.matrix[0][self.num_vars]).abs() < epsilon {
            true // Phase 1 succeeded
        } else {
            false // Phase 1 failed, original problem is infeasible
        }
    }

    pub fn solve(&mut self, original_problem: &LPProblem, standard_objective: &Objective, problem_variables: &[Variable]) -> Result<(Vec<f64>, f64), String> {
        // Step 1: Run Phase 1
        if !self.run_phase1(problem_variables) {
            return Err("Problem is infeasible (Phase 1 failed).".to_string());
        }

        // Step 2: Restore Original Objective (Phase 2 setup)
        // Use the objective coefficients from the standard_problem, which are already adjusted for minimization.
        let mut phase2_objective_coeffs = standard_objective.coefficients.clone();
        phase2_objective_coeffs.resize(self.num_vars, 0.0); // Pad with zeros for aux variables

        // Initialize objective row with these coefficients
        for j in 0..self.num_vars {
            self.matrix[0][j] = phase2_objective_coeffs[j];
        }
        self.matrix[0][self.num_vars] = 0.0; // Initialize RHS of objective row to 0

        // Zero out basic variables in the restored objective row
        for i in 0..self.num_constraints {
            let basic_var_idx = self.basic_variables[i];
            // The coefficient in the original objective function that needs to be eliminated from the objective row
            let basic_var_coeff_in_obj = phase2_objective_coeffs[basic_var_idx]; 

            if (basic_var_coeff_in_obj).abs() > 1e-9 { // Only if coefficient is non-zero
                // Subtract basic_var_coeff_in_obj * constraint_row from objective row
                // This will make the coefficient of basic_var_idx in matrix[0] zero.
                // It also correctly updates the RHS of the objective row (matrix[0][self.num_vars]).
                for j in 0..=self.num_vars {
                    self.matrix[0][j] -= basic_var_coeff_in_obj * self.matrix[i + 1][j];
                }
            }
        }

        // Now, run simplex iterations for Phase 2
        loop {
            if self.is_optimal_phase2() {
                break;
            }

            let pivot_col = match self.find_pivot_column() {
                Some(col) => col,
                None => return Err("Problem is unbounded (no pivot column in Phase 2).".to_string()),
            };

            let pivot_row = match self.find_pivot_row(pivot_col) {
                Some(row) => row,
                None => return Err("Problem is unbounded (no pivot row in Phase 2).".to_string()),
            };

            self.pivot(pivot_row, pivot_col);
        }

        // Step 3: Extract Solution
        let mut solution = vec![0.0; original_problem.variables.len()];
        for i in 0..self.num_constraints {
            let basic_var_idx = self.basic_variables[i];
            // Only consider original decision variables
            if basic_var_idx < original_problem.variables.len() {
                solution[basic_var_idx] = self.matrix[i + 1][self.num_vars];
            }
        }

        // Final objective value
        let mut final_optimal_value = self.matrix[0][self.num_vars];
        if original_problem.objective.objective_type == ObjectiveType::Minimize {
            final_optimal_value *= -1.0;
        }
        

        
        Ok((solution, final_optimal_value))
    }

    /// Finds the pivot column (entering variable) in the tableau.
    /// Returns the index of the pivot column, or None if the tableau is optimal.
    fn find_pivot_column(&self) -> Option<usize> {
        let mut min_coeff = 1e-9; // Use a small epsilon to avoid floating point issues for "negative"
        let mut pivot_col = None;

        for j in 0..self.num_vars { // Iterate through variable columns (excluding RHS)
            if self.matrix[0][j] < min_coeff {
                min_coeff = self.matrix[0][j];
                pivot_col = Some(j);
            }
        }
        pivot_col
    }

    /// Finds the pivot row (leaving variable) using the minimum ratio test.
    /// Returns the index of the pivot row (1-based index for constraint rows), or None if unbounded.
    fn find_pivot_row(&self, pivot_col: usize) -> Option<usize> {
        let mut min_ratio = f64::INFINITY;
        let mut pivot_row = None;
        let epsilon = 1e-9; // For floating point comparisons

        for i in 1..=self.num_constraints { // Iterate through constraint rows (1-based)
            let pivot_val = self.matrix[i][pivot_col];
            let rhs_val = self.matrix[i][self.num_vars];

            if pivot_val > epsilon { // Only consider positive pivot values
                let ratio = rhs_val / pivot_val;
                if ratio < min_ratio {
                    min_ratio = ratio;
                    pivot_row = Some(i);
                }
            }
        }
        pivot_row
    }

    /// Performs the pivot operation on the tableau.
    /// `pivot_row` is 1-based (constraint row index), `pivot_col` is 0-based (variable index).
    fn pivot(&mut self, pivot_row: usize, pivot_col: usize) {
        let pivot_element = self.matrix[pivot_row][pivot_col];

        // Normalize the pivot row
        for j in 0..=self.num_vars { // Iterate through all columns including RHS
            self.matrix[pivot_row][j] /= pivot_element;
        }

        // Zero out other rows in the pivot column
        for i in 0..=self.num_constraints { // Iterate through all rows including objective
            if i != pivot_row {
                let factor = self.matrix[i][pivot_col];
                for j in 0..=self.num_vars { // Iterate through all columns including RHS
                    self.matrix[i][j] -= factor * self.matrix[pivot_row][j];
                }
            }
        }

        // Update basic variables
        self.basic_variables[pivot_row - 1] = pivot_col;
    }

    /// Checks if the Phase 1 auxiliary objective is optimal (all coefficients non-negative).
    fn is_optimal_phase1(&self) -> bool {
        let epsilon = 1e-9;
        for j in 0..self.num_vars {
            if self.matrix[0][j] < -epsilon { // If any coefficient is significantly negative
                return false;
            }
        }
        true
    }

    /// Checks if the Phase 2 objective is optimal (all coefficients non-negative).
    fn is_optimal_phase2(&self) -> bool {
        let epsilon = 1e-9;
        for j in 0..self.num_vars {
            if self.matrix[0][j] < -epsilon { // If any coefficient is significantly negative
                return false;
            }
        }
        true
    }
}



/// Converts an LP problem into standard form.
///
/// Standard form requires:
/// 1. All constraints are equalities.
/// 2. All variables are non-negative.
/// 3. The objective function is to be minimized.
///
/// This function will:
/// - Convert maximize objective to minimize.
/// - Convert >= constraints to <=.
/// - Add slack variables for <= constraints.
/// - Add surplus variables and artificial variables for >= constraints.
/// - Add artificial variables for == constraints.
/// - Handle free variables (if any) by replacing with two non-negative variables (x = x' - x'').
pub fn to_standard_form(problem: &LPProblem) -> (LPProblem, Vec<usize>) {
    let mut new_variables = problem.variables.clone();
    let mut new_constraints = Vec::new();
    let mut new_objective_coefficients = problem.objective.coefficients.clone();
    let mut initial_basic_vars_indices = Vec::new(); // New vector to store indices of initial basic variables

    // Add counters for naming auxiliary variables
    let mut slack_count = 0;
    let mut surplus_count = 0;
    let mut artificial_count = 0;

    // 1. Handle objective function: Convert Maximize to Minimize
    let objective_type = match problem.objective.objective_type {
        ObjectiveType::Maximize => {
            // Multiply objective coefficients by -1
            for coeff in new_objective_coefficients.iter_mut() {
                *coeff *= -1.0;
            }
            ObjectiveType::Minimize
        },
        ObjectiveType::Minimize => ObjectiveType::Minimize,
    };

    // 2. Handle constraints: Convert to equalities and add slack/surplus/artificial variables
    for (_, constraint) in problem.constraints.iter().enumerate() {
        let mut coeffs = constraint.coefficients.clone();
        // Pad coefficients to match the current number of variables before adding new ones
        coeffs.resize(new_variables.len(), 0.0);

        match constraint.constraint_type {
            ConstraintType::Le => {
                // Add slack variable: a_1*x_1 + ... + s_k = rhs
                let slack_var_name = format!("s{}", slack_count);
                let current_var_idx = new_variables.len(); // Index of the new slack variable
                new_variables.push(Variable::new(&slack_var_name, 0.0, f64::INFINITY));
                coeffs.push(1.0); // Coefficient for the slack variable
                new_constraints.push(Constraint::new(coeffs, ConstraintType::Eq, constraint.rhs));
                initial_basic_vars_indices.push(current_var_idx); // This slack variable is basic
                slack_count += 1;
            },
            ConstraintType::Ge => {
                // Add surplus variable and artificial variable: a_1*x_1 + ... - u_k + a_j = rhs
                let surplus_var_name = format!("u{}", surplus_count);
                new_variables.push(Variable::new(&surplus_var_name, 0.0, f64::INFINITY));
                coeffs.push(-1.0); // Coefficient for the surplus variable
                surplus_count += 1;

                let artificial_var_name = format!("a{}", artificial_count);
                let current_var_idx = new_variables.len(); // Index of the new artificial variable
                new_variables.push(Variable::new(&artificial_var_name, 0.0, f64::INFINITY));
                coeffs.push(1.0); // Coefficient for the artificial variable
                new_constraints.push(Constraint::new(coeffs, ConstraintType::Eq, constraint.rhs));
                initial_basic_vars_indices.push(current_var_idx); // This artificial variable is basic
                artificial_count += 1;
            },
            ConstraintType::Eq => {
                // Add artificial variable: a_1*x_1 + ... + a_k = rhs
                let artificial_var_name = format!("a{}", artificial_count);
                let current_var_idx = new_variables.len(); // Index of the new artificial variable
                new_variables.push(Variable::new(&artificial_var_name, 0.0, f64::INFINITY));
                coeffs.push(1.0); // Coefficient for the artificial variable
                new_constraints.push(Constraint::new(coeffs, ConstraintType::Eq, constraint.rhs));
                initial_basic_vars_indices.push(current_var_idx); // This artificial variable is basic
                artificial_count += 1;
            },
        }
    }

    // 3. Handle free variables (variables with negative lower_bound or infinite upper_bound that are not already handled)
    // This implementation assumes all original variables are non-negative by default (lower_bound = 0.0).
    // If a variable has a negative lower_bound, it needs to be transformed.
    // If a variable has f64::NEG_INFINITY as lower_bound and f64::INFINITY as upper_bound, it's a free variable.
    // For now, we assume all variables are non-negative.
    // Future work: implement actual free variable handling (x = x' - x'')

    // Pad objective coefficients for newly added slack/surplus/artificial variables with 0.0
    new_objective_coefficients.resize(new_variables.len(), 0.0);

    (LPProblem::new(
        new_variables,
        new_constraints,
        Objective::new(new_objective_coefficients, objective_type),
    ), initial_basic_vars_indices)
}


#[cfg(test)]
mod tests {
    use super::*;
    use crate::{ConstraintType, LPProblem, ObjectiveType, Variable, Constraint, Objective};

    // Helper to create a simple problem for testing
    fn create_test_lp_problem() -> LPProblem {
        let x = Variable::new("x", 0.0, f64::INFINITY);
        let y = Variable::new("y", 0.0, f64::INFINITY);
        let variables = vec![x, y];

        // x + y <= 10
        let c1 = Constraint::new(vec![1.0, 1.0], ConstraintType::Le, 10.0);
        // 2x - y >= 2
        let c2 = Constraint::new(vec![2.0, -1.0], ConstraintType::Ge, 2.0);
        // x + 2y = 8
        let c3 = Constraint::new(vec![1.0, 2.0], ConstraintType::Eq, 8.0);
        let constraints = vec![c1, c2, c3];

        // Maximize 3x + 2y
        let objective = Objective::new(vec![3.0, 2.0], ObjectiveType::Maximize);

        LPProblem::new(variables, constraints, objective)
    }

    #[test]
    fn test_to_standard_form_objective_conversion() {
        let problem = create_test_lp_problem();
        let (standard_problem, _) = to_standard_form(&problem);

        assert_eq!(standard_problem.objective.objective_type, ObjectiveType::Minimize);
        // Objective coefficients should be negated for Maximize
        assert_eq!(standard_problem.objective.coefficients[0], -3.0);
        assert_eq!(standard_problem.objective.coefficients[1], -2.0);
    }

    #[test]
    fn test_to_standard_form_le_constraint() {
        let problem = create_test_lp_problem();
        let (standard_problem, _) = to_standard_form(&problem);

        // Original: x + y <= 10
        // Expected: x + y + s0 = 10
        let c1_std = &standard_problem.constraints[0];
        assert_eq!(c1_std.constraint_type, ConstraintType::Eq);
        assert_eq!(c1_std.rhs, 10.0);
        assert_eq!(c1_std.coefficients[0], 1.0); // x
        assert_eq!(c1_std.coefficients[1], 1.0); // y
        assert_eq!(c1_std.coefficients[2], 1.0); // s0 (slack)
        assert_eq!(standard_problem.variables[2].name, "s0");
    }

    #[test]
    fn test_to_standard_form_ge_constraint() {
        let problem = create_test_lp_problem();
        let (standard_problem, _) = to_standard_form(&problem);

        // Original: 2x - y >= 2
        // Expected: 2x - y - u0 + a0 = 2 (u0 is surplus, a0 is artificial)
        let c2_std = &standard_problem.constraints[1];
        assert_eq!(c2_std.constraint_type, ConstraintType::Eq);
        assert_eq!(c2_std.rhs, 2.0);
        assert_eq!(c2_std.coefficients[0], 2.0);  // x
        assert_eq!(c2_std.coefficients[1], -1.0); // y
        assert_eq!(c2_std.coefficients[3], -1.0); // u0 (surplus)
        assert_eq!(c2_std.coefficients[4], 1.0);  // a0 (artificial)

        assert_eq!(standard_problem.variables[3].name, "u0");
        assert_eq!(standard_problem.variables[4].name, "a0");
    }

    #[test]
    fn test_to_standard_form_eq_constraint() {
        let problem = create_test_lp_problem();
        let (standard_problem, _) = to_standard_form(&problem);

        // Original: x + 2y = 8
        // Expected: x + 2y + a1 = 8 (a1 is artificial)
        let c3_std = &standard_problem.constraints[2];
        assert_eq!(c3_std.constraint_type, ConstraintType::Eq);
        assert_eq!(c3_std.rhs, 8.0);
        assert_eq!(c3_std.coefficients[0], 1.0);  // x
        assert_eq!(c3_std.coefficients[1], 2.0);  // y
        assert_eq!(c3_std.coefficients[5], 1.0);  // a1 (artificial)

        assert_eq!(standard_problem.variables[5].name, "a1");
    }

    #[test]
    fn test_to_standard_form_variable_counts() {
        let problem = create_test_lp_problem();
        let (standard_problem, _) = to_standard_form(&problem);

        // Original: 2 (x, y)
        // Expected: x, y, s0, u0, a0, a1 = 6 variables
        assert_eq!(standard_problem.variables.len(), 6);
    }

    #[test]
    fn test_objective_coefficients_padded() {
        let problem = create_test_lp_problem();
        let (standard_problem, _) = to_standard_form(&problem);

        // Original objective had 2 coefficients.
        // 4 new variables were added (s0, u0, a0, a1).
        // Objective coefficients should be padded with 0.0s for these new variables.
        assert_eq!(standard_problem.objective.coefficients.len(), 6);
        assert_eq!(standard_problem.objective.coefficients[2], 0.0); // s0
        assert_eq!(standard_problem.objective.coefficients[3], 0.0); // u0
        assert_eq!(standard_problem.objective.coefficients[4], 0.0); // a0
        assert_eq!(standard_problem.objective.coefficients[5], 0.0); // a1
    }

    #[test]
    fn test_simplex_tableau_creation() {
        let problem = create_test_lp_problem();
        let (standard_problem, initial_basic_vars_indices) = to_standard_form(&problem);
        let tableau = SimplexTableau::from_lp_problem(&standard_problem, initial_basic_vars_indices);

        // Assert dimensions
        assert_eq!(tableau.num_constraints, 3);
        assert_eq!(tableau.num_vars, 6);
        assert_eq!(tableau.matrix.len(), 4); // 3 constraints + 1 objective row
        assert_eq!(tableau.matrix[0].len(), 7); // 6 variables + 1 RHS column

        // Assert objective row
        assert_eq!(tableau.matrix[0][0], -3.0);
        assert_eq!(tableau.matrix[0][1], -2.0);
        assert_eq!(tableau.matrix[0][2], 0.0); // s0
        assert_eq!(tableau.matrix[0][3], 0.0); // u0
        assert_eq!(tableau.matrix[0][4], 0.0); // a0
        assert_eq!(tableau.matrix[0][5], 0.0); // a1
        assert_eq!(tableau.matrix[0][6], 0.0); // RHS of objective

        // Assert constraint rows
        // Constraint 1: x + y + s0 = 10
        assert_eq!(tableau.matrix[1][0], 1.0);
        assert_eq!(tableau.matrix[1][1], 1.0);
        assert_eq!(tableau.matrix[1][2], 1.0); // s0
        assert_eq!(tableau.matrix[1][6], 10.0);

        // Constraint 2: 2x - y - u0 + a0 = 2
        assert_eq!(tableau.matrix[2][0], 2.0);
        assert_eq!(tableau.matrix[2][1], -1.0);
        assert_eq!(tableau.matrix[2][3], -1.0); // u0
        assert_eq!(tableau.matrix[2][4], 1.0);  // a0
        assert_eq!(tableau.matrix[2][6], 2.0);

        // Constraint 3: x + 2y + a1 = 8
        assert_eq!(tableau.matrix[3][0], 1.0);
        assert_eq!(tableau.matrix[3][1], 2.0);
        assert_eq!(tableau.matrix[3][5], 1.0);  // a1
        assert_eq!(tableau.matrix[3][6], 8.0);

        // Assert basic variables
        // s0 (index 2) for constraint 0
        // a0 (index 4) for constraint 1
        // a1 (index 5) for constraint 2
        assert_eq!(tableau.basic_variables[0], 2);
        assert_eq!(tableau.basic_variables[1], 4);
        assert_eq!(tableau.basic_variables[2], 5);
    }

    #[test]
    fn test_phase1_successful_completion() {
        let problem = create_test_lp_problem();
        let (standard_problem, initial_basic_vars_indices) = to_standard_form(&problem);
        let mut tableau = SimplexTableau::from_lp_problem(&standard_problem, initial_basic_vars_indices);

        let phase1_result = tableau.run_phase1(&standard_problem.variables);
        assert!(phase1_result, "Phase 1 should succeed for a feasible problem.");

        // After Phase 1, the auxiliary objective value should be 0.0
        let epsilon = 1e-9;
        assert!((tableau.matrix[0][tableau.num_vars]).abs() < epsilon,
                "Expected auxiliary objective value to be 0.0, but got {}", tableau.matrix[0][tableau.num_vars]);
    }

    #[test]
    fn test_simplex_solve_feasible_problem() {
        let original_problem = create_test_lp_problem();
        let (standard_problem, initial_basic_vars_indices) = to_standard_form(&original_problem);
        let mut tableau = SimplexTableau::from_lp_problem(&standard_problem, initial_basic_vars_indices);

        let result = tableau.solve(&original_problem, &standard_problem.objective, &standard_problem.variables);

        assert!(result.is_ok(), "Expected solve to return Ok, but got {:?}", result);

        let (solution, optimal_value) = result.unwrap();
        let epsilon = 1e-9;

        // original problem has 2 variables
        assert_eq!(solution.len(), 2);

        // Expected solution: x=8.0, y=0.0, optimal_value=24.0
        assert!((solution[0] - 8.0).abs() < epsilon, "Expected x to be 8.0, but got {}", solution[0]);
        assert!((solution[1] - 0.0).abs() < epsilon, "Expected y to be 0.0, but got {}", solution[1]);
        assert!((optimal_value - 24.0).abs() < epsilon, "Expected optimal value to be 24.0, but got {}", optimal_value);
    }
}