# Complete PYOMO implementation - Retry Attempt 2

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

def small_bank_optimization():
    """Optimize fund allocation for a small bank using Pyomo"""
    
    # 1. MODEL CREATION
    model = pyo.ConcreteModel()
    
    # 2. DATA SETUP
    # Assuming 3 customers based on the provided data
    n_customers = 3
    
    # Interest rates and total funds from the database
    savings_interest_rates = [0.012, 0.015, 0.018]
    checking_interest_rates = [0.006, 0.008, 0.01]
    total_funds = [5000.0, 10000.0, 15000.0]
    
    # Minimum balance requirements (assumed constants)
    min_savings_balance = [1000.0, 1000.0, 1000.0]
    min_checking_balance = [500.0, 500.0, 500.0]
    
    # CRITICAL: Validate array lengths before indexing
    assert len(savings_interest_rates) == len(checking_interest_rates) == len(total_funds) == n_customers, "Array length mismatch"
    
    # 3. SETS
    model.I = pyo.RangeSet(1, n_customers)  # 1-based indexing for customers
    
    # 4. PARAMETERS
    model.savings_interest_rate = pyo.Param(model.I, initialize={i+1: savings_interest_rates[i] for i in range(n_customers)})
    model.checking_interest_rate = pyo.Param(model.I, initialize={i+1: checking_interest_rates[i] for i in range(n_customers)})
    model.total_funds = pyo.Param(model.I, initialize={i+1: total_funds[i] for i in range(n_customers)})
    model.min_savings_balance = pyo.Param(model.I, initialize={i+1: min_savings_balance[i] for i in range(n_customers)})
    model.min_checking_balance = pyo.Param(model.I, initialize={i+1: min_checking_balance[i] for i in range(n_customers)})
    
    # 5. VARIABLES
    model.savings = pyo.Var(model.I, within=pyo.NonNegativeReals)
    model.checking = pyo.Var(model.I, within=pyo.NonNegativeReals)
    
    # 6. OBJECTIVE FUNCTION
    def obj_rule(model):
        return sum(model.savings_interest_rate[i] * model.savings[i] + model.checking_interest_rate[i] * model.checking[i] for i in model.I)
    model.objective = pyo.Objective(rule=obj_rule, sense=pyo.maximize)
    
    # 7. CONSTRAINTS
    
    # Total Funds Constraint
    def total_funds_rule(model, i):
        return model.savings[i] + model.checking[i] <= model.total_funds[i]
    model.total_funds_constraint = pyo.Constraint(model.I, rule=total_funds_rule)
    
    # Minimum Savings Balance Constraint
    def min_savings_rule(model, i):
        return model.savings[i] >= model.min_savings_balance[i]
    model.min_savings_constraint = pyo.Constraint(model.I, rule=min_savings_rule)
    
    # Minimum Checking Balance Constraint
    def min_checking_rule(model, i):
        return model.checking[i] >= model.min_checking_balance[i]
    model.min_checking_constraint = pyo.Constraint(model.I, rule=min_checking_rule)
    
    # 8. SOLVING WITH GUROBI
    solver = SolverFactory('gurobi')
    
    # Solve the model
    results = solver.solve(model, tee=True)
    
    # 9. RESULT PROCESSING
    if results.solver.termination_condition == pyo.TerminationCondition.optimal:
        print("Optimal solution found!")
        print(f"Optimal value: {pyo.value(model.objective):.2f}")
        
        # Extract variable values
        print("\nVariable values:")
        for i in model.I:
            savings_val = pyo.value(model.savings[i])
            checking_val = pyo.value(model.checking[i])
            print(f"Customer {i}: Savings = {savings_val:.2f}, Checking = {checking_val:.2f}")
        
    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
small_bank_optimization()