################################################################################
# © Copyright 2022 Zapata Computing Inc.
################################################################################
from dataclasses import dataclass
from typing import List, Optional, Protocol

import numpy as np

from ..circuits import Circuit
from ..measurements import ExpectationValues
from ..operators import PauliRepresentation
from .circuit_runner import CircuitRunner


@dataclass(frozen=True)
class EstimationTask:
    """
    Data class defining an estimation problem.

    Args:
        operator: Operator for which we want to calculate the expectation values
        circuit: Circuit used for evaluating the operator
        constraints: Define constraints used in the estimation process,
            e.g. number of shots or target accuracy.
    """

    operator: PauliRepresentation
    circuit: Circuit
    number_of_shots: Optional[int]


class EstimationPreprocessor(Protocol):
    """Protocol defining function which transforms a list of EstimationTasks
    into another list of EstimationTasks.
    """

    def __call__(self, estimation_tasks: List[EstimationTask]) -> List[EstimationTask]:
        ...


class EstimateExpectationValues(Protocol):
    """Protocol defining a function that estimates expectation values for a list of
    estimation tasks. Implementations of this protocol should obey the following rules:

    1. Return one ExpectationValue for each EstimationTask.
    2. The order in which ExpectationValues are returned should match the order
        in which EstimationTasks were provided.
    3. Number of entries in each ExpectationValue is not restricted.
    4. Output ExpectationValue should include coefficients of the terms/operators.
    5. estimation_tasks can include tasks where operator consists of a constant term or
        contains a constant term. The implemented method should include the
        contributions of such constant terms in the return value.
    """

    def __call__(
        self, runner: CircuitRunner, estimation_tasks: List[EstimationTask]
    ) -> List[ExpectationValues]:
        """Estimate expectation values using given runner."""


class EstimationTasksFactory(Protocol):
    """Factory from producing estimation tasks from R^n vectors.

    For instance, this can be used with ansatzes where produced estimation tasks
    are evaluating circuit.
    """

    def __call__(self, parameters: np.ndarray) -> List[EstimationTask]:
        """Produce estimation tasks for given parameters."""
