#!/usr/bin/env python3
"""
DOCPLEX implementation for Loan Allocation Optimization
"""

from docplex.mp.model import Model

def loan_allocation_optimization():
    """Optimize loan allocation across branches to maximize total disbursement"""
    
    # 1. MODEL & DATA SETUP
    mdl = Model(name="loan_allocation")
    
    # Data from the database
    loan_amounts = [5000.0, 10000.0, 7500.0]  # Initial loan amounts
    no_of_customers = [150, 250, 200]         # Number of customers per branch
    credit_scores = [720, 780, 750]           # Credit scores of customers
    
    # Operational parameters
    max_loan_per_customer = 15000
    max_loan_per_credit_score_unit = 600
    max_total_loan_per_customer = 60000
    
    # CRITICAL: Validate array lengths to prevent IndexError
    assert len(loan_amounts) == len(no_of_customers) == len(credit_scores), "Array length mismatch"
    safe_range = range(min(len(loan_amounts), len(no_of_customers), len(credit_scores)))  # Safe indexing
    
    # 2. VARIABLES
    # Decision variable: loan amount to be disbursed for each loan
    amount = {i: mdl.continuous_var(name=f"amount_{i}", lb=0) for i in safe_range}
    
    # 3. OBJECTIVE FUNCTION
    # Maximize the total loan amount disbursed
    total_loan = mdl.sum(amount[i] for i in safe_range)
    mdl.maximize(total_loan)
    
    # 4. CONSTRAINTS
    
    # Constraint 1: Branch customer capacity
    for i in safe_range:
        mdl.add_constraint(amount[i] <= no_of_customers[i] * max_loan_per_customer, ctname=f"branch_capacity_{i}")
    
    # Constraint 2: Customer credit score
    for i in safe_range:
        mdl.add_constraint(amount[i] <= credit_scores[i] * max_loan_per_credit_score_unit, ctname=f"credit_score_{i}")
    
    # Constraint 3: Total loan exposure per customer
    for i in safe_range:
        mdl.add_constraint(amount[i] <= max_total_loan_per_customer, ctname=f"total_exposure_{i}")
    
    # 5. SOLVING & RESULTS
    solution = mdl.solve()
    
    if solution:
        print(f"Optimal value: {solution.objective_value}")
        for i in safe_range:
            value = solution.get_value(amount[i])
            if value > 1e-6:
                print(f"amount[{i}] = {value:.3f}")
    else:
        print("No solution found")
        print(f"Status: {mdl.solve_details.status}")
    
    return mdl

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