#   Copyright 2017 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.cengines._basics.py."""

import types
import pytest
# try:
#     import mock
# except ImportError:
#     from unittest import mock

from projectq import MainEngine
from projectq.types import Qubit
from projectq.cengines import DummyEngine, InstructionFilter
from projectq.meta import DirtyQubitTag
from projectq.ops import (AllocateQubitGate,
                          DeallocateQubitGate,
                          H, FastForwardingGate,
                          ClassicalInstructionGate)

from projectq.cengines import _basics


def test_basic_engine_init():
    eng = _basics.BasicEngine()
    assert eng.main_engine is None
    assert eng.next_engine is None
    assert not eng.is_last_engine


def test_basic_engine_is_available():
    eng = _basics.BasicEngine()
    with pytest.raises(_basics.LastEngineException):
        eng.is_last_engine = True
        eng.is_available("FakeCommand")

    def filter(self, cmd):
        if cmd == "supported":
            return True
        return False

    filter_eng = InstructionFilter(filter)
    eng.next_engine = filter_eng
    eng.is_last_engine = False
    assert eng.is_available("supported")
    assert not eng.is_available("something else")


def test_basic_engine_allocate_and_deallocate_qubit_and_qureg():
    eng = _basics.BasicEngine()
    # custom receive function which checks that main_engine does not send
    # any allocate or deallocate gates
    cmd_sent_by_main_engine = []

    def receive(self, cmd_list): cmd_sent_by_main_engine.append(cmd_list)

    eng.receive = types.MethodType(receive, eng)
    # Create test engines:
    saving_backend = DummyEngine(save_commands=True)
    main_engine = MainEngine(backend=saving_backend,
                             engine_list=[eng, DummyEngine()])
    # Allocate and deallocate qubits
    qubit = eng.allocate_qubit()
    # Try to allocate dirty qubit but it should give a non dirty qubit
    not_dirty_qubit = eng.allocate_qubit(dirty=True)

    # Allocate an actual dirty qubit
    def allow_dirty_qubits(self, meta_tag):
        return meta_tag == DirtyQubitTag

    saving_backend.is_meta_tag_handler = types.MethodType(allow_dirty_qubits,
                                                          saving_backend)
    dirty_qubit = eng.allocate_qubit(dirty=True)
    qureg = eng.allocate_qureg(2)
    # Test qubit allocation
    assert isinstance(qubit, list)
    assert len(qubit) == 1 and isinstance(qubit[0], Qubit)
    assert qubit[0] in main_engine.active_qubits
    assert id(qubit[0].engine) == id(eng)
    # Test non dirty qubit allocation
    assert isinstance(not_dirty_qubit, list)
    assert len(not_dirty_qubit) == 1 and isinstance(not_dirty_qubit[0], Qubit)
    assert not_dirty_qubit[0] in main_engine.active_qubits
    assert id(not_dirty_qubit[0].engine) == id(eng)
    # Test dirty_qubit allocation
    assert isinstance(dirty_qubit, list)
    assert len(dirty_qubit) == 1 and isinstance(dirty_qubit[0], Qubit)
    assert dirty_qubit[0] in main_engine.active_qubits
    assert dirty_qubit[0].id in main_engine.dirty_qubits
    assert id(dirty_qubit[0].engine) == id(eng)
    # Test qureg allocation
    assert isinstance(qureg, list)
    assert len(qureg) == 2
    for tmp_qubit in qureg:
        assert tmp_qubit in main_engine.active_qubits
        assert id(tmp_qubit.engine) == id(eng)
    # Test uniqueness of ids
    assert len(set([qubit[0].id, not_dirty_qubit[0].id, dirty_qubit[0].id,
                    qureg[0].id, qureg[1].id])) == 5
    # Test allocate gates were sent
    assert len(cmd_sent_by_main_engine) == 0
    assert len(saving_backend.received_commands) == 5
    for cmd in saving_backend.received_commands:
        assert cmd.gate == AllocateQubitGate()
    assert saving_backend.received_commands[2].tags == [DirtyQubitTag()]
    # Test deallocate gates were sent
    eng.deallocate_qubit(qubit[0])
    eng.deallocate_qubit(not_dirty_qubit[0])
    eng.deallocate_qubit(dirty_qubit[0])
    eng.deallocate_qubit(qureg[0])
    eng.deallocate_qubit(qureg[1])
    assert len(cmd_sent_by_main_engine) == 0
    assert len(saving_backend.received_commands) == 10
    for cmd in saving_backend.received_commands[5:]:
        assert cmd.gate == DeallocateQubitGate()
    assert saving_backend.received_commands[7].tags == [DirtyQubitTag()]


def test_deallocate_qubit_exception():
    eng = _basics.BasicEngine()
    qubit = Qubit(eng, -1)
    with pytest.raises(ValueError):
        eng.deallocate_qubit(qubit)


def test_basic_engine_is_meta_tag_supported():
    eng = _basics.BasicEngine()
    # BasicEngine needs receive function to function so let's add it:

    def receive(self, cmd_list): self.send(cmd_list)

    eng.receive = types.MethodType(receive, eng)
    backend = DummyEngine()
    engine0 = DummyEngine()
    engine1 = DummyEngine()
    engine2 = DummyEngine()

    def allow_dirty_qubits(self, meta_tag):
        if meta_tag == DirtyQubitTag:
            return True
        return False

    engine2.is_meta_tag_handler = types.MethodType(allow_dirty_qubits,
                                                   engine2)
    main_engine = MainEngine(backend=backend,
                             engine_list=[engine0, engine1, engine2])
    assert not main_engine.is_meta_tag_supported("NotSupported")
    assert main_engine.is_meta_tag_supported(DirtyQubitTag)


def test_forwarder_engine():
    backend = DummyEngine(save_commands=True)
    engine0 = DummyEngine()
    main_engine = MainEngine(backend=backend,
                             engine_list=[engine0])

    def cmd_mod_fun(cmd):
        cmd.tags = "NewTag"
        return cmd

    forwarder_eng = _basics.ForwarderEngine(backend, cmd_mod_fun)
    engine0.next_engine = forwarder_eng
    forwarder_eng2 = _basics.ForwarderEngine(engine0)
    main_engine.next_engine = forwarder_eng2
    qubit = main_engine.allocate_qubit()
    H | qubit
    # Test if H gate was sent through forwarder_eng and tag was added
    received_commands = []
    # Remove Allocate and Deallocate gates
    for cmd in backend.received_commands:
        if not (isinstance(cmd.gate, FastForwardingGate) or
                isinstance(cmd.gate, ClassicalInstructionGate)):
            received_commands.append(cmd)
    for cmd in received_commands:
        print(cmd)
    assert len(received_commands) == 1
    assert received_commands[0].gate == H
    assert received_commands[0].tags == "NewTag"
