# This code is part of Qiskit.
#
# (C) Copyright IBM 2021.
#
# This code is licensed under the Apache License, Version 2.0. You may
# obtain a copy of this license in the LICENSE.txt file in the root directory
# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
#
# Any modifications or derivative works of this code must retain this
# copyright notice, and modified files need to carry a notice indicating
# that they have been altered from the originals.

"""Test RemoveFinalMeasurements pass"""

import unittest

from qiskit import ClassicalRegister, QuantumRegister, QuantumCircuit
from qiskit.circuit.classicalregister import Clbit
from qiskit.transpiler.passes import RemoveFinalMeasurements
from qiskit.converters import circuit_to_dag
from qiskit.test import QiskitTestCase


class TestRemoveFinalMeasurements(QiskitTestCase):
    """Test removing final measurements."""

    def test_multi_bit_register_removed_with_clbits(self):
        """Remove register when all clbits removed."""

        def expected_dag():
            q0 = QuantumRegister(2, "q0")
            qc = QuantumCircuit(q0)
            return circuit_to_dag(qc)

        q0 = QuantumRegister(2, "q0")
        c0 = ClassicalRegister(2, "c0")
        qc = QuantumCircuit(q0, c0)

        # measure into all clbits of c0
        qc.measure(0, 0)
        qc.measure(1, 1)

        dag = circuit_to_dag(qc)
        dag = RemoveFinalMeasurements().run(dag)

        self.assertFalse(dag.cregs)
        self.assertFalse(dag.clbits)
        self.assertEqual(dag, expected_dag())

    def test_register_kept_if_measured_clbit_busy(self):
        """
        A register is kept if the measure destination bit is still
        busy after measure removal.
        """

        def expected_dag():
            q0 = QuantumRegister(1, "q0")
            c0 = ClassicalRegister(1, "c0")
            qc = QuantumCircuit(q0, c0)
            qc.x(0).c_if(c0[0], 0)
            return circuit_to_dag(qc)

        q0 = QuantumRegister(1, "q0")
        c0 = ClassicalRegister(1, "c0")
        qc = QuantumCircuit(q0, c0)

        # make c0 busy
        qc.x(0).c_if(c0[0], 0)

        # measure into c0
        qc.measure(0, c0[0])

        dag = circuit_to_dag(qc)
        dag = RemoveFinalMeasurements().run(dag)

        self.assertListEqual(list(dag.cregs.values()), [c0])
        self.assertListEqual(dag.clbits, list(c0))
        self.assertEqual(dag, expected_dag())

    def test_multi_bit_register_kept_if_not_measured_clbit_busy(self):
        """
        A multi-bit register is kept if it contains a busy bit even if
        the measure destination bit itself is idle.
        """

        def expected_dag():
            q0 = QuantumRegister(1, "q0")
            c0 = ClassicalRegister(2, "c0")
            qc = QuantumCircuit(q0, c0)
            qc.x(q0[0]).c_if(c0[0], 0)
            return circuit_to_dag(qc)

        q0 = QuantumRegister(1, "q0")
        c0 = ClassicalRegister(2, "c0")
        qc = QuantumCircuit(q0, c0)

        # make c0[0] busy
        qc.x(q0[0]).c_if(c0[0], 0)

        # measure into not busy c0[1]
        qc.measure(0, c0[1])

        dag = circuit_to_dag(qc)
        dag = RemoveFinalMeasurements().run(dag)

        # c0 should not be removed because it has busy bit c0[0]
        self.assertListEqual(list(dag.cregs.values()), [c0])

        # note: c0[1] should not be removed even though it is now idle
        # because it is referenced by creg c0.
        self.assertListEqual(dag.clbits, list(c0))
        self.assertEqual(dag, expected_dag())

    def test_overlapping_register_removal(self):
        """Only registers that become idle directly as a result of
        final op removal are removed. In this test, a 5-bit creg
        is implicitly created with its own bits, along with cregs
        ``c0_lower_3`` and ``c0_upper_3`` which reuse those underlying bits.
        ``c0_lower_3`` and ``c0_upper_3`` reference only 1 bit in common.
        A final measure is performed into a bit that exists in ``c0_lower_3``
        but not in ``c0_upper_3``, and subsequently is removed. Consequently,
        both ``c0_lower_3`` and the 5-bit register are removed, because they
        have become unused as a result of the final measure removal.
        ``c0_upper_3`` remains, because it was idle beforehand, not as a
        result of the measure removal, along with all of its bits,
        including the bit shared with ``c0_lower_3``."""

        def expected_dag():
            q0 = QuantumRegister(3, "q0")
            c0 = ClassicalRegister(5, "c0")
            c0_upper_3 = ClassicalRegister(name="c0_upper_3", bits=c0[2:])

            # note c0 is *not* added to circuit!
            qc = QuantumCircuit(q0, c0_upper_3)
            return circuit_to_dag(qc)

        q0 = QuantumRegister(3, "q0")
        c0 = ClassicalRegister(5, "c0")
        qc = QuantumCircuit(q0, c0)

        c0_lower_3 = ClassicalRegister(name="c0_lower_3", bits=c0[:3])
        c0_upper_3 = ClassicalRegister(name="c0_upper_3", bits=c0[2:])
        # Only qc.clbits[2] is shared between the two.

        qc.add_register(c0_lower_3)
        qc.add_register(c0_upper_3)

        qc.measure(0, c0_lower_3[0])

        dag = circuit_to_dag(qc)
        dag = RemoveFinalMeasurements().run(dag)

        self.assertListEqual(list(dag.cregs.values()), [c0_upper_3])
        self.assertListEqual(dag.clbits, list(c0_upper_3))
        self.assertEqual(dag, expected_dag())

    def test_multi_bit_register_removed_if_all_bits_idle(self):
        """A multibit register is removed when all bits are idle."""

        def expected_dag():
            q0 = QuantumRegister(1, "q0")
            qc = QuantumCircuit(q0)
            return circuit_to_dag(qc)

        q0 = QuantumRegister(1, "q0")
        c0 = ClassicalRegister(2, "c0")
        qc = QuantumCircuit(q0, c0)

        # measure into single bit c0[0] of c0
        qc.measure(0, 0)

        dag = circuit_to_dag(qc)
        dag = RemoveFinalMeasurements().run(dag)

        self.assertFalse(dag.cregs)
        self.assertFalse(dag.clbits)
        self.assertEqual(dag, expected_dag())

    def test_multi_reg_shared_bits_removed(self):
        """All registers sharing removed bits should be removed."""

        def expected_dag():
            q0 = QuantumRegister(2, "q0")
            qc = QuantumCircuit(q0)
            return circuit_to_dag(qc)

        q0 = QuantumRegister(2, "q0")
        c0 = ClassicalRegister(2, "c0")
        qc = QuantumCircuit(q0, c0)

        # Create reg with shared bits (same as c0)
        c1 = ClassicalRegister(name="c1", bits=qc.clbits)
        qc.add_register(c1)

        # measure into all clbits of c0
        qc.measure(0, c0[0])
        qc.measure(1, c0[1])

        dag = circuit_to_dag(qc)
        dag = RemoveFinalMeasurements().run(dag)

        self.assertFalse(dag.cregs)
        self.assertFalse(dag.clbits)
        self.assertEqual(dag, expected_dag())

    def test_final_measures_share_dest(self):
        """Multiple final measurements use the same clbit."""

        def expected_dag():
            qc = QuantumCircuit(QuantumRegister(2, "q0"))
            return circuit_to_dag(qc)

        rq = QuantumRegister(2, "q0")
        rc = ClassicalRegister(1, "c0")
        qc = QuantumCircuit(rq, rc)

        qc.measure(0, 0)
        qc.measure(1, 0)

        dag = circuit_to_dag(qc)
        dag = RemoveFinalMeasurements().run(dag)

        self.assertEqual(dag, expected_dag())

    def test_remove_chained_final_measurements(self):
        """Remove successive final measurements."""

        def expected_dag():
            q0 = QuantumRegister(1, "q0")
            q1 = QuantumRegister(1, "q1")
            c0 = ClassicalRegister(1, "c0")
            qc = QuantumCircuit(q0, c0, q1)

            qc.measure(q0, c0)
            qc.measure(q0, c0)
            qc.barrier()
            qc.h(q1)

            return circuit_to_dag(qc)

        q0 = QuantumRegister(1, "q0")
        q1 = QuantumRegister(1, "q1")
        c0 = ClassicalRegister(1, "c0")
        c1 = ClassicalRegister(1, "c1")
        qc = QuantumCircuit(q0, c0, q1, c1)

        qc.measure(q0, c0)
        qc.measure(q0, c0)
        qc.barrier()
        qc.h(q1)
        qc.measure(q1, c1)
        qc.measure(q0, c1)

        dag = circuit_to_dag(qc)
        dag = RemoveFinalMeasurements().run(dag)

        self.assertEqual(dag, expected_dag())

    def test_remove_clbits_without_register(self):
        """clbits of final measurements not in a register are removed."""

        def expected_dag():
            q0 = QuantumRegister(1, "q0")
            qc = QuantumCircuit(q0)
            return circuit_to_dag(qc)

        q0 = QuantumRegister(1, "q0")
        qc = QuantumCircuit(q0)

        # Add clbit without adding register
        qc.add_bits([Clbit()])

        self.assertFalse(qc.cregs)

        # Measure to regless clbit
        qc.measure(0, 0)

        dag = circuit_to_dag(qc)
        dag = RemoveFinalMeasurements().run(dag)

        self.assertFalse(dag.cregs)
        self.assertFalse(dag.clbits)
        self.assertEqual(dag, expected_dag())

    def test_final_barriers_and_measures_complex(self):
        """Test complex final barrier and measure removal."""

        def expected_dag():
            q0 = QuantumRegister(5, "q0")
            c1 = ClassicalRegister(1, "c1")
            qc = QuantumCircuit(q0, c1)
            qc.h(q0[0])
            return circuit_to_dag(qc)

        #       ┌───┐┌─┐ ░     ░ ┌─┐
        # q0_0: ┤ H ├┤M├─░─────░─┤M├───────────────
        #       └┬─┬┘└╥┘ ░     ░ └╥┘┌─┐
        # q0_1: ─┤M├──╫──░─────░──╫─┤M├────────────
        #        └╥┘  ║  ░  ░  ░  ║ └╥┘┌─┐
        # q0_2: ──╫───╫──░──░──░──╫──╫─┤M├─────────
        #         ║   ║  ░  ░  ░  ║  ║ └╥┘┌─┐
        # q0_3: ──╫───╫──░──░──░──╫──╫──╫─┤M├──────
        #         ║   ║  ░  ░  ░  ║  ║  ║ └╥┘┌─┐ ░
        # q0_4: ──╫───╫──░─────░──╫──╫──╫──╫─┤M├─░─
        #         ║   ║  ░     ░  ║  ║  ║  ║ └╥┘ ░
        # c0: 1/══╩═══╩═══════════╬══╬══╬══╬══╬════
        #         0   0           ║  ║  ║  ║  ║
        #                         ║  ║  ║  ║  ║
        # c1: 1/══════════════════╬══╬══╬══╬══╬════
        #                         ║  ║  ║  ║  ║
        # meas: 5/════════════════╩══╩══╩══╩══╩════
        #                         0  1  2  3  4
        q0 = QuantumRegister(5, "q0")
        c0 = ClassicalRegister(1, "c0")
        c1 = ClassicalRegister(1, "c1")
        qc = QuantumCircuit(q0, c0, c1)

        qc.measure(q0[1], c0)
        qc.h(q0[0])
        qc.measure(q0[0], c0[0])
        qc.barrier()
        qc.barrier(q0[2], q0[3])
        qc.measure_all()
        qc.barrier(q0[4])

        dag = circuit_to_dag(qc)
        dag = RemoveFinalMeasurements().run(dag)

        self.assertEqual(dag, expected_dag())


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