#!/usr/bin/env python3

# Propagation_Demo.py
"""
Kapitel CP-SAT: Propagation sichtbar machen.
Vergleicht die Zahl der geprueften Kombinationen bei roher Aufzaehlung
mit der Zahl der Verzweigungen, die CP-SAT tatsaechlich braucht.
"""

import itertools

from ortools.sat.python import cp_model


def rohe_aufzaehlung(n: int, obergrenze: int) -> tuple[int, int]:
    """Zaehlt ALLE Kombinationen und prueft jede einzeln (Brute Force)."""
    geprueft = 0
    loesungen = 0
    for kombination in itertools.product(range(1, obergrenze + 1), repeat=n):
        geprueft += 1
        if len(set(kombination)) == n and sum(kombination) == 3 * n:
            loesungen += 1
    return geprueft, loesungen


def mit_cp_sat(n: int, obergrenze: int) -> tuple[int, int, float]:
    """Dasselbe Problem deklarativ: alle verschieden, Summe = 3n."""
    modell = cp_model.CpModel()
    x = [modell.NewIntVar(1, obergrenze, f"x{i}") for i in range(n)]
    modell.AddAllDifferent(x)                 # globales Constraint
    modell.Add(sum(x) == 3 * n)

    loeser = cp_model.CpSolver()
    loeser.parameters.enumerate_all_solutions = True

    class Zaehler(cp_model.CpSolverSolutionCallback):
        def __init__(self):
            super().__init__()
            self.anzahl = 0

        def on_solution_callback(self):
            self.anzahl += 1

    zaehler = Zaehler()
    loeser.Solve(modell, zaehler)
    return zaehler.anzahl, loeser.NumBranches(), loeser.WallTime()


if __name__ == "__main__":
    print("=" * 84)
    print("  PROPAGATION: WIE VIEL ARBEIT SPART SICH DER SOLVER?")
    print("  Aufgabe: n Zahlen aus 1..G, alle verschieden, Summe = 3n")
    print("=" * 84)
    print(f"{'n':>3} {'G':>4} | {'Brute Force':>14} | {'CP-SAT':>10} | "
          f"{'Ersparnis':>11} | {'Loesungen':>10}")
    print("-" * 84)

    for n, grenze in [(4, 8), (5, 10), (6, 12), (7, 14)]:
        kombis, treffer_bf = rohe_aufzaehlung(n, grenze)
        treffer_cp, verzweigungen, dauer = mit_cp_sat(n, grenze)
        assert treffer_bf == treffer_cp, "Beide Verfahren muessen gleich viele finden!"
        ersparnis = kombis / max(verzweigungen, 1)
        print(f"{n:>3} {grenze:>4} | {kombis:>14,} | {verzweigungen:>10,} | "
              f"{ersparnis:>10.0f}x | {treffer_cp:>10,}")

    print("-" * 84)
    print("'Brute Force' = alle Kombinationen. 'CP-SAT' = tatsaechliche Verzweigungen.")
    print("Der Rest wurde durch Propagation ausgeschlossen, ohne ihn anzusehen.")
    print("=" * 84)