# -*- coding: utf-8 -*-
#   Copyright 2017, 2021 ProjectQ-Framework (www.projectq.ch)
#
#   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.
"""Tests for projectq.ops._command."""

from copy import deepcopy
import sys
import math
import pytest

from projectq import MainEngine
from projectq.cengines import DummyEngine
from projectq.meta import ComputeTag, canonical_ctrl_state
from projectq.ops import BasicGate, Rx, NotMergeable, CtrlAll
from projectq.types import Qubit, Qureg, WeakQubitRef

from projectq.ops import _command


@pytest.fixture
def main_engine():
    return MainEngine(backend=DummyEngine(), engine_list=[DummyEngine()])


def test_command_init(main_engine):
    qureg0 = Qureg([Qubit(main_engine, 0)])
    qureg1 = Qureg([Qubit(main_engine, 1)])
    qureg2 = Qureg([Qubit(main_engine, 2)])
    # qureg3 = Qureg([Qubit(main_engine, 3)])
    # qureg4 = Qureg([Qubit(main_engine, 4)])
    gate = BasicGate()
    cmd = _command.Command(main_engine, gate, (qureg0, qureg1, qureg2))
    assert cmd.gate == gate
    assert cmd.tags == []
    expected_tuple = (qureg0, qureg1, qureg2)
    for cmd_qureg, expected_qureg in zip(cmd.qubits, expected_tuple):
        assert cmd_qureg[0].id == expected_qureg[0].id
        # Testing that Qubits are now WeakQubitRef objects
        assert type(cmd_qureg[0]) == WeakQubitRef
    assert cmd._engine == main_engine
    # Test that quregs are ordered if gate has interchangeable qubits:
    symmetric_gate = BasicGate()
    symmetric_gate.interchangeable_qubit_indices = [[0, 1]]
    symmetric_cmd = _command.Command(main_engine, symmetric_gate, (qureg2, qureg1, qureg0))
    assert cmd.gate == gate
    assert cmd.tags == []
    expected_ordered_tuple = (qureg1, qureg2, qureg0)
    for cmd_qureg, expected_qureg in zip(symmetric_cmd.qubits, expected_ordered_tuple):
        assert cmd_qureg[0].id == expected_qureg[0].id
    assert symmetric_cmd._engine == main_engine


def test_command_deepcopy(main_engine):
    qureg0 = Qureg([Qubit(main_engine, 0)])
    qureg1 = Qureg([Qubit(main_engine, 1)])
    gate = BasicGate()
    cmd = _command.Command(main_engine, gate, (qureg0,))
    cmd.add_control_qubits(qureg1)
    cmd.tags.append("MyTestTag")
    copied_cmd = deepcopy(cmd)
    # Test that deepcopy gives same cmd
    assert copied_cmd.gate == gate
    assert copied_cmd.tags == ["MyTestTag"]
    assert len(copied_cmd.qubits) == 1
    assert copied_cmd.qubits[0][0].id == qureg0[0].id
    assert len(copied_cmd.control_qubits) == 1
    assert copied_cmd.control_qubits[0].id == qureg1[0].id
    # Engine should not be deepcopied but a reference:
    assert id(copied_cmd.engine) == id(main_engine)
    # Test that deepcopy is actually a deepcopy
    cmd.tags = ["ChangedTag"]
    assert copied_cmd.tags == ["MyTestTag"]
    cmd.control_qubits[0].id == 10
    assert copied_cmd.control_qubits[0].id == qureg1[0].id
    cmd.gate = "ChangedGate"
    assert copied_cmd.gate == gate


def test_command_get_inverse(main_engine):
    qubit = main_engine.allocate_qubit()
    ctrl_qubit = main_engine.allocate_qubit()
    cmd = _command.Command(main_engine, Rx(0.5), (qubit,))
    cmd.add_control_qubits(ctrl_qubit)
    cmd.tags = [ComputeTag()]
    inverse_cmd = cmd.get_inverse()
    assert inverse_cmd.gate == Rx(-0.5 + 4 * math.pi)
    assert len(cmd.qubits) == len(inverse_cmd.qubits)
    assert cmd.qubits[0][0].id == inverse_cmd.qubits[0][0].id
    assert id(cmd.qubits[0][0]) != id(inverse_cmd.qubits[0][0])
    assert len(cmd.control_qubits) == len(inverse_cmd.control_qubits)
    assert cmd.control_qubits[0].id == inverse_cmd.control_qubits[0].id
    assert id(cmd.control_qubits[0]) != id(inverse_cmd.control_qubits[0])
    assert cmd.tags == inverse_cmd.tags
    assert id(cmd.tags[0]) != id(inverse_cmd.tags[0])
    assert id(cmd.engine) == id(inverse_cmd.engine)


def test_command_get_merged(main_engine):
    qubit = main_engine.allocate_qubit()
    ctrl_qubit = main_engine.allocate_qubit()
    cmd = _command.Command(main_engine, Rx(0.5), (qubit,))
    cmd.tags = ["TestTag"]
    cmd.add_control_qubits(ctrl_qubit)
    # Merge two commands
    cmd2 = _command.Command(main_engine, Rx(0.5), (qubit,))
    cmd2.add_control_qubits(ctrl_qubit)
    cmd2.tags = ["TestTag"]
    merged_cmd = cmd.get_merged(cmd2)
    expected_cmd = _command.Command(main_engine, Rx(1.0), (qubit,))
    expected_cmd.add_control_qubits(ctrl_qubit)
    expected_cmd.tags = ["TestTag"]
    assert merged_cmd == expected_cmd
    # Don't merge commands as different control qubits
    cmd3 = _command.Command(main_engine, Rx(0.5), (qubit,))
    cmd3.tags = ["TestTag"]
    with pytest.raises(NotMergeable):
        cmd.get_merged(cmd3)
    # Don't merge commands as different tags
    cmd4 = _command.Command(main_engine, Rx(0.5), (qubit,))
    cmd4.add_control_qubits(ctrl_qubit)
    with pytest.raises(NotMergeable):
        cmd.get_merged(cmd4)


def test_command_is_identity(main_engine):
    qubit = main_engine.allocate_qubit()
    qubit2 = main_engine.allocate_qubit()
    cmd = _command.Command(main_engine, Rx(0.0), (qubit,))
    cmd2 = _command.Command(main_engine, Rx(0.5), (qubit2,))
    inverse_cmd = cmd.get_inverse()
    inverse_cmd2 = cmd2.get_inverse()
    assert inverse_cmd.gate.is_identity()
    assert cmd.gate.is_identity()
    assert not inverse_cmd2.gate.is_identity()
    assert not cmd2.gate.is_identity()


def test_command_order_qubits(main_engine):
    qubit0 = Qureg([Qubit(main_engine, 0)])
    qubit1 = Qureg([Qubit(main_engine, 1)])
    qubit2 = Qureg([Qubit(main_engine, 2)])
    qubit3 = Qureg([Qubit(main_engine, 3)])
    qubit4 = Qureg([Qubit(main_engine, 4)])
    qubit5 = Qureg([Qubit(main_engine, 5)])
    gate = BasicGate()
    gate.interchangeable_qubit_indices = [[0, 4, 5], [1, 2]]
    input_tuple = (qubit4, qubit5, qubit3, qubit2, qubit1, qubit0)
    expected_tuple = (qubit0, qubit3, qubit5, qubit2, qubit1, qubit4)
    cmd = _command.Command(main_engine, gate, input_tuple)
    for ordered_qubit, expected_qubit in zip(cmd.qubits, expected_tuple):
        assert ordered_qubit[0].id == expected_qubit[0].id


def test_command_interchangeable_qubit_indices(main_engine):
    gate = BasicGate()
    gate.interchangeable_qubit_indices = [[0, 4, 5], [1, 2]]
    qubit0 = Qureg([Qubit(main_engine, 0)])
    qubit1 = Qureg([Qubit(main_engine, 1)])
    qubit2 = Qureg([Qubit(main_engine, 2)])
    qubit3 = Qureg([Qubit(main_engine, 3)])
    qubit4 = Qureg([Qubit(main_engine, 4)])
    qubit5 = Qureg([Qubit(main_engine, 5)])
    input_tuple = (qubit4, qubit5, qubit3, qubit2, qubit1, qubit0)
    cmd = _command.Command(main_engine, gate, input_tuple)
    assert (
        cmd.interchangeable_qubit_indices
        == [
            [0, 4, 5],
            [1, 2],
        ]
        or cmd.interchangeable_qubit_indices == [[1, 2], [0, 4, 5]]
    )


@pytest.mark.parametrize(
    'state',
    [0, 1, '0', '1', CtrlAll.One, CtrlAll.Zero],
    ids=['int(0)', 'int(1)', 'str(0)', 'str(1)', 'CtrlAll.One', 'CtrlAll.Zero'],
)
def test_commmand_add_control_qubits_one(main_engine, state):
    qubit0 = Qureg([Qubit(main_engine, 0)])
    qubit1 = Qureg([Qubit(main_engine, 1)])
    cmd = _command.Command(main_engine, Rx(0.5), (qubit0,))
    cmd.add_control_qubits(qubit1, state=state)
    assert cmd.control_qubits[0].id == 1
    assert cmd.control_state == canonical_ctrl_state(state, 1)

    with pytest.raises(ValueError):
        cmd.add_control_qubits(qubit0[0])


@pytest.mark.parametrize(
    'state',
    [0, 1, 2, 3, '00', '01', '10', '11', CtrlAll.One, CtrlAll.Zero],
    ids=[
        'int(0)',
        'int(1)',
        'int(2)',
        'int(3)',
        'str(00)',
        'str(01)',
        'str(10)',
        'str(1)',
        'CtrlAll.One',
        'CtrlAll.Zero',
    ],
)
def test_commmand_add_control_qubits_two(main_engine, state):
    qubit0 = Qureg([Qubit(main_engine, 0)])
    qubit1 = Qureg([Qubit(main_engine, 1)])
    qubit2 = Qureg([Qubit(main_engine, 2)])
    qubit3 = Qureg([Qubit(main_engine, 3)])
    cmd = _command.Command(main_engine, Rx(0.5), (qubit0,), qubit1)
    cmd.add_control_qubits(qubit2 + qubit3, state)
    assert cmd.control_qubits[0].id == 1
    assert cmd.control_state == '1' + canonical_ctrl_state(state, 2)


def test_command_all_qubits(main_engine):
    qubit0 = Qureg([Qubit(main_engine, 0)])
    qubit1 = Qureg([Qubit(main_engine, 1)])
    cmd = _command.Command(main_engine, Rx(0.5), (qubit0,))
    cmd.add_control_qubits(qubit1)
    all_qubits = cmd.all_qubits
    assert all_qubits[0][0].id == 1
    assert all_qubits[1][0].id == 0


def test_command_engine(main_engine):
    qubit0 = Qureg([Qubit("fake_engine", 0)])
    qubit1 = Qureg([Qubit("fake_engine", 1)])
    cmd = _command.Command("fake_engine", Rx(0.5), (qubit0,))
    cmd.add_control_qubits(qubit1)
    assert cmd.engine == "fake_engine"
    cmd.engine = main_engine
    assert id(cmd.engine) == id(main_engine)
    assert id(cmd.control_qubits[0].engine) == id(main_engine)
    assert id(cmd.qubits[0][0].engine) == id(main_engine)

    # Avoid raising exception upon Qubit destructions
    qubit0[0].id = -1
    qubit1[0].id = -1


def test_command_comparison(main_engine):
    qubit = Qureg([Qubit(main_engine, 0)])
    ctrl_qubit = Qureg([Qubit(main_engine, 1)])
    cmd1 = _command.Command(main_engine, Rx(0.5), (qubit,))
    cmd1.tags = ["TestTag"]
    cmd1.add_control_qubits(ctrl_qubit)
    # Test equality
    cmd2 = _command.Command(main_engine, Rx(0.5), (qubit,))
    cmd2.tags = ["TestTag"]
    cmd2.add_control_qubits(ctrl_qubit)
    assert cmd2 == cmd1
    # Test not equal because of tags
    cmd3 = _command.Command(main_engine, Rx(0.5), (qubit,))
    cmd3.tags = ["TestTag", "AdditionalTag"]
    cmd3.add_control_qubits(ctrl_qubit)
    assert not cmd3 == cmd1
    # Test not equal because of control qubit
    cmd4 = _command.Command(main_engine, Rx(0.5), (qubit,))
    cmd4.tags = ["TestTag"]
    assert not cmd4 == cmd1
    # Test not equal because of qubit
    qubit2 = Qureg([Qubit(main_engine, 2)])
    cmd5 = _command.Command(main_engine, Rx(0.5), (qubit2,))
    cmd5.tags = ["TestTag"]
    cmd5.add_control_qubits(ctrl_qubit)
    assert cmd5 != cmd1
    # Test not equal because of engine
    cmd6 = _command.Command("FakeEngine", Rx(0.5), (qubit,))
    cmd6.tags = ["TestTag"]
    cmd6.add_control_qubits(ctrl_qubit)
    assert cmd6 != cmd1


def test_command_str(main_engine):
    qubit = Qureg([Qubit(main_engine, 0)])
    ctrl_qubit = Qureg([Qubit(main_engine, 1)])
    cmd = _command.Command(main_engine, Rx(0.5 * math.pi), (qubit,))
    cmd.tags = ["TestTag"]
    cmd.add_control_qubits(ctrl_qubit)
    cmd2 = _command.Command(main_engine, Rx(0.5 * math.pi), (qubit,))
    if sys.version_info.major == 3:
        assert cmd.to_string(symbols=False) == "CRx(1.570796326795) | ( Qureg[1], Qureg[0] )"
        assert str(cmd2) == "Rx(1.570796326795) | Qureg[0]"
    else:
        assert cmd.to_string(symbols=False) == "CRx(1.5707963268) | ( Qureg[1], Qureg[0] )"
        assert str(cmd2) == "Rx(1.5707963268) | Qureg[0]"


def test_command_to_string(main_engine):
    qubit = Qureg([Qubit(main_engine, 0)])
    ctrl_qubit = Qureg([Qubit(main_engine, 1)])
    cmd = _command.Command(main_engine, Rx(0.5 * math.pi), (qubit,))
    cmd.tags = ["TestTag"]
    cmd.add_control_qubits(ctrl_qubit)
    cmd2 = _command.Command(main_engine, Rx(0.5 * math.pi), (qubit,))

    assert cmd.to_string(symbols=True) == u"CRx(0.5π) | ( Qureg[1], Qureg[0] )"
    assert cmd2.to_string(symbols=True) == u"Rx(0.5π) | Qureg[0]"
    if sys.version_info.major == 3:
        assert cmd.to_string(symbols=False) == "CRx(1.570796326795) | ( Qureg[1], Qureg[0] )"
        assert cmd2.to_string(symbols=False) == "Rx(1.570796326795) | Qureg[0]"
    else:
        assert cmd.to_string(symbols=False) == "CRx(1.5707963268) | ( Qureg[1], Qureg[0] )"
        assert cmd2.to_string(symbols=False) == "Rx(1.5707963268) | Qureg[0]"
