# Copyright 2023 Agnostiq Inc.

from typing import Any, Dict, List, Tuple

from pennylane import numpy as np
from pennylane.tape import QuantumScript
from pennylane.transforms import broadcast_expand, map_batch_transform
from pydantic import BaseModel

from ...shared.classes.exceptions import CovalentSDKError
from ..executors import CloudQExecutor
from .qclients import RemoteQClient


class SubmissionRecord(BaseModel):
    """
    A container for useful information about submissions.
    The middleware uses this information to post-process results.
    """

    batch_id: str
    num_qscripts: int
    num_qscripts_expanded: int
    pennylane_active_return: bool

    @property
    def contains_batch(self) -> bool:
        """
        Indicates if the original qscripts were expanded non-trivially.
        """
        return self.num_qscripts != self.num_qscripts_expanded

    @property
    def batch_shape(self) -> Tuple[int, int]:
        """
        Shape of the returned vector.
        """
        return (self.num_qscripts, self.num_qscripts_expanded // self.num_qscripts)


class MiddleWare:
    """
    Middleware object that communicates with the QServer via the Remote QClient.
    """

    def __init__(self) -> None:
        self.qclient = RemoteQClient()

        # Stores records of submission for potential post-processing.
        self._submissions: Dict[str, SubmissionRecord] = {}

    def __new__(cls):
        # Making this a singleton class.
        if not hasattr(cls, 'instance'):
            cls.instance = super(MiddleWare, cls).__new__(cls)
        return cls.instance

    # The following attributes are properties
    # because the qclient might change over time
    # and every time it gets changed, we shouldn't
    # have to set things like:
    # self.database = self.qclient.database
    # Thus, we override the access of these attributes
    # and return/set these dynamically depending upon
    # what the qclient is at that point in time.

    @property
    def selector(self):
        """
        Defers selection to QClient's selector.
        """
        return self.qclient.selector

    @selector.setter
    def selector(self, selector_func):
        self.qclient.selector = selector_func

    @property
    def database(self):
        """
        Points to the QElectron database interface.
        """
        return self.qclient.database

    # Need `QElectronInfo` and `QNodeSpecs` in own module.
    # Can't import these objects here because of circular imports.
    def run_circuits_async(
        self,
        qscripts: List[QuantumScript],
        executors: List[CloudQExecutor],
        qelectron_info,  # NOTE: missing type
        qnode_specs,  # NOTE: missing type
    ) -> str:
        """
        Submit circuits to the QServer via the Remote QClient.

        Args:
            qscripts: List of circuits to be executed.
            executors: :ist of executors to use/
            qelectron_info: QNode device and Global pennylane settings.
            qnode_specs: QNode specifications generated by `qml.specs()`.

        Returns:
            batch_id: Unique ID associated with the submission.
        """
        # Converted batched qscripts into multiple qscripts.
        expanded_qscripts = expand_batch(qscripts)

        # Submit batched qscripts individually.
        batch_id = self.qclient.submit(expanded_qscripts, executors, qelectron_info, qnode_specs)

        # Make record of submission. Post-processing may be required.
        self._record_submission(batch_id, qscripts, expanded_qscripts, qelectron_info)

        return batch_id

    def get_results(self, batch_id: str) -> Any:
        """
        Retrieve results from the QServer via the Remote QClient.

        Args:
            batch_id: Unique ID associated with a submission.

        Returns:
            results: Results of the execution.
        """
        # Get raw results from QServer.
        raw_results = self.qclient.get_results(batch_id)

        # Post-process raw results if needed.
        results = self._post_process(raw_results, batch_id)

        return results

    def _post_process(self, raw_results: List[Any], batch_id: str) -> Any:
        """
        Re-package raw results as necessary. Handle batch submissions.
        """
        # Get submission details.
        submission = self._recall_submission(batch_id)

        if not submission.contains_batch:
            # Post-processing not necessary.
            # Unpack single results. Return multi-return or gradient results as-is.
            return raw_results[0] if len(raw_results) == 1 else raw_results

        # Post-process batched results if needed.
        res = np.asarray(raw_results).squeeze()

        # Adapted from Qiskit device in `devices_base.py` in OS QElectrons.
        shape = submission.batch_shape
        if submission.pennylane_active_return:
            # The case `res.ndim > 1` is for multiple-return circuits.
            return [res] if res.ndim > 1 else list(res.reshape(shape))

        res = res.reshape(-1)
        return [res] if res.ndim > 1 else [res.reshape(shape)]

    def _record_submission(
        self,
        batch_id: str,
        qscripts: List[QuantumScript],
        qscripts_expanded: List[QuantumScript],
        qelectron_info,  # NOTE: missing type
    ) -> None:
        """
        Stores submission details for post-processing.
        """
        if batch_id in self._submissions:
            raise CovalentSDKError(f"Middleware has duplicate batch ID '{batch_id}'")

        self._submissions[batch_id] = SubmissionRecord(
            batch_id=batch_id,
            num_qscripts=len(qscripts),
            num_qscripts_expanded=len(qscripts_expanded),
            pennylane_active_return=qelectron_info.pennylane_active_return,
        )

    def _recall_submission(self, batch_id: str) -> SubmissionRecord:
        """
        Retrieves submission details for post-processing.
        """
        if batch_id not in self._submissions:
            raise CovalentSDKError(f"Middleware missing record for batch ID '{batch_id}'")

        return self._submissions.pop(batch_id)


def expand_batch(circuits: List[QuantumScript]) -> List[QuantumScript]:
    """
    Expand batched circuits into multiple circuits with single parameters.
    """

    if not any(circuit.batch_size for circuit in circuits):
        # No circuit contains a batch. Return the same list of circuits.
        return circuits

    # Expand the circuits.
    expanded_circuits, _ = map_batch_transform(broadcast_expand, circuits)
    return expanded_circuits


middleware = MiddleWare()
