# Complete PYOMO implementation - Retry Attempt 4

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

def optimize_loan_allocation():
    # 1. MODEL CREATION
    model = pyo.ConcreteModel()

    # 2. DATA SETUP
    # Customer data
    customers = [
        {'cust_ID': 1, 'credit_score': 720},
        {'cust_ID': 2, 'credit_score': 680},
        {'cust_ID': 3, 'credit_score': 750}
    ]

    # Loan data
    loans = [
        {'loan_ID': 101, 'amount': 30000, 'branch_ID': 1},
        {'loan_ID': 102, 'amount': 45000, 'branch_ID': 2},
        {'loan_ID': 103, 'amount': 25000, 'branch_ID': 3},
        {'loan_ID': 104, 'amount': 20000, 'branch_ID': 1},
        {'loan_ID': 105, 'amount': 35000, 'branch_ID': 2}
    ]

    # Branch loans data
    branch_loans = [
        {'branch_ID': 1, 'min_loans': 120000},
        {'branch_ID': 2, 'min_loans': 150000},
        {'branch_ID': 3, 'min_loans': 100000}
    ]

    # Validate array lengths
    assert len(customers) == 3, "Customer data length mismatch"
    assert len(loans) == 5, "Loan data length mismatch"
    assert len(branch_loans) == 3, "Branch loans data length mismatch"

    # 3. SETS
    model.C = pyo.Set(initialize=[c['cust_ID'] for c in customers])
    model.L = pyo.Set(initialize=[l['loan_ID'] for l in loans])
    model.B = pyo.Set(initialize=[b['branch_ID'] for b in branch_loans])

    # 4. PARAMETERS
    model.credit_score = pyo.Param(model.C, initialize={c['cust_ID']: c['credit_score'] for c in customers})
    model.min_loans = pyo.Param(model.B, initialize={b['branch_ID']: b['min_loans'] for b in branch_loans})

    # 5. VARIABLES
    model.x = pyo.Var(model.L, within=pyo.NonNegativeReals, bounds=(0, 50000))

    # 6. OBJECTIVE FUNCTION
    def obj_rule(model):
        return sum(model.credit_score[loans[i]['cust_ID']] * model.x[loans[i]['loan_ID']] for i in range(len(loans)))
    model.objective = pyo.Objective(rule=obj_rule, sense=pyo.maximize)

    # 7. CONSTRAINTS
    # Total Budget Constraint
    def total_budget_rule(model):
        return sum(model.x[l['loan_ID']] for l in loans) <= 1000000
    model.total_budget_constraint = pyo.Constraint(rule=total_budget_rule)

    # Branch Minimum Loan Requirement
    def branch_min_loan_rule(model, b):
        return sum(model.x[l['loan_ID']] for l in loans if l['branch_ID'] == b) >= model.min_loans[b]
    model.branch_min_loan_constraint = pyo.Constraint(model.B, rule=branch_min_loan_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)}")

        # Extract variable values
        print("\nLoan allocations:")
        for l in loans:
            loan_val = pyo.value(model.x[l['loan_ID']])
            if loan_val > 1e-6:  # Only print non-zero values
                print(f"Loan ID {l['loan_ID']}: {loan_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
optimize_loan_allocation()