import sys

sys.path.append("..")
from mock_backends.ibm import *

import itertools
from qiskit.circuit.library import CXGate
from clonk.backend_utils.configurable_backend_v2 import ConfigurableFakeBackendV2
import rustworkx as rx
from abc import ABC, abstractmethod


class AbstractModularBackend(ConfigurableFakeBackendV2, ABC):
    """A mock backendv2, tailored for a tree structure with modular connectivity."""

    def __init__(self, module_size=4, children=4, total_levels=3):
        # Validate children for root module
        assert (
            1 < children <= module_size
        ), "Children must be between 2 and the module size, inclusive for the root module."

        self.c_map = []  # To hold all edges
        self.modules = {}  # Maps module ID to its qubits
        self.module_depths = {}  # Maps module ID to its depth
        self.module_size = module_size
        self.children = children  # Max children for the root module
        self.total_levels = total_levels
        self.long_edges = []

        # Construct the modular system
        self._construct_system()

        # Remove duplicate edges and filter to keep only edges from lower to higher qubit IDs
        self.c_map = list(set(self.c_map))
        self.c_map = [edge for edge in self.c_map if edge[0] < edge[1]]

        # Configure gates and measurements based on the generated system
        gate_configuration = {CXGate: [(i, j) for i, j in self.c_map]}
        measurable_qubits = list(range(max(max(pair) for pair in self.c_map) + 1))

        super().__init__(
            name=f"Modular{module_size}-{children}-{total_levels}",
            description="A mock backend with modular architecture",
            n_qubits=len(measurable_qubits),
            gate_configuration=gate_configuration,
            measurable_qubits=measurable_qubits,
            gate_durations={CXGate: 1},
            parameterized_gates={},
        )

    def add_long_edges(self, num_long_edges=1):
        G = rx.PyGraph()
        for node in range(self.num_qubits):
            G.add_node(node)
        for edge in self.c_map:
            G.add_edge(edge[0], edge[1], 1)

        long_edge_endpoints = set()  # Track qubits that are endpoints of long-edges

        for _ in range(num_long_edges):
            all_pairs_lengths = rx.all_pairs_dijkstra_path_lengths(
                G, edge_cost_fn=lambda x: 1
            )

            max_length = 0
            candidates = []
            for source, paths in all_pairs_lengths.items():
                for target, length in paths.items():
                    if source < target and length > max_length:
                        # Check if either qubit is already an endpoint of a long-edge
                        if (
                            source not in long_edge_endpoints
                            and target not in long_edge_endpoints
                        ):
                            max_length = length
                            candidates = [(source, target)]

            # If candidates are found, select the first one as the new long-edge
            if candidates:
                long_edge = candidates[0]
                self.long_edges.append(long_edge)
                G.add_edge(long_edge[0], long_edge[1], 1)
                long_edge_endpoints.update([long_edge[0], long_edge[1]])

    def _find_module_by_qubit(self, qubit):
        """Find the module ID that a given qubit belongs to."""
        for module_id, qubits in self.modules.items():
            if qubit in qubits:
                return module_id
        return None

    @abstractmethod
    def _construct_system(self):
        """
        Subclasses should implement this method to construct
        the modular system based on their specific scheme.
        """
        pass


class TreeModular(AbstractModularBackend):
    def _construct_system(self):
        # Implement the construction logic specific to the Tree scheme
        # Start the recursive construction of the system
        self._recursive_foo(
            qubit_counter=0, current_level=1, parent_id=None, parent_qubit=None, depth=0
        )

    def _recursive_foo(
        self, qubit_counter, current_level, parent_id, parent_qubit, depth
    ):
        module_id = len(self.modules)  # Unique ID for each module
        self.module_depths[module_id] = depth
        qubits = list(range(qubit_counter, qubit_counter + self.module_size))
        self.modules[module_id] = qubits

        # Create all-to-all connectivity within the module
        edges = list(itertools.permutations(qubits, 2))
        self.c_map.extend(edges)

        # Connect this module to its parent module (if applicable)
        if parent_qubit is not None:
            self.c_map.append(
                (parent_qubit, qubits[0])
            )  # Connection from parent to this module
            self.c_map.append((qubits[0], parent_qubit))  # And vice versa

        qubit_counter += self.module_size

        if current_level < self.total_levels:
            child_count = self.children if current_level == 1 else self.children - 1
            for i in range(child_count):
                qubit_counter = self._recursive_foo(
                    qubit_counter,
                    current_level + 1,
                    module_id,
                    qubits[i] if i < child_count - 1 else qubits[-1],
                    depth + 1,
                )

        return qubit_counter


class CorralModular(AbstractModularBackend):
    def __init__(self, num_snails=8, corral_skip_pattern=(0, 0)):
        self.num_snails = num_snails
        self.corral_skip_pattern = corral_skip_pattern
        # Initialize the modules dictionary to track which SNAIL neighborhood (module)
        # each qubit is primarily associated with.
        self.qubit_to_snail_map = {}  # Maps qubits to their primary SNAIL (module)
        super().__init__(
            module_size=4, children=2, total_levels=2
        )  # Example values, adjust as needed

    def _assign_qubit_to_module(self, qubit, snail):
        # Assign a qubit to its first encountered SNAIL neighborhood if it hasn't been assigned yet.
        if qubit not in self.qubit_to_snail_map:
            self.qubit_to_snail_map[qubit] = snail
            # Ensure the module entry exists for the SNAIL
            if snail not in self.modules:
                self.modules[snail] = []
            self.modules[snail].append(qubit)

    def _corral(self, num_snails, skip_pattern):
        num_levels = 2
        assert len(skip_pattern) == num_levels
        snail_edge_list = []
        for snail0, snail1 in zip(range(num_snails), range(1, num_snails + 1)):
            for i in range(num_levels):
                snail_edge_list.append(
                    (snail0, (skip_pattern[i] + snail1) % num_snails)
                )
        return snail_edge_list

    def _snail_to_connectivity(self, snail_edge_list):
        edge_list = []
        for qubit, snail_edge in enumerate(snail_edge_list):
            # Attempt to assign the qubit to a module based on its SNAIL edge
            self._assign_qubit_to_module(qubit, snail_edge[0])
            for temp_qubit, temp_snail_edge in enumerate(snail_edge_list):
                if qubit != temp_qubit and (
                    snail_edge[0] in temp_snail_edge or snail_edge[1] in temp_snail_edge
                ):
                    edge_list.append((qubit, temp_qubit))
        return edge_list

    def _construct_system(self):
        self.snail_edge_list = self._corral(self.num_snails, self.corral_skip_pattern)
        self.c_map = self._snail_to_connectivity(self.snail_edge_list)
        # Each SNAIL is treated as a module for depth purposes, though this could be adjusted
        self.module_depths = {
            snail: snail for snail in range(self.num_snails)
        }  # Example depth assignment


# Example usage
fake_modular = TreeModular(module_size=3, children=3, total_levels=3)
print("Modules:", fake_modular.modules)
print("Coupling Map:", fake_modular.c_map)

import matplotlib.pyplot as plt
import networkx as nx


def pretty_print_with_networkx(pb, label_text):
    G = nx.Graph()

    # Define colors for different levels
    level_colors = [
        "tab:blue",
        "tab:orange",
        "tab:green",
        "tab:red",
        "tab:purple",
        "tab:brown",
        "tab:pink",
        "tab:gray",
    ]

    # Module information for coloring and labeling
    modules, module_depths = pb.modules, pb.module_depths

    # Node colors and labels based on modules
    node_colors = []
    node_labels = {}
    for module_id, qubits in modules.items():
        depth = module_depths[module_id]
        color = level_colors[depth % len(level_colors)]
        for qubit in qubits:
            node_colors.append(color)
            node_labels[qubit] = f"Q{qubit}"

    # Add nodes and edges (regular edges)
    G.add_nodes_from([qubit for _, qubits in modules.items() for qubit in qubits])
    G.add_edges_from(list(pb.coupling_map))

    # Generate layout
    pos = nx.kamada_kawai_layout(G)

    fig, ax = plt.subplots(figsize=(8, 6))  # Adjust figure size as needed

    # Draw the graph without long edges
    nx.draw_networkx_nodes(
        G, pos, node_size=400, node_color=node_colors, alpha=0.6, ax=ax
    )  # Adjust node_size here
    nx.draw_networkx_labels(
        G, pos, labels=node_labels, font_size=8, ax=ax
    )  # Adjust font_size if necessary
    nx.draw_networkx_edges(
        G,
        pos,
        edgelist=list(pb.coupling_map),
        width=2,
        edge_color="gray",
        style="solid",
        ax=ax,
    )

    # Draw long edges on top of the existing layout, if present
    if hasattr(pb, "long_edges"):
        nx.draw_networkx_edges(
            G,
            pos,
            edgelist=pb.long_edges,
            width=5,
            edge_color="red",
            style="dashed",
            ax=ax,
        )

    ax.set_title(label_text)
    ax.axis("off")  # Turn off the axis

    plt.close(
        fig
    )  # Prevent it from displaying immediately in Jupyter or similar environments

    return fig


pretty_print = pretty_print_with_networkx

# Example usage
fake_modular = TreeModular(module_size=4, children=3, total_levels=3)
fake_modular.add_long_edges(num_long_edges=1)  # Example: Add 2 long-edges
print("Modules:", fake_modular.modules)
print("Coupling Map:", fake_modular.c_map)
print("Long Edges:", fake_modular.long_edges)

# Then, visualize
label_text = "Example Visualization with Long-Edges"
image = pretty_print(fake_modular, label_text)

# Display the image
# set image to active figure of plt
display(image)

# Generate the images with varying numbers of long edges
figures = []

# Adjusted loop to include iteration over number of long edges
for i in range(3, 6):  # Module sizes
    for k in range(1, 6):  # Total levels
        for j in range(i, i + 1):  # Children
            if k == 1 and j != 2:
                continue  # Skip configurations that don't make sense

            try:
                # Create a new instance for each configuration of i, j, k
                pb = TreeModular(module_size=i, children=j, total_levels=k)
                num_qubits = pb.num_qubits

                # Determine the maximum number of long edges for this configuration
                max_long_edges = 3 if k > 1 else 0  # At most 4, and 0 if levels=1

                for x in range(max_long_edges + 1):  # Iterate from 0 to max_long_edges
                    if num_qubits <= 64:
                        # Update long edges for the current configuration
                        pb.add_long_edges(num_long_edges=x)

                        label_text = f"Module Size: {i}, Children: {j}, Total Levels: {k}, Num Qubits: {num_qubits}, Long Edges: {x}"
                        print(label_text)

                        # Generate and store the figure for the current configuration
                        fig = pretty_print_with_networkx(pb, label_text)
                        figures.append(fig)

                        # Reset long edges for the next iteration
                        pb.long_edges = []  # Clear existing long edges before next addition
            except AssertionError:
                pass

from PIL import Image
import io

# Create an animation from the figures
# Assuming 'figures' is a list of matplotlib.figure.Figure objects generated from your loop
images = []  # This will hold the PIL Image objects
for fig in figures:
    buf = io.BytesIO()  # Create an in-memory buffer
    fig.savefig(buf, format="png")  # Save the figure to the buffer in PNG format
    buf.seek(0)  # Rewind the buffer to the beginning so we can read from it
    img = Image.open(buf)  # Open the image from the in-memory buffer
    images.append(img)  # Append the PIL Image to the list
    plt.close(fig)  # Close the figure to free memory

gif_path = "networkx_graph_animation.gif"
images[0].save(
    gif_path,
    save_all=True,
    append_images=images[1:],
    optimize=False,
    duration=1500,
    loop=0,
)

# Example usage
corral_modular = CorralModular(num_snails=8, corral_skip_pattern=(0, 0))
corral_modular.add_long_edges(num_long_edges=1)  # Example: Add 2 long-edges
print("Modules:", corral_modular.modules)
print("Coupling Map:", corral_modular.c_map)
print("Long Edges:", corral_modular.long_edges)

# Then, visualize
label_text = "Example Visualization with Long-Edges"
image = pretty_print(corral_modular, label_text)

# Display the image
# set image to active figure of plt
display(image)

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import numpy as np


def pretty_print_snails_3d(corral_modular):
    fig = plt.figure(figsize=(10, 7))
    ax = fig.add_subplot(111, projection="3d")

    num_snails = corral_modular.num_snails
    snail_edge_list = corral_modular._corral(
        num_snails, corral_modular.corral_skip_pattern
    )
    snail_colors = [
        "tab:blue",
        "tab:orange",
        "tab:green",
        "tab:red",
        "tab:purple",
        "tab:brown",
        "tab:pink",
        "tab:gray",
    ]

    # Parameters for the SNAIL posts and slats
    post_height = 0.5  # Height of the SNAIL posts
    post_radius = 0.05  # Radius of the SNAIL posts
    z_offset = post_height / 4
    z_lower = np.zeros(num_snails) + z_offset  # Lower rung positioned at the base
    z_upper = np.ones(num_snails) * 0.25 + z_offset  # Upper rung
    long_edge_height = (z_lower + z_upper) / 2  # Position long edges in between

    # SNAIL positions
    theta = np.linspace(0, 2 * np.pi, num_snails, endpoint=False)
    x = np.cos(theta)
    y = np.sin(theta)

    # Draw the SNAIL posts
    for i in range(num_snails):
        color = snail_colors[corral_modular.module_depths[i] % len(snail_colors)]
        ax.plot(
            [x[i], x[i]],
            [y[i], y[i]],
            [0, post_height],
            color=color,
            lw=6,
            solid_capstyle="round",
        )

    # Function to draw a slat between two points
    def draw_slat(point1, point2, z_height, color="black", linestyle="-"):
        ax.plot(
            [point1[0], point2[0]],
            [point1[1], point2[1]],
            [z_height, z_height],
            color=color,
            lw=4,
            linestyle=linestyle,
            solid_capstyle="round",
        )

    # Divide the edge list into upper and lower rungs
    # Take every second edge starting from the first
    upper_rung_edges = snail_edge_list[::2]
    # Take every second edge starting from the second
    lower_rung_edges = snail_edge_list[1::2]

    # Draw the upper rung slats
    for snail0, snail1 in upper_rung_edges:
        draw_slat((x[snail0], y[snail0]), (x[snail1], y[snail1]), z_upper[snail0])

    # Draw the lower rung slats
    for snail0, snail1 in lower_rung_edges:
        draw_slat((x[snail0], y[snail0]), (x[snail1], y[snail1]), z_lower[snail0])

    # Draw the long edges
    for edge in corral_modular.long_edges:
        # get the snail index for the qubits in the edge
        snail0 = corral_modular._find_module_by_qubit(edge[0])
        snail1 = corral_modular._find_module_by_qubit(edge[1])

        draw_slat(
            (x[snail0], y[snail0]),
            (x[snail1], y[snail1]),
            long_edge_height[edge[0]],
            color="red",
            linestyle="--",
        )

    # Set the aspect ratio and the viewing angle
    ax.set_aspect("auto")
    ax.view_init(elev=30.0, azim=-120)

    # Hide the axes
    ax.set_axis_off()

    return fig


# Example usage:
corral_modular = CorralModular(num_snails=8, corral_skip_pattern=(0, 1))
corral_modular.add_long_edges(num_long_edges=1)  # Example: Add 2 long-edges
corral_modular.module_depths = {i: i for i in range(8)}  # Example module depths
fig = pretty_print_snails_3d(corral_modular)

from PIL import Image


def combine_images(image1, image2):
    # Combine two images side by side
    dst = Image.new(
        "RGB", (image1.width + image2.width, max(image1.height, image2.height))
    )
    dst.paste(image1, (0, 0))
    dst.paste(image2, (image1.width, 0))
    return dst


# Generate the images with varying numbers of long edges for Corral configurations
paired_images = []

for num_snails in range(4, 9, 2):  # Example range for number of snails
    for skip_step in range(0, 3):  # Example range for skip pattern steps
        corral_skip_pattern = (0, skip_step)  # Define skip pattern based on skip_step
        try:
            pb = CorralModular(
                num_snails=num_snails, corral_skip_pattern=corral_skip_pattern
            )
            max_long_edges = min(4, num_snails // 2 - 1)

            for x in range(max_long_edges + 1):
                pb.add_long_edges(num_long_edges=x)
                label_text = f"Num Snails: {num_snails}, Skip Pattern: {corral_skip_pattern}, Num Qubits: {num_qubits}, Long Edges: {x}"
                print(label_text)

                # Generate the figure for the 3D SNAIL visualization
                fig_3d = pretty_print_snails_3d(pb)
                buf_3d = io.BytesIO()
                fig_3d.savefig(buf_3d, format="png")
                buf_3d.seek(0)
                img_3d = Image.open(buf_3d)
                plt.close(fig_3d)

                # Generate the figure for the network graph visualization
                fig_networkx = pretty_print_with_networkx(pb, label_text)
                buf_networkx = io.BytesIO()
                fig_networkx.savefig(buf_networkx, format="png")
                buf_networkx.seek(0)
                img_networkx = Image.open(buf_networkx)
                plt.close(fig_networkx)

                # Combine the two images side by side
                paired_img = combine_images(img_3d, img_networkx)
                paired_images.append(paired_img)

                pb.long_edges = []  # Clear existing long edges before next addition
        except AssertionError:
            pass

# Create the GIF from the paired images
gif_path = "corral_snails_and_networkx.gif"
paired_images[0].save(
    gif_path,
    save_all=True,
    append_images=paired_images[1:],
    optimize=False,
    duration=1500,
    loop=0,
)

print(f"GIF saved to {gif_path}")