#!/usr/bin/env python3

# erzeuge_branch_and_bound.py
"""
Erzeugt den Branch-and-Bound-Suchbaum zum MILP-Kapitel:

    bilder_04/kap_milp_suchbaum.svg / .png

Gezeichnet wird der Baum der Handrechnung des Kapitels:

    max 5 x1 + 4 x2   u.d.N.  6 x1 + 4 x2 <= 24,  x1 + 2 x2 <= 6,  x ganzzahlig

Der Baum wird NICHT abgeschrieben, sondern gerechnet: Eine kleine
Branch-and-Bound-Schleife loest die LP-Relaxationen mit scipy.linprog,
verzweigt ueber die erste gebrochene Variable und kappt nach denselben drei
Regeln, die der Kapiteltext nennt (unzulaessig, Schranke schlechter als der
Incumbent, Loesung bereits ganzzahlig). Am Ende prueft das Skript, dass Optimum
und Zahl der geloesten LPs mit dem Buchtext uebereinstimmen - stimmt etwas
nicht, bricht es ab, statt ein plausibles Bild zu malen.

Warum ueberhaupt neu: Das Vorgaengerbild zeigte einen Baum mit voellig anderen
Zahlen (Wurzel x1 = 2,7, Schranke 150,4, Incumbent 138) zu einem Modell, das
nirgends im Buch steht.

Aufruf (aus dem Repository-Wurzelverzeichnis):
    python3 bilder_04/erzeuge_branch_and_bound.py

Benoetigt: numpy, scipy, graphviz (Python-Paket) und das Werkzeug 'dot'
"""

from __future__ import annotations

import os

import graphviz
import numpy as np
from scipy.optimize import linprog

BASIS = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
BILDER = os.path.join(BASIS, "bilder_04")

INDIGO, CYAN, VIOLETT, GRUEN, AMBER = (
    "#4338ca", "#0891b2", "#7c3aed", "#059669", "#b45309")
TINT = {INDIGO: "#eef2ff", CYAN: "#ecfeff", VIOLETT: "#f5f3ff",
        GRUEN: "#ecfdf5", AMBER: "#fffbeb"}

# Modell der Handrechnung
C = np.array([5.0, 4.0])
A = np.array([[6.0, 4.0], [1.0, 2.0]])
B = np.array([24.0, 6.0])
TOL = 1e-7


def loese(untergrenzen, obergrenzen):
    """LP-Relaxation mit zusaetzlichen Schranken. Liefert (x, Z) oder (None, None)."""
    ergebnis = linprog(c=-C, A_ub=A, b_ub=B,
                       bounds=list(zip(untergrenzen, obergrenzen)),
                       method="highs")
    if not ergebnis.success:
        return None, None
    return np.asarray(ergebnis.x), float(-ergebnis.fun)


def erste_gebrochene(x):
    for j, wert in enumerate(x):
        if abs(wert - round(wert)) > 1e-6:
            return j
    return None


def zahl(wert: float) -> str:
    """Deutsche Schreibweise, hoechstens zwei Nachkommastellen.

    Das '+ 0.0' faengt die negative Null ab: linprog liefert fuer x2 = 0
    gelegentlich -0.0, und '(4; -0)' im Buch waere ein Schoenheitsfehler mit
    Ansage.
    """
    text = f"{wert + 0.0:.2f}".rstrip("0").rstrip(".")
    if text in ("-0", "-"):
        text = "0"
    return text.replace(".", ",")


def branch_and_bound():
    """Baut den Suchbaum und liefert (Knotenliste, Kanten, Optimum, Zielwert).

    Die Reihenfolge - erst der Ast mit der kleineren Obergrenze, dann der
    andere - entspricht der Handrechnung: Ast A (x2 <= 1) wird vor Ast B
    (x2 >= 2) bearbeitet, weshalb der Incumbent aus A stammt und B daran
    scheitert.
    """
    knoten, kanten = [], []
    incumbent, incumbent_x = -np.inf, None
    # TIEFENSUCHE (Stapel, nicht Warteschlange). Das ist keine Kosmetik: Nur so
    # ist Ast A samt seinen Blaettern fertig, bevor Ast B an die Reihe kommt -
    # und nur deshalb scheitert B am Incumbent 20, wie es der Kapiteltext
    # beschreibt. Breitensuche wuerde B zum Incumbent machen und die Aussage
    # des Bildes umdrehen.
    stapel = [("P₀", None, [0.0, 0.0], [np.inf, np.inf], None)]
    zaehler = 0

    while stapel:
        name, elter, unten, oben, kantentext = stapel.pop()
        x, z = loese(unten, oben)
        zaehler += 1
        if elter is not None:
            kanten.append((elter, name, kantentext))

        if x is None:
            knoten.append((name, "unzulässig — gekappt", AMBER, None))
            continue
        if z < incumbent - TOL:
            knoten.append((name, f"Z = {zahl(z)} ≤ Incumbent {zahl(incumbent)}\\n"
                                 "→ gekappt", AMBER, None))
            continue

        j = erste_gebrochene(x)
        if j is None:
            besser = z > incumbent + TOL
            if besser:
                incumbent, incumbent_x = z, x
            knoten.append((name,
                           f"x = ({zahl(x[0])}; {zahl(x[1])}), Z = {zahl(z)}\\n"
                           + ("ganzzahlig — neuer Incumbent" if besser
                              else f"ganzzahlig, aber ≤ Incumbent {zahl(incumbent)}"),
                           GRUEN if besser else AMBER, z))
            continue

        knoten.append((name,
                       f"x = ({zahl(x[0])}; {zahl(x[1])}), Schranke Z ≤ {zahl(z)}\\n"
                       f"x{j + 1} gebrochen — verzweigen", INDIGO if elter is None else CYAN,
                       z))
        grenze = float(np.floor(x[j]))
        links_oben = list(oben); links_oben[j] = grenze
        rechts_unten = list(unten); rechts_unten[j] = grenze + 1
        # Zuerst den ">="-Ast auflegen, damit der "<="-Ast oben liegt und
        # als Erster bearbeitet wird - die Reihenfolge der Handrechnung.
        # Namen wie im Kapiteltext: die Wurzel verzweigt in Ast A und Ast B,
        # deren Kinder heissen A1/A2 bzw. B1/B2.
        links, rechts = ("A", "B") if elter is None else (f"{name}1", f"{name}2")
        stapel.append((rechts, name, rechts_unten, list(oben),
                       f"x{j + 1} ≥ {int(grenze) + 1}"))
        stapel.append((links, name, list(unten), links_oben,
                       f"x{j + 1} ≤ {int(grenze)}"))

    return knoten, kanten, incumbent_x, incumbent, zaehler


def zeichne(knoten, kanten, bestes, zielwert, geloest) -> None:
    punkt = graphviz.Digraph(format="svg")
    punkt.attr(rankdir="TB", splines="polyline", nodesep="0.45", ranksep="0.55",
               fontname="DejaVu Sans", bgcolor="transparent", fontsize="12",
               labelloc="t",
               label=(f"Branch-and-Bound: {geloest} gelöste LPs statt aller "
                      f"ganzzahligen Punkte\\n"
                      f"max 5x₁ + 4x₂  u.d.N.  6x₁+4x₂ ≤ 24,  x₁+2x₂ ≤ 6"))
    punkt.attr("node", fontname="DejaVu Sans", fontsize="10")
    punkt.attr("edge", fontname="DejaVu Sans", fontsize="9", color="#5c6380")

    for name, text, farbe, _ in knoten:
        beschriftung = f"{name}\\n{text}"
        punkt.node(name, beschriftung, shape="box", style="filled,rounded",
                   fillcolor=TINT[farbe], color=farbe, penwidth="2.0",
                   margin="0.16,0.09")
    for von, nach, text in kanten:
        punkt.edge(von, nach, label=text or "")

    punkt.node("ergebnis",
               f"Optimum: x* = ({zahl(bestes[0])}; {zahl(bestes[1])}), "
               f"Z* = {zahl(zielwert)}",
               shape="note", style="filled", fillcolor="#f8fafc",
               color="#94a3b8", fontsize="10", margin="0.18,0.12")

    os.makedirs(BILDER, exist_ok=True)
    for endung in ("svg", "png"):
        punkt.format = endung
        pfad = punkt.render(os.path.join(BILDER, "kap_milp_suchbaum"), cleanup=True)
        print(f"geschrieben: {pfad}")


if __name__ == "__main__":
    knoten, kanten, bestes, zielwert, geloest = branch_and_bound()
    print(f"{len(knoten)} Knoten, {geloest} geloeste LPs")
    for name, text, _, _ in knoten:
        print(f"  {name:5} {text.replace(chr(92) + 'n', ' | ')}")
    print(f"Optimum: ({bestes[0]:g}; {bestes[1]:g}), Z = {zielwert:g}")

    # Gegenprobe gegen die Handrechnung im Kapitel
    if not (np.allclose(bestes, [4.0, 0.0]) and abs(zielwert - 20.0) < TOL):
        raise SystemExit("FEHLER: Ergebnis weicht von der Handrechnung ab "
                         "(erwartet x* = (4; 0), Z* = 20).")
    zeichne(knoten, kanten, bestes, zielwert, geloest)