# Copyright 2019-2020 Cambridge Quantum Computing
#
# Licensed under a Non-Commercial Use Software Licence (the "Licence");
# you may not use this file except in compliance with the Licence.
# You may obtain a copy of the Licence in the LICENCE file accompanying
# these documents or at:
#
#     https://cqcl.github.io/pytket/build/html/licence.html
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the Licence is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the Licence for the specific language governing permissions and
# limitations under the Licence, but note it is strictly for non-commercial use.

"""Methods to allow conversion between pyQuil and t|ket> data types
"""

from collections import namedtuple
import math
from typing import Union

from pyquil import Program
from pyquil.api import QuantumComputer
from pyquil.device import ISA, AbstractDevice, isa_to_graph, specs_from_graph
from pyquil.quilatom import (
    Qubit as Qubit_,
    Expression,
    MemoryReference,
    quil_sin,
    quil_cos,
    Add as Add_,
    Sub,
    Mul as Mul_,
    Div,
    Pow as Pow_,
    Function as Function_,
)
from pyquil.quilbase import Declare, Gate, Halt, Measurement, Pragma
from sympy import pi, Expr, Symbol, sin, cos, Number, Add, Mul, Pow

from pytket.circuit import Circuit, Node, Op, OpType, Qubit
from pytket.device import Device, GateError, GateErrorContainer
from pytket.routing import Architecture

_known_quil_gate = {
    "X": OpType.X,
    "Y": OpType.Y,
    "Z": OpType.Z,
    "H": OpType.H,
    "S": OpType.S,
    "T": OpType.T,
    "RX": OpType.Rx,
    "RY": OpType.Ry,
    "RZ": OpType.Rz,
    "CZ": OpType.CZ,
    "CNOT": OpType.CX,
    "CCNOT": OpType.CCX,
    "CPHASE": OpType.CU1,
    "PHASE": OpType.U1,
    "SWAP": OpType.SWAP,
    "XY": OpType.ISWAP,
}

_known_quil_gate_rev = {v: k for k, v in _known_quil_gate.items()}


def param_to_pyquil(p: Union[float, Expr]) -> Union[float, Expression]:
    ppi = p * pi
    if len(ppi.free_symbols) == 0:
        return float(ppi.evalf())
    else:

        def to_pyquil(e: Expr) -> Union[float, Expression]:
            if isinstance(e, Number):
                return float(e)
            elif isinstance(e, Symbol):
                return MemoryReference(str(e))
            elif isinstance(e, sin):
                return quil_sin(to_pyquil(e))
            elif isinstance(e, cos):
                return quil_cos(to_pyquil(e))
            elif isinstance(e, Add):
                args = [to_pyquil(a) for a in e.args]
                acc = args[0]
                for a in args[1:]:
                    acc += a
                return acc
            elif isinstance(e, Mul):
                args = [to_pyquil(a) for a in e.args]
                acc = args[0]
                for a in args[1:]:
                    acc *= a
                return acc
            elif isinstance(e, Pow):
                args = Pow_(to_pyquil(e.base), to_pyquil(e.exp))
            elif e == pi:
                return math.pi
            else:
                raise NotImplementedError(
                    "Sympy expression could not be converted to a Quil expression: "
                    + str(e)
                )

        return to_pyquil(ppi)


def param_from_pyquil(p: Union[float, Expression]) -> Expr:
    def to_sympy(e: Union[float, Expression]) -> Expr:
        if isinstance(e, float):
            return e
        elif isinstance(e, MemoryReference):
            return Symbol(e.name)
        elif isinstance(e, Function_):
            if e.name == "SIN":
                return sin(to_sympy(e.expression))
            elif e.name == "COS":
                return cos(to_sympy(e.expression))
            else:
                raise NotImplementedError(
                    "Quil expression function "
                    + e.name
                    + " cannot be converted to a sympy expression"
                )
        elif isinstance(e, Add_):
            return to_sympy(e.op1) + to_sympy(e.op2)
        elif isinstance(e, Sub):
            return to_sympy(e.op1) - to_sympy(e.op2)
        elif isinstance(e, Mul_):
            return to_sympy(e.op1) * to_sympy(e.op2)
        elif isinstance(e, Div):
            return to_sympy(e.op1) / to_sympy(e.op2)
        elif isinstance(e, Pow_):
            return to_sympy(e.op1) ** to_sympy(e.op2)
        else:
            raise NotImplementedError(
                "Quil expression could not be converted to a sympy expression: "
                + str(e)
            )

    return to_sympy(p) / pi


def pyquil_to_tk(prog: Program) -> Circuit:
    """
    Convert a :py:class:`pyquil.Program` to a :math:`\\mathrm{t|ket}\\rangle` :py:class:`Circuit` .
    Note that not all pyQuil operations are currently supported by pytket.

    :param prog: A circuit to be converted

    :return: The converted circuit
    """
    qubits = prog.get_qubits()
    tkc = Circuit()
    qmap = {}
    for q in qubits:
        id = Qubit("q", q)
        tkc.add_qubit(id)
        qmap.update({q: id})
    cregmap = {}
    for i in prog.instructions:
        if isinstance(i, Gate):
            try:
                optype = _known_quil_gate[i.name]
            except KeyError as error:
                raise NotImplementedError(
                    "Operation not supported by tket: " + str(i)
                ) from error
            qubits = [qmap[q.index] for q in i.qubits]
            params = [param_from_pyquil(p) for p in i.params]
            tkc.add_gate(optype, params, qubits)
        elif isinstance(i, Measurement):
            qubit = qmap[i.qubit.index]
            reg = cregmap[i.classical_reg.name]
            bit = reg[i.classical_reg.offset]
            tkc.Measure(qubit, bit)
        elif isinstance(i, Declare):
            if i.memory_type == "BIT":
                new_reg = tkc.add_c_register(i.name, i.memory_size)
                cregmap.update({i.name: new_reg})
            elif i.memory_type == "REAL":
                continue
            else:
                raise NotImplementedError(
                    "Cannot handle memory of type " + i.memory_type
                )
        elif isinstance(i, Pragma):
            continue
        elif isinstance(i, Halt):
            return tkc
        else:
            raise NotImplementedError("PyQuil instruction is not a gate: " + str(i))
    return tkc


def tk_to_pyquil(
    tkcirc: Circuit, active_reset: bool = False, return_used_bits: bool = False
) -> Program:
    """
       Convert a :math:`\\mathrm{t|ket}\\rangle` :py:class:`Circuit` to a :py:class:`pyquil.Program` .

    :param tkcirc: A circuit to be converted

    :return: The converted circuit
    """
    p = Program()
    qregs = set()
    for qb in tkcirc.qubits:
        if len(qb.index) != 1:
            raise NotImplementedError("PyQuil registers must use a single index")
        qregs.add(qb.reg_name)
    if len(qregs) > 1:
        raise NotImplementedError(
            "Cannot convert circuit with multiple quantum registers to pyQuil"
        )
    creg_sizes = {}
    for b in tkcirc.bits:
        if len(b.index) != 1:
            raise NotImplementedError("PyQuil registers must use a single index")
        if (b.reg_name not in creg_sizes) or (b.index[0] >= creg_sizes[b.reg_name]):
            creg_sizes.update({b.reg_name: b.index[0] + 1})
    cregmap = {}
    for reg_name, size in creg_sizes.items():
        name = reg_name
        if name == "c":
            name = "ro"
        quil_reg = p.declare(name, "BIT", size)
        cregmap.update({reg_name: quil_reg})
    for sym in tkcirc.free_symbols():
        p.declare(str(sym), "REAL")
    if active_reset:
        p.reset()
    measures = []
    measured_qubits = []
    used_bits = []
    for command in tkcirc:
        op = command.op
        optype = op.type
        if optype == OpType.Measure:
            qb = Qubit_(command.args[0].index[0])
            if qb in measured_qubits:
                raise NotImplementedError(
                    "Cannot apply gate on qubit " + qb.__repr__() + " after measurement"
                )
            bit = command.args[1]
            b = cregmap[bit.reg_name][bit.index[0]]
            measures.append(Measurement(qb, b))
            measured_qubits.append(qb)
            used_bits.append(bit)
            continue
        elif optype == OpType.Barrier:
            continue  # pyQuil cannot handle barriers
        qubits = [Qubit_(qb.index[0]) for qb in command.args]
        for qb in qubits:
            if qb in measured_qubits:
                raise NotImplementedError(
                    "Cannot apply gate on qubit " + qb.__repr__() + " after measurement"
                )
        try:
            gatetype = _known_quil_gate_rev[optype]
        except KeyError as error:
            raise NotImplementedError(
                "Cannot convert tket Op to pyQuil gate: " + op.get_name()
            ) from error
        params = [param_to_pyquil(p) for p in op.params]
        g = Gate(gatetype, params, qubits)
        p += g
    for m in measures:
        p += m
    if return_used_bits:
        return p, used_bits
    return p


def process_device(qc: QuantumComputer) -> Device:
    """Convert a :py:class:`pyquil.api.QuantumComputer` to a :py:class:`Device`.
    
    :param qc: A quantum computer to be converted
    :type qc: QuantumComputer
    :return: The corresponding :math:`\\mathrm{t|ket}\\rangle` :py:class:`Device`
    :rtype: Device
    """
    specs = qc.device.get_specs()

    coupling_map = [
        [n, ni]
        for n, neigh_dict in qc.qubit_topology().adjacency()
        for ni, _ in neigh_dict.items()
    ]

    node_ers_dict = {}
    link_ers_dict = {}

    device_node_ers = specs.f1QRB_std_errs()
    device_link_ers = specs.fCZ_std_errs()
    device_ROs = specs.fROs()
    device_t1s = specs.T1s()
    device_t2s = specs.T2s()

    for index in qc.qubits():
        error_cont = GateErrorContainer({OpType.Rx, OpType.Rz})
        error_cont.add_readout(device_ROs[index])
        error_cont.add_t1_time(device_t1s[index])
        error_cont.add_t2_time(device_t2s[index])
        error_cont.add_error((OpType.Rx, GateError(device_node_ers[index], 1.0)))
        error_cont.add_error((OpType.Rz, GateError(device_node_ers[index], 1.0)))
        node_ers_dict[Node(index)] = error_cont

    for (a, b), err in device_link_ers.items():
        error_cont = GateErrorContainer({OpType.CZ})
        error_cont.add_error((OpType.CZ, GateError(err, 1.0)))
        link_ers_dict[(Node(a), Node(b))] = error_cont
        link_ers_dict[(Node(b), Node(a))] = error_cont

    arc = Architecture(coupling_map)

    device = Device(node_ers_dict, link_ers_dict, arc)
    return device
