#!/usr/bin/env python3
"""
DOCplex implementation for student friendship optimization problem
"""

from docplex.mp.model import Model

def optimize_student_friendship():
    """Optimize student friendships to maximize happiness score"""
    
    # 1. MODEL & DATA SETUP
    mdl = Model(name="student_friendship_optimization")
    
    # Sample data from the problem description
    happiness_scores = [
        (1, 2, 1),
        (2, 3, 1),
        (3, 1, 1),
        (1, 3, 0),
        (2, 1, 0)
    ]
    
    # Extract unique student IDs
    students = sorted(set([s[0] for s in happiness_scores] + [s[1] for s in happiness_scores]))
    
    # Create a dictionary to store happiness scores
    happiness_dict = {(s[0], s[1]): s[2] for s in happiness_scores}
    
    # Business configuration parameters
    min_friends = 2
    max_friends = 5
    
    # 2. VARIABLES
    # Binary decision variables for friendships
    x = {(i, j): mdl.binary_var(name=f"x_{i}_{j}") for i in students for j in students if i != j}
    
    # 3. OBJECTIVE FUNCTION
    # Maximize total happiness score
    objective = mdl.sum(happiness_dict.get((i, j), 0) * x[(i, j)] for i in students for j in students if i != j)
    mdl.maximize(objective)
    
    # 4. CONSTRAINTS
    
    # Minimum friends constraint
    for i in students:
        mdl.add_constraint(mdl.sum(x[(i, j)] for j in students if j != i) >= min_friends, ctname=f"min_friends_{i}")
    
    # Maximum friends constraint
    for i in students:
        mdl.add_constraint(mdl.sum(x[(i, j)] for j in students if j != i) <= max_friends, ctname=f"max_friends_{i}")
    
    # Mutual friendship constraint
    for i in students:
        for j in students:
            if i != j:
                mdl.add_constraint(x[(i, j)] == x[(j, i)], ctname=f"mutual_{i}_{j}")
    
    # 5. SOLVING & RESULTS
    solution = mdl.solve()
    
    if solution:
        print(f"Optimal value: {solution.objective_value}")
        for (i, j), var in x.items():
            if solution.get_value(var) > 0.5:
                print(f"Friendship between {i} and {j} is established.")
    else:
        print("No solution found")
        print(f"Status: {mdl.solve_details.status}")
    
    return mdl

# Run the optimization
if __name__ == "__main__":
    optimize_student_friendship()