import pprint
from collections import namedtuple

import itertools
from ortools.constraint_solver import pywrapcp

class Subject:
    def __init__(self, name: str, periods_per_week: int):
        self.name = name
        self.periods_per_week = periods_per_week
        self.period_names = ['{}-p{}'.format(name, i) for i in range(1, periods_per_week+1)]
        
    def __eq__(self, other):
        return isinstance(other, type(self)) and self.name == other.name

    def __hash__(self):
        return hash(self.name)
        
def solve(students: list, periods_per_week: int):
    """
    Create a timetable for the given number of periods, subjects, and students per subject
    
    Args:
        students: a list of tuples of Subjects, one for each student, containing that student's Subjects
        periods_per_week: The number of periods in the whole timetable
    """
    solver = pywrapcp.Solver("timetable")

    # Flatten the list that contains each student's subjects and remove duplicates
    subjects = set(itertools.chain(*students))

    # Get the list of all the periods for all subjects
    period_names = itertools.chain(*[subject.period_names for subject in subjects])

    # Generate a dict of period_name:period_variable
    # We need a dict to be able to access each period's variable by name
    period_variables = {period_name: solver.IntVar(1, periods_per_week+1, period_name) 
                            for period_name in period_names}

    for student in students:
        # Get the list of all the periods for the student's subjects
        student_period_names = list(itertools.chain(*[subject.period_names for subject in student]))
        
        # Filter the list of all periods to get only those that the student is a part of
        student_period_variables = [
            period_var 
            for period_name, period_var in period_variables.items() 
            if period_name in student_period_names
        ]

        # All of the student's periods must be scheduled at different times
        solver.AddConstraint(solver.AllDifferent(student_period_variables))


    db = solver.Phase(list(period_variables.values()),
                      solver.CHOOSE_MIN_SIZE_LOWEST_MAX,
                      solver.ASSIGN_CENTER_VALUE)
    solver.NewSearch(db)

    while solver.NextSolution():
        yield period_variables

    


# Test data
periods_per_week = 20
HistorySL = Subject("HistorySL", 2)
HistoryHL = Subject("HistoryHL", 3)
MathSL = Subject("MathSL", 2)
MathHL = Subject("MathHL", 3)
BiologySL = Subject("BiologySL", 2)
BiologyHL = Subject("BiologyHL", 3)

students = [
    (HistorySL, MathHL),
    (BiologySL, HistoryHL),
    (MathSL, BiologyHL),
    (HistorySL, BiologyHL)
]

solutions = solve(students, periods_per_week)

pprint.pprint(next(solutions))
print()
pprint.pprint(next(solutions))
print()
pprint.pprint(next(solutions))
print()


@@ -528,21 +528,35 @@
    %0Adef 
-solve
+possible_timetables
 (student
@@ -604,94 +604,33 @@
     
-Create a timetable for the given number of periods, subjects, and students per subject
+Yield possible timetables
 %0A   
@@ -812,24 +812,131 @@
 le timetable
+%0A%0A    Returns:%0A        A dict of %7Bperiod name: the position in the timetable that the period should occupy%7D
 %0A    %22%22%22%0A   
@@ -2399,16 +2399,30 @@
    yield
+ %7Bperiod_name:
  period_
@@ -2433,18 +2433,112 @@
 able
-s%0A%0A    %0A%0A%0A
+.Value() for period_name, period_variable in period_variables.items()%7D%0A%0A%0Aif __name__ == '__main__':%0A    
 # Te
@@ -2545,16 +2545,20 @@
 st data%0A
+    
 periods_
@@ -2571,16 +2571,20 @@
 ek = 20%0A
+    
 HistoryS
@@ -2611,16 +2611,20 @@
 SL%22, 2)%0A
+    
 HistoryH
@@ -2651,16 +2651,20 @@
 HL%22, 3)%0A
+    
 MathSL =
@@ -2685,16 +2685,20 @@
 SL%22, 2)%0A
+    
 MathHL =
@@ -2719,16 +2719,20 @@
 HL%22, 3)%0A
+    
 BiologyS
@@ -2759,16 +2759,20 @@
 SL%22, 2)%0A
+    
 BiologyH
@@ -2800,16 +2800,20 @@
 L%22, 3)%0A%0A
+    
 students
@@ -2817,24 +2817,28 @@
 nts = %5B%0A    
+   
+ 
 (HistorySL, 
@@ -2842,24 +2842,28 @@
 L, MathHL),%0A
+    
     (Biology
@@ -2874,24 +2874,28 @@
 HistoryHL),%0A
+    
     (MathSL,
@@ -2903,24 +2903,28 @@
 BiologyHL),%0A
+    
     (History
@@ -2938,19 +2938,27 @@
 logyHL)%0A
-%5D%0A%0A
+    %5D%0A%0A    
 solution
@@ -2961,21 +2961,35 @@
 tions = 
-solve
+possible_timetables
 (student
@@ -3010,16 +3010,20 @@
 _week)%0A%0A
+    
 pprint.p
@@ -3049,91 +3049,4 @@
 s))%0A
-print()%0Apprint.pprint(next(solutions))%0Aprint()%0Apprint.pprint(next(solutions))%0Aprint()%0A%0A

make the solver's return value easier to use