import numpy as np
from extendedstim.Circuit.Circuit import Circuit
from extendedstim.Code.QuantumCode.MajoranaCSSCode import MajoranaCSSCode
from extendedstim.Code.QuantumCode.PauliCSSCode import PauliCSSCode
from extendedstim.Physics.PauliOperator import PauliOperator


#%%  KEY：将Majorana CSS code转换为现象级噪声下的测试线路
def MajoranaCSSCode2PhenomenologicalCircuit(code,p_noise,p_measure,cycle_number:int):

    ##  ===数据预处理===
    assert isinstance(code,MajoranaCSSCode)
    stabilizers_x=code.generators_x
    stabilizers_z=code.generators_z
    logical_x=code.logical_operators_x
    logical_z=code.logical_operators_z

    ##  获取logical operators
    logical_occupy=[]
    for i in range(len(logical_x)):
        logical_occupy.append(1j*logical_x[i]@logical_z[i])

    majorana_number=code.physical_number
    stabilizer_number = len(stabilizers_x) + len(stabilizers_z)

    ##  ===生成线路===
    circuit_z = Circuit()
    circuit_z.append({"name":"FORCE","majorana_state":np.append(logical_occupy,np.append(stabilizers_x, stabilizers_z)),
                      'pauli_state':[PauliOperator([],[],1)]*(stabilizer_number+len(logical_occupy)),
                      'pauli_number':0,'majorana_number':majorana_number})

    ##  第一轮测量假设完美的初始化
    observable_include= []  # 记录可观测量的索引
    for i,logical_operator in enumerate(logical_occupy):
        circuit_z.append({"name":"MPP","target":logical_operator,"index":i})
        observable_include.append(len(circuit_z.measurements)-1)
    for i,stabilizer in enumerate(stabilizers_x):
        circuit_z.append({"name":"MPP","target":stabilizer,"index":i+len(logical_occupy)})
    for i,stabilizer in enumerate(stabilizers_z):
        circuit_z.append({"name":"MPP","target":stabilizer,"index":i+len(logical_occupy)+len(stabilizers_x)})

    ##  循环多轮，多轮测量错误与量子噪声信道
    for _ in range(cycle_number):
        for i in range(majorana_number):
            circuit_z.append({"name":"FDEPOLARIZE1","target":i,"index":i,"p":p_noise})
        for i,stabilizer in enumerate(stabilizers_x):
            circuit_z.append({"name":"MPP","target":stabilizer,"index":i+len(logical_occupy),'p':p_measure})
        for i,stabilizer in enumerate(stabilizers_z):
            circuit_z.append({"name":"MPP","target":stabilizer,"index":i+len(logical_occupy)+len(stabilizers_x),'p':p_measure})
        for i in range(stabilizer_number):
            circuit_z.append({"name":"DETECTOR","target":[-i - 1, -i - stabilizer_number-1]})

    ##  最后一轮测量假设是没有噪声的
    for i,stabilizer in enumerate(stabilizers_x):
        circuit_z.append({"name":"MPP","target":stabilizer,"index":i+len(logical_occupy)})
    for i,stabilizer in enumerate(stabilizers_z):
        circuit_z.append({"name":"MPP","target":stabilizer,"index":i+len(logical_occupy)+len(stabilizers_x)})

    ##  测量逻辑算符
    for i,logical_operator in enumerate(logical_occupy):
        circuit_z.append({"name":"MPP","target":logical_operator,"index":i})
        circuit_z.append({"name":"OBSERVABLE_INCLUDE","target":[len(circuit_z.measurements) - 1, observable_include[i]]})

    ##  ===返回线路===
    return circuit_z


#%%  KEY：将Pauli CSS code转换为现象级噪声下的测试线路
def PauliCSSCode2PhenomenologicalCircuit(code,p_noise,p_measure,cycle_number:int):
    pass


#%%  KEY：将Majorana CSS code转换为电路级噪声下的测试线路
def MajoranaCSSCode2CircuitLevelCircuit(code,p_noise,p_measure,cycle_number:int):
    pass


#%%  KEY：将Pauli CSS code转换为电路级噪声下的测试线路
def PauliCSSCode2CircuitLevelCircuit(code,p_noise,p_measure,cycle_number:int):
    pass


#%%  KEY：将Majorana CSS code转换为码能力下的测试线路
def MajoranaCSSCode2CodeCapacityCircuit(code,p_noise,p_measure,cycle_number:int):
    pass


#%%  KEY：将Pauli CSS code转换为码能力下的测试线路
def PauliCSSCode2CodeCapacityCircuit(code,p_noise,p_measure,cycle_number:int):
    pass


def Code2Circuit(code,p_noise,p_measure,noise_model,cycle_number):
    if noise_model=='phenomenological':
        if isinstance(code,MajoranaCSSCode):
            return MajoranaCSSCode2PhenomenologicalCircuit(code,p_noise,p_measure,cycle_number)
        elif isinstance(code,PauliCSSCode):
            return PauliCSSCode2PhenomenologicalCircuit(code,p_noise,p_measure,cycle_number)
        else:
            raise NotImplementedError
    elif noise_model=='circuit-level':
        if isinstance(code,MajoranaCSSCode):
            return MajoranaCSSCode2CircuitLevelCircuit(code,p_noise,p_measure,cycle_number)
        elif isinstance(code,PauliCSSCode):
            return PauliCSSCode2CircuitLevelCircuit(code,p_noise,p_measure,cycle_number)
        else:
            raise NotImplementedError
    elif noise_model=='code-capacity':
        if isinstance(code,MajoranaCSSCode):
            return MajoranaCSSCode2CodeCapacityCircuit(code,p_noise,p_measure,cycle_number)
        elif isinstance(code,PauliCSSCode):
            return PauliCSSCode2CodeCapacityCircuit(code,p_noise,p_measure,cycle_number)
        else:
            raise NotImplementedError
    else:
        raise ValueError('noise_model must be phenomenological, circuit-level, or code-capacity')
