# PyZX - Python library for quantum circuit rewriting
#        and optimization using the ZX-calculus
# Copyright (C) 2018 - Aleks Kissinger and John van de Wetering

# Licensed 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.


import unittest
import os
import sys
from types import ModuleType
from typing import Optional

if __name__ == '__main__':
    sys.path.append('..')
    sys.path.append('.')
from pyzx_param.simplify import full_reduce
from pyzx_param.extract import extract_circuit
from pyzx_param.circuit import Circuit
from fractions import Fraction

np: Optional[ModuleType]
try:
    import numpy as np
    from pyzx_param.tensor import compare_tensors
    import math
except ImportError:
    np = None

try:
    from qiskit import quantum_info, transpile
    from qiskit.circuit import QuantumCircuit
    from qiskit.qasm2 import dumps as dumps2
    from qiskit.qasm3 import loads as loads3, dumps as dumps3
except ImportError:
    QuantumCircuit = None


@unittest.skipUnless(np, "numpy needs to be installed for this to run")
class TestQASM(unittest.TestCase):

    def test_parser_state_reset(self):
        from pyzx_param.circuit.qasmparser import QASMParser
        s = """
        OPENQASM 2.0;
        include "qelib1.inc";
        qreg q[1];
        h q[0];
        """
        p = QASMParser()
        c1 = p.parse(s)
        c2 = p.parse(s)
        self.assertEqual(c2.qubits, 1)
        self.assertEqual(len(c2.gates), 1)
        self.assertTrue(c1.verify_equality(c2))

    def test_parse_qasm3(self):
        qasm3 = Circuit.from_qasm("""
        OPENQASM 3;
        include "stdgates.inc";
        qubit[3] q;
        cx q[0], q[1];
        s q[2];
        cx q[2], q[1];
        """)
        c = Circuit(3)
        c.add_gate("CNOT", 0, 1)
        c.add_gate("S", 2)
        c.add_gate("CNOT", 2, 1)
        self.assertEqual(c.qubits, qasm3.qubits)
        self.assertListEqual(c.gates, qasm3.gates)


    def test_parse_qasm3_long_creg(self):
        qasm3 = Circuit.from_qasm("""
        OPENQASM 3;
        include "stdgates.inc";
        qubit[3] q1;
        cx q1[0], q1[1];
        s q1[2];
        cx q1[2], q1[1];
        """)
        c = Circuit(3)
        c.add_gate("CNOT", 0, 1)
        c.add_gate("S", 2)
        c.add_gate("CNOT", 2, 1)
        self.assertEqual(c.qubits, qasm3.qubits)
        self.assertListEqual(c.gates, qasm3.gates)

    def test_load_qasm_from_file(self):
        c = Circuit(1)
        c.add_gate("YPhase", 0, Fraction(1, 4))
        c1 = Circuit.from_qasm_file(os.path.join(os.path.dirname(__file__), "ry.qasm"))
        self.assertEqual(c1.qubits, c.qubits)
        self.assertListEqual(c1.gates,c.gates)

    def test_broadcasting(self):
        """Test that broadcasting is handled correctly.

        If any arguments of a gate are quantum registers instead of qubits, it is a shorthand for broadcasting
        over the qubits of the register."""
        c1 = Circuit.from_qasm("""
        OPENQASM 3;
        include "stdgates.inc";
        qubit[1] q0;
        qubit[3] q1;
        cx q0[0], q1;
        """)
        c2 = Circuit(4)
        c2.add_gate("CNOT", 0, 1)
        c2.add_gate("CNOT", 0, 2)
        c2.add_gate("CNOT", 0, 3)
        self.assertEqual(c1.qubits, c2.qubits)
        self.assertListEqual(c1.gates, c2.gates)

    def test_catch_broadcasting_error(self):
        """Test that all registers are of the same length when broadcasting."""
        with self.assertRaises(TypeError) as context:
            Circuit.from_qasm("""
            OPENQASM 3;
            include "stdgates.inc";
            qubit[1] q0;
            qubit[2] q1;
            qubit[3] q2;
            ccx q0[0], q1, q2;
            """)
        self.assertTrue("Register sizes do not match" in str(context.exception))

    def test_p_same_as_rz(self):
        """Test that the `p` gate is identical to the `rz` gate.

        In the OpenQASM 3 spec, they differ by a global phase. However,
        they are treated in pyzx identically, since the global phase doesn't
        matter.
        """
        p_qasm = Circuit.from_qasm("""
        OPENQASM 2.0;
        include "qelib1.inc";
        qreg q[1];
        p(pi/2) q[0];
        """)
        rz_qasm = Circuit.from_qasm("""
        OPENQASM 2.0;
        include "qelib1.inc";
        qreg q[1];
        rz(pi/2) q[0];
        """)
        self.assertEqual(p_qasm.qubits, rz_qasm.qubits)
        self.assertListEqual(p_qasm.gates, rz_qasm.gates)

    def test_cp_differs_from_crz(self):
        """Test that the `cp` gate and the `crz` gates are different.

        The difference in phase between `p` and `rz` matters when the gates are
        controlled.
        """
        crz_qasm = Circuit.from_qasm("""
        OPENQASM 2.0;
        include "qelib1.inc";
        qreg q[2];
        crz(pi/2) q[0],q[1];
        """)
        cp_qasm = Circuit.from_qasm("""
        OPENQASM 2.0;
        include "qelib1.inc";
        qreg q[2];
        cp(pi/2) q[0],q[1];
        """)
        crz_t = crz_qasm.to_matrix()
        cp_t = cp_qasm.to_matrix()
        self.assertFalse(compare_tensors(crz_t, cp_t, False))

    def test_rzz(self):
        """Regression test for issue #158.

        The qiskit transpiler may rewrite `crz` into `rz` and `rzz` in some
        situations. The following circuits should be equivalent."""
        before = Circuit.from_qasm("""
        OPENQASM 2.0;
        include "qelib1.inc";
        qreg q[2];
        crz(pi/2) q[0],q[1];
        """)

        after = Circuit.from_qasm("""
        OPENQASM 2.0;
        include "qelib1.inc";
        qreg q[2];
        rz(pi/4) q[1];
        rzz(-pi/4) q[0],q[1];
        """)
        self.assertTrue(compare_tensors(
            before.to_matrix(), after.to_matrix(), True))

    def test_custom_gates(self):
        from pyzx_param.circuit.qasmparser import QASMParser
        s1 = """
        OPENQASM 2.0;
        include "qelib1.inc";
        gate majority a,b,c {
            cx c,b;
            cx c,a;
            ccx a,b,c;
        }
        qreg q[3];
        majority q[0],q[1],q[2];
        """
        s2 = """
        OPENQASM 2.0;
        include "qelib1.inc";
        qreg q[3];
        cx q[2],q[1];
        cx q[2],q[0];
        ccx q[0],q[1],q[2];
        """
        p = QASMParser()
        c1 = p.parse(s1)
        c2 = p.parse(s2)
        self.assertEqual(c1.qubits, c2.qubits)
        self.assertListEqual(c1.gates, c2.gates)

    @unittest.skipUnless(QuantumCircuit, "qiskit needs to be installed for this test")
    def test_qasm_qiskit_semantics(self):
        """Verify/document qasm gate semantics when imported into pyzx.

        pyzx's implementation of qasm gates differs from qiskit in global
        phases which are not significant, but which has sometimes led to user
        confusion. This unit test documents those differences.
        """

        def compare_gate_matrix_with_qiskit(gates, num_qubits: int, num_angles: int, qasm_versions=None):
            for gate in gates:
                for qasm_version in qasm_versions if qasm_versions else [2, 3]:
                    header = "OPENQASM 2.0;\ninclude \"qelib1.inc\";\n" if qasm_version == 2 \
                        else "OPENQASM 3;\ninclude \"stdgates.inc\";\n"
                    params = "" if num_angles == 0 else "({})".format(
                        ",".join(["pi/2"] * num_angles))
                    setup = header + f"qreg q[{num_qubits}];\n{gate}{params} "
                    qasm = setup + \
                        ", ".join(
                            [f"q[{i}]" for i in range(num_qubits)]) + ";\n"
                    pyzx_circuit = Circuit.from_qasm(qasm)
                    pyzx_matrix = pyzx_circuit.to_matrix()

                    for g in pyzx_circuit.gates:
                        for b in g.to_basic_gates():
                            self.assertListEqual(b.to_basic_gates(), [b],
                                                 f"\n{gate}.to_basic_gates() contains non-basic gate")

                    # qiskit uses little-endian ordering
                    qiskit_qasm = setup + \
                        ", ".join([f"q[{i}]" for i in reversed(
                            range(num_qubits))]) + ";\n"
                    qc = QuantumCircuit.from_qasm_str(qiskit_qasm) if qasm_version == 2 else loads3(qiskit_qasm)
                    qiskit_matrix = quantum_info.Operator(qc).data

                    # Check that pyzx and qiskit produce the same tensor from the same qasm, modulo qubit endianness.
                    self.assertTrue(compare_tensors(pyzx_matrix, qiskit_matrix, False),
                                    f"Gate: {gate}\nqasm:\n{qasm}\npyzx_matrix:\n{pyzx_matrix}\nqiskit_matrix:\n{qiskit_matrix}")

                    # Check internal round-trip (pyzx to qasm to pyzx) results in the same circuit.
                    qasm_from_pyzx = pyzx_circuit.to_qasm(qasm_version)
                    pyzx_round_trip = Circuit.from_qasm(qasm_from_pyzx)
                    self.assertEqual(pyzx_circuit.qubits, pyzx_round_trip.qubits)
                    self.assertListEqual(pyzx_circuit.gates, pyzx_round_trip.gates)

                    # Check external round-trip (pyzx to qasm to qiskit to qasm to pyzx) results in the same circuit.
                    # Note that the endianness is reversed when going out and again when coming back in, so the overall
                    # result is no change.
                    qiskit_from_qasm = (QuantumCircuit.from_qasm_str(qasm_from_pyzx) if qasm_version == 2
                                        else loads3(qasm_from_pyzx))
                    pyzx_from_qiskit = Circuit.from_qasm(dumps2(qiskit_from_qasm) if qasm_version == 2
                                                         else dumps3(qiskit_from_qasm))
                    self.assertEqual(pyzx_circuit.qubits, pyzx_from_qiskit.qubits)
                    self.assertListEqual(pyzx_circuit.gates, pyzx_from_qiskit.gates)

        # Test standard gates common to both qelib1.inc (OpenQASM 2) and stdgates.inc (OpenQASM 3).
        compare_gate_matrix_with_qiskit(
            ['x', 'y', 'z', 'h', 's', 'sdg', 't', 'tdg', 'sx'], 1, 0)
        compare_gate_matrix_with_qiskit(['u1', 'p', 'rx', 'ry', 'rz'], 1, 1)
        compare_gate_matrix_with_qiskit(['u2'], 1, 2)
        compare_gate_matrix_with_qiskit(['u3'], 1, 3)
        compare_gate_matrix_with_qiskit(
            ['cx', 'CX', 'cy', 'cz', 'ch', 'swap'], 2, 0)
        compare_gate_matrix_with_qiskit(['crx', 'cry', 'crz', 'cp'], 2, 1)
        # 'ccz' not tested because not a standard qasm gate
        compare_gate_matrix_with_qiskit(['ccx', 'cswap'], 3, 0)
        compare_gate_matrix_with_qiskit(['cu'], 2, 4)

        # Test standard gates added to stdgates.inc.
        compare_gate_matrix_with_qiskit(['cphase'], 2, 1, [3])

        # Test standard gates removed from stdgates.inc.
        compare_gate_matrix_with_qiskit(['sxdg'], 1, 0, [2])
        compare_gate_matrix_with_qiskit(['csx'], 2, 0, [2])
        compare_gate_matrix_with_qiskit(['cu1', 'rxx', 'rzz'], 2, 1, [2])
        compare_gate_matrix_with_qiskit(['cu3'], 2, 3, [2])
        compare_gate_matrix_with_qiskit(['u'], 1, 3, [2])

        # Test native OpenQASM 3 gate.
        compare_gate_matrix_with_qiskit(['U'], 1, 3, [3])

    @unittest.skipUnless(QuantumCircuit, "qiskit needs to be installed for this test")
    def test_qiskit_transpile_pyzx_optimization_round_trip(self):
        """Regression test for issue #102.

        Transpile a circuit in qiskit, simplify it using pyzx, and verify that
        it produces the same unitary in qiskit.
        """
        qc = QuantumCircuit(4)
        qc.ccx(2, 1, 0)
        qc.ccz(0, 1, 2)
        qc.h(1)
        qc.ccx(1, 2, 3)
        qc.t(1)
        qc.ccz(0, 1, 2)
        qc.h(1)
        qc.t(0)
        qc.ccz(2, 1, 0)
        qc.s(1)
        qc.ccx(2, 1, 0)
        qc.crz(0.2*np.pi, 0, 1)
        qc.rz(0.8*np.pi, 1)
        qc.cry(0.4*np.pi, 2, 1)
        qc.crx(0.02*np.pi, 2, 0)

        qc1 = transpile(qc)
        t1 = quantum_info.Operator(qc1).data

        # Test round-trip for OpenQASM 2.
        c2 = Circuit.from_qasm(dumps2(qc1))
        g2 = c2.to_graph()
        full_reduce(g2)
        qasm2 = extract_circuit(g2).to_basic_gates().to_qasm(2)

        qc2 = QuantumCircuit().from_qasm_str(qasm2)
        t2 = quantum_info.Operator(qc2).data

        self.assertTrue(compare_tensors(t1, t2))

        # Test round-trip for OpenQASM 3.
        c3 = Circuit.from_qasm(dumps3(qc1))
        g3 = c3.to_graph()
        full_reduce(g3)
        qasm3 = extract_circuit(g3).to_basic_gates().to_qasm(3)

        qc3 = loads3(qasm3)
        t3 = quantum_info.Operator(qc3).data

        self.assertTrue(compare_tensors(t1, t3))


if __name__ == '__main__':
    unittest.main()
