"""
模块: ColorCode
给定边长d，生成对应的color code
------------------------------------------------------------------------------------------------
ColorCode属性：
    self.d：code distance，也就是color code三角形边长
    self.H：check matrix，color code是一个self-dual code

    MajoranaColorCode(MajoranaCSSCode,ColorCode)：
        None

    PauliColorCode(PauliCSSCode,ColorCode)：
        None
------------------------------------------------------------------------------------------------
ColorCode方法：
    __init__（魔法方法）：构造physical sites和checker

        MajoranaColorCode方法：
            __init__（魔法方法，重构自ColorCode和MajoranaCSSCode）：构造physical sites和checker

        PauliColorCode方法：
            __init__（魔法方法，重构自ColorCode和PauliCSSCode）：构造physical sites和checker
"""
import numpy as np
from galois import GF2
from clifforder.Code.QuantumCode import MajoranaCSSCode, PauliCSSCode
from clifforder.Code.Operator import MajoranaOperator, PauliOperator


##  CHAPTER: ====构造color code的check matrix====
class ColorCode:
    __slots__ = ['d','H']

    ##  SECTION: ----构造函数----
    def __init__(self, d: int) -> None:
        """
        input.d：一个奇数，代表了color code的边长
        output：None
        """

        ##  输入检查
        try:
            self.d=int(d)
        except Exception as e:
            raise ValueError('输入的边长不合法') from e

        ##  初始参数
        cell_vector=[np.array([0, 0], dtype=float), np.array([2, 0], dtype=float), np.array([0.5, np.sqrt(3)/2], dtype=float),
                     np.array([1.5, np.sqrt(3)/2], dtype=float), np.array([0.5, -np.sqrt(3)/2], dtype=float),
                     np.array([1.5, -np.sqrt(3)/2], dtype=float)]  # 六边形的六个点
        cells=[]  # 初始化所有的晶格所在的列表
        x_now=0  # 当前的点的x坐标
        y_now=0  # 当前的点的y坐标

        ##  生成所有的六边形
        ##  总共有d//2大层
        for i in range(d//2):

            ##  一大层有3小层
            for j in range(3):

                ##  每一小层含有d//2-i个六边形
                for k in range(d//2-i):
                    x_now_temp=x_now+k*3  # 六边形需要原点平移3个单位得到下一个六边形
                    cells.append([np.array([x_now_temp, y_now], dtype=float)+cell_vector[k] for k in range(6)])  # 六边形内有六个顶点

                ##  每一层的出发点不同
                if j==0:
                    y_now+=np.sqrt(3)/2
                    x_now+=1.5
                elif j==1:
                    y_now+=np.sqrt(3)/2
                    x_now-=1.5
                elif j==2:
                    y_now+=np.sqrt(3)/2
                    x_now+=1.5

        ##  求出处于三角形内的sites
        cells_in_triangle=[]
        for i in range(len(cells)):
            cells_in_triangle.append([])
            for j in range(6):
                if judge_in_triangle(cells[i][j][0], cells[i][j][1],d):
                    cells_in_triangle[i].append(cells[i][j])

        ##  先将最底部的一层site标上序号
        position_list=[]
        stabilizer_list=[]
        for i in range(len(cells_in_triangle)):
            for j in range(len(cells_in_triangle[i])):

                ##  要求现在标的都是最底层的
                if abs(cells_in_triangle[i][j][1])<0.01:
                    flag=True
                    for k in range(len(position_list)):
                        if equal(cells_in_triangle[i][j], position_list[k]):
                            flag=False
                            break
                    if flag:
                        position_list.append(cells_in_triangle[i][j])

        ##  将所有其他的site标上序号
        for i in range(len(cells_in_triangle)):
            stabilizer_temp=[]
            for j in range(len(cells_in_triangle[i])):
                flag=True
                index=0
                for k in range(len(position_list)):
                    if equal(cells_in_triangle[i][j], position_list[k]):
                        flag=False
                        index=k  # 当前的site已经被标记过了
                        break

                ##  如果没标记过就需要标记一下
                if flag:
                    stabilizer_temp.append(len(position_list))
                    position_list.append(cells_in_triangle[i][j])

                ##  如果标记过了就直接加上对应的序号在当前stabilizer里面
                else:
                    stabilizer_temp.append(index)

            ##  将stabilizer加入列表
            stabilizer_list.append(stabilizer_temp)

        ##  生成check matrix
        self.H=GF2(np.zeros((len(stabilizer_list),len(position_list)),dtype=int))
        for i in range(len(stabilizer_list)):
            self.H[i][stabilizer_list[i]]=1

    ##  TODO：----绘制color code的三角图形----
    def plot(self):
        pass


##  CHAPTER: ====Majorana color codes====
class MajoranaColorCode(MajoranaCSSCode,ColorCode):

    ##  SECTION: ----构造函数----
    def __init__(self,d:int):
        ColorCode.__init__(self,d)

        ##  根据校验矩阵生成稳定子生成元
        generators_x=[]
        generators_z=[]
        for i in range(len(self.H)):
            generators_x.append(MajoranaOperator.HermitianOperatorFromOccupy(np.where(self.H[i]!=0)[0],[]))
            generators_z.append(MajoranaOperator.HermitianOperatorFromOccupy([],np.where(self.H[i]!=0)[0]))
        MajoranaCSSCode.__init__(self,generators_x,generators_z,self.H.shape[1])

        ##  先验给出逻辑算符
        self._logical_operators_x=[MajoranaOperator.HermitianOperatorFromOccupy([i for i in range(d)], [])]
        self._logical_operators_z=[MajoranaOperator.HermitianOperatorFromOccupy([], [i for i in range(d)])]


##  CHAPTER: ====Pauli color codes====
class PauliColorCode(PauliCSSCode,ColorCode):

    ##  SECTION: ----构造函数----
    def __init__(self,d:int):

        ##  根据校验矩阵生成稳定子生成元
        ColorCode.__init__(self,d)
        generators_x=[]
        generators_z=[]
        for i in range(len(self.H)):
            generators_x.append(PauliOperator.HermitianOperatorFromOccupy(np.where(self.H[i]!=0)[0],[]))
            generators_z.append(PauliOperator.HermitianOperatorFromOccupy([],np.where(self.H[i]!=0)[0]))
        PauliCSSCode.__init__(self,generators_x,generators_z,self.H.shape[1])

        ##  先验给出逻辑算符
        self._logical_operators_x=[PauliOperator.HermitianOperatorFromOccupy([i for i in range(d)], [])]
        self._logical_operators_z=[PauliOperator.HermitianOperatorFromOccupy([], [i for i in range(d)])]


##  CHAPTER: ====判断当前点在不在规定的三角形中====
def judge_in_triangle(x, y,d) -> bool:
    """
    input.x：点的x坐标
    input.y：点的y坐标
    output：判断结果
    """

    x_1=d//2+d-1
    x_2=x_1/2
    y_2=np.sqrt(3)*x_2

    k_0=0
    b_0=0
    k_1=(y_2)/(x_2)
    b_1=0
    k_2=-k_1
    b_2=2*y_2
    if k_0*x+b_0<y+0.01 and k_1*x+b_1>y-0.01 and k_2*x+b_2>y-0.01:
        return True
    else:
        return False


##  CHAPTER: ====判断两个点是否近似相等====
def equal(a, b) -> bool:
    """
    input.a：a点的的坐标
    input.b：b点的坐标
    output：判断结果
    """
    if abs(a[0]-b[0])+abs(a[1]-b[1])<0.01:
        return True
    else:
        return False
