import json as _json
from dataclasses import dataclass
from enum import Enum
from io import BytesIO
from typing import Tuple, Union

import numpy
import numpy.lib.format as _format
from google.protobuf.json_format import MessageToDict

from qm._QmJobErrors import (
    _handle_job_manager_error,
    MissingElementError,
    ElementWithoutIntermediateFrequencyError,
    ElementWithSingleInputError,
    InvalidElementCorrectionError,
    MissingDigitalInputError,
    InvalidDigitalInputThresholdError,
    InvalidDigitalInputWindowError,
    InvalidDigitalInputPolarityError,
)
from qm._report import ExecutionReport
from qm._results import JobResults
from qm.logger import logger
from qm.pb.frontend_pb2 import (
    IsJobRunningRequest,
    IsJobAcquiringDataRequest,
    IsJobAcquiringDataResponse,
    HaltRequest,
    GetSimulatedQuantumStateRequest,
)
from qm.pb.frontend_pb2 import ResumeRequest, PausedStatusRequest
from qm.pb.inc_qua_pb2 import QuaProgram
from qm.pb.job_manager_pb2 import (
    SetElementCorrectionRequest,
    GetElementCorrectionRequest,
    SetDigitalInputThresholdRequest,
    GetDigitalInputThresholdRequest,
    SetDigitalInputWindowRequest,
    GetDigitalInputWindowRequest,
    SetDigitalInputPolarityRequest,
    ACTIVE_HIGH,
    ACTIVE_LOW,
    GetDigitalInputPolarityRequest,
)
from qm.pb.job_manager_pb2_grpc import JobManagerServiceStub
from qm.pb.job_results_pb2 import PullSimulatorSamplesRequest
from qm.pb.job_results_pb2_grpc import JobResultsServiceStub
from qm.persistence import BinaryAsset
from qm.results.SimulatorSamples import SimulatorSamples
from qm.utils import insert_input_stream_to_job


class AcquiringStatus(Enum):
    AcquireStopped = 0
    NoDataToAcquire = 1
    HasDataToAcquire = 2


@dataclass
class DensityMatrix:
    timestamp: int
    data: numpy.ndarray


class SimulatorOutput:
    def __init__(self, job_id: str, frontend) -> None:
        super().__init__()
        self._id = job_id
        self._frontend = frontend

    def get_quantum_state(self) -> DensityMatrix:
        request = GetSimulatedQuantumStateRequest()
        request.jobId = self._id

        result = self._frontend.GetSimulatedQuantumState(request)
        if result.ok:
            flatten = numpy.array(
                [complex(item.re, item.im) for item in result.state.data]
            )
            n = int(numpy.sqrt(len(flatten)))
            if len(flatten) != (n * n):
                raise RuntimeError("Quantum state matrix is not correct")
            else:
                pass

            matrix = flatten.reshape(n, n)
            timestamp = result.state.timeStamp
            return DensityMatrix(timestamp, matrix)
        else:
            raise RuntimeError("Error while pulling quantum state")


class QmJob:
    def __init__(
        self,
        qmm: "qm.QuantumMachinesManager.QuantumMachinesManager",
        job_id: str,
        execute_response=None,
    ):
        super().__init__()
        self._id = job_id
        self._logger = logger.getChild("job")

        self._simulated_analog_outputs = {"samples": None, "waveforms": None}
        self._simulated_digital_outputs = {"samples": None, "waveforms": None}
        if execute_response is not None and execute_response.HasField("simulated"):
            if execute_response.simulated.HasField("analogOutputs"):
                self._simulated_analog_outputs = MessageToDict(
                    execute_response.simulated.analogOutputs
                )
            if execute_response.simulated.HasField("digitalOutputs"):
                self._simulated_digital_outputs = MessageToDict(
                    execute_response.simulated.digitalOutputs
                )
            if len(execute_response.simulated.errors) > 0:
                raise RuntimeError("\n".join(execute_response.simulated.errors))

        self._qmm = qmm
        self._frontend = qmm._frontend
        self._job_manager_client = JobManagerServiceStub(self.manager._channel)

    def id(self):
        """
        :return: The id of the job
        """
        return self._id

    @property
    def manager(self) -> "qm.QuantumMachinesManager.QuantumMachinesManager":
        """
        The QM object where this job lives

        :return:

        """
        return self._qmm

    def simulated_analog_waveforms(self):
        """
        Return the results of the simulation of quantum elements and analog outputs.

        The returned dictionary has the following keys and entries:

        ``elements``: a dictionary containing the outputs with timestamps and values arranged by quantum elements.

        ``controllers``: a dictionary containing the outputs with timestamps and values arranged by controllers.

            ``ports``: a dictionary containing the outputs with timestamps and values arranged by output ports.

                for each element or output port, the entry is a list of dictionaries with the following information:

                ``timestamp``:

                    The time, in nsec, from the start of the program to the start of the pulse.

                ``samples``:

                    Output information, with ``duration`` given in nsec and ``value`` given normalized OPX output units.


        :return: A dictionary containing output information for the analog outputs of the controller.

        """
        return self._simulated_analog_outputs["waveforms"]

    def simulated_digital_waveforms(self):
        """
        Return the results of the simulation of digital outputs.

        ``controllers``: a dictionary containing the outputs with timestamps and values arranged by controllers.

            ``ports``: a dictionary containing the outputs with timestamps and values arranged by output ports.

                for each element or output port, the entry is a list of dictionaries with the following information:

                ``timestamp``:

                    The time, in nsec, from the start of the program to the start of the pulse.

                ``samples``:

                    A list containing the sequence of outputted values, with ``duration`` given in nsec
                    and ``value`` given as a boolean value.



        :return: A dictionary containing output information for the analog outputs of the controller.
        """
        return self._simulated_digital_outputs["waveforms"]

    def halt(self) -> bool:
        """
        Halts the job on the opx
        """
        request = HaltRequest()
        request.jobId = self._id

        response = self._frontend.Halt(request)
        return response.ok

    def _get_np_simulated_samples(
        self, path=None, include_analog=True, include_digital=True
    ):
        request = PullSimulatorSamplesRequest()
        request.jobId = self._id
        request.includeAnalog = include_analog
        request.includeDigital = include_digital

        writer = BytesIO()
        data_writer = BytesIO()

        for result in self._frontend.PullSimulatorSamples(request):
            if result.ok:
                one_of = result.WhichOneof("body")
                if one_of == "header":
                    _format.write_array_header_2_0(
                        writer,
                        {
                            "descr": _json.loads(result.header.simpleDType),
                            "fortran_order": False,
                            "shape": (result.header.countOfItems,),
                        },
                    )
                elif one_of == "data":
                    data_writer.write(result.data.data)

            else:
                raise RuntimeError("Error while pulling samples")

        data_writer.seek(0)
        for d in data_writer:
            writer.write(d)

        writer.seek(0)
        return numpy.load(writer)

    def get_simulated_samples(self, include_analog=True, include_digital=True):
        """
        Obtain the output samples of a QUA program simulation.

        Samples are returned in an object that holds the controllers in the current simulation,
        where each controller's name will be a property of this object.
        The value of each property of the returned value is an object with the following properties:

        ``analog``:

            holds a dictionary with analog port names as keys and numpy array of samples as values.

        ``digital``:

            holds a dictionary with digital port names as keys and numpy array of samples as values.

        It is also possible to directly plot the outputs using a built-in plot command.

        Example::

            >>> samples = job.get_simulated_samples()
            >>> analog1 = samples.con1.analog["1"]  # obtain analog port 1 of controller "con1"
            >>> digital9 = samples.con1.analog["9"]  # obtain digital port 9 of controller "con1"
            >>> samples.con1.plot()  # Plots all active ports
            >>> samples.con1.plot(analog_ports=['1', '2'], digital_ports=['9'])  # Plots the given output ports

        .. note::

            The simulated digital waveforms are delayed by 136ns relative to the real
            output of the OPX.

        :param include_analog: Should we collect simulated analog samples
        :param include_digital: Should we collect simulated digital samples

        :return: The simulated samples of the job.
        """
        return SimulatorSamples.from_np_array(
            self._get_np_simulated_samples(
                include_analog=include_analog, include_digital=include_digital
            )
        )

    def resume(self):
        """
        Resumes a program that was halted using the pause statement
        """
        request = ResumeRequest()
        request.jobId = self._id
        self._frontend.Resume(request)

    @property
    def result_handles(self) -> JobResults:
        """
        :type: qm._results.JobResults
        :return: The handles that this job generated
        """
        has_job_streaming_state = self.manager._caps.has_job_streaming_state
        if not has_job_streaming_state:
            import qm._results_01

            return qm._results_01.JobResults(
                self._id,
                JobResultsServiceStub(self.manager._channel),
                self.manager.store,
            )
        return JobResults(
            self._id, JobResultsServiceStub(self.manager._channel), self.manager.store
        )

    @property
    def simulator(self) -> SimulatorOutput:
        return SimulatorOutput(self._id, self._frontend)

    def set_element_correction(
        self, element, correction
    ) -> Tuple[float, float, float, float]:
        """
        Sets the correction matrix for correcting gain and phase imbalances
        of an IQ mixer associated with a quantum element.

        Changes will only be done to the current job!

        Values will be rounded to an accuracy of 2^-16.
        Valid values for the correction values are between -2 and (2 - 2^-16).

        Warning - the correction matrix can increase the output voltage which might result in an
        overflow.

        :param element: the name of the element to update the correction for
        :type element: str
        :param correction:

            tuple is of the form (v00, v01, v10, v11) where
            the matrix is
            | v00 v01 |
            | v10 v11 |

        :type correction: tuple
        :return: The correction matrix, after rounding to the OPX resolution.
        """
        request = SetElementCorrectionRequest()
        request.jobId = self._id
        request.qeName = element
        request.correction.v00 = correction[0]
        request.correction.v01 = correction[1]
        request.correction.v10 = correction[2]
        request.correction.v11 = correction[3]

        response = self._job_manager_client.SetElementCorrection(request)
        valid_errors = (
            MissingElementError,
            ElementWithSingleInputError,
            InvalidElementCorrectionError,
            ElementWithoutIntermediateFrequencyError,
        )
        _handle_job_manager_error(request, response, valid_errors)
        return (
            response.correction.v00,
            response.correction.v01,
            response.correction.v10,
            response.correction.v11,
        )

    def get_element_correction(self, element) -> Tuple[float, float, float, float]:
        """
        Gets the correction matrix for correcting gain and phase imbalances
        of an IQ mixer associated with a quantum element.

        :param element: the name of the element to update the correction for
        :type element: str
        :return: The current correction matrix
        """
        request = GetElementCorrectionRequest()
        request.jobId = self._id
        request.qeName = element

        response = self._job_manager_client.GetElementCorrection(request)
        valid_errors = (
            MissingElementError,
            ElementWithSingleInputError,
            ElementWithoutIntermediateFrequencyError,
        )
        _handle_job_manager_error(request, response, valid_errors)
        return (
            response.correction.v00,
            response.correction.v01,
            response.correction.v10,
            response.correction.v11,
        )

    def is_paused(self):
        """
        :return: Returns ``True`` if the job is in a paused state.

        see also:

            ``resume()``
        """
        request = PausedStatusRequest()
        request.jobId = self._id

        response = self._frontend.PausedStatus(request)
        return response.isPaused

    def _is_job_running(self):
        """
        :return: Returns ``True`` if the job is running
        """
        request = IsJobRunningRequest()
        request.jobId = self._id

        response = self._frontend.IsJobRunning(request)
        return response.isRunning

    def _is_data_acquiring(self):
        """
        Returns the data acquiring status.
        The possible statuses are: AcquireStopped, NoDataToAcquire,  HasDataToAcquire

        :return: An AcquiringStatus enum object
        """
        request = IsJobAcquiringDataRequest()
        request.jobId = self._id

        response = self._frontend.IsJobAcquiringData(request)
        result = response.acquiringStatus
        if result == IsJobAcquiringDataResponse.ACQUIRE_STOPPED:
            return AcquiringStatus(0)
        elif result == IsJobAcquiringDataResponse.NO_DATA_TO_ACQUIRE:
            return AcquiringStatus(1)
        elif result == IsJobAcquiringDataResponse.HAS_DATA_TO_ACQUIRE:
            return AcquiringStatus(2)

    def execution_report(self) -> ExecutionReport:
        """
        Get runtime errors report for this job. See :ref:`Guides/error:Runtime Errors`.

        :type: qm._report.ExecutionReport
        :return: An object holding the errors that this job generated.
        """
        return ExecutionReport(self._id, JobResultsServiceStub(self.manager._channel))

    def set_digital_input_threshold(
        self, port: Tuple[str, int], threshold: float
    ) -> float:
        request = SetDigitalInputThresholdRequest()
        request.jobId = self._id
        request.digitalInputPort.SetInParent()
        request.digitalInputPort.controllerName = port[0]
        request.digitalInputPort.portNumber = port[1]
        request.threshold = threshold

        response = self._job_manager_client.SetDigitalInputThreshold(request)
        valid_errors = (MissingDigitalInputError, InvalidDigitalInputThresholdError)
        _handle_job_manager_error(request, response, valid_errors)
        return response.threshold

    def get_digital_input_threshold(self, port: Tuple[str, int]) -> float:
        request = GetDigitalInputThresholdRequest()
        request.jobId = self._id
        request.digitalInputPort.SetInParent()
        request.digitalInputPort.controllerName = port[0]
        request.digitalInputPort.portNumber = port[1]

        response = self._job_manager_client.GetDigitalInputThreshold(request)
        valid_errors = (MissingDigitalInputError,)
        _handle_job_manager_error(request, response, valid_errors)
        return response.threshold

    def set_digital_input_window(self, port: Tuple[str, int], window: int) -> int:
        request = SetDigitalInputWindowRequest()
        request.jobId = self._id
        request.digitalInputPort.SetInParent()
        request.digitalInputPort.controllerName = port[0]
        request.digitalInputPort.portNumber = port[1]
        request.window = window

        response = self._job_manager_client.SetDigitalInputWindow(request)
        valid_errors = (MissingDigitalInputError, InvalidDigitalInputWindowError)
        _handle_job_manager_error(request, response, valid_errors)
        return response.window

    def get_digital_input_window(self, port: Tuple[str, int]) -> int:
        request = GetDigitalInputWindowRequest()
        request.jobId = self._id
        request.digitalInputPort.SetInParent()
        request.digitalInputPort.controllerName = port[0]
        request.digitalInputPort.portNumber = port[1]

        response = self._job_manager_client.GetDigitalInputWindow(request)
        valid_errors = (MissingDigitalInputError,)
        _handle_job_manager_error(request, response, valid_errors)
        return response.window

    def set_digital_input_polarity(self, port: Tuple[str, int], polarity: str) -> str:
        request = SetDigitalInputPolarityRequest()
        request.jobId = self._id
        request.digitalInputPort.SetInParent()
        request.digitalInputPort.controllerName = port[0]
        request.digitalInputPort.portNumber = port[1]
        request.polarity = ACTIVE_HIGH
        if polarity == "HIGH":
            request.polarity = ACTIVE_HIGH
        elif polarity == "LOW":
            request.polarity = ACTIVE_LOW
        else:
            raise InvalidDigitalInputPolarityError(
                f"Invalid value for polarity {polarity}. Valid values are: 'HIGH' "
                f"or 'LOW'"
            )

        response = self._job_manager_client.SetDigitalInputPolarity(request)
        valid_errors = (MissingDigitalInputError, InvalidDigitalInputPolarityError)
        _handle_job_manager_error(request, response, valid_errors)
        return response.polarity

    def get_digital_input_polarity(self, port: Tuple[str, int]) -> str:
        request = GetDigitalInputPolarityRequest()
        request.jobId = self._id
        request.digitalInputPort.SetInParent()
        request.digitalInputPort.controllerName = port[0]
        request.digitalInputPort.portNumber = port[1]

        response = self._job_manager_client.GetDigitalInputPolarity(request)
        valid_errors = (MissingDigitalInputError,)
        _handle_job_manager_error(request, response, valid_errors)
        if response.polarity == ACTIVE_HIGH:
            return "HIGH"
        elif response.polarity == ACTIVE_LOW:
            return "LOW"
        else:
            raise Exception(
                f"Received invalid digital input polarity value {response.polarity} from server."
            )

    def insert_input_stream(
        self,
        input_stream: Union[
            QuaProgram.ArrayVarRefExpression, QuaProgram.AnyScalarExpression
        ],
        data,
    ):
        """
        Insert data to the running job input stream declared in the QUA program.
        The data is then ready to be read by the program using the advance
        input stream QUA statement.
        Note - multiple data entries can be inserted before the data is read by the
        program.

        :param input_stream: The input stream the data is to be inserted to.
        :param data: The data to be inserted.
        """
        insert_input_stream_to_job(self, data, input_stream)


class _SavedResults:
    def __init__(self, asset: BinaryAsset) -> None:
        self._asset = asset
        super().__init__()

    def get_numpy_results(self, *args, **kwargs):
        input = self._asset.for_reading()
        try:
            return numpy.load(input)
        finally:
            input.close()
