"""
Partition the graph in time domain

It takes a decoding graph and position as input.

When used with QEC-Playground outputs, it will likely be able to partition without a problem
"""

import os
import fusion_blossom as fb

# example
# generated by `cargo run --release tool benchmark '[5]' '[5]' '[0.01]' -p10 -m10 --code-type rotated-planar-code --noise-model phenomenological --decoder fusion --decoder-config '{"max_tree_size":0}' --debug-print fusion-blossom-syndrome-file`
# then see output file at `tmp/fusion.syndromes`


EXAMPLE_SYNDROME_FILE = os.path.join(os.path.dirname(
    os.path.abspath(__file__)), "graph_time_partition.syndromes")


def graph_time_partition(initializer: fb.SolverInitializer, positions: list[fb.VisualizePosition]) -> fb.PartitionConfig:
    assert len(positions) > 0
    partition = fb.PartitionConfig(initializer.vertex_num)
    # first check if the time is monotonically increasing
    last_t = positions[0].t
    t_list = [last_t]
    for position in positions:
        assert position.t >= last_t, "t not monotonically increasing, vertex reordering must be performed before calling this"
        if position.t != last_t:
            t_list.append(position.t)
        last_t = position.t
    # pick the t value in the middle to split it
    t_split = t_list[len(t_list) // 2]
    # find the vertices indices
    split_start_index = None
    split_end_index = None
    for vertex_index, position in enumerate(positions):
        if split_start_index is None and position.t == t_split:
            split_start_index = vertex_index
        if position.t == t_split:
            split_end_index = vertex_index + 1
    assert split_start_index is not None
    # partitions are found
    partition.partitions = [
        fb.VertexRange(0, split_start_index),
        fb.VertexRange(split_end_index, len(positions)),
    ]
    partition.fusions = [(0, 1)]
    return partition


def read_syndrome_file(syndrome_filepath):
    reader = fb.ErrorPatternReader(filename=syndrome_filepath)
    return reader.get_initializer(), reader.get_positions()


def main():
    # example
    print(f"reading example from: {EXAMPLE_SYNDROME_FILE}")
    initializer, positions = read_syndrome_file(EXAMPLE_SYNDROME_FILE)
    partition = graph_time_partition(initializer, positions)
    print(partition)
    print(partition.to_json())


if __name__ == "__main__":
    main()