"""
模块: FermionicLatticeSurgery
定义了费米子晶格Surgery，生成辅助测量joint observable的subsystem code
------------------------------------------------------------------------------------------------
属性：
    SubsystemCode的属性：
        self.code：中心子码
        self.bare_logical_operators：选定的logical operator
-------------------------------------------------------------------------
方法：
    SubsystemCode的方法：
        __init__（魔法方法）：构造一个subsystem code
       logical_operators（属性方法）：返回已经给定的bare logical operators
       distance（属性方法）：返回subsystem code distance
"""
import copy
import networkx as nx
import numpy as np
from clifforder.Code.SubsystemCode import SubsystemCode
from clifforder.Code.QuantumCode import MajoranaCSSCode, MajoranaCode
from clifforder.Code.ColorCode import MajoranaColorCode
from clifforder.Code.Operator import MajoranaOperator


##  CHAPTER：====Fermionic lattice surgery====
class FermionicLatticeSurgery:

    ##  SECTION: ----构造函数----
    def __init__(self, code_A: MajoranaCSSCode, index_A: int, code_B: MajoranaCSSCode | None = None, index_B: int | None = None) -> None:
        """
        input.code_A: 某一个Majorana CSS Code
        input.index_A: 选取其中一个logical Majorana operator作为测量目标
        input.code_B: Majorana CSS Code，不选代表是color code
        input.index_B: 选取其中一个logical Majorana operator作为测量目标，不选默认color code的第一个
        output: None
        """

        ##  PART: ----初始化----
        self.related_stabilizers_A_x = []  # 与joint observable有overlap的A中的X型stabilizers
        self.related_stabilizers_B_x = []  # 与joint observable有overlap的B中的X型stabilizers
        self.related_stabilizers_A_z = []  # 与joint observable有overlap的A中的Z型stabilizers
        self.related_stabilizers_B_z = []  # 与joint observable有overlap的B中的Z型stabilizers
        self.irrelated_stabilizers_A_z = []  # 与joint observable没有overlap的A中的Z型stabilizers
        self.irrelated_stabilizers_B_z = []  # 与joint observable没有overlap的B中的Z型stabilizers
        self.irrelated_stabilizers_A_x = []  # 与joint observable没有overlap的A中的X型stabilizers
        self.irrelated_stabilizers_B_x = []  # 与joint observable没有overlap的B中的X型stabilizers
        self.ancilla_stabilizers = []  # 一开始就作用在ancilla上的stabilizers
        self.modify_stabilizers_A = []  # 修改A中与joint observable有overlap的stabilizers
        self.modify_stabilizers_B = []  # 修改B中与joint observable有overlap的stabilizers
        self.gauge_stabilizers = []  # 作用在ancilla上的stabilizers
        self.measurement_stabilizers = []  # 联合两边的stabilizers
        self.ancilla_number = 0  # Ancilla sites的数目
        self.observable: None | MajoranaOperator = None  # 待测量的可观测量

        ##  PART: ----区分general fermionic lattice surgery和Zechuan lattice surgery----
        if index_B is None or code_B is None:
            self.code_A: MajoranaCSSCode = code_A
            self.index_A: int = index_A
            self.code_B: MajoranaColorCode = MajoranaColorCode(code_A.logical_operators_x[self.index_A].weight)
            self.index_B: int = 0
            self._ZechuanLatticeSurgery()
        else:

            ##  总是让weight较大的放在A
            if code_A.logical_operators_x[index_A].weight > code_B.logical_operators_x[index_B].weight:
                self.code_A: MajoranaCSSCode = code_A.copy()
                self.code_B: MajoranaCSSCode = code_B.copy()
                self.index_A: int = index_A
                self.index_B: int = index_B
            else:
                self.code_A: MajoranaCSSCode = code_B.copy()
                self.code_B: MajoranaCSSCode = code_A.copy()
                self.index_A: int = index_B
                self.index_B: int = index_A
            self._GeneralFermionicLatticeSurgery()

        ##  PART: ----根据结果赋值----
        self.physical_number = code_A.physical_number + self.code_B.physical_number + self.ancilla_number  # 总共用到的sites数目

        ##  规定其余剩下的logical operators是bare logical operators
        self.bare_logical_operators_x = []
        self.bare_logical_operators_z = []
        for i in range(len(self.code_A.logical_operators_x)):
            if i != self.index_A:
                self.bare_logical_operators_z.append(self.code_A.logical_operators_z[i])
                self.bare_logical_operators_x.append(self.code_A.logical_operators_x[i])
        for i in range(len(self.code_B.logical_operators_x)):
            if i != self.index_B:
                self.bare_logical_operators_z.append(MajoranaOperator.HermitianOperatorFromOccupy([], self.code_B.logical_operators_z[i].occupy_z + self.code_A.physical_number))
                self.bare_logical_operators_x.append(MajoranaOperator.HermitianOperatorFromOccupy(self.code_B.logical_operators_x[i].occupy_x + self.code_A.physical_number, []))

        ##  规定index_A和index_B对应的Z logical operators 是bare logical operator
        self.bare_logical_operators_z.append(MajoranaOperator.HermitianOperatorFromOccupy(
            [], np.hstack([self.code_A.logical_operators_z[self.index_A].occupy_z, self.code_B.logical_operators_z[self.index_B].occupy_z + self.code_A.physical_number])))

        ##  生成前中后三种code
        self.code = SubsystemCode(MajoranaCode(self.irrelated_stabilizers_A_z + self.irrelated_stabilizers_B_z + self.irrelated_stabilizers_A_x + self.irrelated_stabilizers_B_x +
                                 self.related_stabilizers_A_z + self.related_stabilizers_B_z + self.modify_stabilizers_A + self.modify_stabilizers_B + self.gauge_stabilizers,
                                               self.physical_number), self.bare_logical_operators_x + self.bare_logical_operators_z)
        self.code_before_surgery = MajoranaCSSCode([MajoranaOperator.HermitianOperatorFromOccupy(temp.occupy_x, []) for temp in self.code_A.generators_x] +
                                                   [MajoranaOperator.HermitianOperatorFromOccupy(temp.occupy_x + code_A.physical_number, []) for temp in self.code_B.generators_x],
                                                   [MajoranaOperator.HermitianOperatorFromOccupy([], temp.occupy_z) for temp in self.code_A.generators_z] +
                                                   [MajoranaOperator.HermitianOperatorFromOccupy([], temp.occupy_z + code_A.physical_number) for temp in self.code_B.generators_z],
                                                   self.code_A.physical_number + self.code_B.physical_number)
        self.code_after_surgery = MajoranaCode(self.irrelated_stabilizers_A_x + self.irrelated_stabilizers_B_x + self.irrelated_stabilizers_A_z + self.irrelated_stabilizers_B_x +
                                               self.related_stabilizers_A_z + self.related_stabilizers_B_z + self.modify_stabilizers_A + self.modify_stabilizers_B +
                                               self.measurement_stabilizers + self.gauge_stabilizers, self.physical_number)


    ##  CHAPTER: ====通用的fermionic lattice surgery====
    def _GeneralFermionicLatticeSurgery(self):

        ##  PART: ----预处理----
        ##  标度对齐
        code_A: MajoranaCSSCode = self.code_A.copy()
        code_B: MajoranaCSSCode = self.code_B.copy()
        majorana_number_A = code_A.physical_number
        majorana_number_B = code_B.physical_number
        logical_operator_A = code_A.logical_operators_x[self.index_A]
        logical_operator_B = code_B.logical_operators_x[self.index_B]
        logical_operator_B.occupy_x = logical_operator_B.occupy_x + majorana_number_A
        code_A.index_map(np.arange(majorana_number_A), majorana_number_A + majorana_number_B)
        code_B.index_map(np.arange(majorana_number_A, majorana_number_A + majorana_number_B), majorana_number_A + majorana_number_B)
        self.observable = 1j * logical_operator_A @ logical_operator_B

        ##  提取stabilizers
        self.related_stabilizers_A_x = copy.deepcopy([temp for temp in code_A.generators if len(set(temp.occupy_x) & set(logical_operator_A.occupy_x)) > 0])
        self.related_stabilizers_B_x = copy.deepcopy([temp for temp in code_B.generators if len(set(temp.occupy_x) & set(logical_operator_B.occupy_x)) > 0])
        self.related_stabilizers_A_z = copy.deepcopy([temp for temp in code_A.generators_z if len(set(temp.occupy_z) & set(logical_operator_A.occupy_x)) > 0])
        self.related_stabilizers_B_z = copy.deepcopy([temp for temp in code_B.generators_z if len(set(temp.occupy_z) & set(logical_operator_B.occupy_x)) > 0])
        self.irrelated_stabilizers_A_z = copy.deepcopy([temp for temp in code_A.generators_z if len(set(temp.occupy_z) & set(logical_operator_A.occupy_x)) == 0])
        self.irrelated_stabilizers_B_z = copy.deepcopy([temp for temp in code_B.generators_z if len(set(temp.occupy_z) & set(logical_operator_B.occupy_x)) == 0])
        self.irrelated_stabilizers_A_x = copy.deepcopy([temp for temp in code_A.generators_x if len(set(temp.occupy_x) & set(logical_operator_A.occupy_x)) == 0])
        self.irrelated_stabilizers_B_x = copy.deepcopy([temp for temp in code_B.generators_x if len(set(temp.occupy_x) & set(logical_operator_B.occupy_x)) == 0])

        ##  PART: ----执行lattice surgery----
        ancilla_mode_number = 0  # 初始化ancilla majorana modes数目

        ##  每一个related stabilizers添加ancilla Majorana modes数目等于overlap大小
        for i in range(len(self.related_stabilizers_A_x)):
            result_temp = [ancilla_mode_number + temp for temp in range(len(set(self.related_stabilizers_A_x[i].occupy_x) & set(logical_operator_A.occupy_x)))]
            self.modify_stabilizers_A.append([self.related_stabilizers_A_x[i].occupy_x.tolist(), result_temp])
            ancilla_mode_number += len(result_temp)
        for i in range(len(self.related_stabilizers_B_x)):
            result_temp = [ancilla_mode_number + temp for temp in range(len(set(self.related_stabilizers_B_x[i].occupy_x) & set(logical_operator_B.occupy_x)))]
            self.modify_stabilizers_B.append([self.related_stabilizers_B_x[i].occupy_x.tolist(), result_temp])
            ancilla_mode_number += len(result_temp)

        ##  Measurement stabilizers首先要把两边的logical operators成对添加，然后再处理A侧落单的
        for i in range(len(logical_operator_B.occupy_x)):
            self.measurement_stabilizers.append([[logical_operator_A.occupy_x[i], logical_operator_B.occupy_x[i]], []])
        for i in range((len(logical_operator_A.occupy_x) - len(logical_operator_B.occupy_x)) // 2):
            self.measurement_stabilizers.append([[logical_operator_A.occupy_x[2 * i + len(logical_operator_B.occupy_x)], logical_operator_A.occupy_x[2 * i + len(logical_operator_B.occupy_x) + 1]], []])

        ##  Modified stabilizer和measurement stabilizer在joint observable上有多少overlap，就把modified stabilizer上的多少ancilla modes拨给measurement stabilizer，每个ancilla modes只拨一次
        for i in range(len(self.modify_stabilizers_A)):
            flag_temp = 0
            for j in range(len(self.measurement_stabilizers)):
                number = len(set(self.modify_stabilizers_A[i][0]) & set(self.measurement_stabilizers[j][0]))
                self.measurement_stabilizers[j][1] = self.measurement_stabilizers[j][1] + [self.modify_stabilizers_A[i][1][flag_temp + k] for k in range(number)]
                flag_temp += number
        for i in range(len(self.modify_stabilizers_B)):
            flag_temp = 0
            for j in range(len(self.measurement_stabilizers)):
                number = len(set(self.modify_stabilizers_B[i][0]) & set(self.measurement_stabilizers[j][0]))
                self.measurement_stabilizers[j][1] = self.measurement_stabilizers[j][1] + [self.modify_stabilizers_B[i][1][flag_temp + k] for k in range(number)]
                flag_temp += number

        ##  如果measurement stabilizer是odd-weight，那么还需要添加配平
        single_mode_list = []
        for i in range(len(self.measurement_stabilizers)):
            if len(self.measurement_stabilizers[i][1]) % 2 != 0:
                self.measurement_stabilizers[i][1].append(ancilla_mode_number)
                single_mode_list.append(ancilla_mode_number)
                ancilla_mode_number += 1

        ##   添加gauge stabilizers需要图论计算
        mode_vertices = ['Mode' + str(temp) for temp in range(ancilla_mode_number)]  # ancilla modes
        A_vertices = []  # A中的modified stabilizer nodes
        B_vertices = []  # B中的modified stabilizer nodes
        measurement_vertices = []  # measurement stabilizer nodes
        edges = []  # 图中的边

        ## 如果一个modified stabilizer node只关联两个ancilla modes，那么就直接添加到A vertex中，反之需要拆分A vertex到多个子nodes
        for i in range(len(self.modify_stabilizers_A)):
            if len(self.modify_stabilizers_A[i][1]) // 2 == 1:
                A_vertices.append('A' + str(i))
                edges.append(('A' + str(i), 'Mode' + str(self.modify_stabilizers_A[i][1][0])))
                edges.append(('A' + str(i), 'Mode' + str(self.modify_stabilizers_A[i][1][1])))
            else:
                for j in range(len(self.modify_stabilizers_A[i][1]) // 2):
                    A_vertices.append('A' + str(i) + '_' + str(j))
                    edges.append(('A' + str(i), 'Mode' + str(self.modify_stabilizers_A[i][1][j * 2])))
                    edges.append(('A' + str(i), 'Mode' + str(self.modify_stabilizers_A[i][1][j * 2 + 1])))
        for i in range(len(self.modify_stabilizers_B)):
            if len(self.modify_stabilizers_B[i][1]) // 2 == 1:
                B_vertices.append('B' + str(i))
                edges.append(('B' + str(i), 'Mode' + str(self.modify_stabilizers_B[i][1][0])))
                edges.append(('B' + str(i), 'Mode' + str(self.modify_stabilizers_B[i][1][1])))
            else:
                for j in range(len(self.modify_stabilizers_B[i][1]) // 2):
                    B_vertices.append('B' + str(i) + '_' + str(j))
                    edges.append(('B' + str(i), 'Mode' + str(self.modify_stabilizers_B[i][1][j * 2])))
                    edges.append(('B' + str(i), 'Mode' + str(self.modify_stabilizers_B[i][1][j * 2 + 1])))

        ##  Measurement stabilizer作为节点
        for i in range(len(self.measurement_stabilizers)):
            measurement_vertices.append('Measure' + str(i))
            for j in range(len(self.measurement_stabilizers[i][1])):
                edges.append(('Measure' + str(i), 'Mode' + str(self.measurement_stabilizers[i][1][j])))

        ##  单个用于配平的ancilla modes需要两两添加虚拟边
        virtual_vertices = []
        for i in range(len(single_mode_list) // 2):
            virtual_vertices.append('Virtual' + str(i))
            edges.append(('Virtual' + str(i), 'Mode' + str(single_mode_list[i])))
            edges.append(('Virtual' + str(i), 'Mode' + str(single_mode_list[i])))

        ##  构建图论模型
        graph = nx.Graph()
        graph.add_nodes_from(mode_vertices)
        graph.add_nodes_from(A_vertices)
        graph.add_nodes_from(B_vertices)
        graph.add_nodes_from(measurement_vertices)
        graph.add_nodes_from(virtual_vertices)
        graph.add_edges_from(edges)

        ##  用圈基作为gauge stabilizers
        cyclic_basis = nx.minimum_cycle_basis(graph)
        for i in range(len(cyclic_basis)):
            modes_temp = []
            for j in range(len(cyclic_basis[i])):
                if 'Mode' in cyclic_basis[i][j]:
                    modes_temp.append(int(cyclic_basis[i][j][4:]))
            self.gauge_stabilizers.append([[], modes_temp])

        ##  PART: ----将stabilizers转换为MajoranaOperator格式----
        alls = [self.modify_stabilizers_A, self.modify_stabilizers_B, self.measurement_stabilizers, self.gauge_stabilizers]
        for i in range(len(alls)):
            for j in range(len(alls[i])):
                occupy_x_temp = alls[i][j][0] + [temp // 2 + majorana_number_A + majorana_number_B for temp in alls[i][j][1] if temp % 2 == 0]
                occupy_z_temp = [temp // 2 + majorana_number_A + majorana_number_B for temp in alls[i][j][1] if temp % 2 == 1]
                alls[i][j] = MajoranaOperator.HermitianOperatorFromOccupy(occupy_x_temp, occupy_z_temp)
        self.ancilla_number = ancilla_mode_number // 2  # ancilla sites数目


    ##  CHAPTER: ====刘泽川提出的fermionic lattice surgery方案====
    def _ZechuanLatticeSurgery(self):

        ##  PART: ----预处理----
        ##  标度对齐
        code_A: MajoranaCSSCode = self.code_A.copy()
        code_B: MajoranaColorCode = self.code_B.copy()
        majorana_number_A = code_A.physical_number
        majorana_number_B = code_B.physical_number
        logical_operator_A = code_A.logical_operators_x[self.index_A]
        logical_operator_B = code_B.logical_operators_x[0]
        logical_operator_B.occupy_x = logical_operator_B.occupy_x + majorana_number_A
        code_A.index_map(np.arange(majorana_number_A), majorana_number_A + majorana_number_B)
        code_B.index_map(np.arange(majorana_number_A, majorana_number_A + majorana_number_B), majorana_number_A + majorana_number_B)
        self.observable = 1j * logical_operator_A @ logical_operator_B

        ##  提取stabilizers
        self.related_stabilizers_A_x = copy.deepcopy([temp for temp in code_A.generators if len(set(temp.occupy_x) & set(logical_operator_A.occupy_x)) > 0])
        self.related_stabilizers_B_x = copy.deepcopy([temp for temp in code_B.generators if len(set(temp.occupy_x) & set(logical_operator_B.occupy_x)) > 0])
        self.irrelated_stabilizers_A_z = copy.deepcopy([temp for temp in code_A.generators_z if len(set(temp.occupy_z) & set(logical_operator_A.occupy_x)) == 0])
        self.irrelated_stabilizers_B_z = copy.deepcopy([temp for temp in code_B.generators_z if len(set(temp.occupy_z) & set(logical_operator_B.occupy_x)) == 0])
        self.irrelated_stabilizers_A_x = copy.deepcopy([temp for temp in code_A.generators_x if len(set(temp.occupy_x) & set(logical_operator_A.occupy_x)) == 0])
        self.irrelated_stabilizers_B_x = copy.deepcopy([temp for temp in code_B.generators_x if len(set(temp.occupy_x) & set(logical_operator_B.occupy_x)) == 0])

        ##  PART: ----执行lattice surgery----
        ##  初始化
        ancilla_mode_number = 0  # ancilla Majorana modes数目初始化为0
        edge_site_bin_A: list[list[int]] = [[] for _ in range(code_B.d)]  # 需要添加到A边上的modes
        edge_site_bin_B: list[list[int]] = [[] for _ in range(code_B.d)]  # 需要添加到B边上的modes
        used_flags = [False for _ in range(code_B.d - 1)]  # 检测哪个B边上的stabilizer还没有被用过

        ##  对A上每一个stabilizer处理：有多少overlap就添加多少ancilla modes，然后将ancilla modes两两分配给B上的stabilizers
        for i in range(len(self.related_stabilizers_A_x)):

            ##  计算A上stabilizer与logical operator的overlap，得到需要两两配对多少
            overlaps = list(set(self.related_stabilizers_A_x[i].occupy_x) & set(logical_operator_A.occupy_x))
            modification_A=[]

            ##  分配给A和B上的sites并生成B的modification
            for j in range(len(overlaps) // 2):
                index_0 = int(np.where(logical_operator_A.occupy_x == overlaps[2 * j])[0][0])
                index_1 = int(np.where(logical_operator_A.occupy_x == overlaps[2 * j + 1])[0][0])
                edge_site_bin_A[index_0].append(ancilla_mode_number)
                edge_site_bin_A[index_1].append(ancilla_mode_number + 1)
                modification_A.append(ancilla_mode_number)
                modification_A.append(ancilla_mode_number + 1)
                edge_site_bin_B[index_0].append(ancilla_mode_number + 2)
                edge_site_bin_B[index_1].append(ancilla_mode_number + 3)
                B_temp = get_between(code_B, index_0, index_1)
                self.related_stabilizers_B_z.append(MajoranaOperator.HermitianOperatorFromOccupy([], [temp for temp in B_temp.occupy_x]))
                self.modify_stabilizers_B.append([[temp for temp in B_temp.occupy_x], [ancilla_mode_number + 2, ancilla_mode_number + 3]])

                ##  标记一下哪里用过了
                small = min(index_0, index_1)
                big = max(index_0, index_1)
                used_flags[small] = True
                used_flags[big - 1] = True

                ##  直接生成gauge stabilizers
                self.gauge_stabilizers.append([[], [ancilla_mode_number + temp for temp in range(4)]])
                ancilla_mode_number += 4

            ##  生成A的modification
            self.modify_stabilizers_A.append([self.related_stabilizers_A_x[i].occupy_x.tolist(), modification_A])
            self.related_stabilizers_A_z.append(MajoranaOperator.HermitianOperatorFromOccupy([], self.related_stabilizers_A_x[i].occupy_x.tolist()))

        ##  处理B上没有用过的stabilizers
        for i in range(len(used_flags)):
            if not used_flags[i]:
                edge_site_bin_B[i].append(ancilla_mode_number)
                edge_site_bin_B[i].append(ancilla_mode_number + 1)
                edge_site_bin_B[i + 1].append(ancilla_mode_number + 2)
                edge_site_bin_B[i + 1].append(ancilla_mode_number + 3)
                self.modify_stabilizers_B.append([get_between(code_B, i, i + 1).occupy_x.tolist(), [ancilla_mode_number, ancilla_mode_number + 2]])
                self.gauge_stabilizers.append([[], [ancilla_mode_number + temp for temp in range(4)]])
                ancilla_mode_number += 4

        ##  生成measurement stabilizers
        for i in range(len(edge_site_bin_A)):
            self.measurement_stabilizers.append(([logical_operator_A.occupy_x[i], logical_operator_B.occupy_x[i]], edge_site_bin_A[i] + edge_site_bin_B[i]))

        ##  PART: ----将stabilizers转换为MajoranaOperator格式----
        alls = [self.modify_stabilizers_A, self.modify_stabilizers_B, self.measurement_stabilizers, self.gauge_stabilizers]
        for i in range(len(alls)):
            for j in range(len(alls[i])):
                occupy_x_temp = alls[i][j][0] + [temp // 2 + majorana_number_A + majorana_number_B for temp in alls[i][j][1] if temp % 2 == 0]
                occupy_z_temp = [temp // 2 + majorana_number_A + majorana_number_B for temp in alls[i][j][1] if temp % 2 == 1]
                alls[i][j] = MajoranaOperator.HermitianOperatorFromOccupy(occupy_x_temp, occupy_z_temp)
        self.ancilla_number = ancilla_mode_number // 2  # ancilla sites数目


##  CHAPTER: ====计算B上stabilizer之间的乘积====
def get_between(color_code: MajoranaColorCode, index_0: int, index_1: int):
    """
    input.color_code: MajoranaColorCode
    input.index_0: site索引
    input.index_1: site索引
    output: 两个site之间的stabilizers，遵循Zechuan lattice surgery的规则
    """

    ##  对边上的stabilizers重排序
    edge_stabilizers = []
    for i in range(color_code.d // 2):
        edge_stabilizers.append(color_code.generators_x[0:color_code.d][i])
        edge_stabilizers.append(color_code.generators_x[0:color_code.d][i + color_code.d // 2])

    ##  重新编织大小
    if index_0 > index_1:
        aa = index_1
        bb = index_0
    else:
        aa = index_0
        bb = index_1

    ##  计算stabilizer的乘积或单个的stabilizer
    if aa == bb - 1:
        return edge_stabilizers[aa].copy()
    else:
        return edge_stabilizers[aa] @ edge_stabilizers[bb - 1]

