use crate::errors::{ProblemError, ProblemResult};
#[cfg(feature = "serde")]
use serde::{Deserialize, Serialize};
use std::fmt;
use approx::AbsDiffEq;

/// Represents the type of the objective function
#[derive(Debug, PartialEq, Clone, Copy)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub enum ObjectiveType {
    /// Minimize the objective function
    Minimize,
    /// Maximize the objective function
    Maximize,
}

impl fmt::Display for ObjectiveType {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            ObjectiveType::Minimize => write!(f, "minimize"),
            ObjectiveType::Maximize => write!(f, "maximize"),
        }
    }
}

/// Represents the objective function `c_1*x_1 + c_2*x_2 + ...`
#[derive(Debug, Clone, PartialEq)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct Objective {
    /// Coefficients of the variables in the objective function
    pub coefficients: Vec<f64>,
    /// Type of optimization (minimize or maximize)
    pub objective_type: ObjectiveType,
    /// Constant term added to the objective value
    pub constant: f64,
    /// Optional name for the objective
    pub name: Option<String>,
}

impl Objective {
    /// Creates a new objective function
    pub fn new(coefficients: Vec<f64>, objective_type: ObjectiveType) -> ProblemResult<Self> {
        if coefficients.is_empty() {
            return Err(ProblemError::InvalidObjective {
                reason: "Objective must have at least one coefficient".to_string(),
            });
        }

        if coefficients.iter().any(|&c| !c.is_finite()) {
            return Err(ProblemError::InvalidObjective {
                reason: "All coefficients must be finite".to_string(),
            });
        }

        Ok(Objective {
            coefficients,
            objective_type,
            constant: 0.0,
            name: None,
        })
    }

    /// Creates a new objective with a constant term
    pub fn new_with_constant(
        coefficients: Vec<f64>,
        objective_type: ObjectiveType,
        constant: f64,
    ) -> ProblemResult<Self> {
        if !constant.is_finite() {
            return Err(ProblemError::InvalidObjective {
                reason: "Constant term must be finite".to_string(),
            });
        }

        let mut objective = Self::new(coefficients, objective_type)?;
        objective.constant = constant;
        Ok(objective)
    }

    /// Creates a new objective with a name
    pub fn new_with_name(
        name: impl Into<String>,
        coefficients: Vec<f64>,
        objective_type: ObjectiveType,
    ) -> ProblemResult<Self> {
        let mut objective = Self::new(coefficients, objective_type)?;
        objective.name = Some(name.into());
        Ok(objective)
    }

    /// Sets the name of the objective
    pub fn with_name(mut self, name: impl Into<String>) -> Self {
        self.name = Some(name.into());
        self
    }

    /// Sets the constant term
    pub fn with_constant(mut self, constant: f64) -> Self {
        self.constant = constant;
        self
    }

    /// Gets the number of variables in this objective
    pub fn num_variables(&self) -> usize {
        self.coefficients.len()
    }

    /// Evaluates the objective function for a given variable assignment
    pub fn evaluate(&self, variables: &[f64]) -> ProblemResult<f64> {
        if variables.len() != self.coefficients.len() {
            return Err(ProblemError::DimensionMismatch {
                expected: self.coefficients.len(),
                actual: variables.len(),
            });
        }

        let linear_part: f64 = self.coefficients
            .iter()
            .zip(variables.iter())
            .map(|(&coeff, &var)| coeff * var)
            .sum();

        Ok(linear_part + self.constant)
    }

    /// Converts the objective to minimization form
    /// If already minimizing, returns a clone. If maximizing, negates coefficients and constant.
    pub fn to_minimization(&self) -> Self {
        match self.objective_type {
            ObjectiveType::Minimize => self.clone(),
            ObjectiveType::Maximize => Objective {
                coefficients: self.coefficients.iter().map(|&c| -c).collect(),
                objective_type: ObjectiveType::Minimize,
                constant: -self.constant,
                name: self.name.clone(),
            },
        }
    }

    /// Converts the objective to maximization form
    /// If already maximizing, returns a clone. If minimizing, negates coefficients and constant.
    pub fn to_maximization(&self) -> Self {
        match self.objective_type {
            ObjectiveType::Maximize => self.clone(),
            ObjectiveType::Minimize => Objective {
                coefficients: self.coefficients.iter().map(|&c| -c).collect(),
                objective_type: ObjectiveType::Maximize,
                constant: -self.constant,
                name: self.name.clone(),
            },
        }
    }

    /// Multiplies the objective by a scalar
    pub fn multiply_by(&mut self, scalar: f64) {
        if scalar < 0.0 {
            // Flip the objective type when multiplying by negative
            self.objective_type = match self.objective_type {
                ObjectiveType::Minimize => ObjectiveType::Maximize,
                ObjectiveType::Maximize => ObjectiveType::Minimize,
            };
        }
        
        for coeff in &mut self.coefficients {
            *coeff *= scalar;
        }
        self.constant *= scalar;
    }

    /// Checks if the objective is linear (all finite coefficients)
    pub fn is_linear(&self) -> bool {
        self.coefficients.iter().all(|&c| c.is_finite()) && self.constant.is_finite()
    }

    /// Returns the gradient of the objective (which is just the coefficients for linear objectives)
    pub fn gradient(&self) -> &[f64] {
        &self.coefficients
    }

    /// Computes the directional derivative in a given direction
    pub fn directional_derivative(&self, direction: &[f64]) -> ProblemResult<f64> {
        if direction.len() != self.coefficients.len() {
            return Err(ProblemError::DimensionMismatch {
                expected: self.coefficients.len(),
                actual: direction.len(),
            });
        }

        Ok(self.coefficients
            .iter()
            .zip(direction.iter())
            .map(|(&coeff, &dir)| coeff * dir)
            .sum())
    }
}

impl fmt::Display for Objective {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{} ", self.objective_type)?;

        if let Some(ref name) = self.name {
            write!(f, "{}: ", name)?;
        }

        for (i, &coeff) in self.coefficients.iter().enumerate() {
            if i == 0 {
                if coeff == 1.0 {
                    write!(f, "x{}", i)?;
                } else if coeff == -1.0 {
                    write!(f, "-x{}", i)?;
                } else {
                    write!(f, "{:.3}x{}", coeff, i)?;
                }
            } else if coeff > 0.0 {
                if coeff == 1.0 {
                    write!(f, " + x{}", i)?;
                } else {
                    write!(f, " + {:.3}x{}", coeff, i)?;
                }
            } else if coeff < 0.0 {
                if coeff == -1.0 {
                    write!(f, " - x{}", i)?;
                } else {
                    write!(f, " - {:.3}x{}", -coeff, i)?;
                }
            }
        }

        if self.constant != 0.0 {
            if self.constant > 0.0 {
                write!(f, " + {:.3}", self.constant)?;
            } else {
                write!(f, " - {:.3}", -self.constant)?;
            }
        }

        Ok(())
    }
}

impl AbsDiffEq for Objective {
    type Epsilon = f64;

    fn default_epsilon() -> Self::Epsilon {
        f64::EPSILON
    }

    fn abs_diff_eq(&self, other: &Self, epsilon: Self::Epsilon) -> bool {
        self.objective_type == other.objective_type &&
        self.coefficients.len() == other.coefficients.len() &&
        self.coefficients.iter()
            .zip(other.coefficients.iter())
            .all(|(&a, &b)| (a - b).abs() <= epsilon) &&
        (self.constant - other.constant).abs() <= epsilon
    }
}

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

    #[test]
    fn test_objective_creation() {
        let objective = Objective::new(vec![5.0, 7.0], ObjectiveType::Maximize).unwrap();
        assert_eq!(objective.coefficients, vec![5.0, 7.0]);
        assert_eq!(objective.objective_type, ObjectiveType::Maximize);
        assert_eq!(objective.constant, 0.0);
        assert!(objective.name.is_none());
    }

    #[test]
    fn test_objective_with_constant() {
        let objective = Objective::new_with_constant(
            vec![1.0, 2.0], 
            ObjectiveType::Minimize, 
            10.0
        ).unwrap();
        assert_eq!(objective.constant, 10.0);
    }

    #[test]
    fn test_invalid_objective() {
        // Empty coefficients
        assert!(Objective::new(vec![], ObjectiveType::Minimize).is_err());
        
        // Infinite coefficient
        assert!(Objective::new(vec![f64::INFINITY], ObjectiveType::Minimize).is_err());
    }

    #[test]
    fn test_objective_evaluation() {
        let objective = Objective::new_with_constant(
            vec![3.0, 2.0],
            ObjectiveType::Maximize,
            5.0
        ).unwrap();
        
        let value = objective.evaluate(&[2.0, 1.5]).unwrap();
        assert_abs_diff_eq!(value, 3.0 * 2.0 + 2.0 * 1.5 + 5.0);
    }

    #[test]
    fn test_objective_conversion() {
        let max_obj = Objective::new(vec![1.0, -2.0], ObjectiveType::Maximize).unwrap();
        let min_obj = max_obj.to_minimization();
        
        assert_eq!(min_obj.objective_type, ObjectiveType::Minimize);
        assert_eq!(min_obj.coefficients, vec![-1.0, 2.0]);
        
        let restored = min_obj.to_maximization();
        assert_abs_diff_eq!(restored, max_obj, epsilon = 1e-10);
    }

    #[test]
    fn test_objective_multiplication() {
        let mut objective = Objective::new_with_constant(
            vec![1.0, 2.0],
            ObjectiveType::Minimize,
            5.0
        ).unwrap();
        
        objective.multiply_by(-2.0);
        assert_eq!(objective.coefficients, vec![-2.0, -4.0]);
        assert_eq!(objective.constant, -10.0);
        assert_eq!(objective.objective_type, ObjectiveType::Maximize);
    }

    #[test]
    fn test_directional_derivative() {
        let objective = Objective::new(vec![3.0, -1.0, 2.0], ObjectiveType::Minimize).unwrap();
        let direction = vec![1.0, -1.0, 0.5];
        
        let derivative = objective.directional_derivative(&direction).unwrap();
        assert_abs_diff_eq!(derivative, 3.0 * 1.0 + (-1.0) * (-1.0) + 2.0 * 0.5);
    }
}