# Complete PYOMO implementation

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

def optimize_transaction_fees():
    """Optimize transaction fees by allocating transactions across card types."""
    
    # 1. MODEL CREATION
    model = pyo.ConcreteModel()
    
    # 2. DATA SETUP
    customers = [1, 2, 3]
    card_types = ['VISA', 'MASTERCARD', 'AMEX']
    
    # Fees for each card type
    card_fees = {
        'VISA': 0.02,
        'MASTERCARD': 0.03,
        'AMEX': 0.035
    }
    
    # Account limits for each customer
    account_limits = {
        1: 1000,
        2: 1500,
        3: 2000
    }
    
    # CRITICAL: Validate array lengths before indexing
    assert len(customers) == len(account_limits), "Array length mismatch"
    
    # 3. SETS (Pyomo way to define indices)
    model.Customers = pyo.Set(initialize=customers)
    model.CardTypes = pyo.Set(initialize=card_types)
    
    # 4. PARAMETERS (data containers)
    model.CardFees = pyo.Param(model.CardTypes, initialize=card_fees)
    model.AccountLimits = pyo.Param(model.Customers, initialize=account_limits)
    
    # 5. VARIABLES
    # Transaction amounts allocated to each card type for each customer
    model.TransactionAmounts = pyo.Var(model.Customers, model.CardTypes, within=pyo.NonNegativeReals)
    
    # 6. OBJECTIVE FUNCTION
    # Minimize the total transaction fees
    def obj_rule(model):
        return sum(model.CardFees[card] * model.TransactionAmounts[cust, card] 
                   for cust in model.Customers for card in model.CardTypes)
    model.Objective = pyo.Objective(rule=obj_rule, sense=pyo.minimize)
    
    # 7. CONSTRAINTS
    # Ensure the total transaction amount for each customer does not exceed their account limit
    def account_limit_rule(model, cust):
        return sum(model.TransactionAmounts[cust, card] for card in model.CardTypes) <= model.AccountLimits[cust]
    model.AccountLimitConstraint = pyo.Constraint(model.Customers, rule=account_limit_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("\nTransaction Amounts:")
        for cust in model.Customers:
            for card in model.CardTypes:
                amount = pyo.value(model.TransactionAmounts[cust, card])
                if amount > 1e-6:  # Only print non-zero values
                    print(f"Customer {cust}, Card {card}: {amount:.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

# Execute the optimization
if __name__ == "__main__":
    optimize_transaction_fees()