#!/usr/bin/env python3
"""
DOCPLEX implementation for small_bank_1 optimization problem
"""

from docplex.mp.model import Model

def small_bank_optimization():
    """Optimize fund allocation between savings and checking accounts"""
    
    # 1. MODEL & DATA SETUP
    mdl = Model(name="small_bank_1")
    
    # Data from CUSTOMER_FUNDS table
    customer_ids = [1, 2, 3]
    total_funds = [1000.0, 1500.0, 2000.0]
    
    # Business configuration parameters
    savings_interest_rate = 0.03
    checking_interest_rate = 0.01
    min_savings_balance = 100.0
    min_checking_balance = 50.0
    max_savings_balance = 50000.0
    max_checking_balance = 30000.0
    total_bank_funds = 1000000.0
    
    # CRITICAL: Validate array lengths to prevent IndexError
    assert len(customer_ids) == len(total_funds), "Array length mismatch"
    safe_range = range(min(len(customer_ids), len(total_funds)))  # Safe indexing
    
    # 2. VARIABLES
    savings = {i: mdl.continuous_var(name=f"savings_{i}", lb=0) for i in safe_range}
    checking = {i: mdl.continuous_var(name=f"checking_{i}", lb=0) for i in safe_range}
    
    # 3. OBJECTIVE FUNCTION
    total_interest = mdl.sum(savings_interest_rate * savings[i] + checking_interest_rate * checking[i] for i in safe_range)
    mdl.maximize(total_interest)
    
    # 4. CONSTRAINTS
    
    # Total Funds Constraint
    for i in safe_range:
        mdl.add_constraint(savings[i] + checking[i] <= total_funds[i], ctname=f"total_funds_{i}")
    
    # Minimum Savings Balance Constraint
    for i in safe_range:
        mdl.add_constraint(savings[i] >= min_savings_balance, ctname=f"min_savings_{i}")
    
    # Minimum Checking Balance Constraint
    for i in safe_range:
        mdl.add_constraint(checking[i] >= min_checking_balance, ctname=f"min_checking_{i}")
    
    # Maximum Savings Balance Constraint
    for i in safe_range:
        mdl.add_constraint(savings[i] <= max_savings_balance, ctname=f"max_savings_{i}")
    
    # Maximum Checking Balance Constraint
    for i in safe_range:
        mdl.add_constraint(checking[i] <= max_checking_balance, ctname=f"max_checking_{i}")
    
    # Total Bank Funds Constraint
    total_allocated = mdl.sum(savings[i] + checking[i] for i in safe_range)
    mdl.add_constraint(total_allocated <= total_bank_funds, ctname="total_bank_funds")
    
    # 5. SOLVING & RESULTS
    solution = mdl.solve()
    
    if solution:
        print(f"Optimal value: {solution.objective_value}")
        for i in safe_range:
            savings_value = solution.get_value(savings[i])
            checking_value = solution.get_value(checking[i])
            print(f"Customer {customer_ids[i]}: Savings = {savings_value:.2f}, Checking = {checking_value:.2f}")
    else:
        print("No solution found")
        print(f"Status: {mdl.solve_details.status}")
    
    return mdl

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