from qiskit import QuantumCircuit
from qiskit.circuit import CircuitInstruction
from qiskit.quantum_info import Operator
from ..mps import MatrixProductState
from ..mpo import MatrixProductOperator
from ..tn import TensorNetwork
import numpy as np
from numpy import ndarray
from scipy.linalg import qr
import copy

class NoiseData:

    def __init__(self, calibration_data: dict):
        """Constructor
        
        Args:
            calibration_data: Calibration data for the quantum device"
        """
        self.calibration_data = calibration_data

class OperatorSum:

    def __init__(self, operators: list[ndarray], weights: list[float]):
        """Constructor
        
        Args:
            operators: Matrix representation of the operators in the Kraus channel
            weights: Weights for each operator
        """
        assert len(operators) == len(weights)
        self.num_ops = len(operators)
        self.dimension = len(operators[0])
        self.operators = operators
        self.weights = weights
        self.op_dict = {}
        for idx in range(len(operators)):
            op = operators[idx]
            weight = weights[idx]
            weighted_op = np.sqrt(weight) * op 
            self.op_dict[str(idx)] = weighted_op


    def to_superoperator(self) -> ndarray:
        """Convert a Kraus channel to a superoperator"""
        super_op = np.zeros((self.dimension**2, self.dimension**2))
        for op in self.op_dict.values():
            op_conj = op.conj()
            temp_super_op = np.kron(op_conj, op)
            super_op = super_op + temp_super_op
        return super_op


class TNQuantumErrorMitigation:

    def __init__(self, circuit: QuantumCircuit, max_bond_dimension: int, calibration_data: dict):
        """Constructor
        
        Args:
            circuit: The input quantum circuit
            max_bond_dimension: Maximum allowed bond dimension for TNQEM MPO
            calibration_data: The calibration data for the quantum device
        """
        self.circuit = circuit
        self.num_qubits = circuit.num_qubits
        self.max_bond_dimension = max_bond_dimension
        self.calibration_data = calibration_data
        self.noise_model = NoiseData(calibration_data)

    
    def bitstring_to_vectorised_dm_mps(self, bitstring: str) -> MatrixProductState:
        """Convert a bitstring into an MPS representing the vectorised density matrix
        
        Args:
            bitstring: Input bitstring

        Returns:
            A MatrixProductState
        """
        if len(bitstring) == 1:
            if bitstring[0] == "0":
                first_array = np.array([1,0,0,0], dtype=complex).reshape((4,))
            else:
                first_array = np.array([0,0,0,1], dtype=complex).reshape((4,))
            mps = MatrixProductState.from_arrays([first_array])
            return mps

        arrays = []

        first_bit = bitstring[0]
        if first_bit == "0":
            first_array = np.array([1,0,0,0], dtype=complex).reshape((1,4))
        else:
            first_array = np.array([0,0,0,1], dtype=complex).reshape((1,4))
        arrays.append(first_array)
        
        for bit in bitstring[1:-1]:
            if bit == "0":
                array = np.array([1,0,0,0], dtype=complex).reshape((1,1,4))
            else:
                array = np.array([0,0,0,1], dtype=complex).reshape((1,1,4))
            arrays.append(array)
        
        last_bit = bitstring[-1]
        if last_bit == "0":
            last_array = np.array([1,0,0,0], dtype=complex).reshape((1,4))
        else:
            last_array = np.array([0,0,0,1], dtype=complex).reshape((1,4))
        arrays.append(last_array)

        mps = MatrixProductState.from_arrays(arrays)
        return mps 
    
    def circuit_to_supercircuit(self) -> list[tuple[ndarray, list[int]]]:
        """Convert a set of circuit instructions to a list of superoperators"""
        supercircuit = []
        for inst in self.circuit.data:
            qidxs = [inst.qubits[idx]._index for idx in range(len(inst.qubits))]
            op = Operator(inst.operation).reverse_qargs().data
            operator_sum = OperatorSum([op], [1.0])
            superop = operator_sum.to_superoperator()
            supercircuit.append((superop, qidxs))
        return supercircuit
    
    def gate_to_noisegate(self, gate: CircuitInstruction) -> ndarray: # type: ignore
        """Convert a circuit instruction into a superoperator representing the noisy gate"""
        noisegates = []
        num_qubits = gate.operation.num_qubits
        if num_qubits == 1:
            op = Operator(gate.operation).reverse_qargs().data
            operator_sum = OperatorSum([op], [1.0])
            superop = operator_sum.to_superoperator()
            noisegates.append(superop)
        elif num_qubits == 2:
            # The expected gate
            op = Operator(gate.operation).reverse_qargs().data
            operator_sum = OperatorSum([op], [1.0])
            superop = operator_sum.to_superoperator()
            noisegates.append(superop)

            # Depolarising noise

        # Combine
        noisegates = noisegates[::-1]
        for idx in range(len(noisegates)-1):
            noisegates[idx+1] = noisegates[idx] @ noisegates[idx+1]
        noisegate = noisegates[-1]

        return noisegate
        
    def circuit_to_noisecircuit(self) -> list[tuple[ndarray,list[int]]]:
        """Convert a set of circuit instructions to a dictionary of superoperators for the noisy circuit"""
        noisecircuit = []
        for inst in self.circuit.data:
            qidxs = [inst.qubits[idx]._index for idx in range(len(inst.qubits))]
            noisegate = self.gate_to_noisegate(inst)
            noisecircuit.append((noisegate, qidxs))
        return noisecircuit

    def invert_noisecircuit(self, noisecircuit: list[tuple[ndarray,list[int]]]) -> list[tuple[ndarray,list[int]]]:
        """Invert a noisy circuit"""
        inv_noisecircuit = []
        for noisegate, qidxs in noisecircuit:
            inv_noisegate = np.linalg.inv(noisegate)
            inv_noisecircuit.append((inv_noisegate, qidxs))
        return inv_noisecircuit
    
    def superoperator_to_submpo(self, superop: ndarray):
        """Convert a superoperator to an MPO"""
        dim = len(superop)
        if dim == 4:
            mpo = MatrixProductOperator.from_arrays([superop])
        elif dim == 16:
            superop = np.reshape(superop, (4,4,4,4))
            superop = np.moveaxis(superop, [0,1,2,3], [0,2,1,3])
            superop = np.reshape(superop, (16,16))
            q, r = qr(superop)
            first_array = np.reshape(q, (4,4,16))
            first_array = np.moveaxis(first_array, [0,1,2], [1,2,0])
            second_array = np.reshape(r, (16,4,4))
            mpo = MatrixProductOperator.from_arrays([first_array, second_array])
        return mpo
    
    def build_identity_mpo(self) -> MatrixProductOperator:
        """Build identity MPO with physical dim 4"""
        id_array = np.eye(4)
        first_array = np.reshape(id_array, (1,4,4))
        middle_arrays = []
        for _ in range(self.num_qubits - 2):
            middle_array = np.reshape(id_array, (1,1,4,4))
            middle_arrays.append(middle_array)
        last_array = np.reshape(id_array, (1,4,4))
        all_arrays = [first_array] + middle_arrays + [last_array]
        mpo = MatrixProductOperator.from_arrays(all_arrays)
        return mpo
    
    def contract_one_layer(self, mpo: MatrixProductOperator, submpo_right: MatrixProductOperator, submpo_left: MatrixProductOperator, qidxs: list[int]) -> MatrixProductOperator:
        """Contract one layer of the TNQEM MPO construction"""
        mpo_copy = copy.deepcopy(mpo)
        submpo_right_copy = copy.deepcopy(submpo_right)
        submpo_left_copy = copy.deepcopy(submpo_left)

        if len(qidxs) == 1:
            submpo_left_copy.set_default_indices("B", "A", "C")
            mpo_copy.set_default_indices("D", "C", "E")
            submpo_right_copy.set_default_indices("F", "E", "G")

            site_tensor = mpo_copy.tensors[qidxs[0]-1]
            tn = TensorNetwork([site_tensor] + submpo_left_copy.tensors + submpo_right_copy.tensors)
            updated_site_tensor = tn.contract_entire_network()
            if qidxs[0] == 1:
                index_order = ["D1", "G1", "A1"]
            elif qidxs[0] == mpo.num_sites:
                index_order = [f"D{mpo.num_sites-1}", f"G{mpo.num_sites}", f"A{mpo.num_sites}"]
            else:
                index_order = [f"D{qidxs[0]-1}", f"D{qidxs[0]}", f"G{qidxs[0]}", f"A{qidxs[0]}"]
            updated_site_tensor.reorder_indices(index_order)
            updated_array = updated_site_tensor.data 
            new_arrays = [mpo_copy.tensors[idx].data for idx in range(qidxs[0]-1)] + [updated_array] + [mpo_copy.tensors[idx].data for idx in range(qidxs[0], mpo.num_sites)]
            new_mpo = MatrixProductOperator.from_arrays(new_arrays)
            return new_mpo

        elif len(qidxs) == 2:
            submpo_left_copy.set_default_indices("B", "A", "C")
            mpo_copy.set_default_indices("D", "C", "E")
            submpo_right_copy.set_default_indices("F", "E", "G")

            
            pass

        return


    def build_tnqem_mpo(self) -> MatrixProductOperator:
        """Build the TNQEM MPO"""
        supercircuit = self.circuit_to_supercircuit()
        noisecircuit = self.circuit_to_noisecircuit()
        inv_noisecircuit = self.invert_noisecircuit(noisecircuit)
        num_gates = len(supercircuit)

        mpo = self.build_identity_mpo()
        for idx in range(num_gates):
            gate, qidxs = supercircuit[idx]
            inv_noisegate, qidxs2 = inv_noisecircuit[idx]
            assert qidxs == qidxs2

            qidxs = [qidx+1 for qidx in qidxs]
            gate_mpo = self.superoperator_to_submpo(gate)
            inv_noisegate_mpo = self.superoperator_to_submpo(inv_noisegate)
            # mpo = self.contract_one_layer(mpo, gate_mpo, inv_noisegate_mpo, qidxs)
            mpo = mpo.contract_sub_mpo(gate_mpo, qidxs, None)
            mpo = mpo.contract_sub_mpo(inv_noisegate_mpo, qidxs, None, False)
            mpo.compress(self.max_bond_dimension)
            mpo.update_bond_information()
        
        return mpo
    
    def sample_from_vectorised_dm_mps(self, mps: MatrixProductState, num_bitstrings: int = 1) -> list[str]:
        """Sample from the output MPS of the TNQEM scheme"""
        samples = {}
        prefix_prob_dict = {}  # {prefix : (prob0, prob1)}
        sample_prob_dict = {}  # {bitstring : probabiity}

        samples_collected = 0
        mps_copy = copy.deepcopy(mps)
        while samples_collected < num_bitstrings:
            prob_existing_sample = np.sum(list(sample_prob_dict.values()))
            choose_existing_sample = np.random.choice(
                ["y", "n"], p=[prob_existing_sample, 1 - prob_existing_sample]
            )
            if choose_existing_sample == "y":
                probs_norm = np.sum(list(sample_prob_dict.values()))
                normalised_probs = [x / probs_norm for x in sample_prob_dict.values()]
                bitstring = np.random.choice(
                    list(sample_prob_dict.keys()), p=normalised_probs
                )
                samples[bitstring] += 1
                samples_collected += 1
                continue
            bitstring = ""
            current_mps = copy.deepcopy(mps_copy)
            current_mpo = current_mps.form_density_operator()
            total_prob = 1.0
            temp_prob = 1.0
            for site in range(1, mps.num_sites + 1):
                if bitstring in prefix_prob_dict:
                    prob0, prob1 = prefix_prob_dict[bitstring]
                else:
                    num_sites_to_trace = len(bitstring) - (
                        mps.num_sites - current_mpo.num_sites
                    )
                    if len(bitstring) > 0 and num_sites_to_trace > 0:
                        sub_mpo_bitstring = bitstring[-num_sites_to_trace:]
                        sub_mps = self.bitstring_to_vectorised_dm_mps(sub_mpo_bitstring)
                        sub_mpo = sub_mps.form_density_operator()
                        current_mpo = current_mpo.contract_sub_mpo(
                            sub_mpo, list(range(1, num_sites_to_trace + 1))
                        )
                        current_mpo = current_mpo.partial_trace(
                            list(range(1, num_sites_to_trace + 1)),
                            set_default_indices=True,
                        )
                        current_mpo.multiply_by_constant(1 / temp_prob**2)
                        current_mpo.indices = current_mpo.get_all_indices()
                        temp_prob = 1.0
                    if site != mps.num_sites:
                        site_rdm = (
                            current_mpo.partial_trace(
                                list(range(2, current_mpo.num_sites + 1))
                            )
                            .tensors[0]
                            .data.todense()
                        )
                    else:
                        site_rdm = current_mpo.tensors[0].data.todense()
                    prob0 = np.sqrt(min(site_rdm[0, 0].real, 1.0))
                    prob1 = 1.0 - prob0

                site_bit = np.random.choice(["0", "1"], p=[prob0, prob1])
                bitstring += site_bit
                if site_bit == "0":
                    temp_prob *= prob0
                    total_prob *= prob0
                else:
                    temp_prob *= prob1
                    total_prob *= prob1
                if bitstring[:-1] not in prefix_prob_dict:
                    prefix_prob_dict[bitstring[:-1]] = (prob0, prob1)
            if bitstring not in sample_prob_dict:
                if bitstring in samples:
                    samples[bitstring] += 1
                else:
                    samples[bitstring] = 1
                samples_collected += 1
                sample_prob_dict[bitstring] = total_prob
        
        return samples

    def run_single_shot_tnqem(self, input_bitstring: str, num_samples: int=1) -> list[str]:
        """Perform single shot TNQEM"""
        tnqem_mpo = self.build_tnqem_mpo()
        input_mps = self.bitstring_to_vectorised_dm_mps(input_bitstring)
        output_mps = input_mps.apply_mpo(tnqem_mpo)
        samples = self.sample_from_vectorised_dm_mps(output_mps, num_samples)
        return samples
