#!/usr/bin/env python3
"""
Pyomo 6.9.2 Cheat Sheet - Common Patterns & Syntax
Use as in-context example for LLM code generation
Compatible with Gurobi solver (your available solver)
"""

import pyomo.environ as pyo
from pyomo.opt import SolverFactory

def optimization_example():
    """Essential Pyomo patterns - covers 90% of use cases"""
    
    # 1. MODEL CREATION
    model = pyo.ConcreteModel()
    
    # 2. DATA SETUP
    n_items = 5
    costs = [10, 15, 12, 8, 25]
    demands = [20, 30, 15, 40, 35]
    
    # CRITICAL: Validate array lengths before indexing
    assert len(costs) == len(demands) == n_items, "Array length mismatch"
    
    # 3. SETS (Pyomo way to define indices)
    model.I = pyo.RangeSet(1, n_items)  # 1-based indexing
    model.J = pyo.Set(initialize=[1, 2, 3])  # Custom set
    
    # 4. PARAMETERS (data containers)
    model.cost = pyo.Param(model.I, initialize={i+1: costs[i] for i in range(n_items)})
    model.demand = pyo.Param(model.I, initialize={i+1: demands[i] for i in range(n_items)})
    
    # 5. VARIABLES
    # Continuous variables
    model.x = pyo.Var(model.I, within=pyo.NonNegativeReals, bounds=(0, 100))
    
    # Binary variables
    model.y = pyo.Var(model.I, within=pyo.Binary)
    
    # Integer variables
    model.z = pyo.Var(model.I, within=pyo.NonNegativeIntegers, bounds=(0, 10))
    
    # Single variable
    model.total = pyo.Var(within=pyo.NonNegativeReals)
    
    # 6. OBJECTIVE FUNCTION
    # Method A: Using summation (most common)
    def obj_rule(model):
        return sum(model.cost[i] * model.x[i] for i in model.I)
    model.objective = pyo.Objective(rule=obj_rule, sense=pyo.minimize)
    
    # Method B: Direct expression (alternative)
    # model.objective = pyo.Objective(expr=pyo.summation(model.cost, model.x), sense=pyo.minimize)
    
    # 7. CONSTRAINTS - CORRECT SYNTAX PATTERNS
    
    # Pattern A: Individual constraint with rule function (RECOMMENDED)
    def capacity_rule(model):
        return sum(model.x[i] for i in model.I) <= 100
    model.capacity_constraint = pyo.Constraint(rule=capacity_rule)
    
    # Pattern B: Indexed constraints with rule function
    def linking_rule(model, i):
        return model.x[i] <= model.demand[i] * model.y[i]
    model.linking_constraint = pyo.Constraint(model.I, rule=linking_rule)
    
    # Pattern C: Direct expression constraints
    model.total_constraint = pyo.Constraint(
        expr=model.total == sum(model.cost[i] * model.x[i] for i in model.I)
    )
    
    # Pattern D: Constraint list for dynamic constraints
    model.cuts = pyo.ConstraintList()
    for i in model.I:
        model.cuts.add(model.x[i] >= 0.1 * model.demand[i])
    
    # Pattern E: Range constraints (both bounds)
    def bounds_rule(model, i):
        return (0.05 * model.demand[i], model.x[i], 0.95 * model.demand[i])
    model.bounds_constraint = pyo.Constraint(model.I, rule=bounds_rule)
    
    # 8. SOLVING WITH GUROBI (your available solver)
    solver = SolverFactory('gurobi')
    
    # Optional: Set solver options
    solver.options['TimeLimit'] = 300  # 5 minutes
    solver.options['MIPGap'] = 0.01    # 1% gap
    
    # Solve the model
    results = solver.solve(model, tee=True)  # tee=True shows solver output
    
    # 9. RESULT PROCESSING
    # Check solver status
    if results.solver.termination_condition == pyo.TerminationCondition.optimal:
        print("Optimal solution found!")
        print(f"Optimal value: {pyo.value(model.objective)}")
        
        # Extract variable values
        print("\nVariable values:")
        for i in model.I:
            x_val = pyo.value(model.x[i])
            y_val = pyo.value(model.y[i])
            if x_val > 1e-6:  # Only print non-zero values
                print(f"x[{i}] = {x_val:.3f}, y[{i}] = {int(y_val)}")
        
        print(f"Total = {pyo.value(model.total):.3f}")
        
    elif results.solver.termination_condition == pyo.TerminationCondition.infeasible:
        print("Problem is infeasible")
    elif results.solver.termination_condition == pyo.TerminationCondition.unbounded:
        print("Problem is unbounded")
    else:
        print(f"Solver terminated with condition: {results.solver.termination_condition}")
    
    return model

# SYNTAX RULES TO FOLLOW:
"""
 CORRECT PYOMO PATTERNS:

**Model Creation:**
model = pyo.ConcreteModel()

**Sets & Parameters:**
model.I = pyo.RangeSet(1, n)
model.param = pyo.Param(model.I, initialize=data_dict)

**Variables:**
model.x = pyo.Var(model.I, within=pyo.NonNegativeReals, bounds=(0, 100))
model.y = pyo.Var(model.I, within=pyo.Binary)

**Constraints (Rule Functions - RECOMMENDED):**
def constraint_rule(model, i):
    return model.x[i] <= model.demand[i]
model.constraint = pyo.Constraint(model.I, rule=constraint_rule)

**Objective:**
def obj_rule(model):
    return sum(model.cost[i] * model.x[i] for i in model.I)
model.objective = pyo.Objective(rule=obj_rule, sense=pyo.minimize)

**Summations:**
sum(model.x[i] for i in model.I)  # Python sum (recommended)
pyo.summation(model.cost, model.x)  # Pyomo summation
pyo.quicksum(model.x[i] for i in model.I)  # Pyomo quicksum (if available)

**Solving:**
solver = SolverFactory('gurobi')
results = solver.solve(model, tee=True)

**Value Extraction:**
optimal_value = pyo.value(model.objective)
variable_value = pyo.value(model.x[i])

 COMMON MISTAKES TO AVOID:
- Forgetting to validate array lengths before indexing
- Using 0-based indexing with Pyomo Sets (Pyomo is typically 1-based)
- Not checking solver termination condition before extracting values
- Using model.x[i].value instead of pyo.value(model.x[i])
- Missing rule functions for indexed constraints
"""

# ADVANCED PATTERNS
def advanced_patterns_example():
    """More complex Pyomo patterns"""
    
    model = pyo.ConcreteModel()
    
    # Multi-dimensional sets
    model.I = pyo.RangeSet(1, 3)
    model.J = pyo.RangeSet(1, 4)
    
    # Multi-dimensional variables
    model.x = pyo.Var(model.I, model.J, within=pyo.Binary)
    
    # Complex constraints with multiple indices
    def complex_rule(model, i):
        return sum(model.x[i, j] for j in model.J) <= 1
    model.complex_constraint = pyo.Constraint(model.I, rule=complex_rule)
    
    # Conditional constraints
    def conditional_rule(model, i, j):
        if i <= 2:
            return model.x[i, j] <= 0.5
        else:
            return pyo.Constraint.Skip  # Skip this constraint
    model.conditional = pyo.Constraint(model.I, model.J, rule=conditional_rule)
    
    return model
