"""
模块: Frame
生成量子计算机的Pauli帧，并可以执行Clifford gate和Pauli or Majorana measurement
------------------------------------------------------------------------------------------------
Frame的属性：
    self.pauli_number：Qubits数目
    self.majorana_number：Fermionic sites数目
    self.frame：平台当前的噪声情况
------------------------------------------------------------------------------------------------

"""
import galois
import numpy as np
from clifforder.Code.Operator import MajoranaOperator, PauliOperator


class Frame:

    GF = galois.GF(2)

    ##  SECTION: ----构造函数----
    def __init__(self):
        """
        influence: 初始化一个空帧
        """
        self.pauli_number = 0
        self.majorana_number = 0
        self.frame = None

    ##  SECTION: ----初始化帧----
    def trap(self, majorana_number: int, pauli_number: int):
        """
        input.majorana_number：Fermionic sites的数目
        input.pauli_number：Qubits的数目
        influence: 随机初始化帧，模拟投影到随机+1/-1本征态
        """

        ##  输入检查
        try:
            self.pauli_number = int(pauli_number)
            self.majorana_number = int(majorana_number)
        except Exception as e:
            raise ValueError("当前输入的参数不合法") from e

        ##  帧变量放到空态并随机初始化
        self.frame = self.GF.Zeros(2 * pauli_number + 2 * majorana_number)
        for i in range(majorana_number):
            if np.random.rand() > 0.5:
                temp = self.GF(np.zeros(2 * majorana_number, dtype=int))
                temp[i * 2] = 1
                temp[i * 2 + 1] = 1
                self.frame[0:self.majorana_number * 2] += temp
        for i in range(pauli_number):
            if np.random.rand() > 0.5:
                temp = self.GF(np.zeros(2 * pauli_number, dtype=int))
                temp[i * 2 + 1] = 1
                self.frame[self.majorana_number * 2:] += temp

    ##  SECTION: ----测量某一个string operator----
    def measure(self, op: MajoranaOperator | PauliOperator, reference_value: int) -> int:
        """
        input.op: 待测量的厄米算符
        input.reference_value: 参考值，在无噪声情况下的测量结果
        output: int, 实际测量结果 (+1 or -1).
        influence: 根据当前帧与待测量算符的 commute/anticommutate 关系，翻转参考值，并以50%概率更新帧。
        """

        ##  Computing final state
        if isinstance(op, MajoranaOperator):
            v0 = op.get_vector(self.majorana_number)
            v1 = self.frame[0:self.majorana_number * 2]
            overlap_number = np.dot(v0, v1)
            weight_0 = np.dot(v0, v0)
            weight_1 = np.dot(v1, v1)
            if overlap_number + weight_0 * weight_1 == 0:
                result = reference_value
            else:
                result = -reference_value
            if np.random.rand() > 0.5:
                self.frame[0:self.majorana_number * 2] += v0
            return result
        elif isinstance(op, PauliOperator):
            v0 = op.get_vector(self.pauli_number)
            v0_x = v0[0::2]
            v0_z = v0[1::2]
            v1 = self.frame[self.majorana_number * 2:]
            v1_x = v1[0::2]
            v1_z = v1[1::2]
            overlap_number = np.dot(v0_x, v1_z) + np.dot(v0_z, v1_x)
            if overlap_number == 0:
                result = reference_value
            else:
                result = -reference_value
            if np.random.rand() > 0.5:
                self.frame[self.majorana_number * 2:] += v0
            return result
        else:
            raise ValueError

    ##  SECTION: ----Pauli X gate什么也不做----
    def x(self, qubit_index: int):
        """
        input.qubit_index: Target qubit的索引
        influence: 什么也不做
        """
        pass

    ##  SECTION: ----Pauli Y gate什么也不做----
    def y(self, qubit_index: int):
        """
        input.qubit_index: Target qubit的索引
        influence: 什么也不做
        """
        pass

    ##  SECTION: ----Pauli Z gate什么也不做----
    def z(self, qubit_index: int):
        """
        input.qubit_index: Target qubit的索引
        influence: 什么也不做
        """
        pass

    ##  SECTION: ----Hadamard gate----
    def h(self, qubit_index: int):
        """
        input.qubit_index: Target qubit的索引
        influence: Swap X和Z组件
        """

        qubit_index_x = self.majorana_number * 2 + qubit_index * 2
        qubit_index_z = self.majorana_number * 2 + qubit_index * 2 + 1

        ##  Computing final state
        cache = self.frame[qubit_index_x].copy()
        self.frame[qubit_index_x] = self.frame[qubit_index_z]
        self.frame[qubit_index_z] = cache


    ##  SECTION: ----Pauli phase gate----
    def s(self, pauli_index: int):
        """
        input.pauli_index: Target qubit的索引
        influence: Add X组件到Z组件中（Z -> ZX）
        """

        qubit_index_x = self.majorana_number * 2 + pauli_index * 2
        qubit_index_z = self.majorana_number * 2 + pauli_index * 2 + 1

        ##  Computing final state
        self.frame[qubit_index_z] += self.frame[qubit_index_x]


    ##  SECTION: ----Majorana U gate什么也不做----
    def u(self, majorana_index: int):
        """
        input.majorana_index: Target fermionic site的索引
        influence: 什么也不做
        """
        pass


    ##  SECTION: ----Majorana V gate什么也不做----
    def v(self, majorana_index: int):
        """
        input.majorana_index: Target fermionic site的索引
        influence: 什么也不做
        """
        pass


    ##  SECTION: ----Majorana N gate什么也不做----
    def n(self, majorana_index: int):
        """
        input.majorana_index: Target fermionic site的索引
        influence: 什么也不做
        """
        pass

    ##  SECTION: ----Majorana phase gate----
    def p(self, majorana_index: int):
        """
        input.majorana_index: Target fermionic site的索引
        influence: Swap X和Z组件
        """

        ##  Check inputs
        majorana_index_x = majorana_index * 2
        majorana_index_z = majorana_index * 2 + 1

        ##  Computing fianl state
        cache = self.frame[majorana_index_x].copy()
        self.frame[majorana_index_x] = self.frame[majorana_index_z]
        self.frame[majorana_index_z] = cache

    ##  SECTION: ----Pauli control X gate----
    def cx(self, control_index: int, target_index: int):
        """
        input.control_index: Control qubit的索引
        input.target_index: Target qubit的索引
        influence: 更新Pauli frame: target_X += control_X, control_Z += target_Z.
        """

        ##  Check inputs
        control_qubit_index_x = self.majorana_number * 2 + control_index * 2
        control_qubit_index_z = self.majorana_number * 2 + 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

        ##  Computing final state
        self.frame[target_qubit_index_x] += self.frame[control_qubit_index_x]
        self.frame[control_qubit_index_z] += self.frame[target_qubit_index_z]


    ##  SECTION: ----执行Fermionic control X gate----
    def cnx(self, control_index: int, target_index: int):
        """
        input.control_index: Fermionic control site的索引
        input.target_index: Target qubit的索引
        influence: 更新 fermion-qubit frame 基于混合门规则.
        """

        ##  索引对齐
        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

        ##  Computing final state
        target_x = self.frame[target_qubit_index_x] + self.frame[control_majorana_index_x] + self.frame[control_majorana_index_z]
        control_x = self.frame[control_majorana_index_x] + self.frame[target_qubit_index_z]
        control_z = self.frame[control_majorana_index_z] + self.frame[target_qubit_index_z]
        self.frame[target_qubit_index_x] = target_x
        self.frame[control_majorana_index_x] = control_x
        self.frame[control_majorana_index_z] = control_z


    ##  SECTION: ----执行Fermionic 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: 更新fermion-fermionic frame基于混合门规则.
        """

        ##  索引对齐
        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

        ##  Computing final state
        control_x = self.frame[control_majorana_index_x] + self.frame[target_majorana_index_x] + self.frame[target_majorana_index_z]
        control_z = self.frame[control_majorana_index_z] + self.frame[target_majorana_index_x] + self.frame[target_majorana_index_z]
        target_x = self.frame[target_majorana_index_x] + self.frame[control_majorana_index_x] + self.frame[control_majorana_index_z]
        target_z = self.frame[target_majorana_index_z] + self.frame[control_majorana_index_x] + self.frame[control_majorana_index_z]
        self.frame[control_majorana_index_x] = control_x
        self.frame[control_majorana_index_z] = control_z
        self.frame[target_majorana_index_x] = target_x
        self.frame[target_majorana_index_z] = target_z


    ##  SECTION: ----执行Fermionic braid gate----
    def braid(self, control_index: int, target_index: int, *args):
        """
        input.control_index: Fermionic site的索引
        input.target_index: Fermionic site的索引
        influence: 更新frame.
        """

        ##  索引对齐
        control_majorana_index_z = control_index * 2 + 1
        target_majorana_index_x = target_index * 2

        ##  Computing final state
        cache = self.frame[control_majorana_index_z].copy()
        self.frame[control_majorana_index_z] = self.frame[target_majorana_index_x]
        self.frame[target_majorana_index_x] = cache


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

        ##  Compute
        if np.random.rand() < p:
            string_0 = error_pattern[0]
            string_1 = error_pattern[1]
            if string_0 == 'X':
                self.frame[index_0 * 2 + self.majorana_number * 2] += self.GF(1)
            elif string_0 == 'Y':
                self.frame[index_0 * 2 + self.majorana_number * 2] += self.GF(1)
                self.frame[index_0 * 2 + self.majorana_number * 2 + 1] += self.GF(1)
            elif string_0 == 'Z':
                self.frame[index_0 * 2 + self.majorana_number * 2 + 1] += self.GF(1)
            elif string_0 == 'U':
                self.frame[index_0 * 2] += self.GF(1)
            elif string_0 == 'V':
                self.frame[index_0 * 2 + 1] += self.GF(1)
            elif string_0 == 'N':
                self.frame[index_0 * 2] += self.GF(1)
                self.frame[index_0 * 2 + 1] += self.GF(1)
            else:
                raise ValueError(f"string_0={string_0} is not valid.")
            if string_1 == 'X':
                self.frame[index_1 * 2 + self.majorana_number * 2] += self.GF(1)
            elif string_1 == 'Y':
                self.frame[index_1 * 2 + self.majorana_number * 2] += self.GF(1)
                self.frame[index_1 * 2 + self.majorana_number * 2 + 1] += self.GF(1)
            elif string_1 == 'Z':
                self.frame[index_1 * 2 + self.majorana_number * 2 + 1] += self.GF(1)
            elif string_1 == 'U':
                self.frame[index_1 * 2] += self.GF(1)
            elif string_1 == 'V':
                self.frame[index_1 * 2 + 1] += self.GF(1)
            elif string_1 == 'N':
                self.frame[index_1 * 2] += self.GF(1)
                self.frame[index_1 * 2 + 1] += self.GF(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: 重置目标qubit的状态，并随机赋值
        """

        ##  索引对齐
        qubit_index_x = self.majorana_number * 2 + pauli_index * 2
        qubit_index_z = self.majorana_number * 2 + pauli_index * 2 + 1

        ##  Compute
        self.frame[qubit_index_x] = 0
        self.frame[qubit_index_z] = 0
        if np.random.rand() < 0.5:
            self.frame[qubit_index_z] = self.GF(1)
