# -*- coding: utf-8 -*-

"""
    Licensed to the Apache Software Foundation (ASF) under one
    or more contributor license agreements.  See the NOTICE file
    distributed with this work for additional information
    regarding copyright ownership.  The ASF licenses this file
    to you under the Apache License, Version 2.0 (the
    "License"); you may not use this file except in compliance
    with the License.  You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

    Unless required by applicable law or agreed to in writing,
    software distributed under the License is distributed on an
    "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
    KIND, either express or implied.  See the License for the
    specific language governing permissions and limitations
    under the License.
"""

from qiskit.quantum_info.operators.channel import Choi, PTM, Kraus, Chi, SuperOp

import numpy as np

from qat.comm.quops.ttypes import QuantumChannel, RepresentationType
from qat.comm.datamodel.ttypes import Matrix, ComplexNumber


def array_to_matrix(array):
    """
    Transform a two dimmentional numpy array to a myqlm Matrix.

    Args:
        array: (ndarray) a two dimmentional numpy array

    Returns:
        (Matrix): a myqlm Matrix
    """

    assert len(array.shape) == 2, "The array must be two dimmentional"
    data = []
    for arr in array:
        for elem in arr:
            data.append(ComplexNumber(np.real(elem), np.imag(elem)))
    matri = Matrix(array.shape[0], array.shape[1], data)
    return matri


def qiskit_to_qchannel(representation):
    """
    Create a myqlm representation of quantum channel from a qiskit representation
    of a quantum channel.

    Args:
        representation: (Kraus|Choi|Chi|SuperOp|PTM) qiskit representation of a quantum channel.

    Returns:
        (QuantumChannel): myqlm representation of a quantum channel.
    """

    qchannel = None
    qiskit_data = representation.data
    # Find what representation it is.
    # Then create the corresponding matrix (kraus_ops|basis|matrix)from the data
    # of the representation.
    # Finally, create the QuantumChannel with the RepresentationType, the arity
    # (got from the qiskit representation) and the matrix.
    if isinstance(representation, Kraus):
        kraus_ops = []
        for arr in qiskit_data:
            kraus_ops.append(array_to_matrix(arr))
        qchannel = QuantumChannel(
            representation=RepresentationType.KRAUS,
            arity=representation.num_qubits,
            kraus_ops=kraus_ops)
    elif isinstance(representation, Chi):
        basis = []
        basis.append(array_to_matrix(qiskit_data))
        qchannel = QuantumChannel(
            representation=RepresentationType.CHI,
            arity=representation.num_qubits,
            basis=basis)
    elif isinstance(representation, SuperOp):
        basis = []
        basis.append(array_to_matrix(qiskit_data))
        qchannel = QuantumChannel(
            representation=RepresentationType.SUPEROP,
            arity=representation.num_qubits,
            basis=basis)
    elif isinstance(representation, PTM):
        matri = array_to_matrix(qiskit_data)
        qchannel = QuantumChannel(
            representation=RepresentationType.PTM,
            arity=representation.num_qubits,
            matrix=matri)
    elif isinstance(representation, Choi):
        matri = array_to_matrix(qiskit_data)
        qchannel = QuantumChannel(
            representation=RepresentationType.CHOI,
            arity=representation.num_qubits,
            matrix=matri)
    return qchannel


def qchannel_to_qiskit(representation):
    """
    Create a qiskit representation of quantum channel from a myqlm representation
    of a quantum channel.

    Args:
        representation: (QuantumChannel) myqlm representation of a quantum channel.

    Returns:
        (Kraus|Choi|Chi|SuperOp|PTM): qiskit representation of a quantum channel.
    """

    rep = representation.representation
    # Find what representation it is.
    # Then create the corresponding matrix and shape it like qiskit is expecting it.
    # Finally, create the qiskit representation from that matrix.
    if rep in (RepresentationType.PTM, RepresentationType.CHOI):
        matri = representation.matrix
        data_re = []
        data_im = []
        for i in range(matri.nRows):
            for j in range(matri.nCols):
                data_re.append(matri.data[i * matri.nRows + j].re + 0.j)
                data_im.append(matri.data[i * matri.nRows + j].im)
        data = np.array(data_re)
        data.imag = np.array(data_im)
        data = data.reshape((matri.nRows, matri.nCols))
        return PTM(data) if (rep == RepresentationType.PTM) else Choi(data)
    if rep in (RepresentationType.CHI, RepresentationType.SUPEROP):
        final_data = []
        for matri in representation.basis:
            data_re = []
            data_im = []
            for i in range(matri.nRows):
                for j in range(matri.nCols):
                    data_re.append(matri.data[i * matri.nRows + j].re + 0.j)
                    data_im.append(matri.data[i * matri.nRows + j].im)
            data = np.array(data_re)
            data.imag = np.array(data_im)
            data = data.reshape((matri.nRows, matri.nCols))
            final_data.append(data)
        if rep == RepresentationType.CHI:
            return Chi(final_data) if len(final_data) > 1 else Chi(final_data[0])
        return SuperOp(final_data) if len(final_data) > 1 else SuperOp(final_data[0])
    if rep == RepresentationType.KRAUS:
        final_data = []
        for matri in representation.kraus_ops:
            data_re = []
            data_im = []
            for i in range(matri.nRows):
                for j in range(matri.nCols):
                    data_re.append(matri.data[i * matri.nRows + j].re + 0.j)
                    data_im.append(matri.data[i * matri.nRows + j].im)
            data = np.array(data_re)
            data.imag = np.array(data_im)
            data = data.reshape((matri.nRows, matri.nCols))
            final_data.append(data)
        return Kraus(final_data)
    return None

[((46, 12, 46, 56), 'qat.comm.datamodel.ttypes.Matrix', 'Matrix', ({(46, 19, 46, 33): 'array.shape[0]', (46, 35, 46, 49): 'array.shape[1]', (46, 51, 46, 55): 'data'}, {}), '(array.shape[0], array.shape[1], data)', False, 'from qat.comm.datamodel.ttypes import Matrix, ComplexNumber\n'), ((73, 19, 76, 32), 'qat.comm.quops.ttypes.QuantumChannel', 'QuantumChannel', (), '', False, 'from qat.comm.quops.ttypes import QuantumChannel, RepresentationType\n'), ((130, 15, 130, 32), 'numpy.array', 'np.array', ({(130, 24, 130, 31): 'data_re'}, {}), '(data_re)', True, 'import numpy as np\n'), ((131, 20, 131, 37), 'numpy.array', 'np.array', ({(131, 29, 131, 36): 'data_im'}, {}), '(data_im)', True, 'import numpy as np\n'), ((163, 15, 163, 32), 'qiskit.quantum_info.operators.channel.Kraus', 'Kraus', ({(163, 21, 163, 31): 'final_data'}, {}), '(final_data)', False, 'from qiskit.quantum_info.operators.channel import Choi, PTM, Kraus, Chi, SuperOp\n'), ((80, 19, 83, 24), 'qat.comm.quops.ttypes.QuantumChannel', 'QuantumChannel', (), '', False, 'from qat.comm.quops.ttypes import QuantumChannel, RepresentationType\n'), ((133, 15, 133, 24), 'qiskit.quantum_info.operators.channel.PTM', 'PTM', ({(133, 19, 133, 23): 'data'}, {}), '(data)', False, 'from qiskit.quantum_info.operators.channel import Choi, PTM, Kraus, Chi, SuperOp\n'), ((133, 65, 133, 75), 'qiskit.quantum_info.operators.channel.Choi', 'Choi', ({(133, 70, 133, 74): 'data'}, {}), '(data)', False, 'from qiskit.quantum_info.operators.channel import Choi, PTM, Kraus, Chi, SuperOp\n'), ((143, 19, 143, 36), 'numpy.array', 'np.array', ({(143, 28, 143, 35): 'data_re'}, {}), '(data_re)', True, 'import numpy as np\n'), ((144, 24, 144, 41), 'numpy.array', 'np.array', ({(144, 33, 144, 40): 'data_im'}, {}), '(data_im)', True, 'import numpy as np\n'), ((149, 15, 149, 34), 'qiskit.quantum_info.operators.channel.SuperOp', 'SuperOp', ({(149, 23, 149, 33): 'final_data'}, {}), '(final_data)', False, 'from qiskit.quantum_info.operators.channel import Choi, PTM, Kraus, Chi, SuperOp\n'), ((149, 63, 149, 85), 'qiskit.quantum_info.operators.channel.SuperOp', 'SuperOp', ({(149, 71, 149, 84): 'final_data[0]'}, {}), '(final_data[0])', False, 'from qiskit.quantum_info.operators.channel import Choi, PTM, Kraus, Chi, SuperOp\n'), ((159, 19, 159, 36), 'numpy.array', 'np.array', ({(159, 28, 159, 35): 'data_re'}, {}), '(data_re)', True, 'import numpy as np\n'), ((160, 24, 160, 41), 'numpy.array', 'np.array', ({(160, 33, 160, 40): 'data_im'}, {}), '(data_im)', True, 'import numpy as np\n'), ((87, 19, 90, 24), 'qat.comm.quops.ttypes.QuantumChannel', 'QuantumChannel', (), '', False, 'from qat.comm.quops.ttypes import QuantumChannel, RepresentationType\n'), ((148, 19, 148, 34), 'qiskit.quantum_info.operators.channel.Chi', 'Chi', ({(148, 23, 148, 33): 'final_data'}, {}), '(final_data)', False, 'from qiskit.quantum_info.operators.channel import Choi, PTM, Kraus, Chi, SuperOp\n'), ((148, 63, 148, 81), 'qiskit.quantum_info.operators.channel.Chi', 'Chi', ({(148, 67, 148, 80): 'final_data[0]'}, {}), '(final_data[0])', False, 'from qiskit.quantum_info.operators.channel import Choi, PTM, Kraus, Chi, SuperOp\n'), ((45, 38, 45, 51), 'numpy.real', 'np.real', ({(45, 46, 45, 50): 'elem'}, {}), '(elem)', True, 'import numpy as np\n'), ((45, 53, 45, 66), 'numpy.imag', 'np.imag', ({(45, 61, 45, 65): 'elem'}, {}), '(elem)', True, 'import numpy as np\n'), ((93, 19, 96, 25), 'qat.comm.quops.ttypes.QuantumChannel', 'QuantumChannel', (), '', False, 'from qat.comm.quops.ttypes import QuantumChannel, RepresentationType\n'), ((99, 19, 102, 25), 'qat.comm.quops.ttypes.QuantumChannel', 'QuantumChannel', (), '', False, 'from qat.comm.quops.ttypes import QuantumChannel, RepresentationType\n')]