# 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 collections import Counter
from unittest.mock import patch

import numpy as np
import pytest

from braket.circuits import Observable, ResultType
from braket.ir import jaqcd, openqasm
from braket.task_result import (
    AdditionalMetadata,
    GateModelTaskResult,
    ResultTypeValue,
    TaskMetadata,
)
from braket.task_result.iqm_metadata_v1 import IqmMetadata
from braket.task_result.oqc_metadata_v1 import OqcMetadata
from braket.task_result.rigetti_metadata_v1 import RigettiMetadata
from braket.tasks import GateModelQuantumTaskResult


@pytest.fixture
def task_metadata_shots():
    return TaskMetadata(**{"id": "task_arn", "deviceId": "default", "shots": 100})


@pytest.fixture
def task_metadata_zero_shots():
    return TaskMetadata(**{"id": "task_arn", "deviceId": "default", "shots": 0})


@pytest.fixture
def additional_metadata():
    program = jaqcd.Program(instructions=[jaqcd.CNot(control=0, target=1)])
    return AdditionalMetadata(action=program)


@pytest.fixture
def additional_metadata_openqasm():
    program = openqasm.Program(
        source="""
        qubit[2] q;
        bit[2] c;

        h q[0];
        cnot q[0], q[1];

        c = measure q;
        """
    )
    return AdditionalMetadata(action=program)


@pytest.fixture
def quil_program():
    return """
PRAGMA INITIAL_REWIRING "NAIVE"
RESET
DECLARE ro BIT[2]
PRAGMA PRESERVE_BLOCK
RX(1.5707963267948966) 0
RX(1.5707963267948966) 7
RZ(1.5707963267948966) 0
RZ(1.5707963267948966) 7
RX(-1.5707963267948966) 7
CZ 0 7
RZ(3.141592653589793) 7
RX(1.5707963267948966) 7
RZ(1.5707963267948966) 7
RX(-1.5707963267948966) 7
PRAGMA END_PRESERVE_BLOCK
MEASURE 0 ro[0]
MEASURE 7 ro[1]
"""


@pytest.fixture
def additional_metadata_rigetti(quil_program):
    program = openqasm.Program(
        source="""
        OPENQASM 3.0;
        bit[2] b;
        h $0;
        cnot $0, $7;
        b[0] = measure $0;
        b[1] = measure $7;
        """
    )
    rigetti_metadata = RigettiMetadata(compiledProgram=quil_program)

    return AdditionalMetadata(action=program, rigettiMetadata=rigetti_metadata)


@pytest.fixture
def additional_metadata_iqm():
    source = """
    OPENQASM 3.0;
    bit[2] b;
    h $0;
    cnot $0, $7;
    b[0] = measure $0;
    b[1] = measure $7;
    """
    program = openqasm.Program(source=source)
    iqm_metadata = IqmMetadata(compiledProgram=source)

    return AdditionalMetadata(action=program, iqmMetadata=iqm_metadata)


@pytest.fixture
def qasm2_program():
    return """
    OPENQASM 2.0;
    include "qelib1.inc";

    qreg node[8];
    creg b[2];
    u3(0.5*pi,0.0*pi,1.0*pi) node[0];
    cx node[0],node[1];
    measure node[0] -> b[0];
    measure node[1] -> b[1];
    """


@pytest.fixture
def additional_metadata_oqc(qasm2_program):
    program = openqasm.Program(
        source="""
        OPENQASM 3.0;
        bit[2] b;
        qubit[8] q;
        h q[0];
        cnot q[0], q[7];
        b[0] = measure q[0];
        b[1] = measure q[7];
        """
    )
    oqc_metadata = OqcMetadata(compiledProgram=qasm2_program)

    return AdditionalMetadata(action=program, oqcMetadata=oqc_metadata)


@pytest.fixture
def result_obj_1(task_metadata_shots, additional_metadata):
    return GateModelTaskResult(
        measurements=[[0, 0], [0, 1], [0, 1], [0, 1]],
        measuredQubits=[0, 1],
        taskMetadata=task_metadata_shots,
        additionalMetadata=additional_metadata,
    )


@pytest.fixture
def result_rigetti(result_obj_1, additional_metadata_rigetti):
    result = GateModelQuantumTaskResult.from_object(result_obj_1)
    result.additional_metadata = additional_metadata_rigetti
    return result


@pytest.fixture
def result_oqc(result_obj_1, additional_metadata_oqc):
    result = GateModelQuantumTaskResult.from_object(result_obj_1)
    result.additional_metadata = additional_metadata_oqc
    return result


@pytest.fixture
def result_iqm(result_obj_1, additional_metadata_iqm):
    result = GateModelQuantumTaskResult.from_object(result_obj_1)
    result.additional_metadata = additional_metadata_iqm
    return result


@pytest.fixture
def result_str_1(result_obj_1):
    return result_obj_1.json()


@pytest.fixture
def result_str_2(result_obj_1):
    result_obj_1.taskMetadata.id = "task_arn_2"
    return result_obj_1.json()


@pytest.fixture
def result_str_3(task_metadata_shots, additional_metadata):
    return GateModelTaskResult(
        measurementProbabilities={"011000": 0.9999999999999982},
        measuredQubits=list(range(6)),
        taskMetadata=task_metadata_shots,
        additionalMetadata=additional_metadata,
    ).json()


@pytest.fixture
def result_obj_4(task_metadata_zero_shots, additional_metadata):
    return GateModelTaskResult(
        resultTypes=[
            ResultTypeValue.construct(
                type=jaqcd.Probability(targets=[0]), value=np.array([0.5, 0.5])
            ),
            ResultTypeValue.construct(
                type=jaqcd.StateVector(),
                value=np.array([complex(0.70710678, 0), 0, 0, complex(0.70710678, 0)]),
            ),
            ResultTypeValue.construct(
                type=jaqcd.Expectation(observable=["y"], targets=[0]), value=0.0
            ),
            ResultTypeValue.construct(
                type=jaqcd.Variance(observable=["y"], targets=[0]), value=0.1
            ),
            ResultTypeValue.construct(
                type=jaqcd.Amplitude(states=["00"]), value={"00": complex(0.70710678, 0)}
            ),
        ],
        measuredQubits=list(range(2)),
        taskMetadata=task_metadata_zero_shots,
        additionalMetadata=additional_metadata,
    )


@pytest.fixture
def result_str_4(task_metadata_zero_shots, additional_metadata):
    result = GateModelTaskResult(
        resultTypes=[
            ResultTypeValue(type=jaqcd.Probability(targets=[0]), value=[0.5, 0.5]),
            ResultTypeValue(
                type=jaqcd.StateVector(), value=[(0.70710678, 0), (0, 0), (0, 0), (0.70710678, 0)]
            ),
            ResultTypeValue(type=jaqcd.Expectation(observable=["y"], targets=[0]), value=0.0),
            ResultTypeValue(type=jaqcd.Variance(observable=["y"], targets=[0]), value=0.1),
            ResultTypeValue(type=jaqcd.Amplitude(states=["00"]), value={"00": (0.70710678, 0)}),
        ],
        measuredQubits=list(range(2)),
        taskMetadata=task_metadata_zero_shots,
        additionalMetadata=additional_metadata,
    )
    return result.json()


@pytest.fixture
def result_obj_5(task_metadata_shots):
    return GateModelTaskResult(
        taskMetadata=task_metadata_shots,
        additionalMetadata=AdditionalMetadata(
            action=jaqcd.Program(
                instructions=[jaqcd.CNot(control=0, target=1), jaqcd.CNot(control=2, target=3)],
                results=[jaqcd.Probability(targets=[1]), jaqcd.Expectation(observable=["z"])],
            )
        ),
        measurements=[
            [0, 0, 1, 0],
            [1, 1, 1, 1],
            [1, 0, 0, 1],
            [0, 0, 1, 0],
            [1, 1, 1, 1],
            [0, 1, 1, 1],
            [0, 0, 0, 1],
            [0, 1, 1, 1],
            [0, 0, 0, 0],
            [0, 0, 0, 1],
        ],
        measuredQubits=[0, 1, 2, 3],
    )


@pytest.fixture
def malformatted_results_1(task_metadata_shots, additional_metadata):
    return GateModelTaskResult(
        measuredQubits=list(range(6)),
        taskMetadata=task_metadata_shots,
        additionalMetadata=additional_metadata,
    ).json()


@pytest.fixture
def openqasm_result_obj_shots(task_metadata_shots, additional_metadata_openqasm):
    return GateModelTaskResult.construct(
        measurements=[[0, 0], [0, 1], [0, 1], [0, 1]],
        measuredQubits=[0, 1],
        taskMetadata=task_metadata_shots,
        additionalMetadata=additional_metadata_openqasm,
        resultTypes=[jaqcd.Probability()],
    )


def test_openqasm_shots_calculate_result_types(openqasm_result_obj_shots):
    result = GateModelQuantumTaskResult._from_object_internal_computational_basis_sampling(
        openqasm_result_obj_shots
    )
    assert result.result_types[0].type == jaqcd.Probability()
    assert np.array_equal(result.result_types[0].value, [0.25, 0.75, 0, 0])


test_ir_results = [
    (jaqcd.Probability(targets=[1]), np.array([0.6, 0.4])),
    (jaqcd.Probability(targets=[1, 2]), np.array([0.4, 0.2, 0.0, 0.4])),
    (
        jaqcd.Probability(),
        np.array([0.1, 0.2, 0.2, 0.0, 0.0, 0.0, 0.0, 0.2, 0.0, 0.1, 0.0, 0.0, 0.0, 0.0, 0.0, 0.2]),
    ),
    (jaqcd.Sample(targets=[1], observable=["z"]), np.array([1, -1, 1, 1, -1, -1, 1, -1, 1, 1])),
    (
        jaqcd.Sample(targets=[1, 2], observable=["x", "y"]),
        np.array([-1, 1, 1, -1, 1, 1, 1, 1, 1, 1]),
    ),
    (
        jaqcd.Sample(observable=["z"]),
        [
            np.array([1, -1, -1, 1, -1, 1, 1, 1, 1, 1]),
            np.array([1, -1, 1, 1, -1, -1, 1, -1, 1, 1]),
            np.array([-1, -1, 1, -1, -1, -1, 1, -1, 1, 1]),
            np.array([1, -1, -1, 1, -1, -1, -1, -1, 1, -1]),
        ],
    ),
    (jaqcd.Expectation(targets=[1], observable=["z"]), 0.2),
    (jaqcd.Expectation(targets=[1], observable=[[[[-1, 0], [0, 0]], [[0, 0], [1, 0]]]]), -0.2),
    (jaqcd.Expectation(targets=[1, 2], observable=["z", "y"]), 0.6),
    (jaqcd.Expectation(observable=["z"]), [0.4, 0.2, -0.2, -0.4]),
    (jaqcd.Variance(targets=[1], observable=["z"]), 0.96),
    (jaqcd.Variance(targets=[1], observable=[[[[-1, 0], [0, 0]], [[0, 0], [1, 0]]]]), 0.96),
    (jaqcd.Variance(targets=[1, 2], observable=["z", "y"]), 0.64),
    (jaqcd.Variance(observable=["z"]), [0.84, 0.96, 0.96, 0.84]),
]


def test_get_compiled_circuit_rigetti(result_rigetti, quil_program):
    """Test get_compiled_circuit method."""
    assert result_rigetti.get_compiled_circuit() == quil_program


def test_get_compiled_circuit_oqc(result_oqc, qasm2_program):
    """Test get_compiled_circuit method."""
    assert result_oqc.get_compiled_circuit() == qasm2_program


def test_get_compiled_circuit_iqm(result_iqm):
    """Test get_compiled_circuit method."""
    assert result_iqm.get_compiled_circuit() == result_iqm.additional_metadata.action.source


def test_get_compiled_circuit_no_qhp_metadata(result_obj_1):
    """Test get_compiled_circuit method."""
    result = GateModelQuantumTaskResult.from_object(result_obj_1)
    assert result.get_compiled_circuit() is None


def test_get_compiled_circuit_no_metadata(result_obj_1):
    """Test that the method does not raise an error if metadata is missing."""
    result = GateModelQuantumTaskResult.from_object(result_obj_1)
    result.additional_metadata = None
    assert result.get_compiled_circuit() is None


def test_measurement_counts_from_measurements():
    measurements: np.ndarray = np.array([
        [1, 0, 1, 0],
        [0, 0, 0, 0],
        [1, 0, 1, 0],
        [1, 0, 0, 0],
        [1, 0, 0, 0],
        [1, 0, 1, 0],
    ])
    measurement_counts = GateModelQuantumTaskResult.measurement_counts_from_measurements(
        measurements
    )
    expected_counts: Counter = {"1010": 3, "0000": 1, "1000": 2}
    assert expected_counts == measurement_counts


def test_measurement_probabilities_from_measurement_counts():
    counts = {"00": 1, "01": 1, "10": 1, "11": 97}
    probabilities = {"00": 0.01, "01": 0.01, "10": 0.01, "11": 0.97}

    m_probabilities = GateModelQuantumTaskResult.measurement_probabilities_from_measurement_counts(
        counts
    )

    assert m_probabilities == probabilities


def test_measurements_from_measurement_probabilities():
    shots = 5
    probabilities = {"00": 0.2, "01": 0.2, "10": 0.2, "11": 0.4}
    measurements_list = [["0", "0"], ["0", "1"], ["1", "0"], ["1", "1"], ["1", "1"]]
    expected_results = np.asarray(measurements_list, dtype=int)

    measurements = GateModelQuantumTaskResult.measurements_from_measurement_probabilities(
        probabilities, shots
    )

    assert np.allclose(measurements, expected_results)


def test_from_string_measurements(result_str_1):
    result_obj = GateModelTaskResult.parse_raw(result_str_1)
    task_result = GateModelQuantumTaskResult.from_string(result_str_1)
    expected_measurements = np.asarray(result_obj.measurements, dtype=int)
    assert task_result.task_metadata == result_obj.taskMetadata
    assert task_result.additional_metadata == result_obj.additionalMetadata
    assert np.array2string(task_result.measurements) == np.array2string(expected_measurements)
    assert not task_result.measurement_counts_copied_from_device
    assert not task_result.measurement_probabilities_copied_from_device
    assert task_result.measurements_copied_from_device
    assert task_result.measured_qubits == result_obj.measuredQubits
    assert task_result.values == []
    assert task_result.result_types == []


def test_from_object_result_types(result_obj_5):
    result_obj = result_obj_5
    task_result = GateModelQuantumTaskResult.from_object(result_obj)
    expected_measurements = np.asarray(result_obj.measurements, dtype=int)
    assert np.array2string(task_result.measurements) == np.array2string(expected_measurements)
    assert np.allclose(task_result.values[0], np.array([0.6, 0.4]))
    assert task_result.values[1] == [0.4, 0.2, -0.2, -0.4]
    assert task_result.result_types[0].type == jaqcd.Probability(targets=[1])
    assert task_result.result_types[1].type == jaqcd.Expectation(observable=["z"])


def test_from_string_measurement_probabilities(result_str_3):
    result_obj = GateModelTaskResult.parse_raw(result_str_3)
    task_result = GateModelQuantumTaskResult.from_string(result_str_3)
    assert task_result.measurement_probabilities == result_obj.measurementProbabilities
    shots = 100
    measurement_list = [list("011000") for _ in range(shots)]
    expected_measurements = np.asarray(measurement_list, dtype=int)
    assert np.allclose(task_result.measurements, expected_measurements)
    assert task_result.measurement_counts == Counter(["011000" for _ in range(shots)])
    assert not task_result.measurement_counts_copied_from_device
    assert task_result.measurement_probabilities_copied_from_device
    assert not task_result.measurements_copied_from_device


def test_from_object_equal_to_from_string(result_obj_1, result_str_1, result_str_3):
    assert GateModelQuantumTaskResult.from_object(
        result_obj_1
    ) == GateModelQuantumTaskResult.from_string(result_str_1)
    assert GateModelQuantumTaskResult.from_object(
        GateModelTaskResult.parse_raw(result_str_3)
    ) == GateModelQuantumTaskResult.from_string(result_str_3)


def test_equality(result_str_1, result_str_2):
    result_1 = GateModelQuantumTaskResult.from_string(result_str_1)
    result_2 = GateModelQuantumTaskResult.from_string(result_str_1)
    other_result = GateModelQuantumTaskResult.from_string(result_str_2)
    non_result = "not a quantum task result"

    assert result_1 == result_2
    assert result_1 is not result_2
    assert result_1 != other_result
    assert result_1 != non_result


def test_from_string_simulator_only(result_obj_4, result_str_4):
    result_obj = result_obj_4
    result = GateModelQuantumTaskResult.from_string(result_str_4)
    assert len(result.result_types) == len(result_obj.resultTypes)
    for i in range(len(result.result_types)):
        rt = result.result_types[i]
        expected_rt = result_obj.resultTypes[i]
        assert rt.type == expected_rt.type
        if isinstance(rt.value, np.ndarray):
            assert np.allclose(rt.value, expected_rt.value)
        else:
            assert rt.value == expected_rt.value


def test_get_value_by_result_type(result_obj_4):
    result = GateModelQuantumTaskResult.from_object(result_obj_4)
    assert np.allclose(
        result.get_value_by_result_type(ResultType.Probability(target=0)), result.values[0]
    )
    assert np.allclose(result.get_value_by_result_type(ResultType.StateVector()), result.values[1])
    assert (
        result.get_value_by_result_type(ResultType.Expectation(observable=Observable.Y(), target=0))
        == result.values[2]
    )
    assert (
        result.get_value_by_result_type(ResultType.Variance(observable=Observable.Y(), target=0))
        == result.values[3]
    )
    assert result.get_value_by_result_type(ResultType.Amplitude(state=["00"])) == result.values[4]


@pytest.mark.xfail(raises=ValueError)
def test_get_value_by_result_type_value_error(result_obj_4):
    result = GateModelQuantumTaskResult.from_object(result_obj_4)
    result.get_value_by_result_type(ResultType.Probability(target=[0, 1]))


@pytest.mark.xfail(raises=ValueError)
def test_shots_no_measurements_no_measurement_probs(malformatted_results_1):
    GateModelQuantumTaskResult.from_string(malformatted_results_1)


@pytest.mark.parametrize("ir_result,expected_result", test_ir_results)
def test_calculate_ir_results(ir_result, expected_result):
    ir_string = jaqcd.Program(
        instructions=[jaqcd.H(target=i) for i in range(4)], results=[ir_result]
    ).json()
    measured_qubits = [0, 1, 2, 3]
    measurements = np.array([
        [0, 0, 1, 0],
        [1, 1, 1, 1],
        [1, 0, 0, 1],
        [0, 0, 1, 0],
        [1, 1, 1, 1],
        [0, 1, 1, 1],
        [0, 0, 0, 1],
        [0, 1, 1, 1],
        [0, 0, 0, 0],
        [0, 0, 0, 1],
    ])
    result_types = GateModelQuantumTaskResult._calculate_result_types(
        ir_string, measurements, measured_qubits
    )
    assert len(result_types) == 1
    assert result_types[0].type == ir_result
    assert np.allclose(result_types[0].value, expected_result)


@pytest.mark.xfail(raises=ValueError)
def test_calculate_ir_results_value_error():
    ir_string = json.dumps({"results": [{"type": "foo"}]})
    measured_qubits = [0]
    measurements = np.array([[0]])
    GateModelQuantumTaskResult._calculate_result_types(ir_string, measurements, measured_qubits)


@pytest.mark.parametrize(
    "observable_1, observable_2",
    [
        (
            jaqcd.Expectation(targets=[1, 2], observable=["x"]),
            jaqcd.Expectation(observable=["x"], targets=[1, 2]),
        ),
        pytest.param(
            jaqcd.Expectation(observable=["x"], targets=[1, 2]),
            jaqcd.Expectation(observable=["x"], targets=[2, 1]),
            marks=pytest.mark.xfail,
        ),
        pytest.param(
            jaqcd.Expectation(observable=["x"], targets=[1, 2]),
            jaqcd.Sample(observable=["x"], targets=[2, 1]),
            marks=pytest.mark.xfail,
        ),
        (
            jaqcd.Expectation(
                observable=[
                    [[[0, 0], [0.512345, 0]], [[0.543215, 0], [0, 0]]],
                    [[[1, 0], [1, 0]], [[1, 0], [-1, 0]]],
                ],
                targets=[1, 2],
            ),
            jaqcd.Expectation(
                observable=[
                    [[[0, 0], [0.512345, 0]], [[0.543215, 0], [0, 0]]],
                    [[[1, 0], [1, 0]], [[1, 0], [-1, 0]]],
                ],
                targets=[1, 2],
            ),
        ),
        (
            jaqcd.Expectation(observable=["y", "z"], targets=[1, 2]),
            jaqcd.Expectation(observable=["y", "z"], targets=[1, 2]),
        ),
    ],
)
def test_hash_result_types(observable_1, observable_2):
    assert GateModelQuantumTaskResult._result_type_hash(
        observable_1
    ) == GateModelQuantumTaskResult._result_type_hash(observable_2)


@patch(
    "braket.tasks.gate_model_quantum_task_result.GateModelQuantumTaskResult._calculate_result_types"
)
def test_result_type_skips_computation_already_populated(calculate_result_types_mocked):
    result_str = json.dumps({
        "braketSchemaHeader": {
            "name": "braket.task_result.gate_model_task_result",
            "version": "1",
        },
        "measurements": [[0]],
        "resultTypes": [
            {"type": {"observable": ["z"], "targets": [0], "type": "variance"}, "value": 12.0}
        ],
        "measuredQubits": [0],
        "taskMetadata": {
            "braketSchemaHeader": {"name": "braket.task_result.task_metadata", "version": "1"},
            "id": "arn:aws:braket:us-east-1:1234567890:quantum-task/22a238b2-ae96",
            "shots": 1,
            "deviceId": "arn:aws:braket:::device/quantum-simulator/amazon/dm1",
            "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": 1,
                    "disableQubitRewiring": False,
                },
            },
            "createdAt": "2022-01-12T06:05:22.633Z",
            "endedAt": "2022-01-12T06:05:24.136Z",
            "status": "COMPLETED",
        },
        "additionalMetadata": {
            "action": {
                "braketSchemaHeader": {"name": "braket.ir.openqasm.program", "version": "1"},
                "source": "\nqubit[1] q;\nh q[0];\n#pragma braket result variance z(q[0])\n",
            },
            "simulatorMetadata": {
                "braketSchemaHeader": {
                    "name": "braket.task_result.simulator_metadata",
                    "version": "1",
                },
                "executionDuration": 16,
            },
        },
    })
    res = GateModelQuantumTaskResult.from_string(result_str)
    assert (
        res.get_value_by_result_type(ResultType.Variance(observable=Observable.Z(), target=[0]))
        == 12
    )
    calculate_result_types_mocked.assert_not_called()
