#!/usr/bin/env python3
"""
Gurobipy 12.0.2 Implementation for Wrestler Selection Optimization
"""

import gurobipy as gp
from gurobipy import GRB

def wrestler_selection_optimization():
    """Optimize wrestler selection to maximize total days held while adhering to constraints."""
    
    # 1. MODEL & DATA SETUP
    model = gp.Model("wrestler_selection")
    
    # Data from the database
    wrestler_data = [
        {'Days_held': 150, 'selection_status': False},
        {'Days_held': 250, 'selection_status': True},
        {'Days_held': 100, 'selection_status': False}
    ]
    
    wrestler_team_data = [
        {'wrestler_id': 1, 'team': 'Team A'},
        {'wrestler_id': 2, 'team': 'Team B'},
        {'wrestler_id': 3, 'team': 'Team A'}
    ]
    
    wrestler_eliminations_data = [
        {'wrestler_id': 1, 'eliminations': 8},
        {'wrestler_id': 2, 'eliminations': 12},
        {'wrestler_id': 3, 'eliminations': 5}
    ]
    
    # Extract relevant data
    days_held = [w['Days_held'] for w in wrestler_data]
    eliminations = [e['eliminations'] for e in wrestler_eliminations_data]
    team_assignments = [t['team'] for t in wrestler_team_data]
    
    n_wrestlers = len(wrestler_data)
    
    # CRITICAL: Validate array lengths before loops
    assert len(days_held) == len(eliminations) == len(team_assignments) == n_wrestlers, "Array length mismatch"
    
    # 2. VARIABLES
    x = {i: model.addVar(vtype=GRB.BINARY, name=f"x_{i}") for i in range(n_wrestlers)}
    
    # 3. OBJECTIVE FUNCTION
    model.setObjective(gp.quicksum(days_held[i] * x[i] for i in range(n_wrestlers)), GRB.MAXIMIZE)
    
    # 4. CONSTRAINTS
    
    # Total Wrestlers Constraint
    model.addConstr(gp.quicksum(x[i] for i in range(n_wrestlers)) <= 6, name="total_wrestlers_limit")
    
    # Total Eliminations Constraint
    model.addConstr(gp.quicksum(eliminations[i] * x[i] for i in range(n_wrestlers)) <= 30, name="total_eliminations_limit")
    
    # Team Balance Constraint
    team_a_indices = [i for i, team in enumerate(team_assignments) if team == 'Team A']
    team_b_indices = [i for i, team in enumerate(team_assignments) if team == 'Team B']
    
    model.addConstr(gp.quicksum(x[i] for i in team_a_indices) == gp.quicksum(x[i] for i in team_b_indices), name="team_balance")
    
    # 5. SOLVING & RESULTS
    model.optimize()
    
    if model.status == GRB.OPTIMAL:
        print(f"Optimal value: {model.objVal}")
        for i in range(n_wrestlers):
            if x[i].x > 1e-6:
                print(f"Wrestler {i+1} selected: {x[i].x:.0f}")
    elif model.status == GRB.INFEASIBLE:
        print("Problem is infeasible")
    elif model.status == GRB.UNBOUNDED:
        print("Problem is unbounded")
    
    return model

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