"""
模块: Platform
生成量子计算机，并可以执行Clifford gate和Pauli or Majorana measurement
------------------------------------------------------------------------------------------------
Platform的属性：
    self.pauli_number：Qubits数目
    self.majorana_number：Fermionic sites数目
    self.stabilizers：平台当前的量子态
    self.coffs = 平台当前量子态的系数
------------------------------------------------------------------------------------------------

"""
import galois
import numpy as np
from galois import GF2
from clifforder.Circuit.tools import solve
from clifforder.Code.Operator import MajoranaOperator, PauliOperator


##  CHAPTER: ====量子计算平台====
class Platform:

    GF = galois.GF(2)

    ##  SECTION: ----构造函数----
    def __init__(self):
        """
        influence：初始化一个空的量子计算平台
        """

        self.pauli_number = 0  # Qubits数目
        self.majorana_number = 0  # Fermionic sites数目
        self.stabilizers = None  # Stabilizer state
        self.coffs = None  # Stabilizer state的系数


    ##  SECTION: ----对平台初始化----
    def trap(self, majorana_number: int, pauli_number: int):
        """
        input.majorana_number: Number of fermionic sites.
        input.pauli_number: Number of qubits.
        influence: Initialize stabilizers for the given number of sites/qubits.
        """

        ##  输入检查
        try:
            majorana_number = int(majorana_number)
            pauli_number = int(pauli_number)
            assert majorana_number >= 0 and pauli_number >= 0, "输入的fermionic site数目或qubit数目不能为负数"
        except Exception as e:
            raise ValueError("当前输入的不合法") from e

        ##  初始化赋值
        self.pauli_number = pauli_number
        self.majorana_number = majorana_number
        self.stabilizers = self.GF(np.zeros((majorana_number + pauli_number, 2 * majorana_number + 2 * pauli_number), dtype=int))
        self.coffs = np.ones(majorana_number + pauli_number, dtype=complex)

        ##  将平台放在全0态上
        for i in range(majorana_number):
            self.stabilizers[i, 2 * i] = 1
            self.stabilizers[i, 2 * i + 1] = 1
            self.coffs[i] = 1j
        for i in range(pauli_number):
            self.stabilizers[i + majorana_number, 2 * i + 1 + 2 * majorana_number] = 1
            self.coffs[i + majorana_number] = 1

    ##  SECTION: ----测量某个string operator----
    def measure(self, op: MajoranaOperator | PauliOperator) -> int:
        """
        input.op: 待测量的厄米算符
        output: int, +1 or -1 测量结果
        influence: 更新stabilizer state
        """

        ##  输入检查
        if isinstance(op, MajoranaOperator):
            assert op.is_hermitian, "Majorana operator必须是厄米算符"
            vector_op = np.append(op.get_vector(self.majorana_number), self.GF.Zeros(self.pauli_number * 2))
        elif isinstance(op, PauliOperator):
            assert op.is_hermitian, "Pauli operator必须是厄米算符"
            vector_op = np.append(self.GF.Zeros(self.majorana_number * 2), op.get_vector(self.pauli_number))
        else:
            raise TypeError("输入的厄米算符必须是MajoranaOperator或PauliOperator")

        ##  执行测量
        first_index = None
        for i in range(len(self.stabilizers)):
            pauli_commute = (np.dot(self.stabilizers[i][self.majorana_number * 2::2], vector_op[self.majorana_number * 2 + 1::2]) +
                             np.dot(self.stabilizers[i][self.majorana_number * 2 + 1::2], vector_op[self.majorana_number * 2::2]))
            majorana_commute = (np.dot(self.stabilizers[i][0:self.majorana_number * 2], vector_op[0:self.majorana_number * 2]) +
                                np.sum(vector_op[0:self.majorana_number * 2]) * np.sum(self.stabilizers[i][0:self.majorana_number * 2]))
            if majorana_commute + pauli_commute == 1 and first_index is None:
                first_index = i
            elif majorana_commute + pauli_commute == 1 and first_index is not None:
                factor = np.sum([np.sum(self.stabilizers[i][temp + 1:2 * self.majorana_number]) for temp in range(2 * self.majorana_number) if
                                 self.stabilizers[first_index][temp] == 1])
                self.stabilizers[i] += self.stabilizers[first_index]
                if np.mod(factor, 2) == 0:
                    factor = 1
                else:
                    factor = -1
                self.coffs[i] = self.coffs[i] * self.coffs[first_index] * factor
            else:
                pass

        ##  当state在测量算符的本征空间内
        if first_index is not None:
            if np.random.rand() < 0.5:
                self.stabilizers[first_index] = vector_op
                self.coffs[first_index] = op.coff
                return 1
            else:
                self.stabilizers[first_index] = vector_op
                self.coffs[first_index] = -op.coff
                return -1

        ##  当state不在测量算符的本征空间内
        else:
            solution = solve(self.stabilizers, vector_op)
            coff = complex(np.prod([self.coffs[i] for i in range(len(self.stabilizers)) if solution[i] == 1]))
            vector_now, factor = majorana_factor(self.majorana_number,[self.stabilizers[i] for i in range(len(self.stabilizers)) if solution[i] == 1])
            coff = coff * factor
            if coff == op.coff:
                return 1
            else:
                return -1


    ##  SECTION: ----执行Pauli X gate----
    def x(self, qubit_index: int):
        """
        input.qubit_index: 目标qubit的索引
        influence: 对目标qubit执行Pauli X门，即对该qubit的Z算符取反
        """

        ##  Compute final state
        indices = np.where(self.stabilizers[:, self.majorana_number * 2 + qubit_index * 2 + 1] == 1)[0]
        self.coffs[indices] = -self.coffs[indices]


    ##  SECTION: ----执行Pauli Y gate----
    def y(self, qubit_index: int):
        """
        input.qubit_index: 目标qubit的索引
        influence: 对目标qubit执行Pauli Y门，即对该qubit的X和Z算符取反
        """

        ##  Compute final state
        indices_x = np.where(self.stabilizers[:, self.majorana_number * 2 + qubit_index * 2 + 1] == 1)[0]
        self.coffs[indices_x] = -self.coffs[indices_x]
        indices_z = np.where(self.stabilizers[:, self.majorana_number * 2 + qubit_index * 2] == 1)[0]
        self.coffs[indices_z] = -self.coffs[indices_z]


    ##  SECTION: ----执行Pauli Z gate----
    def z(self, qubit_index: int):
        """
        input.qubit_index: 目标qubit的索引
        influence: 对目标qubit执行Pauli Z门，即对该qubit的X算符取反
        """

        ##  Compute final state
        indices = np.where(self.stabilizers[:, self.majorana_number * 2 + qubit_index * 2] == 1)[0]
        self.coffs[indices] = -self.coffs[indices]


    ##  SECTION: ----执行Hadamard gate----
    def h(self, qubit_index: int):
        """
        input.qubit_index: 目标qubit的索引
        influence: 对目标qubit执行Hadamard gate，即交换该qubit的X和Z算符支持并更新相位
        """

        ##  Compute final state
        indices = np.where(np.logical_and(self.stabilizers[:, self.majorana_number * 2 + qubit_index * 2 + 1] == 1,self.stabilizers[:, self.majorana_number * 2 + qubit_index * 2] == 1))[0]
        self.coffs[indices] = -self.coffs[indices]
        caches = self.stabilizers[:, self.majorana_number * 2 + qubit_index * 2].copy()
        self.stabilizers[:, self.majorana_number * 2 + qubit_index * 2] = self.stabilizers[:, self.majorana_number * 2 + qubit_index * 2 + 1]
        self.stabilizers[:, self.majorana_number * 2 + qubit_index * 2 + 1] = caches


    ##  SECTION: ----执行Pauli phase gate----
    def s(self, pauli_index: int):
        """
        input.pauli_index: 目标qubit的索引
        influence: 对目标qubit执行Pauli S门，即对该qubit的Z算符取反并乘以i
        """

        ##  Compute final state
        indices = np.where(self.stabilizers[:, self.majorana_number * 2 + pauli_index * 2] == 1)[0]
        self.coffs[indices] = 1j * self.coffs[indices]
        self.stabilizers[:, self.majorana_number * 2 + pauli_index * 2 + 1] += self.stabilizers[:, self.majorana_number * 2 + pauli_index * 2]


    ##  SECTION: ----执行Majorana U gate----
    def u(self, majorana_index: int):
        """
        input.majorana_index: 目标fermionic site的索引
        influence: 对目标fermionic site执行Majorana U门，即根据重叠权重翻转相位
        """

        ##  Compute final state
        weights = np.sum(self.stabilizers[:, 0:2 * self.majorana_number], axis=1)
        overlaps = self.GF(np.where(self.stabilizers[:, majorana_index * 2] == 1, 1, 0))
        indices = np.where(weights + overlaps == 0)[0]
        self.coffs[indices] = -self.coffs[indices]


    ##  SECTION: ----执行Majorana V gate----
    def v(self, majorana_index: int):
        """
        input.majorana_index: 目标fermionic site的索引
        influence: 对目标fermionic site执行Majorana V门，即根据Z重叠权重翻转相位
        """

        ##  Compute final state
        weights = np.sum(self.stabilizers[:, 0:2 * self.majorana_number], axis=1)
        overlaps = self.GF(np.where(self.stabilizers[:, majorana_index * 2 + 1] == 1, 1, 0))
        indices = np.where(weights + overlaps == 0)[0]
        self.coffs[indices] = -self.coffs[indices]


    ##  SECTION: ----执行Majorana N gate----
    def n(self, majorana_index: int):
        """
        input.majorana_index: 目标fermionic site的索引
        influence: 对目标fermionic site执行Majorana N门，即根据X或Z奇偶性翻转相位
        """

        ##  Compute final state
        weights = np.sum(self.stabilizers[:, 0:2 * self.majorana_number], axis=1)
        overlaps = self.GF(
            np.where(np.logical_xor(self.stabilizers[:, majorana_index * 2] == 1, self.stabilizers[:, majorana_index * 2 + 1] == 1), 1, 0))
        indices = np.where(weights + overlaps == 1)[0]
        self.coffs[indices] = -self.coffs[indices]


    ##  SECTION: ----执行Majorana phase gate----
    def p(self, majorana_index: int):
        """
        input.majorana_index: 目标fermionic site的索引
        influence: 对目标fermionic site执行Majorana P门，即交换该site的X和Z算符支持
        """

        ##  Compute final state
        indices = np.where(np.logical_and(self.stabilizers[:, majorana_index * 2 + 1] == 1, self.stabilizers[:, majorana_index * 2] == 0))[0]
        self.coffs[indices] = -self.coffs[indices]
        caches = self.stabilizers[:, majorana_index * 2].copy()
        self.stabilizers[:, majorana_index * 2] = self.stabilizers[:, majorana_index * 2 + 1]
        self.stabilizers[:, majorana_index * 2 + 1] = caches


    ##  SECTION: ----执行Pauli CNOT gate----
    def cx(self, control_index: int, target_index: int):
        """
        input.control_index: Control qubit的索引
        input.target_index: Target qubit的索引
        influence: 对目标qubit执行Pauli CNOT门，即根据control qubit的X算符支持更新target qubit的X和Z算符支持
        """

        ##  Compute final state
        control_qubit_index_x = control_index * 2 + self.majorana_number * 2
        control_qubit_index_z = control_index * 2 + 1 + self.majorana_number * 2
        target_qubit_index_x = target_index * 2 + self.majorana_number * 2
        target_qubit_index_z = target_index * 2 + 1 + self.majorana_number * 2
        targets_x = self.stabilizers[:, control_qubit_index_x] + self.stabilizers[:, target_qubit_index_x]
        control_z = self.stabilizers[:, control_qubit_index_z] + self.stabilizers[:, target_qubit_index_z]
        self.stabilizers[:, target_qubit_index_x] = targets_x
        self.stabilizers[:, control_qubit_index_z] = control_z


    ##  SECTION: ----执行Majorana CNOT gate----
    def cnx(self, control_index: int, target_index: int):
        """
        input.control_index: Control fermionic site的索引
        input.target_index: Target qubit的索引
        influence: 对目标qubit执行Majorana CNOT门，即根据control fermionic site的X算符支持更新target qubit的X和Z算符支持
        """

        ##  Compute final state
        control_majorana_index_x = control_index * 2
        control_majorana_index_z = control_index * 2 + 1
        target_qubit_index_x = self.majorana_number * 2 + target_index * 2
        target_qubit_index_z = self.majorana_number * 2 + target_index * 2 + 1
        factor = [1j if self.stabilizers[temp, target_qubit_index_z] == 1 else 1 for temp in range(len(self.stabilizers))]
        indices = np.where(np.logical_and(self.stabilizers[:, target_qubit_index_z] == 1, self.stabilizers[:, control_majorana_index_z] == 1))[0]
        self.coffs = self.coffs * np.array(factor)
        self.coffs[indices] = -self.coffs[indices]
        targets_x = self.stabilizers[:, target_qubit_index_x] + self.stabilizers[:, control_majorana_index_x] + self.stabilizers[
            :, control_majorana_index_z]
        controls_x = self.stabilizers[:, control_majorana_index_x] + self.stabilizers[:, target_qubit_index_z]
        controls_z = self.stabilizers[:, control_majorana_index_z] + self.stabilizers[:, target_qubit_index_z]
        self.stabilizers[:, target_qubit_index_x] = targets_x
        self.stabilizers[:, control_majorana_index_x] = controls_x
        self.stabilizers[:, control_majorana_index_z] = controls_z


    ##  SECTION: ----执行Majorana control N gate----
    def cnn(self, control_index: int, target_index: int):
        """
        input.control_index: Control fermionic site的索引
        input.target_index: Target fermionic site的索引
        influence: 对目标fermionic site执行Majorana control N门，即根据control fermionic site的X算符支持更新target fermionic site的X和Z算符支持
        """

        if target_index < control_index:
            cache = target_index
            target_index = control_index
            control_index = cache

        ##  Compute final state
        control_majorana_index_x = control_index * 2
        control_majorana_index_z = control_index * 2 + 1
        target_majorana_index_x = target_index * 2
        target_majorana_index_z = target_index * 2 + 1
        temp = self.stabilizers[:, (control_majorana_index_x, control_majorana_index_z, target_majorana_index_x, target_majorana_index_z)]
        indices_1j = np.where(np.logical_or(np.logical_or(np.all(temp == GF2([1, 0, 0, 0]), axis=1), np.all(temp == GF2([0, 1, 0, 0]), axis=1)),
                                            np.logical_or(np.all(temp == GF2([0, 0, 1, 0]), axis=1), np.all(temp == GF2([0, 0, 0, 1]), axis=1))))[0]
        indices_minus_1 = np.where(np.logical_or(np.all(temp == GF2([1, 0, 0, 1]), axis=1), np.all(temp == GF2([0, 1, 1, 0]), axis=1)))[0]
        indices_minus_1j = np.where(np.logical_or(np.logical_or(np.all(temp == GF2([1, 1, 1, 0]), axis=1), np.all(temp == GF2([1, 1, 0, 1]), axis=1)),
                                                  np.logical_or(np.all(temp == GF2([1, 0, 1, 1]), axis=1),
                                                                np.all(temp == GF2([0, 1, 1, 1]), axis=1))))[0]
        self.coffs[indices_1j] = self.coffs[indices_1j] * 1j
        self.coffs[indices_minus_1] = -self.coffs[indices_minus_1]
        self.coffs[indices_minus_1j] = self.coffs[indices_minus_1j] * (-1j)
        control_x = self.stabilizers[:, control_majorana_index_x] + self.stabilizers[:, target_majorana_index_x] + self.stabilizers[
            :, target_majorana_index_z]
        control_z = self.stabilizers[:, control_majorana_index_z] + self.stabilizers[:, target_majorana_index_x] + self.stabilizers[
            :, target_majorana_index_z]
        target_x = self.stabilizers[:, target_majorana_index_x] + self.stabilizers[:, control_majorana_index_x] + self.stabilizers[
            :, control_majorana_index_z]
        target_z = self.stabilizers[:, target_majorana_index_z] + self.stabilizers[:, control_majorana_index_x] + self.stabilizers[
            :, control_majorana_index_z]
        self.stabilizers[:, control_majorana_index_x] = control_x
        self.stabilizers[:, control_majorana_index_z] = control_z
        self.stabilizers[:, target_majorana_index_x] = target_x
        self.stabilizers[:, target_majorana_index_z] = target_z


    ##  SECTION: ----执行Majorana BRAID gate----
    def braid(self, control_index: int, target_index: int):
        """
        input.control_index: Control fermionic site的索引
        input.target_index: Target fermionic site的索引
        influence: 对目标fermionic site执行Majorana BRAID门，即根据control fermionic site的Z算符支持更新target fermionic site的X和Z算符支持
        """

        control_majorana_index_z = control_index * 2 + 1
        target_majorana_index_x = target_index * 2
        if target_index > control_index:
            numbers_of_mid = np.sum(self.stabilizers[:, control_majorana_index_z + 1:target_majorana_index_x], axis=1)
        else:
            numbers_of_mid = np.sum(self.stabilizers[:, target_majorana_index_x + 1:control_majorana_index_z], axis=1)

        ##  Compute final state
        factors_0 = self.stabilizers[:, control_majorana_index_z] * numbers_of_mid
        factors_1 = self.stabilizers[:, target_majorana_index_x] * numbers_of_mid
        factors_2 = self.stabilizers[:, control_majorana_index_z] * self.stabilizers[:, target_majorana_index_x]
        factors_4 = self.stabilizers[:, target_majorana_index_x]
        judge = factors_0 + factors_1 + factors_2 + factors_4
        factors = [-1 if judge[temp] == 1 else 1 for temp in range(len(self.stabilizers))]
        self.coffs = self.coffs * np.array(factors)
        caches = self.stabilizers[:, control_majorana_index_z].copy()
        self.stabilizers[:, control_majorana_index_z] = self.stabilizers[:, target_majorana_index_x]
        self.stabilizers[:, target_majorana_index_x] = caches


    ##  SECTION: ----执行 ERROR gate----
    def error(self, p: float, error_pattern: str, index_0: int, index_1: int = None):
        """
        input.p: Error probability.
        input.error_pattern: Error pattern string.
        input.index_0: Target fermionic site的索引 0.
        input.index_1: Target fermionic site的索引 1.
        influence: Apply an error with probability p.
        """

        if np.random.rand() < p:
            string_0 = error_pattern[0]
            string_1 = error_pattern[1]
            if string_0 == 'X':
                self.x(index_0)
            elif string_0 == 'Y':
                self.y(index_0)
            elif string_0 == 'Z':
                self.z(index_0)
            elif string_0 == 'U':
                self.u(index_0)
            elif string_0 == 'V':
                self.v(index_0)
            elif string_0 == 'N':
                self.n(index_0)
            else:
                raise ValueError(f"string_0={string_0} is not valid.")
            if string_1 == 'X':
                self.x(index_1)
            elif string_1 == 'Y':
                self.y(index_1)
            elif string_1 == 'Z':
                self.z(index_1)
            elif string_1 == 'U':
                self.u(index_1)
            elif string_1 == 'V':
                self.v(index_1)
            elif string_1 == 'N':
                self.n(index_1)
            elif string_1 == '_':
                pass
            else:
                raise ValueError(f"string_0={string_0} or string_1={string_1} is not valid.")


    ##  SECTION: ----重置qubit----
    def reset(self, pauli_index: int):
        """
        input.pauli_index: Target qubit的索引.
        influence: Set corresponding stabilizer row to Z=+1.
        """

        ##  将对应qubit放到0态上
        vector_op = np.append(self.GF.Zeros(self.majorana_number * 2), PauliOperator([], [pauli_index], 1).get_vector(self.pauli_number))
        first_index = None
        for i in range(len(self.stabilizers)):
            pauli_commute = (np.dot(self.stabilizers[i][self.majorana_number * 2::2], vector_op[self.majorana_number * 2 + 1::2]) +
                             np.dot(self.stabilizers[i][self.majorana_number * 2 + 1::2], vector_op[self.majorana_number * 2::2]))
            majorana_commute = (np.dot(self.stabilizers[i][0:self.majorana_number * 2], vector_op[0:self.majorana_number * 2]) +
                                np.sum(vector_op[0:self.majorana_number * 2]) * np.sum(self.stabilizers[i][0:self.majorana_number * 2]))

            if majorana_commute + pauli_commute == 1 and first_index is None:
                first_index = i
            elif majorana_commute + pauli_commute == 1 and first_index is not None:
                factor = majorana_factor(self.majorana_number, [self.stabilizers[i], self.stabilizers[first_index]])
                self.stabilizers[i] += self.stabilizers[first_index]
                if np.mod(factor, 2) == 0:
                    factor = 1
                else:
                    factor = -1
                self.coffs[i] = self.coffs[i] * self.coffs[first_index] * factor
            else:
                pass
        if first_index is not None:
            self.stabilizers[first_index] = vector_op
            self.coffs[first_index] = 1
        else:
            solution = solve(self.stabilizers, vector_op)
            coff = complex(np.prod([self.coffs[i] for i in range(len(self.stabilizers)) if solution[i] == 1]))
            vector_now, factor = majorana_factor(self.majorana_number,
                                                 [self.stabilizers[i] for i in range(len(self.stabilizers)) if solution[i] == 1])
            coff = coff * factor
            if coff == -1:
                self.x(pauli_index)


    ##  SECTION: ----检测state是否在算符本征子空间----
    def detect(self, op: MajoranaOperator | PauliOperator) -> int | None:
        """
        input.op: 待检测的MajoranaOperator或PauliOperator.
        output: int or None, +1/-1 if deterministic, None if random.
        influence: None.
        """

        ##  Check inputs
        if isinstance(op, MajoranaOperator):
            vector_op = np.append(op.get_vector(self.majorana_number), self.GF.Zeros(self.pauli_number * 2))
        else:
            vector_op = np.append(self.GF.Zeros(self.majorana_number * 2), op.get_vector(self.pauli_number))

        ##  Compute final state
        solution = solve(self.stabilizers, vector_op)
        if solution is None:
            return None
        coff = complex(np.prod([self.coffs[i] for i in range(len(self.stabilizers)) if solution[i] == 1]))
        vector_now, factor = majorana_factor(self.majorana_number,
                                             [vector_op] + [self.stabilizers[i] for i in range(len(self.stabilizers)) if solution[i] == 1])
        coff = coff * factor
        if coff == op.coff:
            return 1
        else:
            return -1


##  CHAPTER: ====计算Majorana factor====
def majorana_factor(majorana_number: int, args: list) -> tuple[np.ndarray, int]:
    """
    input.majorana_number: Fermionic sites的数目
    input.args: vectors的列表
    output: tuple, 包括resultant vector和一个phase factor
    """

    first_vector = args[0]
    vector_now = first_vector
    factor = 0
    for vector in args[1:]:
        factor = factor + np.sum(
            [np.count_nonzero(vector_now[temp + 1:2 * majorana_number]) for temp in range(2 * majorana_number) if vector[temp] == 1])
        vector_now = vector_now + vector
    if np.mod(factor, 2) == 0:
        factor = 1
    elif np.mod(factor, 2) == 1:
        factor = -1
    return vector_now, factor
