from bnbvisual import *

raw_nodes = {
    "0": ("Node 0", 250.0, (0.5, 1.0, 1.0, 0.5, 0.0, 1.0, 1.0, 0.2), {0: "root"}),
    "1": ("Node 1", 245.5, (0.0, 1.0, 0.8, 0.5, 0.0, 1.0, 1.0, 0.2), {1: "fractional"}),
    "2": ("Node 2", 248.2, (1.0, 0.9, 1.0, 0.5, 0.0, 1.0, 1.0, 0.2), {1: "fractional"}),
    "3": ("Node 3", 240.0, (0.0, 1.0, 0.0, 0.5, 0.0, 1.0, 1.0, 0.0), {2: "fractional"}),
    "4": ("Node 4", "N/A", None, {2: "infeasible"}),
    "5": ("Node 5", 242.1, (1.0, 0.0, 1.0, 0.5, 0.0, 1.0, 1.0, 0.2), {2: "fractional"}),
    "6": ("Node 6", 247.9, (1.0, 1.0, 0.8, 0.5, 0.0, 1.0, 1.0, 0.2), {2: "fractional"}),
    "7": (
        "Node 7",
        230.0,
        (0.0, 1.0, 0.0, 0.0, 0.0, 1.0, 1.0, 0.0),
        {3: "incumbent", 10: "pruned"},
    ),
    "8": ("Node 8", 238.5, (0.0, 1.0, 0.0, 1.0, 0.0, 0.5, 1.0, 0.0), {3: "fractional"}),
    "9": ("Node 9", 235.0, (1.0, 0.0, 1.0, 0.0, 0.0, 1.0, 1.0, 0.2), {4: "fractional"}),
    "10": (
        "Node 10",
        241.0,
        (1.0, 0.0, 1.0, 1.0, 0.0, 1.0, 0.6, 0.2),
        {4: "fractional"},
    ),
    "11": (
        "Node 11",
        245.0,
        (1.0, 1.0, 0.0, 0.5, 0.0, 1.0, 1.0, 0.2),
        {5: "fractional"},
    ),
    "12": (
        "Node 12",
        246.8,
        (1.0, 1.0, 1.0, 0.5, 0.0, 1.0, 1.0, 0.2),
        {5: "fractional"},
    ),
    "13": ("Node 13", 228.0, (0.0, 1.0, 0.0, 1.0, 0.0, 0.0, 1.0, 0.0), {6: "pruned"}),
    "14": ("Node 14", "N/A", None, {6: "infeasible"}),
    "15": (
        "Node 15",
        243.5,
        (1.0, 1.0, 1.0, 0.0, 0.0, 1.0, 1.0, 0.2),
        {7: "fractional"},
    ),
    "16": (
        "Node 16",
        245.2,
        (1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 0.6),
        {7: "fractional"},
    ),
    "17": (
        "Node 17",
        240.0,
        (1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 0.0),
        {8: "incumbent", 11: "pruned"},
    ),
    "18": (
        "Node 18",
        244.1,
        (1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 0.4, 1.0),
        {8: "fractional"},
    ),
    "19": ("Node 19", "N/A", None, {9: "infeasible"}),
    "20": (
        "Node 20",
        243.8,
        (1.0, 1.0, 1.0, 1.0, 0.0, 0.5, 1.0, 1.0),
        {9: "fractional"},
    ),
    "21": (
        "Node 21",
        241.5,
        (1.0, 1.0, 1.0, 1.0, 0.0, 0.0, 1.0, 1.0),
        {10: "incumbent"},
    ),
    "22": (
        "Node 22",
        242.9,
        (1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0),
        {10: "incumbent"},
    ),
}

edges = [
    ("0", "1", r"$x_{1} = 0$"),
    ("0", "2", r"$x_{1} = 1$"),
    ("1", "3", r"$x_{3} = 0$"),
    ("1", "4", r"$x_{3} = 1$"),
    ("2", "5", r"$x_{2} = 0$"),
    ("2", "6", r"$x_{2} = 1$"),
    ("3", "7", r"$x_{4} = 0$"),
    ("3", "8", r"$x_{4} = 1$"),
    ("5", "9", r"$x_{4} = 0$"),
    ("5", "10", r"$x_{4} = 1$"),
    ("6", "11", r"$x_{3} = 0$"),
    ("6", "12", r"$x_{3} = 1$"),
    ("8", "13", r"$x_{6} = 0$"),
    ("8", "14", r"$x_{6} = 1$"),
    ("12", "15", r"$x_{4} = 0$"),
    ("12", "16", r"$x_{4} = 1$"),
    ("16", "17", r"$x_{8} = 0$"),
    ("16", "18", r"$x_{8} = 1$"),
    ("18", "19", r"$x_{7} = 0$"),
    ("18", "20", r"$x_{7} = 1$"),
    ("20", "21", r"$x_{6} = 0$"),
    ("20", "22", r"$x_{6} = 1$"),
]

reveal_sequence = [
    ["0"],
    ["1", "2"],
    ["3", "4", "5", "6"],
    ["7", "8"],
    ["9", "10"],
    ["11", "12"],
    ["13", "14"],
    ["15", "16"],
    ["17", "18"],
    ["19", "20"],
    ["21", "22"],
]

nodes = generate_tree_layout(raw_nodes, edges)

draw_bnb_incremental(
    nodes, edges, reveal_sequence, filename="bnb_massive_depth", showplt=False
)