#!/usr/bin/env python3
"""
Gurobipy Implementation for Student Friendship Optimization
"""

import gurobipy as gp
from gurobipy import GRB

def optimize_student_friendship():
    """Optimize student friendships to maximize happiness score."""
    
    # 1. MODEL & DATA SETUP
    model = gp.Model("student_friendship_optimization")
    
    # Sample data from the problem context
    happiness_scores = {
        (1, 2): 1,
        (2, 3): 1,
        (3, 1): 1,
        (1, 3): 0,
        (2, 1): 0
    }
    
    students = {1, 2, 3}
    min_friends = 2
    max_friends = 5
    
    # CRITICAL: Validate data consistency
    assert all(i in students and j in students for (i, j) in happiness_scores), "Invalid student IDs in happiness scores"
    
    # 2. VARIABLES
    x = model.addVars(happiness_scores.keys(), vtype=GRB.BINARY, name="x")
    
    # 3. OBJECTIVE FUNCTION
    model.setObjective(gp.quicksum(happiness_scores[i, j] * x[i, j] for (i, j) in happiness_scores), GRB.MAXIMIZE)
    
    # 4. CONSTRAINTS
    
    # Minimum Friends Constraint
    for i in students:
        model.addConstr(gp.quicksum(x[i, j] for j in students if (i, j) in happiness_scores) >= min_friends, name=f"min_friends_{i}")
    
    # Maximum Friends Constraint
    for i in students:
        model.addConstr(gp.quicksum(x[i, j] for j in students if (i, j) in happiness_scores) <= max_friends, name=f"max_friends_{i}")
    
    # Mutual Friendship Constraint
    for (i, j) in happiness_scores:
        if i < j:  # Avoid duplicate constraints
            model.addConstr(x[i, j] == x[j, i], name=f"mutual_{i}_{j}")
    
    # 5. SOLVING & RESULTS
    model.optimize()
    
    if model.status == GRB.OPTIMAL:
        print(f"Optimal value: {model.objVal}")
        for (i, j) in happiness_scores:
            if x[i, j].x > 0.5:
                print(f"Friendship between student {i} and student {j} is established.")
    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__":
    optimize_student_friendship()