import json
from pprint import pprint
from random import randrange, randint

league_count = 4
tabelle = []

with open('data.json') as f:
    data = json.load(f)

# Augment example orte with Raucher flag
ident = 0
for o in data["orte"]:
    o["id"] = ident
    o["raucher"] = not randrange(2)
    o["teams"] = []
    ident += 1

# Augment example teams with weekday, league and Raucher flag
ident = 0
for t in data["teams"]:

    t["id"] = ident
    t["tag"] = randrange(7)
    t["liga"] = randrange(league_count)
    t["nichtraucher"] = not randrange(4)

    team_nr = t["nichtraucher"]

    found = False
    while not found:
        heimort = randrange(len(data["orte"]))
        ersatzort = randrange(len(data["orte"]))

        if data["orte"][ersatzort]["raucher"]:
            continue

        ort_nr = not data["orte"][heimort]["raucher"]

        if team_nr != ort_nr:
            continue

        if team_nr:
            ersatzort = -1

        found = True

    t["ersatzort"] = ersatzort
    t["heimort"] = heimort
    data["orte"][heimort]["teams"].append(ident)

    ident += 1

# Count necessary spieltage
max_spieltage = 0
for league in range(league_count):
    pairings_count = len([t for t in data["teams"] if t["liga"] == league])
    if pairings_count % 2:
        pairings_count += 1
    max_spieltage = max(2 * (pairings_count - 1), max_spieltage);

print( "Max Spieltage: ", max_spieltage)

# Generate all pairings with our demo data, making up wildcards as we go
ident = 0
for league in range(league_count):
    league_teams = [t for t in data["teams"] if t["liga"] == league]

    for team_a in league_teams:
        game_count = 0

        for team_b in league_teams:
            if team_a["id"] == team_b["id"]: continue

            # If Heimspiel-Team is smokers and the Gastteam is not, the alternative
            # location needs to be chosen
            ort = next((ort for ort in data["orte"] if ort["id"] == team_a["heimort"]), None)
            if team_b["nichtraucher"] and ort["raucher"]:
                ort = next((ort for ort in data["orte"] if ort["id"] == team_a["ersatzort"]), None)

            tabelle.append({"team_a": team_a["id"], "team_b": team_b["id"], "liga": league, "ort": ort["id"], "tag": team_a["tag"], "id": ident})
            ident += 1
            game_count += 1

        # For leagues with fewer Spieltage, fill up the remainder with wildcard games
        while game_count < max_spieltage / 2:
            tabelle.append({"team_a": team_a["id"], "team_b": -1, "liga": league, "ort": -1, "tag": -1, "id": ident})
            ident += 1
            tabelle.append({"team_a": -1, "team_b": team_a["id"], "liga": league, "ort": -1, "tag": -1, "id": ident})
            ident += 1
            game_count += 1

# To ensure that Rueckspiele only happen in Rueckrunde, this is the threshold
rueckrundenstart = max_spieltage // 2

from ortools.sat.python import cp_model
model = cp_model.CpModel()

# Create variables in our solver's model to hold all the possible Spieltage
variables = {game["id"]: model.new_int_var(0, max_spieltage - 1, f"game{game['id']}") for game in tabelle}

# Make sure each team only plays once each spieltag
for team in data["teams"]:
    # team_games = [game["id"] for game in tabelle if game["team_a"] == team["id"] or game["team_b"] == team["id"]]
    # print (len(team_games), team_games, team)
    model.add_all_different([variables[game["id"]] for game in tabelle if game["team_a"] == team["id"] or game["team_b"] == team["id"]])

# Make sure each ort is only used once per spieltag+week
for ort in data["orte"]:
    ort_games = [game for game in tabelle if game["ort"] == ort["id"]]
    for tag in list({game["tag"] for game in ort_games}):
        model.add_all_different([variables[game["id"]] for game in ort_games if game["tag"] == tag ])

# Make sure that Rueckrundenspiele only happen at Rueckrunde
for game in tabelle:
    # Don't force this constraint on wildcard games
    if game["ort"] == -1: continue

    # Find Rueckspiel
    rueckspiel = next((candidate for candidate in tabelle if candidate["team_a"] == game["team_b"] and candidate["team_b"] == game["team_a"]), None)

    va = variables[game["id"]]
    vb = variables[rueckspiel["id"]]

    runde_a = model.new_int_var(0, 1, "tmpa")
    runde_b = model.new_int_var(0, 1, "tmpb")

    model.add_division_equality(runde_a, va, rueckrundenstart)
    model.add_division_equality(runde_b, vb, rueckrundenstart)
    model.add(runde_a != runde_b)

# Make sure that Heimspiele and Gastspiele alternate
violations = []
for team in data["teams"]:
    heimspiele = [candidate for candidate in tabelle if candidate["team_a"] == team["id"] ]

    for heimspiel in heimspiele:
        for conflict in heimspiele:
            if heimspiel["id"] == conflict["id"]: continue

            va = variables[heimspiel["id"]]
            vb = variables[conflict["id"]]

            tmpdiff_a = model.new_int_var(-max_spieltage, max_spieltage, "tmpa")
            tmpdiff_b = model.new_int_var(-max_spieltage, max_spieltage, "tmpb")

            model.add(tmpdiff_a == va - vb)
            model.add(tmpdiff_b == vb - va)

            tmpconsequtive_a = model.new_bool_var("tmpa")
            tmpconsequtive_b = model.new_bool_var("tmpb")

            model.add(tmpdiff_a == 1).only_enforce_if(tmpconsequtive_a)
            model.add(tmpdiff_a != 1).only_enforce_if(tmpconsequtive_a.Not())

            model.add(tmpdiff_b == 1).only_enforce_if(tmpconsequtive_b)
            model.add(tmpdiff_b != 1).only_enforce_if(tmpconsequtive_b.Not())

            is_consecutive = model.new_bool_var('is_consecutive')

            model.add_bool_or([tmpconsequtive_a, tmpconsequtive_b]).OnlyEnforceIf(is_consecutive)
            model.add_bool_and([tmpconsequtive_a.Not(), tmpconsequtive_b.Not()]).OnlyEnforceIf(is_consecutive.Not())

            violations.append(is_consecutive)

# Most probably this constraint can't be fully satisfied, so just
# ask the solver to minimize clusters of Heimspiele for a team
model.minimize(sum(violations))

print ("All set, solving")

solver = cp_model.CpSolver()
solver.parameters.max_time_in_seconds = 30.0
solver.parameters.random_seed = randint(0, 1_000_000)
status = solver.solve(model)
print(status, solver.status_name())

from sys import exit
if status == cp_model.INFEASIBLE:
    exit(-1)

# Distribute Spieltage back to the games so we can sort them
for game in tabelle:
    game["spieltag"] = solver.value(variables[game["id"]])

# tabelle.sort(key=lambda game: game["spieltag"])

# PRINT OUT RESULTS AS TABLE

# helper table for quicker lookup of teams and orte by their id
all_teams = {team["id"]: team for team in data["teams"]}
all_teams[-1] = {"name": "*"}
all_orte = {ort["id"]: ort for ort in data["orte"]}
all_orte[-1] = {"name": "*"}

for league in range(league_count):
    league_tabelle = [game for game in tabelle if game["liga"] == league]
    league_tabelle.sort(key=lambda game: game["spieltag"])

    for game in league_tabelle:
        print(game["spieltag"], all_teams[game["team_a"]]["name"], "::", all_teams[game["team_b"]]["name"], "::", all_orte[game["ort"]]["name"])

#print( json.dumps(data))
#print( json.dumps(tabelle))