"""
模块: PolynomialCode
给定多项式，生成对应的polynomial code，同时实现了表示多项式环的类
------------------------------------------------------------------------------------------------
Polynomial的属性：
    self.number：多项式内变量个数
    self.polynomial：多项式每一项的每个变量的幂
    self.lengths：多项式每个维度的循环长度

PolynomialRing的属性：
    self.number：多项式内变量个数
    self.lengths：多项式每个维度的循环长度

PolynomialCode的属性：
    self.lengths：多项式每个维度的循环长度
    self.cell_size：每个cell的大小，也就是每个cell内sites数目
    self.polynomials_x：每个cell内sites对应的位移多项式，对应X checker
    self.polynomials_z：每个cell内sites对应的位移多项式，对应Z checker
    self.H_X：X checker的check matrix
    self.H_Z：Z checker的check matrix

    MajoranaPolynomialCode(MajoranaCSSCode,PolynomialCode)：
        None

    PauliPolynomialCode(PauliCSSCode,PolynomialCode)：
        None
------------------------------------------------------------------------------------------------
Polynomial的方法：
    __init__（魔法方法）：构造多项式
    __str__（魔法方法）：返回多项式的字符串表示
    __add__（魔法方法）：多项式加法
    __radd__（魔法方法）：多项式加法的右结合
    __mul__（魔法方法）：多项式乘法
    __rmul__（魔法方法）：多项式乘法的右结合
    __pow__（魔法方法）：多项式取幂
    matrix（属性方法）：多项式对应的矩阵
    T（属性方法）：多项式对应的转置

PolynomialRing的方法：
    __init__（魔法方法）：构造多项式环
    element：返回某一个变量的一次项

PolynomialCode的方法：
    __init__（魔法方法）：构造quantum polynomial code
    __str__（魔法方法）：返回polynomial code的字符串表示

    MajoranaPolynomialCode(MajoranaCSSCode,PolynomialCode)的方法：
        __init__（魔法方法，重构自MajoranaCSSCode,PolynomialCode）：构造Majorana polynomial code

    PauliPolynomialCode(PauliCSSCode,PolynomialCode)的方法：
        __init__（魔法方法，重构自PauliCSSCode,PolynomialCode）：构造Pauli polynomial code

------------------------------------------------------------------------------------------------
多项式例子：
    $x_0x_1x_2x_3+2x_0x_1x_2x_3+3x_0x_1x_2x_3$
    element = Polynomial(4,[[0,0,0,1],[0,0,1,1],[0,1,1,1]],[2,2,2,2])
"""
import copy
from typing import Iterable
import galois
from galois import GF2
from clifforder.Code.QuantumCode import MajoranaCSSCode, PauliCSSCode
from clifforder.Code.Operator import MajoranaOperator, PauliOperator
import numpy as np


##  CHAPTER: ====作为一个多项式环的元素类====
class Polynomial:

    ##  SECTION: ----构造函数----
    def __init__(self,number:int,polynomial,lengths):
        """
        input.number：多项式内变量个数
        input.polynomial：多项式每一项的每个变量的幂
        lengths：多项式每个维度的循环长度
        """

        ##  输入检查
        try:
            number = int(number)
            polynomial = np.array(polynomial,dtype=int)
            assert polynomial.shape[1]==number
        except Exception as e:
            raise ValueError('输入的参量不合法') from e

        ##  确保合并同类项，并将重复项相加湮灭
        for i in range(polynomial.shape[1]):
            polynomial[:,i]=np.mod(polynomial[:,i],lengths[i])
        row_dtype = np.dtype([("f", polynomial.dtype, (polynomial.shape[1],))])
        x_view = polynomial.view(row_dtype).ravel()
        uniq, counts = np.unique(x_view, return_counts=True)
        polynomial = uniq[counts == 1].view(polynomial.dtype).reshape(-1, polynomial.shape[1])

        ##  赋值
        self.number = number
        self.lengths = np.array(lengths)
        self.polynomial = polynomial


    ##  SECTION: ----加法----
    def __add__(self, other: 'Polynomial|int'):
        """
        input.other：0，1或一个多项式
        output：计算结果多项式，就是把两个多项式的项放在一起合并同类项为0
        """

        ##  当对方等于0的情况
        if other==0:
            return self.copy()

        ##  对方等于1的情况
        elif other==1:
            return Polynomial(self.number, np.vstack([np.array([0] * self.number, dtype=int), self.polynomial]), self.lengths)

        ##  对方是一般的多项式的情况
        elif isinstance(other, Polynomial):
            assert self.number == other.number and all([self.lengths[i] == other.lengths[i] for i in range(len(self.lengths))]),"两个多项式的空间必须相同"
            return Polynomial(self.number, np.vstack([self.polynomial, other.polynomial]), self.lengths)

        ##  其他输入抛出错误
        else:
            raise TypeError('不支持当前元素的环加法')


    ##  SECTION: ----右加法----
    def __radd__(self, other: 'Polynomial|int'):
        """
        input.other：0，1或一个多项式
        output：计算结果多项式，就是把两个多项式的项放在一起合并同类项为0
        """

        return self.__add__(other)

    ##  SECTION: ----乘法----
    def __mul__(self, other: 'Polynomial|int'):
        """
        input.other：0，1或一个多项式
        output：计算结果多项式，就是把两个多项式单项式两两相加然后合并同类项为0
        """

        ##  当对方等于0的情况
        if other==1:
            return self.copy()

        ##  对方等于1的情况
        elif other==0:
            return 0

        ##  对方是一般的多项式的情况
        elif isinstance(other, Polynomial):
            x = []
            for i in range(len(self.polynomial)):
                for j in range(len(other.polynomial)):
                    x.append(self.polynomial[i]+other.polynomial[j])
            return Polynomial(self.number, np.array(x, dtype=int), self.lengths)

        ##  其他输入抛出错误
        else:
            raise TypeError('不支持当前元素的环加法')


    ##  SECTION: ----右乘法----
    def __rmul__(self, other: 'Polynomial|int'):
        """
        input.other：0，1或一个多项式
        output：和一般的乘法一样，因为这是交换环
        """

        return self.__mul__(other)


    ##  SECTION: ----幂----
    def __pow__(self,other:int):

        ##  检查输入
        try:
            other = int(other)
        except Exception as e:
            raise TypeError("输入的幂不合法") from e

        ##  计算并返回结果
        temp=self

        ##  幂为0的时候恒等于1
        if other==0:
            return Polynomial(self.number, np.array([[0] * self.number, ], dtype=int), self.lengths)

        ##  其他情况需要做other-1次乘法
        else:
            for i in range(other-1):
                temp=temp*self
            return temp


    ##  SECTION: ----字符串输出----
    def __str__(self):
        """
        output：多项式的字符串表示
        """

        poly_str=""  # 多项式的字符串表示初始化
        for q in range(len(self.polynomial)):

            ##  第一项前面没有加号
            if q==0:

                ##  遍历每一个变量的幂
                for p in range(len(self.polynomial[q])):

                    ##  如果都是0幂，则给出1
                    if np.all(self.polynomial[q,p]==0):
                        poly_str+="1"
                    else:

                        ##  如果某一项为0则不写出来
                        if self.polynomial[q][p]==0:
                            continue

                        ##  如果某一项为1则不写幂
                        elif self.polynomial[q][p]==1:
                            poly_str=poly_str+f"x_{q}"

                        ##  其他情况写上变量及其幂
                        else:
                            poly_str=poly_str+f"x_{q}^{self.polynomial[q][p]}"

            ##  其他项前面有加号
            else:
                poly_str = poly_str + f"+"

                ##  规则同上
                for p in range(len(self.polynomial[q])):
                    if np.all(self.polynomial[q,p]==0):
                        poly_str+="1"
                    else:
                        if self.polynomial[q][p]==0:
                            continue
                        elif self.polynomial[q][p]==1:
                            poly_str=poly_str+f"x_{q}"
                        else:
                            poly_str=poly_str+f"x_{q}^{self.polynomial[q][p]}"

        ##  返回结果
        return poly_str


    ##  SECTION: ----复制----
    def copy(self):
        """
        output：当前多项式的一个副本
        """

        ##  直接返回结果
        return Polynomial(self.number, copy.deepcopy(self.polynomial), self.lengths)


    ##  SECTION: ----生成多项式对应的矩阵----
    def matrix(self)->galois.GF2:
        """
        output：多项式对应的01矩阵
        """

        result=None  # 矩阵初始化

        ##  多项式的单项式之间相加，多项式单项式内部做张量积
        for i in range(len(self.polynomial)):
            temp=1
            for j in range(len(self.polynomial[i])):
                temp=np.kron(temp,shift(self.lengths[j],self.polynomial[i][j]))
            if i==0:
                result=GF2(temp)
            else:
                result=result+GF2(temp)

        ##  返回结果
        return result


    ##  SECTION: ----求多项式的转置----
    @property
    def T(self):
        """
        output：转置之后的多项式，也就是多项式里面每一项都将幂取负号
        """

        ##  复制当前多项式
        result=self.copy()

        ##  当前多项式的单项式内变量的幂全部取负号
        for i in range(result.polynomial.shape[1]):
            result.polynomial[:,i]=self.lengths[i]-result.polynomial[:,i]

        ##  返回结果
        return result


##  CHAPTER: ====多项式环类====
class PolynomialRing:

    ##  SECTION: ----构造函数----
    def __init__(self,number,lengths):
        """
        number: 多项式的变量个数
        lengths: 每个变量的长度
        """

        ##  输入检查
        try:
            number=int(number)
            length=np.array([int(length) for length in lengths],dtype=int)
            assert np.all(length>0)
            assert len(length)==number
        except Exception as e:
            raise ValueError("当前多项式环不合法") from e
        self.number=number
        self.lengths=lengths

    ##  SECTION: ----生成变量的一次项----
    def element(self,index):
        """
        input.index：标明是第几个变量
        output：返回第index个变量的一次项
        """

        ##  检查输入
        try:
            index=int(index)
        except Exception as e:
            raise ValueError("当前输入的变量索引不合法") from e

        ##  构造对应变量的一次项
        temp=[0]*self.number
        temp[index]=1

        ##  返回结果
        return Polynomial(self.number, np.array([temp, ], dtype=int), self.lengths)


##  CHAPTER: ====Polynomial code的check matrix====
class PolynomialCode:

    ##  SECTION: ----构造函数----
    def __init__(self, poly_x_list:Iterable[Polynomial], poly_z_list:Iterable[Polynomial]):
        """
        lengths: Lengths of each dimension
        cell_size: Number of qubits in each unit cell
        polynomials: PolynomialRing description, structure A[q][i][j]
        q: Represents the qth qubit in the unit cell (0 <= q < cell_size)
        i: Represents the ith position of the qubit in the stabilizer (i.e., the ith term in the polynomial)
        j: Represents the dimension indices of the offset vector [x, y, z...]
        """

        ##  输入检查
        try:
            assert np.all([isinstance(polynomial,Polynomial) for polynomial in poly_x_list]),"输入的X多项式列表中包含非多项式对象"
            assert np.all([isinstance(polynomial,Polynomial) for polynomial in poly_z_list]),"输入的Z多项式列表中包含非多项式对象"
            poly_x_list=[polynomial.copy() for polynomial in poly_x_list]
            poly_z_list=[polynomial.copy() for polynomial in poly_z_list]
            lengths=poly_x_list[0].lengths.copy()
            cell_size=len(poly_x_list)
            number=poly_x_list[0].number
            assert len(poly_z_list)==cell_size,"晶格原胞内sites数目必须相等"
            assert np.all([temp.number==number for temp in poly_x_list+poly_z_list]),"多项式的变量数必须相等"
            assert np.all([np.all(temp.lengths==lengths) for temp in poly_x_list+poly_z_list]),"每个多项式的变量长度必须相等"
        except Exception as e:
            raise ValueError("当前输入的参数不合法") from e

        ##  赋值
        self.lengths = lengths
        self.cell_size = cell_size
        self.polynomials_x = poly_x_list
        self.polynomials_z = poly_z_list

        ##  生成多项式对应的矩阵
        matrix_list_x = []
        matrix_list_z = []
        for i in range(self.cell_size):
            matrix_list_x.append(self.polynomials_x[i].matrix())
            matrix_list_z.append(self.polynomials_z[i].matrix())
        self.H_X=np.hstack(matrix_list_x)
        self.H_Z=np.hstack(matrix_list_z)


    ##  SECTION: ----打印对象----
    def __str__(self)->str:
        """
        output：表示多项式
        influence：打印多项式
        """

        ##  构造X多项式字符串，每一项占向量的一个
        x_str = "("  # 多项式字符串初始化
        for i in range(self.cell_size):
            if i==0:
                x_str=x_str+f"{self.polynomials_x[i]}"
            else:
                x_str=x_str+f",{self.polynomials_x[i]}"
        x_str=x_str+"|0,...,0)^T"  # X多项式字符串结束

        # 构造Z多项式字符串，每一项占向量的一个
        z_str = "(0,...,0|"  # 多项式字符串初始化
        for i in range(self.cell_size):
            if i == 0:
                z_str = z_str + f"{self.polynomials_z[i]}"
            else:
                z_str = z_str + f",{self.polynomials_z[i]}"
        z_str = z_str + ")^T"  # Z多项式字符串结束

        return f"$H_X={x_str}\nH_Z={z_str}$"


##  CHAPTER: ====Majorana polynomial Code====
class MajoranaPolynomialCode(MajoranaCSSCode,PolynomialCode):

    ##  SECTION: ----构造函数----
    def __init__(self,poly_x_list:Iterable[Polynomial], poly_z_list:Iterable[Polynomial]):

        ##  构造校验矩阵
        PolynomialCode.__init__(self,poly_x_list,poly_z_list)

        ##  生成稳定子生成元
        generators_x=[]
        generators_z=[]
        for i in range(len(self.H_X)):
            generators_x.append(MajoranaOperator.HermitianOperatorFromOccupy(np.where(self.H_X[i]!=0)[0],[]))
        for i in range(len(self.H_Z)):
            generators_z.append(MajoranaOperator.HermitianOperatorFromOccupy([],np.where(self.H_Z[i]!=0)[0]))
        MajoranaCSSCode.__init__(self,generators_x,generators_z,self.H_X.shape[1])


##  CHAPTER: ====Pauli polynomial Code====
class PauliPolynomialCode(PauliCSSCode,PolynomialCode):

    ##  SECTION: ----构造函数----
    def __init__(self, poly_x_list:Iterable[Polynomial], poly_z_list:Iterable[Polynomial]):

        ##  构造校验矩阵
        PolynomialCode.__init__(self,poly_x_list,poly_z_list)

        ##  生成稳定子生成元
        generators_x=[]
        generators_z=[]
        for i in range(len(self.H_X)):
            generators_x.append(PauliOperator.HermitianOperatorFromOccupy(np.where(self.H_X[i]!=0)[0],[]))
        for i in range(len(self.H_Z)):
            generators_z.append(PauliOperator.HermitianOperatorFromOccupy([],np.where(self.H_Z[i]!=0)[0]))
        PauliCSSCode.__init__(self,generators_x,generators_z,self.H_X.shape[1])


##  CHAPTER: ====位移矩阵====
def shift(number:int, shift_number:int):
    """
    input.number: Matrix dimension
    input.shift_number: Cyclic shift_number step length
    output: GF(2) matrix, permutation matrix
    """

    ##  初始化矩阵
    S = np.zeros((number, number), dtype=int)

    ##  矩阵在第i行第i+shift_number列为1，其余为0
    for i in range(number):
        S[i, (i + shift_number) % number] = 1

    ##  转换为GF(2)矩阵
    return galois.GF(2)(S)
