import galois
from tqdm.contrib.itertools import product
from extendedstim.Code.PrimitiveCode.MajoranaCSSCode import MajoranaCSSCode
from extendedstim.Code.PrimitiveCode.PauliCSSCode import PauliCSSCode
from extendedstim.Physics.MajoranaOperator import MajoranaOperator
from extendedstim.Physics.PauliOperator import PauliOperator
from extendedstim.tools.TypingTools import islist, isinteger
import numpy as np


class PolynomialMatrix:
    def __init__(self, dimension,lengths,cell_size:int,polynomials):
        """""
        input：
            dimension: 维度
            lengths: 每个维度的长度
            cell_size: 每个unit cell中的qubit数量
            polynomials: 多项式描述，结构为 A[q][i][j]
                q: 代表unit cell中的第q个qubit (0 <= q < cell_size)
                i: 代表该qubit在stabilizer中出现的第i个位置 (即多项式中的第i项)
                j: 代表偏移向量的维度索引 [x, y, z...]
        """""
        assert len(polynomials)==cell_size
        assert islist(lengths) and all([isinteger(length) for length in lengths])
        
        # 规范化多项式中的偏移量 (mod lengths)
        for q in range(len(polynomials)): # q: qubit index within cell
            for i in range(len(polynomials[q])): # i: term index in polynomial
                polynomials[q][i]=np.mod(polynomials[q][i],lengths)
                
        self.polynomial = polynomials
        self.dimension = dimension
        self.cell_size = cell_size
        self.lengths = lengths
        self.physical_number=cell_size*np.prod(lengths)
        
        stabilizers=[]
        # 遍历整个晶格中的每一个位置 vector (作为stabilizer的中心位置)
        for vector in product(*[range(length) for length in lengths]):
            stabilizer_temp = np.zeros((cell_size,)+tuple(lengths))
            # index 对应 q (qubit type in cell)
            # polynomial 对应 A[q] (list of offsets for this qubit type)
            for q, offsets in enumerate(self.polynomial):
                for pos in offsets: # pos is one offset vector [x, y, z]
                    # Stabilizer中心在 vector，连接到的qubit位置是 vector + pos
                    vector_temp=np.array(vector,dtype=int)+np.array(pos,dtype=int)
                    # 标记对应的qubit (第q种qubit，位置为 vector_temp % lengths)
                    stabilizer_temp[q][*np.mod(vector_temp,lengths)]=1
            stabilizers.append(stabilizer_temp)
        stabilizers=np.array(stabilizers,dtype=int)
        stabilizers=galois.GF2(stabilizers.reshape((self.physical_number//cell_size,self.physical_number)))
        self.H=stabilizers

    #%%  SECTION：----重载函数----
    def __str__(self)->str:
        """""
        input：None
        output：字符串表示
        influence：返回多项式矩阵的字符串描述
        """""
        if self.dimension <= 3:
            vars = ['x', 'y', 'z']
        else:
            vars = [f'x_{i}' for i in range(self.dimension)]
            
        poly_strs = []
        for q in range(len(self.polynomial)):
            terms = []
            offsets = self.polynomial[q]
            if not offsets:
                terms.append("0")
            else:
                for pos in offsets:
                    term_parts = []
                    for dim_idx, p in enumerate(pos):
                        if p != 0:
                            if p == 1:
                                term_parts.append(f"{vars[dim_idx]}")
                            else:
                                term_parts.append(f"{vars[dim_idx]}^{p}")
                    if not term_parts:
                        terms.append("1")
                    else:
                        terms.append("".join(term_parts))
            poly_strs.append(" + ".join(terms) if terms else "0")
            
        return f"({', '.join(poly_strs)})"


class PolynomialCode:
    def __init__(self, h_x: PolynomialMatrix, h_z: PolynomialMatrix):
        assert h_x.dimension == h_z.dimension
        assert h_x.lengths == h_z.lengths
        assert h_x.cell_size == h_z.cell_size
        self.h_x = h_x
        self.h_z = h_z
        self.dimension = h_x.dimension
        self.lengths = h_x.lengths
        self.cell_size = h_x.cell_size
        self.H_X=h_x.H
        self.H_Z=h_z.H

    def _get_visual_pos(self, target, center, length):
        """Helper to get the nearest visual position accounting for PBC."""
        diff = target - center
        # Map diff to [-length/2, length/2]
        diff = (diff + length/2) % length - length/2
        return center + diff

    def _get_qubit_offset(self, q, cell_size):
        """
        计算cell中第q个qubit的视觉偏移量。
        排列规则：从左往右，从下往上 (Grid layout)
        """
        if cell_size == 1:
            return 0, 0
        
        # Grid layout dimensions
        nx = int(np.ceil(np.sqrt(cell_size)))
        ny = int(np.ceil(cell_size / nx))
        
        # Grid coordinates (col, row)
        col = q % nx
        row = q // nx
        
        # Map to visual offsets centered at (0,0)
        # spacing = 0.2
        spacing = 0.2
        
        # Center the grid
        # x range: [-(nx-1)/2, (nx-1)/2] * spacing
        dx = (col - (nx - 1) / 2) * spacing
        dy = (row - (ny - 1) / 2) * spacing
        
        return dx, dy

    def _get_qubit_style(self, q, cell_size):
        """
        获取qubit的颜色样式。
        使用低饱和度高级配色。
        """
        # 配色表 (fill, edge)
        # 0: Green (#86B257, #74A742)
        # 1: Blue (#87CEEB, #4682B4) - SkyBlue / SteelBlue
        # 2: Purple (#D8BFD8, #DA70D6) - Thistle / Orchid
        # 3: Orange (#FFDAB9, #FFA07A) - PeachPuff / LightSalmon
        
        colors = [
            ('#86B257', '#74A742'), # Green
            ('#87CEEB', '#4682B4'), # Blue
            ('#D8BFD8', '#DA70D6'), # Purple
            ('#FFDAB9', '#FFA07A'), # Orange
            ('#F0E68C', '#BDB76B'), # Khaki
            ('#E0FFFF', '#00CED1')  # LightCyan / DarkTurquoise
        ]
        
        fill, edge = colors[q % len(colors)]
        return fill, edge

    #%%  SECTION：----重载函数----
    def __str__(self)->str:
        """""
        input：None
        output：字符串表示
        influence：返回PolynomialCode的详细描述
        """""
        info = []
        info.append(f"Polynomial Code (Dimension: {self.dimension})")
        info.append(f"Lengths: {self.lengths}")
        info.append(f"Cell Size: {self.cell_size}")
        info.append(f"H_X Polynomials: {self.h_x}")
        info.append(f"H_Z Polynomials: {self.h_z}")
        return "\n".join(info)

    def _get_visual_pos_3d(self, tx, ty, tz, cx, cy, cz, Lx, Ly, Lz):
        dx = (tx - cx + Lx/2) % Lx - Lx/2
        dy = (ty - cy + Ly/2) % Ly - Ly/2
        dz = (tz - cz + Lz/2) % Lz - Lz/2
        return cx+dx, cy+dy, cz+dz

    #%%  SECTION：----绘图函数----
    def draw(self)->None:
        """""
        input：None
        output：None
        influence：绘制量子码的结构图，包括qubit的位置、正方晶格、以及X和Z stabilizer的结构。
        坐标系说明：
        - polynomial中的[x, y, z]分别对应[横坐标, 纵坐标, 层坐标]
        - 绘图时，x轴为水平向右，y轴为垂直向上(2D)或深度(3D)，z轴为垂直堆叠(3D)
        - Stabilizer绘制在坐标原点(0,0,...)所在的晶格中心
        """""
        try:
            import matplotlib.pyplot as plt
            import matplotlib.patches as patches
            from mpl_toolkits.mplot3d import Axes3D
        except ImportError:
            print("Matplotlib not found.")
            return
        
        # Stabilizer Colors
        color_x = '#DD705D'
        color_z = '#5874B5'

        ##  PART：----一维情况----
        if self.dimension == 1:
            fig, ax = plt.subplots(figsize=(max(6, self.lengths[0]), 3))
            L = self.lengths[0]
            
            ##  绘制正方晶格（一维为线）
            ax.plot([-0.5, L-0.5], [0, 0], color='lightgray', linestyle='-', linewidth=2, zorder=0)
            ax.plot([L-1, L-0.5], [0, 0], color='lightgray', linestyle='--', zorder=0)
            ax.plot([-0.5, 0], [0, 0], color='lightgray', linestyle='--', zorder=0)

            ##  绘制Cell标识和Qubit
            for x in range(L):
                # 绘制Cell椭圆
                ellipse = patches.Ellipse((x, 0), width=0.6, height=0.3, 
                                        edgecolor='gray', facecolor='none', 
                                        linestyle='--', alpha=0.3, zorder=1)
                ax.add_patch(ellipse)
                
                for i in range(self.cell_size):
                    dx, dy = self._get_qubit_offset(i, self.cell_size)
                    fill_c, edge_c = self._get_qubit_style(i, self.cell_size)
                    # 缩小Qubit圆圈 (radius 0.1 -> 0.05?)
                    circle = plt.Circle((x + dx, dy), 0.06, facecolor=fill_c, edgecolor=edge_c, linewidth=1.5, zorder=10)
                    ax.add_patch(circle)

            ##  绘制Stabilizers
            cx = 0
            
            # X Stabilizer (Red)
            check_vis_x = cx + 0.5
            check_vis_y = 0.2
            ax.scatter([check_vis_x], [check_vis_y], marker='s', color=color_x, s=60, label='X Check', zorder=15)

            for q, offsets in enumerate(self.h_x.polynomial):
                for pos in offsets:
                    tx = (cx + pos[0]) % L
                    qubit_offset_x, qubit_offset_y = self._get_qubit_offset(q, self.cell_size)
                    vis_tx = self._get_visual_pos(tx, cx, L)
                    ax.plot([check_vis_x, vis_tx + qubit_offset_x], [check_vis_y, qubit_offset_y], color=color_x, linestyle='-', alpha=0.6, zorder=5)
            
            # Z Stabilizer (Blue)
            check_vis_x = cx + 0.5
            check_vis_y = -0.2
            ax.scatter([check_vis_x], [check_vis_y], marker='s', color=color_z, s=60, label='Z Check', zorder=15)

            for q, offsets in enumerate(self.h_z.polynomial):
                for pos in offsets:
                    tx = (cx + pos[0]) % L
                    qubit_offset_x, qubit_offset_y = self._get_qubit_offset(q, self.cell_size)
                    vis_tx = self._get_visual_pos(tx, cx, L)
                    ax.plot([check_vis_x, vis_tx + qubit_offset_x], [check_vis_y, qubit_offset_y], color=color_z, linestyle='-', alpha=0.6, zorder=5)

            ax.set_xlim(-1, L)
            ax.set_ylim(-1, 1)
            ax.set_aspect('equal')
            ax.axis('off')
            plt.title(f"1D Polynomial Code (PBC), L={L}")
            plt.legend(loc='upper right')
            plt.show()

        ##  PART：----二维情况----
        elif self.dimension == 2:
            fig, ax = plt.subplots(figsize=(6, 6))
            Lx, Ly = self.lengths
            
            ##  绘制正方晶格
            for y in range(Ly):
                ax.plot([0, Lx-1], [y, y], color='lightgray', linestyle='-', zorder=0)
            for x in range(Lx):
                ax.plot([x, x], [0, Ly-1], color='lightgray', linestyle='-', zorder=0)

            ##  绘制Cell标识和Qubit
            for x in range(Lx):
                for y in range(Ly):
                    # 绘制Cell椭圆 (圆形)
                    ellipse = patches.Ellipse((x, y), width=0.6, height=0.6, 
                                            edgecolor='gray', facecolor='none', 
                                            linestyle='--', alpha=0.3, zorder=1)
                    ax.add_patch(ellipse)
                    
                    for i in range(self.cell_size):
                        dx, dy = self._get_qubit_offset(i, self.cell_size)
                        fill_c, edge_c = self._get_qubit_style(i, self.cell_size)
                        circle = plt.Circle((x+dx, y+dy), 0.06, facecolor=fill_c, edgecolor=edge_c, linewidth=1.5, zorder=10)
                        ax.add_patch(circle)

            ##  绘制Stabilizers
            cx, cy = 0, 0
            
            # X Lines (Red)
            # One check marker at center for X, slightly offset
            check_vis_x, check_vis_y = cx + 0.5 - 0.1, cy + 0.5 - 0.1
            ax.scatter([check_vis_x], [check_vis_y], marker='s', color=color_x, s=60, label='X Check', zorder=15)

            for q, offsets in enumerate(self.h_x.polynomial):
                for pos in offsets:
                    # pos[0] -> x (horizontal), pos[1] -> y (vertical)
                    tx, ty = (cx + pos[0]) % Lx, (cy + pos[1]) % Ly
                    
                    # Qubit offset
                    dx, dy = self._get_qubit_offset(q, self.cell_size)
                    
                    vis_tx = self._get_visual_pos(tx, cx, Lx)
                    vis_ty = self._get_visual_pos(ty, cy, Ly)
                    
                    ax.plot([check_vis_x, vis_tx+dx], [check_vis_y, vis_ty+dy], color=color_x, linestyle='-', alpha=0.6, linewidth=1.5, zorder=5)

            # Z Lines (Blue)
            # One check marker at center for Z, slightly offset
            check_vis_x, check_vis_y = cx + 0.5 + 0.1, cy + 0.5 + 0.1
            ax.scatter([check_vis_x], [check_vis_y], marker='s', color=color_z, s=60, label='Z Check', zorder=15)

            for q, offsets in enumerate(self.h_z.polynomial):
                for pos in offsets:
                    tx, ty = (cx + pos[0]) % Lx, (cy + pos[1]) % Ly
                    
                    dx, dy = self._get_qubit_offset(q, self.cell_size)
                    
                    vis_tx = self._get_visual_pos(tx, cx, Lx)
                    vis_ty = self._get_visual_pos(ty, cy, Ly)
                    
                    ax.plot([check_vis_x, vis_tx+dx], [check_vis_y, vis_ty+dy], color=color_z, linestyle='-', alpha=0.6, linewidth=1.5, zorder=5)


            ##  体现PBC
            for y in range(Ly):
                ax.plot([Lx-1, Lx-0.5], [y, y], color='lightgray', linestyle='--')
                ax.plot([-0.5, 0], [y, y], color='lightgray', linestyle='--')
            for x in range(Lx):
                ax.plot([x, x], [Ly-1, Ly-0.5], color='lightgray', linestyle='--')
                ax.plot([x, x], [-0.5, 0], color='lightgray', linestyle='--')

            ax.set_aspect('equal')
            ax.axis('off')
            plt.title(f"2D Polynomial Code (PBC), Size={Lx}x{Ly}")
            plt.legend(loc='upper right')
            plt.show()

        ##  PART：----三维情况----
        elif self.dimension == 3:
            Lx, Ly, Lz = self.lengths
            # Create subplots: one for each z-layer
            fig, axes = plt.subplots(1, Lz, figsize=(4*Lz, 5), constrained_layout=True)
            if Lz == 1:
                axes = [axes]
            
            # Stabilizer center at (0,0,0)
            cx, cy, cz = 0, 0, 0
            
            # Visual positions for X and Z checks (slightly offset)
            check_vis_x_X, check_vis_y_X = cx + 0.5 - 0.1, cy + 0.5 - 0.1
            check_vis_x_Z, check_vis_y_Z = cx + 0.5 + 0.1, cy + 0.5 + 0.1

            for z_layer in range(Lz):
                ax = axes[z_layer]
                ax.set_title(f"Layer z={z_layer}")
                
                ##  绘制正方晶格 (2D Grid for this layer)
                for y in range(Ly):
                    ax.plot([0, Lx-1], [y, y], color='lightgray', linestyle='-', zorder=0)
                for x in range(Lx):
                    ax.plot([x, x], [0, Ly-1], color='lightgray', linestyle='-', zorder=0)

                ##  绘制Cell标识和Qubit
                for x in range(Lx):
                    for y in range(Ly):
                        # Cell ellipse
                        ellipse = patches.Ellipse((x, y), width=0.6, height=0.6, 
                                                edgecolor='gray', facecolor='none', 
                                                linestyle='--', alpha=0.3, zorder=1)
                        ax.add_patch(ellipse)
                        
                        for i in range(self.cell_size):
                            dx, dy = self._get_qubit_offset(i, self.cell_size)
                            fill_c, edge_c = self._get_qubit_style(i, self.cell_size)
                            circle = plt.Circle((x+dx, y+dy), 0.06, facecolor=fill_c, edgecolor=edge_c, linewidth=1.5, zorder=10)
                            ax.add_patch(circle)

                ##  绘制Stabilizers
                
                # Check Markers
                # Only solid if on this layer, else ghost
                if z_layer == cz:
                    ax.scatter([check_vis_x_X], [check_vis_y_X], marker='s', color=color_x, s=60, label='X Check', zorder=15)
                    ax.scatter([check_vis_x_Z], [check_vis_y_Z], marker='s', color=color_z, s=60, label='Z Check', zorder=15)
                else:
                    # Ghost markers
                    ax.scatter([check_vis_x_X], [check_vis_y_X], marker='s', color=color_x, s=60, alpha=0.15, zorder=5)
                    ax.scatter([check_vis_x_Z], [check_vis_y_Z], marker='s', color=color_z, s=60, alpha=0.15, zorder=5)

                # X Connections
                for q, offsets in enumerate(self.h_x.polynomial):
                    for pos in offsets:
                        tx, ty, tz = (cx + pos[0]) % Lx, (cy + pos[1]) % Ly, (cz + pos[2]) % Lz
                        
                        if tz == z_layer:
                            dx, dy = self._get_qubit_offset(q, self.cell_size)
                            vis_tx = self._get_visual_pos(tx, cx, Lx)
                            vis_ty = self._get_visual_pos(ty, cy, Ly)
                            ax.plot([check_vis_x_X, vis_tx+dx], [check_vis_y_X, vis_ty+dy], color=color_x, linestyle='-', alpha=0.6, linewidth=1.5, zorder=10)

                # Z Connections
                for q, offsets in enumerate(self.h_z.polynomial):
                    for pos in offsets:
                        tx, ty, tz = (cx + pos[0]) % Lx, (cy + pos[1]) % Ly, (cz + pos[2]) % Lz
                        
                        if tz == z_layer:
                            dx, dy = self._get_qubit_offset(q, self.cell_size)
                            vis_tx = self._get_visual_pos(tx, cx, Lx)
                            vis_ty = self._get_visual_pos(ty, cy, Ly)
                            ax.plot([check_vis_x_Z, vis_tx+dx], [check_vis_y_Z, vis_ty+dy], color=color_z, linestyle='-', alpha=0.6, linewidth=1.5, zorder=10)

                ##  体现PBC
                for y in range(Ly):
                    ax.plot([Lx-1, Lx-0.5], [y, y], color='lightgray', linestyle='--')
                    ax.plot([-0.5, 0], [y, y], color='lightgray', linestyle='--')
                for x in range(Lx):
                    ax.plot([x, x], [Ly-1, Ly-0.5], color='lightgray', linestyle='--')
                    ax.plot([x, x], [-0.5, 0], color='lightgray', linestyle='--')

                ax.set_aspect('equal')
                ax.axis('off')
                
            # Global Legend (taking handles from first plot if available, or creating custom)
            handles = [
                plt.Line2D([0], [0], marker='s', color='w', markerfacecolor=color_x, markersize=10, label='X Check'),
                plt.Line2D([0], [0], marker='s', color='w', markerfacecolor=color_z, markersize=10, label='Z Check')
            ]
            fig.legend(handles=handles, loc='upper right')
            plt.suptitle(f"3D Polynomial Code Layers (L={Lx}x{Ly}x{Lz})")
            plt.show()


# %%  CHAPTER：====Majorana finite projective geometry code====
class MajoranaPolynomialCode(MajoranaCSSCode,PolynomialCode):

    ##  SECTION：----构造方法----
    def __init__(self,h_x:PolynomialMatrix,h_z:PolynomialMatrix):
        PolynomialCode.__init__(self,h_x,h_z)
        generators_x=[]
        generators_z=[]
        for i in range(len(self.H)):
            generators_x.append(MajoranaOperator.HermitianOperatorFromOccupy(np.where(self.H_X[i]!=0)[0],[]))
            generators_z.append(MajoranaOperator.HermitianOperatorFromOccupy([],np.where(self.H_Z[i]!=0)[0]))
        MajoranaCSSCode.__init__(self,generators_x,generators_z,self.H.shape[1])


# %%  CHAPTER：====Pauli finite projective geometry code====
class PauliPolynomialCode(PauliCSSCode,PolynomialCode):

    ##  SECTION：----构造方法----
    def __init__(self,h_x:PolynomialMatrix,h_z:PolynomialMatrix):
        PolynomialCode.__init__(self,h_x,h_z)
        generators_x=[]
        generators_z=[]
        for i in range(len(self.H)):
            generators_x.append(PauliOperator.HermitianOperatorFromOccupy(np.where(self.H_X[i]!=0)[0],[]))
            generators_z.append(PauliOperator.HermitianOperatorFromOccupy([],np.where(self.H_Z[i]!=0)[0]))
        PauliCSSCode.__init__(self,generators_x,generators_z,self.H.shape[1])