#!/usr/bin/python3

from dataclasses import dataclass
import random
import typing
import copy
import re
import os


@dataclass
class Semaphore:
    z: int
    c: int
    startPhase: int

    @property
    def period(self) -> int:
        return self.z + self.c

    def print(self, file):
        assert self.z >= 1 and self.z <= 10**9
        assert self.c >= 0 and self.c <= 10**9
        startState = "Z" if self.startPhase % self.period < self.z else "C"
        startPhase = self.startPhase % self.period
        f = self.z - startPhase if startPhase < self.z else self.period - startPhase
        assert f > 0
        assert f <= (self.z if startState == "Z" else self.c)
        print(self.z, self.c, startState, f, file=file)


@dataclass
class Edge:
    a: int
    b: int
    t: int
    semaphore: Semaphore

    def print(self, file):
        assert self.a != self.b
        assert self.t >= 1 and self.t <= 10**9
        print(self.a + 1, self.b + 1, self.t, end=" ", file=file)
        self.semaphore.print(file=file)


@dataclass
class Graph:
    n: int
    edges: list[Edge]

    def relabel_swap(self, a: int, b: int) -> "Graph":
        for edge in self.edges:
            if edge.a in [a, b]:
                edge.a = a + b - edge.a
            if edge.b in [a, b]:
                edge.b = a + b - edge.b
        return self

    def relabel_insert(self, a: int, b: int) -> "Graph":
        """Inserts a before b."""
        for edge in self.edges:
            # first, simulate removing a
            marker = a + self.n
            if edge.a == a:
                edge.a = marker
            elif edge.a > a:
                edge.a -= 1
            if edge.b == a:
                edge.b = marker
            elif edge.b > a:
                edge.b -= 1
            # then, simulate reinserting before b
            # make sure it's idempotent when a == b
            insert_before_what = b if a >= b else b - 1
            if edge.a == marker:
                edge.a = b
            elif edge.a >= insert_before_what:
                edge.a += 1
            if edge.b == marker:
                edge.b = b
            elif edge.b >= insert_before_what:
                edge.b += 1
        return self

    def shuffled(self) -> "Graph":
        """Shuffles all vertices except 0 and n-1."""
        for edge in self.edges:
            assert edge.a >= 0 and edge.a < self.n
            assert edge.b >= 0 and edge.b < self.n
        perm = list(range(1, self.n - 1))
        random.shuffle(perm)
        perm = [0] + perm + [self.n - 1]
        for edge in self.edges:
            edge.a = perm[edge.a]
            edge.b = perm[edge.b]
        return self

    def concatenate(self, other: "Graph") -> "Graph":
        """
        Mutates and returns self
        Note: this results in a disconnected graph
        """
        for edge in copy.deepcopy(other.edges):
            edge.a += self.n
            self.b += self.n
            self.edges.append(edge)
        self.n += other.n
        return self

    def print(self, file):
        assert self.n <= 10**5
        assert len(self.edges) <= 2 * 10**5
        print(self.n, len(self.edges), file=file)
        for edge in self.edges:
            edge.print(file)


@dataclass
class TestSuite:
    number: int
    tests: list[Graph]

    def gen(self):
        fs = os.listdir("test")
        for f in fs:
            if re.match(f"{self.number:02d}\\..\\...", f) is not None:
                os.remove(f"test/{f}")
        for i, test in enumerate(self.tests):
            ch = chr(ord("a") + i)
            f = f"test/{self.number:02d}.{ch}.in"
            print("Creating", f)
            with open(f, "w") as fout:
                test.print(fout)


def gen_dense_layered(
    n: int,
    layers: int,
    gen_base: typing.Callable[[int, int], Edge],
) -> Graph:
    """m <= 2 * (n/layers)**2 * (layers - 1)"""
    # nodes split into layers; complete bipartite edges between consecutive layers
    # later edges give slightly smaller weights to force repeated improvements
    L = [list(range(i * n // layers, (i + 1) * n // layers)) for i in range(layers)]
    heavy_edges: list[Edge] = []
    light_edges: list[Edge] = []
    for i in range(layers - 1):
        A, B = L[i], L[i + 1]
        # adversarial order: add heavier edges first, then slightly lighter ones
        for u in A:
            order = B[:]
            random.shuffle(order)
            # first pass: heavy edges
            hes = []
            for v in order:
                edge = gen_base(u, v)
                edge.t = max(edge.t, 2)
                hes.append(edge)
            # second pass: slightly lighter so distances improve again
            les = copy.deepcopy(hes)
            for edge in les:
                edge.t -= 1
            random.shuffle(les)
            heavy_edges.extend(hes)
            light_edges.extend(les)
    edges = heavy_edges = light_edges
    assert len(edges) <= 2 * (n / layers) ** 2 * (layers - 1)
    return Graph(n=n, edges=edges)


def gen_grid_equal_paths(
    k: int, gen_edge: typing.Callable[[tuple[int, int], tuple[int, int]], Edge]
):
    """
    k: grid dimension
    Edges generated through gen_edge have their a, b fields ignored (those are set based on position within grid)
    m <= 4 * k**2
    """
    # nodes numbered r*k + c
    edges: list[Edge] = []
    n = k * k

    def id(r: int, c: int) -> int:
        return r * k + c

    for r in range(k):
        for c in range(k):
            if r + 1 < k:
                edge = gen_edge((r, c), (r + 1, c))
                edge.a = id(r, c)
                edge.b = id(r + 1, c)
                edges.append(edge)
            if r > 0:
                edge = gen_edge((r, c), (r - 1, c))
                edge.a = id(r, c)
                edge.b = id(r - 1, c)
                edges.append(edge)
            if c + 1 < k:
                edge = gen_edge((r, c), (r, c + 1))
                edge.a = id(r, c)
                edge.b = id(r, c + 1)
                edges.append(edge)
            if c > 0:
                edge = gen_edge((r, c), (r, c - 1))
                edge.a = id(r, c)
                edge.b = id(r, c - 1)
                edges.append(edge)
    assert len(edges) <= 4 * k**2
    return Graph(n=n, edges=edges)


def gen_near_ties(
    depth: int,
    gen_edge: typing.Callable[
        [int, int, typing.Literal["big", "small", "medium"]], Edge
    ],
) -> Graph:
    """n <= 2*depth, m <= 4*depth"""
    # a ladder: top path edges weight=1, bottom path weight=1+eps, many cross-links
    edges: list[Edge] = []
    n = 2 * depth
    top = [i for i in range(0, n, 2)]
    bot = [i for i in range(1, n, 2)]
    for i in range(depth - 1):
        edges.append(gen_edge(top[i], top[i + 1], "small"))
        edges.append(gen_edge(bot[i], bot[i + 1], "big"))
        # cross links that make many almost-equal combinations
        edges.append(gen_edge(top[i], bot[i + 1], "medium"))
        edges.append(gen_edge(bot[i], top[i + 1], "medium"))
    assert n <= 2 * depth
    assert len(edges) <= 4 * depth
    return Graph(n=n, edges=edges)


def gen_path_graph(n: int, gen_edge: typing.Callable[[int, int], Edge]) -> Graph:
    """m = n - 1"""
    edges = [gen_edge(i, i + 1) for i in range(n - 1)]
    assert len(edges) == n - 1
    return Graph(n=n, edges=edges)


def gen_star_with_spokes(n: int, gen_edge: typing.Callable[[int, int], Edge]) -> Graph:
    """m <= 3 * n"""
    edges: list[Edge] = []
    for v in range(n - 1):
        edges.append(gen_edge(v, n - 1))  # star
        edges.append(gen_edge(n - 1, v))
    # a few long heavy spokes from leaves
    for v in range(n // 2):
        edges.append(gen_edge(v, v + 1))
        edges.append(gen_edge(v + 1, v))
    assert len(edges) <= 3 * n
    return Graph(n=n, edges=edges).relabel_swap(n - 2, n - 1)


def gen_multigraph_bits(
    n: int, density: int, gen_edge: typing.Callable[[int, int], Edge]
) -> Graph:
    """
    m <= 3 * n * density
    Can +1 edge length generated from gen_edge
    """
    edges: list[Edge] = []
    for _ in range(n * density):
        u = random.randrange(n)
        v = u
        while v == u:
            v = random.randrange(n)
        edge = gen_edge(u, v)
        edge.t += 1
        edges.append(edge)
        if random.random() < 0.2:
            edge2 = copy.deepcopy(edge)
            edge2.t -= 1
            edges.append(edge2)  # parallel with slightly better weight
        if random.random() < 0.05:
            edge3 = copy.deepcopy(edge)
            edge3.t = 1
            edges.append(edge3)
    assert len(edges) <= 3 * n * density
    return Graph(n=n, edges=edges)


def gen_preferential_attachment(
    n: int, density: int, gen_edge: typing.Callable[[int, int], Edge]
):
    """
    m <= n * 3 * density
    Always generates edges both directions (though the edge params may differ for the two directions)
    Among vertices in 1, ..., n-2, those with small ID are more likely to have a large degree
    """
    edges: list[Edge] = []
    targets = [0, 1]  # seed
    for v in range(2, n):
        # connect to m nodes chosen proportional to degree
        chosen = set()
        num_choose = random.randint(density // 2, density + density // 2)
        for _ in range(3 * num_choose):
            if len(chosen) >= min(num_choose, len(targets)):
                break
            chosen.add(random.choice(targets))
        for u in chosen:
            edges.append(gen_edge(v, u))
            edges.append(gen_edge(u, v))
            targets.append(u)
            targets.append(v)
    assert len(edges) <= n * 3 * density
    return Graph(n=n, edges=edges).relabel_insert(n - 1, 0)


def gen_random(
    n: int, m: int, reach: int, gen_edge: typing.Callable[[int, int], Edge]
) -> Graph:
    """reach: when generating an edge with 1 endpoint in x, can reach only vertices x+1, ..., x+reach"""
    assert m >= n
    edges: list[Edge] = []
    for u in range(n - 1):
        v = u
        while v == u:
            v = random.randint(u + 1, min(n - 1, u + reach))
        if u > v:
            tmp = u
            u = v
            v = tmp
        edges.append(gen_edge(u, v))
    for _ in range(m - n + 1):
        u = random.randrange(n)
        v = u
        while v == u:
            v = random.randint(max(0, u - reach), min(n - 1, u + reach))
        if u > v:
            tmp = u
            u = v
            v = tmp
        edges.append(gen_edge(u, v))
    assert len(edges) == m
    return Graph(n=n, edges=edges)


def gen_edge_subtask1(a: int, b: int) -> Edge:
    return Edge(
        a,
        b,
        t=random.randint(1, 10),
        semaphore=Semaphore(z=1, c=1, startPhase=random.randint(0, 1)),
    )


random.seed(564453543)

t1 = TestSuite(
    1,
    [
        gen_dense_layered(n=3000, layers=100, gen_base=gen_edge_subtask1),
        gen_grid_equal_paths(
            k=200, gen_edge=lambda _a, _b: gen_edge_subtask1(0, 0)
        ).shuffled(),
        gen_multigraph_bits(
            n=5000,
            density=10,
            gen_edge=lambda a, b: Edge(
                a,
                b,
                t=random.randint(1, 9),
                semaphore=Semaphore(z=1, c=1, startPhase=random.randint(0, 1)),
            ),
        ).shuffled(),
        gen_path_graph(n=100000, gen_edge=gen_edge_subtask1).shuffled(),
        gen_preferential_attachment(
            n=10000, density=5, gen_edge=gen_edge_subtask1
        ).shuffled(),
        gen_preferential_attachment(
            n=5000, density=10, gen_edge=gen_edge_subtask1
        ).shuffled(),
        gen_random(
            n=100000,
            m=200000,
            reach=100000,
            gen_edge=lambda a, b: Edge(
                a,
                b,
                t=1
                + random.randint(
                    ((abs(a - b) - 1) // 10000) // 2, ((abs(a - b) - 1) // 10000)
                ),
                semaphore=Semaphore(z=1, c=1, startPhase=random.randint(0, 1)),
            ),
        ).shuffled(),
        gen_random(
            n=100000,
            m=200000,
            reach=100,
            gen_edge=lambda a, b: Edge(
                a,
                b,
                t=random.randint(1, 1 + (abs(a - b) - 1) // 10),
                semaphore=Semaphore(z=1, c=1, startPhase=random.randint(0, 1)),
            ),
        ).shuffled(),
        gen_star_with_spokes(n=60000, gen_edge=gen_edge_subtask1).shuffled(),
    ],
)


def gen_edge_subtask2(a: int, b: int) -> Edge:
    z = random.randint(1, 500)
    c = 1000 - z
    return Edge(
        a,
        b,
        t=1,
        semaphore=Semaphore(z=z, c=c, startPhase=random.randint(0, 999)),
    )


def gen_edge_subtask2_hard(a: int, b: int) -> Edge:
    return Edge(
        a,
        b,
        t=1,
        semaphore=Semaphore(z=1, c=999, startPhase=random.randint(0, 999)),
    )


t2 = TestSuite(
    2,
    [
        gen_dense_layered(n=350, layers=50, gen_base=gen_edge_subtask2_hard),
        gen_dense_layered(n=150, layers=10, gen_base=gen_edge_subtask2_hard),
        gen_grid_equal_paths(
            k=32, gen_edge=lambda _a, _b: gen_edge_subtask2(0, 0)
        ).shuffled(),
        gen_path_graph(n=5000, gen_edge=gen_edge_subtask2).shuffled(),
        gen_preferential_attachment(
            n=750, density=2, gen_edge=gen_edge_subtask2_hard
        ).shuffled(),
        gen_preferential_attachment(
            n=750, density=2, gen_edge=gen_edge_subtask2
        ).shuffled(),
        gen_random(
            n=1000,
            m=5000,
            reach=5,
            gen_edge=gen_edge_subtask2,
        ).shuffled(),
        gen_random(
            n=1000,
            m=5000,
            reach=5,
            gen_edge=gen_edge_subtask2_hard,
        ).shuffled(),
        gen_star_with_spokes(n=1600, gen_edge=gen_edge_subtask2_hard).shuffled(),
    ],
)


def gen_edge_subtask3(a: int, b: int) -> Edge:
    max_period = random.randint(2, 7)
    z = random.randint(1, max_period)
    c = max_period - z
    return Edge(
        a,
        b,
        t=1,
        semaphore=Semaphore(z=z, c=c, startPhase=random.randint(0, 6)),
    )


def gen_edge_subtask3_hard(a: int, b: int) -> Edge:
    max_period = random.randint(2, 7)
    z = 1
    c = max_period - z
    return Edge(
        a,
        b,
        t=1,
        semaphore=Semaphore(z=z, c=c, startPhase=random.randint(0, 6)),
    )


t3 = TestSuite(
    3,
    [
        gen_dense_layered(n=350, layers=50, gen_base=gen_edge_subtask3_hard),
        gen_dense_layered(n=150, layers=10, gen_base=gen_edge_subtask3_hard),
        gen_grid_equal_paths(
            k=32, gen_edge=lambda _a, _b: gen_edge_subtask3(0, 0)
        ).shuffled(),
        gen_path_graph(n=5000, gen_edge=gen_edge_subtask3).shuffled(),
        gen_preferential_attachment(
            n=750, density=2, gen_edge=gen_edge_subtask3_hard
        ).shuffled(),
        gen_preferential_attachment(
            n=750, density=2, gen_edge=gen_edge_subtask3
        ).shuffled(),
        gen_random(
            n=1000,
            m=5000,
            reach=5,
            gen_edge=gen_edge_subtask3,
        ).shuffled(),
        gen_random(
            n=1000,
            m=5000,
            reach=5,
            gen_edge=gen_edge_subtask3_hard,
        ).shuffled(),
        gen_star_with_spokes(n=1600, gen_edge=gen_edge_subtask3_hard).shuffled(),
    ],
)


def gen_edge_subtask4(a: int, b: int) -> Edge:
    return Edge(
        a,
        b,
        t=random.randint(1, 10**9),
        semaphore=Semaphore(z=1, c=0, startPhase=0),
    )


def gen_edge_subtask4_near_ties(
    a: int, b: int, edge_type: typing.Literal["small", "medium", "big"]
) -> Edge:
    t = 0
    if edge_type == "small":
        t = 9 * 10**8 + random.randint(0, 3)
    elif edge_type == "medium":
        t = 9 * 10**8 + random.randint(3, 6)
    else:
        t = 9 * 10**8 + random.randint(6, 9)
    return Edge(
        a,
        b,
        t=t,
        semaphore=Semaphore(z=1, c=0, startPhase=0),
    )


t4 = TestSuite(
    4,
    [
        gen_dense_layered(n=10000, layers=2000, gen_base=gen_edge_subtask4),
        gen_grid_equal_paths(
            k=200, gen_edge=lambda _a, _b: gen_edge_subtask4(0, 0)
        ).shuffled(),
        gen_preferential_attachment(
            n=33000,
            density=2,
            gen_edge=lambda a, b: Edge(
                a,
                b,
                t=random.randint(abs(a - b) * 10**3, abs(a - b) * 10**4),
                semaphore=Semaphore(z=1, c=0, startPhase=0),
            ),
        ).shuffled(),
        gen_preferential_attachment(
            n=12000,
            density=5,
            gen_edge=lambda a, b: Edge(
                a,
                b,
                t=random.randint(abs(a - b) * 10**3, abs(a - b) * 10**4),
                semaphore=Semaphore(z=1, c=0, startPhase=0),
            ),
        ).shuffled(),
        gen_random(
            n=100000,
            m=200000,
            reach=5,
            gen_edge=gen_edge_subtask4,
        ).shuffled(),
        gen_random(
            n=100000,
            m=200000,
            reach=100,
            gen_edge=lambda a, b: Edge(
                a,
                b,
                t=random.randint(abs(a - b) * 10**5, abs(a - b) * 10**6),
                semaphore=Semaphore(z=1, c=0, startPhase=0),
            ),
        ).shuffled(),
        gen_random(
            n=10000,
            m=200000,
            reach=10,
            gen_edge=lambda a, b: Edge(
                a,
                b,
                t=random.randint(abs(a - b) * 10**6, abs(a - b) * 10**7),
                semaphore=Semaphore(z=1, c=0, startPhase=0),
            ),
        ).shuffled(),
        gen_near_ties(depth=50000, gen_edge=gen_edge_subtask4_near_ties),
    ],
)


def make_gen_edge_subtask5(
    max_z: int,
    c_min_max: tuple[int, int],
    gen_edge_length: typing.Callable[[int, int], int],
):
    z = random.randint(1, max_z)
    c = random.randint(*c_min_max)

    def gen_edge(a: int, b: int) -> Edge:
        return Edge(
            a,
            b,
            t=gen_edge_length(a, b),
            semaphore=Semaphore(z=z, c=c, startPhase=0),
        )

    return gen_edge


make_gen_edge_subtask5_default = lambda: make_gen_edge_subtask5(
    100, (10**7, 10**8), lambda a, b: random.randint(1, 10**9)
)
gen_edge_subtask5_grid = make_gen_edge_subtask5_default()


def make_gen_edge_subtask5_near_ties(
    max_z: int,
    c_min_max: tuple[int, int],
):
    z = random.randint(1, max_z)
    c = random.randint(*c_min_max)

    def fn(a: int, b: int, edge_type: typing.Literal["small", "medium", "big"]) -> Edge:
        t = 0
        if edge_type == "small":
            t = 9 * 10**8 + random.randint(0, 3)
        elif edge_type == "medium":
            t = 9 * 10**8 + random.randint(3, 6)
        else:
            t = 9 * 10**8 + random.randint(6, 9)
        return Edge(
            a,
            b,
            t=t,
            semaphore=Semaphore(z=z, c=c, startPhase=0),
        )

    return fn


t5 = TestSuite(
    5,
    [
        gen_dense_layered(
            n=10000,
            layers=2000,
            gen_base=make_gen_edge_subtask5_default(),
        ),
        gen_grid_equal_paths(
            k=200, gen_edge=lambda _a, _b: gen_edge_subtask5_grid(0, 0)
        ).shuffled(),
        gen_preferential_attachment(
            n=33000,
            density=2,
            gen_edge=make_gen_edge_subtask5(
                100,
                (10**7, 10**8),
                lambda a, b: random.randint(abs(a - b) * 10**3, abs(a - b) * 10**4),
            ),
        ).shuffled(),
        gen_preferential_attachment(
            n=12000,
            density=5,
            gen_edge=make_gen_edge_subtask5(
                100,
                (10**7, 10**8),
                lambda a, b: random.randint(abs(a - b) * 10**3, abs(a - b) * 10**4),
            ),
        ).shuffled(),
        gen_random(
            n=100000,
            m=200000,
            reach=10,
            gen_edge=make_gen_edge_subtask5_default(),
        ).shuffled(),
        gen_random(
            n=100000,
            m=200000,
            reach=50,
            gen_edge=make_gen_edge_subtask5(
                100,
                (10**7, 10**8),
                lambda a, b: random.randint(abs(a - b) * 10**5, abs(a - b) * 10**6),
            ),
        ).shuffled(),
        gen_random(
            n=10000,
            m=200000,
            reach=20,
            gen_edge=make_gen_edge_subtask5(
                100,
                (10**7, 10**8),
                lambda a, b: random.randint(abs(a - b) * 10**6, abs(a - b) * 10**7),
            ),
        ).shuffled(),
        gen_near_ties(
            depth=50000, gen_edge=make_gen_edge_subtask5_near_ties(100, (10**7, 10**8))
        ),
    ],
)


def gen_edge_subtask6(a: int, b: int) -> Edge:
    z = random.randint(1, 10**3)
    c = random.randint(0, 10**8)
    startPhase = random.randint(0, z + c - 1)
    return Edge(
        a,
        b,
        t=random.randint(1, 10**9),
        semaphore=Semaphore(z=z, c=c, startPhase=startPhase),
    )


def gen_edge_subtask6_scaling_dist(a: int, b: int, k: tuple[int, int]) -> Edge:
    z = random.randint(1, 10**3)
    c = random.randint(0, 10**8)
    startPhase = random.randint(0, z + c - 1)
    return Edge(
        a,
        b,
        t=random.randint(abs(a - b) * k[0], abs(a - b) * k[1]),
        semaphore=Semaphore(z=z, c=c, startPhase=startPhase),
    )


def gen_edge_subtask6_near_ties(
    a: int, b: int, edge_type: typing.Literal["small", "medium", "big"]
) -> Edge:
    z = random.randint(1, 10**3)
    c = random.randint(0, 10**8)
    startPhase = random.randint(0, z + c - 1)
    t = 0
    if edge_type == "small":
        t = 9 * 10**8 + random.randint(0, 3)
    elif edge_type == "medium":
        t = 9 * 10**8 + random.randint(3, 6)
    else:
        t = 9 * 10**8 + random.randint(6, 9)
    return Edge(
        a,
        b,
        t=t,
        semaphore=Semaphore(z=z, c=c, startPhase=startPhase),
    )


t6 = TestSuite(
    6,
    [
        gen_dense_layered(n=10000, layers=2000, gen_base=gen_edge_subtask6),
        gen_grid_equal_paths(
            k=200, gen_edge=lambda _a, _b: gen_edge_subtask6(0, 0)
        ).shuffled(),
        gen_preferential_attachment(
            n=33000,
            density=2,
            gen_edge=lambda a, b: gen_edge_subtask6_scaling_dist(a, b, (10**3, 10**4)),
        ).shuffled(),
        gen_preferential_attachment(
            n=12000,
            density=5,
            gen_edge=lambda a, b: gen_edge_subtask6_scaling_dist(a, b, (10**3, 10**4)),
        ).shuffled(),
        gen_random(
            n=100000,
            m=200000,
            reach=15,
            gen_edge=gen_edge_subtask6,
        ).shuffled(),
        gen_random(
            n=100000,
            m=200000,
            reach=200,
            gen_edge=lambda a, b: gen_edge_subtask6_scaling_dist(a, b, (10**5, 10**6)),
        ).shuffled(),
        gen_random(
            n=10000,
            m=200000,
            reach=15,
            gen_edge=lambda a, b: gen_edge_subtask6_scaling_dist(a, b, (10**6, 10**7)),
        ).shuffled(),
        gen_near_ties(depth=50000, gen_edge=gen_edge_subtask6_near_ties),
    ],
)

t1.gen()
t2.gen()
t3.gen()
t4.gen()
t5.gen()
t6.gen()
