"""Move and Rotation
=================

This example demonstrates moving and rotating the spatial boundary of a logical qubit.
"""

# ruff: noqa: E402

#!/usr/bin/env python
# coding: utf-8

# # Move Rotation
#
# This notebook demonstrates moving and rotating the spatial boundary of a logical
# qubit. Rotating the boundary types of a logical qubit is a crucial operation in
# lattice surgery. For instance, merging two logical qubits requires their boundary
# types to be aligned. Therefore, performing such boundary-type rotations is often
# essential to facilitate seamless lattice merging.
# ### Construction
#
# `tqec` provides the builtin function `tqec.gallery.move_rotation` to construct it.

# %%

from tqec import Basis
from tqec.gallery import move_rotation

graph = move_rotation()
# %%

graph.view_as_html()


# This operation rotates the orientation of the logical observable through a
# spatial L-shape junction. As shown below, the correlation surface initially
# aligns with the Y-axis and finally aligns with the X-axis.
# %%


correlation_surfaces = graph.find_correlation_surfaces()
# %%

graph.view_as_html(
    pop_faces_at_directions=("-Y",),
    show_correlation_surface=correlation_surfaces[0],
)


# ### Example Circuit
#
# Here we show an example circuit of move rotation with $d=3$ surface code that is initialized and measured in the $X$ basis. You can download the circuit :download:`here <../media/gallery/move_rotation/circuit.stim>` or view it in `Crumble <https://algassert.com/crumble#circuit=Q(0,4>`_0;Q(0,8)1;Q(0,12)2;Q(1,1)3;Q(1,3)4;Q(1,5)5;Q(1,7)6;Q(1,9)7;Q(1,11)8;Q(1,13)9;Q(2,0)10;Q(2,2)11;Q(2,4)12;Q(2,6)13;Q(2,8)14;Q(2,10)15;Q(2,12)16;Q(2,14)17;Q(3,1)18;Q(3,3)19;Q(3,5)20;Q(3,7)21;Q(3,9)22;Q(3,11)23;Q(3,13)24;Q(4,2)25;Q(4,4)26;Q(4,6)27;Q(4,8)28;Q(4,10)29;Q(4,12)30;Q(5,1)31;Q(5,3)32;Q(5,5)33;Q(5,7)34;Q(5,9)35;Q(5,11)36;Q(5,13)37;Q(6,2)38;Q(6,6)39;Q(6,10)40;Q(6,12)41;Q(6,14)42;Q(7,9)43;Q(7,11)44;Q(7,13)45;Q(8,8)46;Q(8,10)47;Q(8,12)48;Q(9,9)49;Q(9,11)50;Q(9,13)51;Q(10,10)52;Q(10,12)53;Q(10,14)54;Q(11,9)55;Q(11,11)56;Q(11,13)57;Q(12,8)58;Q(12,10)59;Q(12,12)60;Q(13,9)61;Q(13,11)62;Q(13,13)63;Q(14,12)64;RX_0_3_4_5_10_11_12_18_19_20_25_26_27_31_32_33_38;TICK;CX_12_4_25_18_27_20;CZ_11_3_26_19_38_31;TICK;CX_12_19_25_31_27_33;CZ_0_4;TICK;CX_10_3_12_5_25_19;CZ_11_4_26_20_38_32;TICK;CZ_11_18_26_32;TICK;CX_10_18_12_20_25_32;CZ_0_5_11_19_26_33;TICK;MX_0_10_11_12_25_26_27_38;DT(2,0,0)rec[-7];DT(2,4,0)rec[-5];DT(4,2,0)rec[-4];DT(4,6,0)rec[-2];TICK;RX_0_10_11_12_25_26_27_38;TICK;CX_12_4_25_18_27_20;CZ_11_3_26_19_38_31;TICK;CX_12_19_25_31_27_33;CZ_0_4;TICK;CX_10_3_12_5_25_19;CZ_11_4_26_20_38_32;TICK;CZ_11_18_26_32;TICK;CX_10_18_12_20_25_32;CZ_0_5_11_19_26_33;TICK;MX_0_10_11_12_25_26_27_38;DT(0,4,1)rec[-8]_rec[-16];DT(2,0,1)rec[-7]_rec[-15];DT(2,2,1)rec[-6]_rec[-14];DT(2,4,1)rec[-5]_rec[-13];DT(4,2,1)rec[-4]_rec[-12];DT(4,4,1)rec[-3]_rec[-11];DT(4,6,1)rec[-2]_rec[-10];DT(6,2,1)rec[-1]_rec[-9];TICK;RX_0_10_11_12_25_26_27_38;TICK;CX_12_4_25_18_27_20;CZ_11_3_26_19_38_31;TICK;CX_12_19_25_31_27_33;CZ_0_4;TICK;CX_10_3_12_5_25_19;CZ_11_4_26_20_38_32;TICK;CZ_11_18_26_32;TICK;CX_10_18_12_20_25_32;CZ_0_5_11_19_26_33;TICK;MX_0_10_11_12_25_26_27_38;DT(0,4,2)rec[-8]_rec[-16];DT(2,0,2)rec[-7]_rec[-15];DT(2,2,2)rec[-6]_rec[-14];DT(2,4,2)rec[-5]_rec[-13];DT(4,2,2)rec[-4]_rec[-12];DT(4,4,2)rec[-3]_rec[-11];DT(4,6,2)rec[-2]_rec[-10];DT(6,2,2)rec[-1]_rec[-9];TICK;RX_0_1_2_6_7_8_9_10_11_12_13_14_15_16_17_21_22_23_24_25_26_27_28_29_30_34_35_36_37_38_39_40_41_42_43_44_45_46_47_48_49_50_51_52_53_54_55_56_57_58_59_60_61_62_63_64;TICK;CX_12_4_14_6_16_8_25_18_27_20_29_22_41_36_47_43_53_50_59_55_64_62;CZ_11_3_13_5_15_7_17_9_26_19_28_21_30_23_38_31_39_33_40_35_42_37_48_44_52_49_54_51_60_56;TICK;CX_12_19_14_21_25_31_27_33_29_35;CZ_0_4_1_6_2_8_17_24_30_36_40_43_42_45_48_50_52_55_54_57_60_62;TICK;CX_10_3_12_5_14_7_16_9_25_19_27_21_29_23_41_37_47_44_53_51_59_56_64_63;CZ_11_4_13_6_15_8_26_20_28_22_30_24_38_32_39_34_40_36_46_43_48_45_52_50_58_55_60_57;TICK;CX_16_23_41_44_47_49_53_56_59_61;CZ_11_18_13_20_15_22_26_32_28_34;TICK;CX_10_18_12_20_14_22_16_24_25_32_27_34_29_36_41_45_47_50_53_57_59_62;CZ_0_5_1_7_2_9_11_19_13_21_15_23_26_33_28_35_30_37_40_44_46_49_48_51_52_56_58_61_60_63;TICK;MX_0_1_2_10_11_12_13_14_15_16_17_25_26_27_28_29_30_38_39_40_41_42_46_47_48_52_53_54_58_59_60_64;DT(0,4,3)rec[-32]_rec[-40];DT(2,0,3)rec[-29]_rec[-39];DT(2,2,3)rec[-28]_rec[-38];DT(2,4,3)rec[-27]_rec[-37];DT(2,8,3)rec[-25];DT(2,12,3)rec[-23];DT(4,2,3)rec[-21]_rec[-36];DT(4,4,3)rec[-20]_rec[-35];DT(4,6,3)rec[-19]_rec[-34];DT(4,10,3)rec[-17];DT(6,2,3)rec[-15]_rec[-33];DT(6,12,3)rec[-12];DT(8,10,3)rec[-9];DT(10,12,3)rec[-6];DT(12,10,3)rec[-3];DT(14,12,3)rec[-1];TICK;RX_0_1_2_10_11_12_13_14_15_16_17_25_26_27_28_29_30_38_39_40_41_42_46_47_48_52_53_54_58_59_60_64;TICK;CX_12_4_14_6_16_8_25_18_27_20_29_22_41_36_47_43_53_50_59_55_64_62;CZ_11_3_13_5_15_7_17_9_26_19_28_21_30_23_38_31_39_33_40_35_42_37_48_44_52_49_54_51_60_56;TICK;CX_12_19_14_21_25_31_27_33_29_35;CZ_0_4_1_6_2_8_17_24_30_36_40_43_42_45_48_50_52_55_54_57_60_62;TICK;CX_10_3_12_5_14_7_16_9_25_19_27_21_29_23_41_37_47_44_53_51_59_56_64_63;CZ_11_4_13_6_15_8_26_20_28_22_30_24_38_32_39_34_40_36_46_43_48_45_52_50_58_55_60_57;TICK;CX_16_23_41_44_47_49_53_56_59_61;CZ_11_18_13_20_15_22_26_32_28_34;TICK;CX_10_18_12_20_14_22_16_24_25_32_27_34_29_36_41_45_47_50_53_57_59_62;CZ_0_5_1_7_2_9_11_19_13_21_15_23_26_33_28_35_30_37_40_44_46_49_48_51_52_56_58_61_60_63;TICK;MX_0_1_2_10_11_12_13_14_15_16_17_25_26_27_28_29_30_38_39_40_41_42_46_47_48_52_53_54_58_59_60_64;DT(0,4,4)rec[-32]_rec[-64];DT(0,8,4)rec[-31]_rec[-63];DT(0,12,4)rec[-30]_rec[-62];DT(2,0,4)rec[-29]_rec[-61];DT(2,2,4)rec[-28]_rec[-60];DT(2,4,4)rec[-27]_rec[-59];DT(2,6,4)rec[-26]_rec[-58];DT(2,8,4)rec[-25]_rec[-57];DT(2,10,4)rec[-24]_rec[-56];DT(2,12,4)rec[-23]_rec[-55];DT(2,14,4)rec[-22]_rec[-54];DT(4,2,4)rec[-21]_rec[-53];DT(4,4,4)rec[-20]_rec[-52];DT(4,6,4)rec[-19]_rec[-51];DT(4,8,4)rec[-18]_rec[-50];DT(4,10,4)rec[-17]_rec[-49];DT(4,12,4)rec[-16]_rec[-48];DT(6,2,4)rec[-15]_rec[-47];DT(6,6,4)rec[-14]_rec[-46];DT(6,10,4)rec[-13]_rec[-45];DT(6,12,4)rec[-12]_rec[-44];DT(6,14,4)rec[-11]_rec[-43];DT(8,8,4)rec[-10]_rec[-42];DT(8,10,4)rec[-9]_rec[-41];DT(8,12,4)rec[-8]_rec[-40];DT(10,10,4)rec[-7]_rec[-39];DT(10,12,4)rec[-6]_rec[-38];DT(10,14,4)rec[-5]_rec[-37];DT(12,8,4)rec[-4]_rec[-36];DT(12,10,4)rec[-3]_rec[-35];DT(12,12,4)rec[-2]_rec[-34];DT(14,12,4)rec[-1]_rec[-33];TICK;RX_0_1_2_10_11_12_13_14_15_16_17_25_26_27_28_29_30_38_39_40_41_42_46_47_48_52_53_54_58_59_60_64;TICK;CX_12_4_14_6_16_8_25_18_27_20_29_22_41_36_47_43_53_50_59_55_64_62;CZ_11_3_13_5_15_7_17_9_26_19_28_21_30_23_38_31_39_33_40_35_42_37_48_44_52_49_54_51_60_56;TICK;CX_12_19_14_21_25_31_27_33_29_35;CZ_0_4_1_6_2_8_17_24_30_36_40_43_42_45_48_50_52_55_54_57_60_62;TICK;CX_10_3_12_5_14_7_16_9_25_19_27_21_29_23_41_37_47_44_53_51_59_56_64_63;CZ_11_4_13_6_15_8_26_20_28_22_30_24_38_32_39_34_40_36_46_43_48_45_52_50_58_55_60_57;TICK;CX_16_23_41_44_47_49_53_56_59_61;CZ_11_18_13_20_15_22_26_32_28_34;TICK;CX_10_18_12_20_14_22_16_24_25_32_27_34_29_36_41_45_47_50_53_57_59_62;CZ_0_5_1_7_2_9_11_19_13_21_15_23_26_33_28_35_30_37_40_44_46_49_48_51_52_56_58_61_60_63;TICK;MX_0_1_2_3_4_5_6_7_8_9_10_11_12_13_14_15_16_17_18_19_20_21_22_23_24_25_26_27_28_29_30_31_32_33_34_35_36_37_38_39_40_41_42_43_44_45_46_47_48_52_53_54_58_59_60_64;DT(0,4,5)rec[-56]_rec[-88];DT(0,8,5)rec[-55]_rec[-87];DT(0,12,5)rec[-54]_rec[-86];DT(2,0,5)rec[-46]_rec[-85];DT(2,2,5)rec[-45]_rec[-84];DT(2,4,5)rec[-44]_rec[-83];DT(2,6,5)rec[-43]_rec[-82];DT(2,8,5)rec[-42]_rec[-81];DT(2,10,5)rec[-41]_rec[-80];DT(2,12,5)rec[-40]_rec[-79];DT(2,14,5)rec[-39]_rec[-78];DT(3,1,5)rec[-38]_rec[-46]_rec[-53];DT(3,5,5)rec[-36]_rec[-37]_rec[-44]_rec[-51]_rec[-52];DT(3,9,5)rec[-34]_rec[-35]_rec[-42]_rec[-49]_rec[-50];DT(1,11,5)rec[-32]_rec[-33]_rec[-40]_rec[-47]_rec[-48];DT(4,2,5)rec[-31]_rec[-77];DT(4,4,5)rec[-30]_rec[-76];DT(4,6,5)rec[-29]_rec[-75];DT(4,8,5)rec[-28]_rec[-74];DT(4,10,5)rec[-27]_rec[-73];DT(4,12,5)rec[-26]_rec[-72];DT(5,3,5)rec[-24]_rec[-25]_rec[-31]_rec[-37]_rec[-38];DT(5,5,5)rec[-22]_rec[-23]_rec[-29]_rec[-35]_rec[-36];DT(3,11,5)rec[-20]_rec[-21]_rec[-27]_rec[-33]_rec[-34];DT(6,2,5)rec[-18]_rec[-71];DT(6,6,5)rec[-17]_rec[-70];DT(6,10,5)rec[-16]_rec[-69];DT(6,12,5)rec[-15]_rec[-68];DT(6,14,5)rec[-14]_rec[-67];DT(5,11,5)rec[-11]_rec[-12]_rec[-15]_rec[-19]_rec[-20];DT(8,8,5)rec[-10]_rec[-66];DT(8,10,5)rec[-9]_rec[-65];DT(8,12,5)rec[-8]_rec[-64];DT(10,10,5)rec[-7]_rec[-63];DT(10,12,5)rec[-6]_rec[-62];DT(10,14,5)rec[-5]_rec[-61];DT(12,8,5)rec[-4]_rec[-60];DT(12,10,5)rec[-3]_rec[-59];DT(12,12,5)rec[-2]_rec[-58];DT(14,12,5)rec[-1]_rec[-57];TICK;RX_47_52_53_54_58_59_60_64;TICK;CX_53_50_59_55_64_62;CZ_52_49_54_51_60_56;TICK;CX_47_49;CZ_52_55_54_57_60_62;TICK;CX_53_51_59_56_64_63;CZ_52_50_58_55_60_57;TICK;CX_53_56_59_61;TICK;CX_47_50_53_57_59_62;CZ_52_56_58_61_60_63;TICK;MX_47_52_53_54_58_59_60_64;DT(8,10,6)rec[-8]_rec[-17]_rec[-20]_rec[-21];DT(10,10,6)rec[-7]_rec[-15];DT(10,12,6)rec[-6]_rec[-14];DT(10,14,6)rec[-5]_rec[-13];DT(12,8,6)rec[-4]_rec[-12];DT(12,10,6)rec[-3]_rec[-11];DT(12,12,6)rec[-2]_rec[-10];DT(14,12,6)rec[-1]_rec[-9];TICK;RX_47_52_53_54_58_59_60_64;TICK;CX_53_50_59_55_64_62;CZ_52_49_54_51_60_56;TICK;CX_47_49;CZ_52_55_54_57_60_62;TICK;CX_53_51_59_56_64_63;CZ_52_50_58_55_60_57;TICK;CX_53_56_59_61;TICK;CX_47_50_53_57_59_62;CZ_52_56_58_61_60_63;TICK;MX_47_52_53_54_58_59_60_64;DT(8,10,7)rec[-8]_rec[-16];DT(10,10,7)rec[-7]_rec[-15];DT(10,12,7)rec[-6]_rec[-14];DT(10,14,7)rec[-5]_rec[-13];DT(12,8,7)rec[-4]_rec[-12];DT(12,10,7)rec[-3]_rec[-11];DT(12,12,7)rec[-2]_rec[-10];DT(14,12,7)rec[-1]_rec[-9];TICK;RX_47_52_53_54_58_59_60_64;TICK;CX_53_50_59_55_64_62;CZ_52_49_54_51_60_56;TICK;CX_47_49;CZ_52_55_54_57_60_62;TICK;CX_53_51_59_56_64_63;CZ_52_50_58_55_60_57;TICK;CX_53_56_59_61;TICK;CX_47_50_53_57_59_62;CZ_52_56_58_61_60_63;TICK;MX_47_49_50_51_52_53_54_55_56_57_58_59_60_61_62_63_64;DT(8,10,8)rec[-17]_rec[-25];DT(9,9,8)rec[-15]_rec[-16]_rec[-17];DT(10,10,8)rec[-13]_rec[-24];DT(10,12,8)rec[-12]_rec[-23];DT(10,14,8)rec[-11]_rec[-22];DT(11,13,8)rec[-8]_rec[-9]_rec[-12]_rec[-14]_rec[-15];DT(12,8,8)rec[-7]_rec[-21];DT(12,10,8)rec[-6]_rec[-20];DT(12,12,8)rec[-5]_rec[-19];DT(13,9,8)rec[-3]_rec[-4]_rec[-6]_rec[-9]_rec[-10];DT(13,11,8)rec[-1]_rec[-2]_rec[-3];DT(14,12,8)rec[-1]_rec[-18];OI(0)rec[-3]_rec[-9]_rec[-15]_rec[-45]_rec[-53]_rec[-66]_rec[-67]_rec[-68]_rec[-69]_rec[-70]_rec[-71]). # noqa: E501
# %%


from tqec import NoiseModel, compile_block_graph

graph = move_rotation(Basis.X)
compiled_graph = compile_block_graph(graph)
circuit = compiled_graph.generate_stim_circuit(
    k=1, noise_model=NoiseModel.uniform_depolarizing(p=0.001)
)


# ### Simulation
#
# Here we show the simulation results of both $X$-basis and $Z$-basis
# experiments under a **uniform depolarizing** noise model.
#
# <details><summary>Click to show the full code used for simulation</summary>
#
# ```py
# from multiprocessing import cpu_count
# from pathlib import Path
#
# import matplotlib.pyplot as plt
# import numpy
# import sinter
#
# from tqec.gallery.memory import memory
# from tqec.gallery.move_rotation import move_rotation
# from tqec import NoiseModel
# from tqec.simulation.plotting.inset import plot_observable_as_inset
# from tqec.simulation.simulation import start_simulation_using_sinter
# from tqec.utils.enums import Basis
#
# SAVE_DIR = Path("results")
#
#
# def generate_graphs(support_observable_basis: Basis) -> None:
#     block_graph = move_rotation(support_observable_basis)
#     zx_graph = block_graph.to_zx_graph()
#
#     correlation_surfaces = block_graph.find_correlation_surfaces()
#
#     stats = start_simulation_using_sinter(
#         block_graph,
#         range(1, 4),
#         list(numpy.logspace(-4, -1, 10)),
#         NoiseModel.uniform_depolarizing,
#         manhattan_radius=2,
#         observables=correlation_surfaces,
#         num_workers=cpu_count(),
#         max_shots=1_000_000,
#         max_errors=5_000,
#         decoders=["pymatching"],
#         print_progress=True,
#         save_resume_filepath=Path(
#             f"../_examples_database/move_rotation_stats_{support_observable_basis.value}.csv"
#         ),
#         database_path=Path("../_examples_database/database.pkl"),
#     )
#
#     for i, stat in enumerate(stats):
#         _, ax = plt.subplots()
#         sinter.plot_error_rate(
#             ax=ax,
#             stats=stat,
#             x_func=lambda stat: stat.json_metadata["p"],
#             failure_units_per_shot_func=lambda stat: stat.json_metadata["d"],
#             group_func=lambda stat: stat.json_metadata["d"],
#         )
#         plot_observable_as_inset(ax, zx_graph, correlation_surfaces[i])
#         ax.grid(axis="both")
#         ax.legend()
#         ax.loglog()
#         ax.set_title("Move Rotation Error Rate")
#         ax.set_xlabel("Physical Error Rate")
#         ax.set_ylabel("Logical Error Rate(per round)")
#         fig.savefig(
#             SAVE_DIR
#             / f"move_rotation_result_{support_observable_basis}_observable_{i}.png"
#         )
#
#
# def main():
#     SAVE_DIR.mkdir(exist_ok=True)
#     generate_graphs(Basis.Z)
#     generate_graphs(Basis.X)
#
#
# if __name__ == "__main__":
#     main()
#
#
# ```
#
# </details>
#

# %%


from multiprocessing import cpu_count
from pathlib import Path

import matplotlib.pyplot as plt
import numpy
import sinter

from tqec import NoiseModel
from tqec.gallery.move_rotation import move_rotation
from tqec.simulation.plotting.inset import plot_observable_as_inset
from tqec.simulation.simulation import start_simulation_using_sinter
from tqec.utils.enums import Basis


def generate_graphs(support_observable_basis: Basis) -> None:
    """Generate the logical error-rate graphs corresponding to the provided basis."""
    block_graph = move_rotation(support_observable_basis)
    zx_graph = block_graph.to_zx_graph()

    correlation_surfaces = block_graph.find_correlation_surfaces()

    stats = start_simulation_using_sinter(
        block_graph,
        range(1, 4),
        list(numpy.logspace(-4, -1, 10)),
        NoiseModel.uniform_depolarizing,
        manhattan_radius=2,
        observables=correlation_surfaces,
        num_workers=cpu_count(),
        max_shots=1_000_000,
        max_errors=5_000,
        decoders=["pymatching"],
        # note that save_resume_filepath and database_path can help reduce the time taken
        # by the simulation after the database and result statistics have been saved to
        # the chosen path
        save_resume_filepath=Path(
            f"../_examples_database/move_rotation_stats_{support_observable_basis.value}.csv"
        ),
        database_path=Path("../_examples_database/database.pkl"),
    )

    for i, stat in enumerate(stats):
        _, ax = plt.subplots()
        sinter.plot_error_rate(
            ax=ax,
            stats=stat,
            x_func=lambda stat: stat.json_metadata["p"],
            failure_units_per_shot_func=lambda stat: stat.json_metadata["d"],
            group_func=lambda stat: stat.json_metadata["d"],
        )
        plot_observable_as_inset(ax, zx_graph, correlation_surfaces[i])
        ax.grid(axis="both")
        ax.legend()
        ax.loglog()
        ax.set_title("Move Rotation Error Rate")
        ax.set_xlabel("Physical Error Rate")
        ax.set_ylabel("Logical Error Rate(per round)")


# #### Z Basis

# %%


generate_graphs(Basis.Z)


# #### X Basis

# %%


generate_graphs(Basis.X)