#   Copyright 2022 Entropica Labs
#
#   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 os
import pytest
import unittest
import networkx as nw
from braket.jobs.local import LocalQuantumJob
from unittest.mock import MagicMock

from openqaoa.problems import MinimumVertexCover
from openqaoa.algorithms import AWSJobs
from openqaoa.algorithms import QAOA, RQAOA
from openqaoa.backends import create_device


class TestingAwsJobs(unittest.TestCase):
    def setUp(self):
        # the input data directory opt/braket/input/data
        os.environ["AMZN_BRAKET_INPUT_DIR"] = "./tests/jobs_test_input/"
        # the output directory opt/braket/model to write ob results to
        os.environ["AMZN_BRAKET_JOB_RESULTS_DIR"] = "/oq_release_tests/testing_jobs/"
        # the name of the job
        os.environ["AMZN_BRAKET_JOB_NAME"] = "oq_release_test"
        # the checkpoint directory
        os.environ["AMZN_BRAKET_CHECKPOINT_DIR"] = "oq_test_suite_checkpoint"
        # the hyperparameter
        os.environ["AMZN_BRAKET_HP_FILE"] = ""
        # the device ARN (AWS Resource Number)
        os.environ[
            "AMZN_BRAKET_DEVICE_ARN"
        ] = "arn:aws:braket:::device/quantum-simulator/amazon/sv1"
        # the output S3 bucket, as specified in the CreateJob request’s OutputDataConfig
        os.environ["AMZN_BRAKET_OUT_S3_BUCKET"] = "amazon-braket-us-east-1-oq-testing"
        # the entry point as specified in the CreateJob request’s ScriptModeConfig
        os.environ["AMZN_BRAKET_SCRIPT_ENTRY_POINT"] = ""
        # the compression type as specified in the CreateJob request’s ScriptModeConfig
        os.environ["AMZN_BRAKET_SCRIPT_COMPRESSION_TYPE"] = ""
        # the S3 location of the user’s script as specified in the CreateJob request’s ScriptModeConfig
        os.environ["AMZN_BRAKET_SCRIPT_S3_URI"] = ""
        # the S3 location where the SDK would store the task results by default for the job
        os.environ["AMZN_BRAKET_TASK_RESULTS_S3_URI"] = ""
        # the S3 location where the job results would be stored, as specified in CreateJob request’s OutputDataConfig
        os.environ["AMZN_BRAKET_JOB_RESULTS_S3_PATH"] = ""
        # the string that should be passed to CreateQuantumTask’s jobToken parameter for quantum tasks created in the job container
        # os.environ["AMZN_BRAKET_JOB_TOKEN"] = ''

        self.n_qubits = 10

        self.vc = MinimumVertexCover(
            nw.circulant_graph(self.n_qubits, [1]), field=1.0, penalty=10
        ).qubo

    def testOsEnvironAssignment(self):

        qaoa_workflow = AWSJobs(algorithm="QaoA")
        assert qaoa_workflow.algorithm == "qaoa"
        assert qaoa_workflow.input_dir == "./tests/jobs_test_input/"
        assert qaoa_workflow.device.device_name == os.environ["AMZN_BRAKET_DEVICE_ARN"]

        rqaoa_workflow = AWSJobs(algorithm="rqAoa")
        assert rqaoa_workflow.algorithm == "rqaoa"
        assert rqaoa_workflow.device.device_name == os.environ["AMZN_BRAKET_DEVICE_ARN"]

    @pytest.mark.docker_aws
    def testLocalJob(self):
        """Test an end-to-end qaoa running on a local docker instance"""

        input_data_path = os.path.join(
            os.environ["AMZN_BRAKET_INPUT_DIR"], "input_data/"
        )

        # Create the qubo and the qaoa
        q = QAOA()
        q.set_classical_optimizer(maxiter='5')
        q.set_device(
            create_device("aws", "arn:aws:braket:::device/quantum-simulator/amazon/sv1")
        )
        
        ### The following lines are needed to fool the github actions into correctly executing q.compile() !!
        q.device.check_connection = MagicMock(return_value = True)
        q.device.qpu_connected = True
        q.device.provider_connected = True
        q.device.n_qubits = self.n_qubits
        q.device.backend_device = ''
        q.device.aws_region = 'us-east-1'

        q.compile(self.vc)
        q.dump(
            file_name="openqaoa_params.json",
            file_path=input_data_path,
            prepend_id=False,
            overwrite=True,
        )

        job = LocalQuantumJob.create(
            device="arn:aws:braket:::device/quantum-simulator/amazon/sv1",
            source_module="./tests/jobs_test_input/aws_braket_source_module/openqaoa_qaoa_script.py",
            image_uri="amazon-braket-oq-dev",
            input_data={"input_data": input_data_path},
        )

        assert (job.state() == "COMPLETED") and (job.result() != None) == True


    @pytest.mark.docker_aws
    def testLocalJobRQAOA(self):
        """Test an end-to-end rqaoa running on a local docker instance"""

        input_data_path = os.path.join(
            os.environ["AMZN_BRAKET_INPUT_DIR"], "input_data/"
        )

        # Create the rqaoa
        r = RQAOA()
        r.set_rqaoa_parameters(n_cutoff=6)
        r.set_classical_optimizer(maxiter=3, save_intermediate=False)
        r.set_device(
            create_device("aws", "arn:aws:braket:::device/quantum-simulator/amazon/sv1")
        )
        
        ### The following lines are needed to fool the github actions into correctly executing q.compile() !!
        r.device.check_connection = MagicMock(return_value = True)
        r.device.qpu_connected = True
        r.device.provider_connected = True
        r.device.n_qubits = self.n_qubits
        r.device.backend_device = ''
        r.device.aws_region = 'us-east-1'

        r.compile(self.vc)
        r.dump(
            file_name="openqaoa_params.json",
            file_path=input_data_path,
            prepend_id=False,
            overwrite=True,
        )

        job = LocalQuantumJob.create(
            device="arn:aws:braket:::device/quantum-simulator/amazon/sv1",
            source_module="./tests/jobs_test_input/aws_braket_source_module/openqaoa_rqaoa_script.py",
            image_uri="amazon-braket-oq-dev",
            input_data={"input_data": input_data_path},
        )

        assert (job.state() == "COMPLETED") and (job.result() != None) == True


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