theory_core 0.1.0

//! Constraint Intermediate Representation (IR)
//!
//! Provides a typed, versioned intermediate representation for constraints
//! that can be consumed by both neural networks and SMT solvers.

use indexmap::IndexMap;
use serde::{Deserialize, Serialize};
use smallvec::SmallVec;
use std::fmt;

/// Versioned constraint IR to enable evolution
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ConstraintIR {
    pub version: u32,
    pub variables: IndexMap<VarId, Variable>,
    pub constraints: Vec<Constraint>,
    pub objective: Option<Objective>,
    pub theory_tags: SmallVec<[TheoryTag; 4]>,
}

impl ConstraintIR {
    pub fn new() -> Self {
        Self {
            version: 1,
            variables: IndexMap::new(),
            constraints: Vec::new(),
            objective: None,
            theory_tags: SmallVec::new(),
        }
    }

    pub fn add_variable(&mut self, var: Variable) -> VarId {
        let id = VarId(self.variables.len());
        self.variables.insert(id, var);
        id
    }

    pub fn add_constraint(&mut self, constraint: Constraint) {
        self.constraints.push(constraint);
    }

    pub fn set_objective(&mut self, objective: Objective) {
        self.objective = Some(objective);
    }

    pub fn add_theory_tag(&mut self, tag: TheoryTag) {
        if !self.theory_tags.contains(&tag) {
            self.theory_tags.push(tag);
        }
    }

    /// Create a simple test constraint IR for demonstration purposes
    pub fn new_simple_test() -> Self {
        let mut ir = Self::new();

        // Add variables
        let x = ir.add_variable(Variable {
            name: "x".to_string(),
            domain: Domain::Integer {
                min: Some(0),
                max: Some(100),
            },
            metadata: VariableMetadata::default(),
        });

        let y = ir.add_variable(Variable {
            name: "y".to_string(),
            domain: Domain::Integer {
                min: Some(0),
                max: Some(100),
            },
            metadata: VariableMetadata::default(),
        });

        // Add constraints: x + y = 10
        ir.add_constraint(Constraint::Equal {
            lhs: Expr::Binary {
                op: BinaryOp::Add,
                lhs: Box::new(Expr::Var(x)),
                rhs: Box::new(Expr::Var(y)),
            },
            rhs: Expr::Const(ConstValue::Integer(10)),
        });

        // Add constraint: x >= 2*y (represented as NOT(x < 2*y))
        ir.add_constraint(Constraint::Not {
            constraint: Box::new(Constraint::Less {
                lhs: Expr::Var(x),
                rhs: Expr::Binary {
                    op: BinaryOp::Mul,
                    lhs: Box::new(Expr::Const(ConstValue::Integer(2))),
                    rhs: Box::new(Expr::Var(y)),
                },
            }),
        });

        // Add theory tag for integer linear programming
        ir.add_theory_tag(TheoryTag::LIA);

        ir
    }
}

/// Variable identifier
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct VarId(pub usize);

impl fmt::Display for VarId {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "v{}", self.0)
    }
}

/// Variable in the constraint system
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Variable {
    pub name: String,
    pub domain: Domain,
    pub metadata: VariableMetadata,
}

/// Domain specification for variables
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum Domain {
    /// Boolean variable (true/false)
    Boolean,
    /// Integer with optional bounds
    Integer { min: Option<i64>, max: Option<i64> },
    /// Real number with optional bounds
    Real { min: Option<f64>, max: Option<f64> },
    /// Bit-vector of specified width
    BitVector { width: u32 },
    /// String with optional length bounds
    String { max_length: Option<usize> },
    /// Array of elements
    Array {
        element_type: Box<Domain>,
        size: Option<usize>,
    },
    /// Enumeration of possible values
    Enum { values: Vec<String> },
}

/// Metadata for variables (optimization hints, etc.)
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct VariableMetadata {
    pub description: Option<String>,
    pub importance: f32,       // 0.0 to 1.0
    pub computation_cost: f32, // estimated cost to compute
}

/// Constraint in the system
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum Constraint {
    /// Equality constraint: lhs = rhs
    Equal { lhs: Expr, rhs: Expr },
    /// Inequality constraint: lhs < rhs
    Less { lhs: Expr, rhs: Expr },
    /// Less than or equal: lhs <= rhs
    LessEqual { lhs: Expr, rhs: Expr },
    /// Logical AND of sub-constraints
    And { constraints: Vec<Constraint> },
    /// Logical OR of sub-constraints
    Or { constraints: Vec<Constraint> },
    /// Negation of constraint
    Not { constraint: Box<Constraint> },
    /// Implication: if antecedent then consequent
    Implies {
        antecedent: Box<Constraint>,
        consequent: Box<Constraint>,
    },
    /// Custom constraint with solver-specific encoding
    Custom {
        name: String,
        params: IndexMap<String, Expr>,
    },
}

/// Expression in constraints
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum Expr {
    /// Variable reference
    Var(VarId),
    /// Constant values
    Const(ConstValue),
    /// Binary operations
    Binary {
        op: BinaryOp,
        lhs: Box<Expr>,
        rhs: Box<Expr>,
    },
    /// Unary operations
    Unary { op: UnaryOp, operand: Box<Expr> },
    /// Function call
    Call { function: String, args: Vec<Expr> },
}

/// Constant values in expressions
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum ConstValue {
    Bool(bool),
    Integer(i64),
    Real(f64),
    String(String),
    BitVector { value: u64, width: u32 },
}

/// Binary operators
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum BinaryOp {
    Add,
    Sub,
    Mul,
    Div,
    Mod,
    Power, // Exponentiation: a^b
    And,
    Or,
    Xor,
    Equal,
    NotEqual,
    Less,
    LessEqual,
    Greater,
    GreaterEqual,
}

/// Unary operators
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum UnaryOp {
    Not,
    Negate, // Arithmetic negation (-x)
    Sqrt,   // Square root
    Abs,    // Absolute value
}

/// Optimization objective
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Objective {
    pub goal: OptimizationGoal,
    pub expr: Expr,
    pub weight: f64,
}

/// Optimization goals
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum OptimizationGoal {
    Minimize,
    Maximize,
    Satisfy, // Just find a satisfying assignment
}

/// Theory tags to route to appropriate solvers
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum TheoryTag {
    /// Linear integer arithmetic
    LIA,
    /// Linear real arithmetic  
    LRA,
    /// Non-linear arithmetic
    NLA,
    /// Bit-vectors
    BV,
    /// Arrays
    Array,
    /// Strings
    String,
    /// Floating-point arithmetic
    FP,
    /// Uninterpreted functions
    UF,
    /// Quantified formulas
    Quantified,
}

impl Default for ConstraintIR {
    fn default() -> Self {
        Self::new()
    }
}