# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"). You
# may not use this file except in compliance with the License. A copy of
# the License is located at
#
#     http://aws.amazon.com/apache2.0/
#
# or in the "license" file accompanying this file. This file 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 json
from unittest.mock import Mock

from braket.aws import AwsQuantumTaskBatch

DWAVE_ARN = "arn:aws:braket:::device/qpu/d-wave/Advantage_system1"
RIGETTI_ARN = "arn:aws:braket:::device/qpu/rigetti/Ankaa-2"
IONQ_ARN = "arn:aws:braket:us-east-1::device/qpu/ionq/Harmony"
OQC_ARN = "arn:aws:braket:eu-west-2::device/qpu/oqc/Lucy"
SV1_ARN = "arn:aws:braket:::device/quantum-simulator/amazon/sv1"
DM1_ARN = "arn:aws:braket:::device/quantum-simulator/amazon/dm1"
TN1_ARN = "arn:aws:braket:::device/quantum-simulator/amazon/tn1"
XANADU_ARN = "arn:aws:braket:us-east-1::device/qpu/xanadu/Borealis"

RIGETTI_REGION = "us-west-1"


class MockS3:
    MOCK_TASK_METADATA = {
        "braketSchemaHeader": {"name": "braket.task_result.task_metadata", "version": "1"},
        "id": "task_arn",
        "shots": 100,
        "deviceId": "default",
    }

    MOCK_S3_RESULT_GATE_MODEL = json.dumps({
        "braketSchemaHeader": {
            "name": "braket.task_result.gate_model_task_result",
            "version": "1",
        },
        "measurements": [[0, 0], [0, 0], [0, 0], [1, 1]],
        "measuredQubits": [0, 1],
        "taskMetadata": MOCK_TASK_METADATA,
        "additionalMetadata": {
            "action": {
                "braketSchemaHeader": {"name": "braket.ir.jaqcd.program", "version": "1"},
                "instructions": [{"control": 0, "target": 1, "type": "cnot"}],
            },
        },
    })

    MOCK_S3_RESULT_GATE_MODEL_WITH_RESULT_TYPES = json.dumps({
        "braketSchemaHeader": {
            "name": "braket.task_result.gate_model_task_result",
            "version": "1",
        },
        "measurements": [[0, 0], [0, 0], [0, 0], [1, 1]],
        "measuredQubits": [0, 1],
        "resultTypes": [
            {
                "type": {"observable": ["h", "x"], "targets": [0, 1], "type": "expectation"},
                "value": 0.7071067811865474,
            },
            {
                "type": {"states": ["01", "10", "00", "11"], "type": "amplitude"},
                "value": {
                    "01": [0.0, 0.0],
                    "10": [0.0, 0.0],
                    "00": [0.7071067811865475, 0.0],
                    "11": [0.7071067811865475, 0.0],
                },
            },
        ],
        "taskMetadata": MOCK_TASK_METADATA,
        "additionalMetadata": {
            "action": {
                "braketSchemaHeader": {"name": "braket.ir.jaqcd.program", "version": "1"},
                "instructions": [{"control": 0, "target": 1, "type": "cnot"}],
            }
        },
    })

    MOCK_S3_RESULT_PROGRAM_SET = json.dumps({
        "braketSchemaHeader": {
            "name": "braket.task_result.program_set_task_result",
            "version": "1",
        },
        "programResults": [
            {
                "braketSchemaHeader": {
                    "name": "braket.task_result.program_result",
                    "version": "1",
                },
                "executableResults": [
                    {
                        "braketSchemaHeader": {
                            "name": "braket.task_result.program_set_executable_result",
                            "version": "1",
                        },
                        "measurements": [
                            [0, 0],
                            [0, 0],
                            [1, 1],
                            [0, 0],
                            [1, 1],
                            [0, 0],
                            [1, 1],
                            [0, 0],
                            [1, 1],
                            [0, 0],
                            [1, 1],
                            [0, 0],
                            [0, 0],
                            [0, 0],
                            [1, 1],
                            [1, 1],
                            [1, 1],
                            [0, 0],
                            [1, 1],
                            [0, 0],
                        ],
                        "measuredQubits": [0, 1],
                        "inputsIndex": 0,
                    },
                    {
                        "braketSchemaHeader": {
                            "name": "braket.task_result.program_set_executable_failure",
                            "version": "1",
                        },
                        "inputsIndex": 0,
                        "failureMetadata": {
                            "failureReason": "QPU was sick, should be good again after getting some sleep",
                            "retryable": True,
                            "category": "DEVICE",
                        },
                    },
                ],
                "source": {
                    "braketSchemaHeader": {
                        "name": "braket.ir.openqasm.program",
                        "version": "1",
                    },
                    "source": "OPENQASM 3.0;\nbit[2] b;\nqubit[2] q;\nh q[0];\ncnot q[0], q[1];\nb[0] = measure q[0];\nb[1] = measure q[1];",  # noqa
                    "inputs": {"theta": [0.12, 2.1]},
                },
                "additionalMetadata": {
                    "simulatorMetadata": {
                        "braketSchemaHeader": {
                            "name": "braket.task_result.simulator_metadata",
                            "version": "1",
                        },
                        "executionDuration": 50,
                    }
                },
            }
        ],
        "taskMetadata": {
            "braketSchemaHeader": {
                "name": "braket.task_result.program_set_task_metadata",
                "version": "1",
            },
            "id": "arn:aws:braket:us-west-2:667256736152:quantum-task/bfebc86f-e4ed-4d6f-8131-addd1a49d6dc",  # noqa
            "deviceId": "arn:aws:braket:::device/quantum-simulator/amazon/sv1",
            "requestedShots": 120,
            "successfulShots": 100,
            "programMetadata": [{"executables": [{}]}],
            "deviceParameters": {
                "braketSchemaHeader": {
                    "name": "braket.device_schema.simulators.gate_model_simulator_device_parameters",
                    "version": "1",
                },
                "paradigmParameters": {
                    "braketSchemaHeader": {
                        "name": "braket.device_schema.gate_model_parameters",
                        "version": "1",
                    },
                    "qubitCount": 5,
                    "disableQubitRewiring": False,
                },
            },
            "createdAt": "2024-10-15T19:06:58.986Z",
            "endedAt": "2024-10-15T19:07:00.382Z",
            "status": "COMPLETED",
            "totalFailedExecutables": 1,
        },
    })

    MOCK_S3_RESULT_ANNEALING = json.dumps(
        {
            "braketSchemaHeader": {
                "name": "braket.task_result.annealing_task_result",
                "version": "1",
            },
        },
    )

    MOCK_S3_RESULT_ANNEALING = json.dumps({
        "braketSchemaHeader": {
            "name": "braket.task_result.annealing_task_result",
            "version": "1",
        },
        "solutions": [[-1, -1, -1, -1], [1, -1, 1, 1], [1, -1, -1, 1]],
        "solutionCounts": [3, 2, 4],
        "values": [0.0, 1.0, 2.0],
        "variableCount": 4,
        "taskMetadata": {
            "id": "task_arn",
            "shots": 100,
            "deviceId": DWAVE_ARN,
        },
        "additionalMetadata": {
            "action": {
                "type": "ISING",
                "linear": {"0": 0.3333, "1": -0.333, "4": -0.333, "5": 0.333},
                "quadratic": {"0,4": 0.667, "0,5": -1.0, "1,4": 0.667, "1,5": 0.667},
            },
            "dwaveMetadata": {
                "activeVariables": [0],
                "timing": {
                    "qpuSamplingTime": 100,
                    "qpuAnnealTimePerSample": 20,
                    "qpuAccessTime": 10917,
                    "qpuAccessOverheadTime": 3382,
                    "qpuReadoutTimePerSample": 274,
                    "qpuProgrammingTime": 9342,
                    "qpuDelayTimePerSample": 21,
                    "postProcessingOverheadTime": 117,
                    "totalPostProcessingTime": 117,
                    "totalRealTime": 10917,
                    "runTimeChip": 1575,
                    "annealTimePerRun": 20,
                    "readoutTimePerRun": 274,
                },
            },
        },
    })

    MOCK_S3_RESULT_PHOTONIC_MODEL = json.dumps({
        "braketSchemaHeader": {
            "name": "braket.task_result.photonic_model_task_result",
            "version": "1",
        },
        "measurements": [[[1, 2, 3, 4]], [[4, 3, 2, 1]], [[0, 0, 0, 0]]],
        "taskMetadata": {
            "id": "task_arn",
            "shots": 3,
            "deviceId": XANADU_ARN,
        },
        "additionalMetadata": {
            "action": {
                "source": "Vac | q[0]",
            },
            "xanaduMetadata": {
                "compiledProgram": "DECLARE ro BIT[2];",
            },
        },
    })

    MOCK_S3_RESULT_ANALOG_HAMILTONIAN_SIMULTION = json.dumps({
        "braketSchemaHeader": {
            "name": "braket.task_result.analog_hamiltonian_simulation_task_result",
            "version": "1",
        },
        "taskMetadata": {
            "id": "task_arn",
            "shots": 3,
            "deviceId": "mock_arn",
        },
        "measurements": [
            {
                "shotMetadata": {"shotStatus": "Success"},
                "shotResult": {
                    "preSequence": [1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1],
                    "postSequence": [0, 0, 0, 1, 0, 1, 1, 0, 0, 1, 0],
                },
            },
            {
                "shotMetadata": {"shotStatus": "Partial Success"},
                "shotResult": {
                    "preSequence": [1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1],
                    "postSequence": None,
                },
            },
            {
                "shotMetadata": {"shotStatus": "Failure"},
                "shotResult": {"preSequence": None, "postSequence": None},
            },
        ],
    })


def run_and_assert(
    aws_quantum_task_mock,
    device,
    default_s3_folder,
    default_shots,
    default_poll_timeout,
    default_poll_interval,
    circuit,
    s3_destination_folder,  # Treated as positional arg
    shots,  # Treated as positional arg
    poll_timeout_seconds,  # Treated as positional arg
    poll_interval_seconds,  # Treated as positional arg
    inputs,  # Treated as positional arg
    gate_definitions,  # Treated as positional arg
    reservation_arn,  # Treated as positional arg
    extra_args,
    extra_kwargs,
):
    task_mock = Mock()
    aws_quantum_task_mock.return_value = task_mock

    run_args = []
    if s3_destination_folder is not None:
        run_args.append(s3_destination_folder)
    if shots is not None:
        run_args.append(shots)
    if poll_timeout_seconds is not None:
        run_args.append(poll_timeout_seconds)
    if poll_interval_seconds is not None:
        run_args.append(poll_interval_seconds)
    if inputs is not None:
        run_args.append(inputs)
    if gate_definitions is not None:
        run_args.append(gate_definitions)
    run_args += extra_args or []
    run_kwargs = extra_kwargs or {}
    if reservation_arn:
        run_kwargs.update({"reservation_arn": reservation_arn})

    task = device.run(circuit, *run_args, **run_kwargs)
    assert task == task_mock

    create_args, create_kwargs = _create_task_args_and_kwargs(
        default_s3_folder,
        default_shots,
        default_poll_timeout,
        default_poll_interval,
        s3_destination_folder,
        shots,
        poll_timeout_seconds,
        poll_interval_seconds,
        inputs,
        gate_definitions,
        reservation_arn,
        extra_args,
        extra_kwargs,
    )

    aws_quantum_task_mock.assert_called_with(
        device._aws_session, device.arn, circuit, *create_args, **create_kwargs
    )


def run_batch_and_assert(
    aws_quantum_task_mock,
    aws_session_mock,
    device,
    default_s3_folder,
    default_shots,
    default_poll_timeout,
    default_poll_interval,
    circuits,
    s3_destination_folder,
    shots,
    max_parallel,
    max_connections,
    poll_timeout_seconds,
    poll_interval_seconds,
    inputs,
    gate_definitions,
    reservation_arn,
    extra_args,
    extra_kwargs,
):
    task_mock = Mock()
    task_mock.state.return_value = "COMPLETED"
    aws_quantum_task_mock.return_value = task_mock
    new_session_mock = Mock()
    aws_session_mock.return_value = new_session_mock

    run_args = []
    if s3_destination_folder is not None:
        run_args.append(s3_destination_folder)
    if shots is not None:
        run_args.append(shots)
    if max_parallel is not None:
        run_args.append(max_parallel)
    if max_connections is not None:
        run_args.append(max_connections)
    if poll_timeout_seconds is not None:
        run_args.append(poll_timeout_seconds)
    if poll_interval_seconds is not None:
        run_args.append(poll_interval_seconds)
    if inputs is not None:
        run_args.append(inputs)
    if gate_definitions is not None:
        run_args.append(gate_definitions)
    run_args += extra_args or []
    run_kwargs = extra_kwargs or {}
    if reservation_arn:
        run_kwargs.update({"reservation_arn": reservation_arn})

    batch = device.run_batch(circuits, *run_args, **run_kwargs)
    assert batch.tasks == [task_mock for _ in range(len(circuits))]

    create_args, create_kwargs = _create_task_args_and_kwargs(
        default_s3_folder,
        default_shots,
        default_poll_timeout,
        default_poll_interval,
        s3_destination_folder,
        shots,
        poll_timeout_seconds,
        poll_interval_seconds,
        inputs,
        gate_definitions,
        reservation_arn,
        extra_args,
        extra_kwargs,
    )

    max_pool_connections = max_connections or AwsQuantumTaskBatch.MAX_CONNECTIONS_DEFAULT

    # aws_session_mock.call_args.kwargs syntax is newer than Python 3.7
    assert aws_session_mock.call_args[1]["config"].max_pool_connections == max_pool_connections
    aws_quantum_task_mock.assert_any_call(
        new_session_mock, device.arn, circuits[0], *create_args, **create_kwargs
    )
    aws_quantum_task_mock.assert_any_call(
        new_session_mock, device.arn, circuits[1], *create_args, **create_kwargs
    )


def _create_task_args_and_kwargs(
    default_s3_folder,
    default_shots,
    default_poll_timeout,
    default_poll_interval,
    s3_folder,
    shots,
    poll_timeout_seconds,
    poll_interval_seconds,
    inputs,
    gate_definitions,
    reservation_arn,
    extra_args,
    extra_kwargs,
):
    create_args = [
        s3_folder if s3_folder is not None else default_s3_folder,
        shots if shots is not None else default_shots,
    ]
    create_args += extra_args or []
    create_kwargs = extra_kwargs or {}
    create_kwargs.update({
        "poll_timeout_seconds": (
            poll_timeout_seconds if poll_timeout_seconds is not None else default_poll_timeout
        ),
        "poll_interval_seconds": (
            poll_interval_seconds if poll_interval_seconds is not None else default_poll_interval
        ),
        "inputs": inputs,
        "gate_definitions": gate_definitions,
        "reservation_arn": reservation_arn,
    })
    return create_args, create_kwargs
