from __future__ import annotations

import logging
from typing import TYPE_CHECKING

import numpy as np
from qiskit.primitives import BaseEstimator, EstimatorResult, PrimitiveJob

from .provider import BlueQubitBackendV2

if TYPE_CHECKING:
    from qiskit.circuit import QuantumCircuit
    from qiskit.quantum_info import SparsePauliOp

logger = logging.getLogger("bluequbit-python-sdk")


class BlueQubitEstimatorV1(BaseEstimator):
    """Evaluates expectation values for provided quantum circuit and observable combinations.

    Qiskit Estimator V1 interface to BlueQubit cloud computation. Compatible with
    ``qiskit_algorithms.VQE`` and other primitives that require the V1 estimator interface.

    .. deprecated::
        ``BaseEstimator`` (V1) is deprecated in Qiskit and will be removed in a future
        release. Prefer :class:`BlueQubitEstimatorV2` for new code.
    """

    def __init__(
        self,
        *,
        backend: BlueQubitBackendV2 | None = None,
        options: dict | None = None,
    ) -> None:
        """Initialize BlueQubitEstimatorV1.

        Args:
            backend: The BlueQubit backend to use. If ``None``, a new backend is
                created using environment variables or the provided options.
            options: Optional dict with keys ``"backend_options"`` (passed to
                ``BlueQubitBackendV2``) and ``"run_options"`` (passed to ``backend.run``).
        """
        super().__init__()
        self._custom_options = options if options is not None else {}
        if backend is None:
            backend = BlueQubitBackendV2(
                **self._custom_options.get("backend_options", {})
            )
        self._run_options = self._custom_options.get("run_options", {})
        self._backend = backend

    def _run(
        self,
        circuits: tuple[QuantumCircuit, ...],
        observables: tuple[SparsePauliOp, ...],
        parameter_values: tuple[tuple[float, ...], ...],
        **run_options: object,
    ) -> PrimitiveJob:
        job = PrimitiveJob(
            self._compute, circuits, observables, parameter_values, run_options
        )
        job._submit()  # noqa: SLF001
        return job

    def _compute(
        self,
        circuits: tuple[QuantumCircuit, ...],
        observables: tuple[SparsePauliOp, ...],
        parameter_values: tuple[tuple[float, ...], ...],
        run_options: dict,
    ) -> EstimatorResult:
        merged_run_options = {**self._run_options, **run_options}
        values = []
        metadata = []
        for circuit, observable, params in zip(
            circuits, observables, parameter_values, strict=False
        ):
            bound_circuit = circuit.assign_parameters(
                dict(zip(circuit.parameters, params, strict=True))
            )
            pauli_sums = list(zip(observable.paulis, observable.coeffs, strict=True))
            bq_job = self._backend.run(
                bound_circuit,
                pauli_sum=pauli_sums,  # type: ignore[arg-type]
                **merged_run_options,
            )
            result = bq_job.result()
            ev = np.real_if_close(result.expectation_value)
            values.append(ev)
            metadata.append({"circuit_metadata": circuit.metadata})
        return EstimatorResult(np.array(values, dtype=np.float64), metadata)
