#!/usr/bin/env python3
"""
DOCPLEX implementation for dormitory assignment problem
"""

from docplex.mp.model import Model

def dorm_assignment_optimization():
    """Optimize dormitory assignments to minimize total distance traveled by students."""
    
    # 1. MODEL & DATA SETUP
    mdl = Model(name="dorm_assignment")
    
    # Sample data from the problem description
    students = [1, 2, 3]
    dorms = [1, 2]
    
    # DistanceMatrix data
    distance_matrix = {
        (1, 1): 12.5,
        (1, 2): 18.3,
        (2, 1): 15.0,
        (2, 2): 10.2,
        (3, 1): 20.1,
        (3, 2): 14.7
    }
    
    # GenderInfo data
    gender_info = {
        (1, 1): 'Male',
        (2, 2): 'Female',
        (3, 1): 'Male'
    }
    
    # Dorm data
    dorm_capacity = {
        1: 100,
        2: 150
    }
    
    dorm_gender = {
        1: 'Male',
        2: 'Female'
    }
    
    # CRITICAL: Validate array lengths to prevent IndexError
    assert len(students) == 3, "Students array length mismatch"
    assert len(dorms) == 2, "Dorms array length mismatch"
    
    # 2. VARIABLES
    # Binary decision variables: x[s, d] = 1 if student s is assigned to dorm d
    x = {(s, d): mdl.binary_var(name=f"x_{s}_{d}") for s in students for d in dorms}
    
    # 3. OBJECTIVE FUNCTION
    # Minimize total distance traveled by all students
    objective = mdl.sum(distance_matrix[(s, d)] * x[(s, d)] for s in students for d in dorms)
    mdl.minimize(objective)
    
    # 4. CONSTRAINTS
    
    # Single Assignment: Each student must be assigned to exactly one dorm
    for s in students:
        mdl.add_constraint(mdl.sum(x[(s, d)] for d in dorms) == 1, ctname=f"single_assignment_{s}")
    
    # Dorm Capacity: The number of students assigned to a dorm cannot exceed its capacity
    for d in dorms:
        mdl.add_constraint(mdl.sum(x[(s, d)] for s in students) <= dorm_capacity[d], ctname=f"capacity_{d}")
    
    # Gender Matching: A student can only be assigned to a dorm if the dorm's gender matches the student's gender
    for s in students:
        for d in dorms:
            if (s, d) in gender_info and gender_info[(s, d)] != dorm_gender[d]:
                mdl.add_constraint(x[(s, d)] == 0, ctname=f"gender_matching_{s}_{d}")
    
    # 5. SOLVING & RESULTS
    solution = mdl.solve()
    
    if solution:
        print(f"Optimal value: {solution.objective_value}")
        for s in students:
            for d in dorms:
                if solution.get_value(x[(s, d)]) > 0.5:
                    print(f"Student {s} assigned to Dorm {d}")
    else:
        print("No solution found")
        print(f"Status: {mdl.solve_details.status}")
    
    return mdl

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