# Copyright 2022-2025 Qiskit on IQM developers
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License 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.
"""Circuit execution jobs."""

from __future__ import annotations

from collections import Counter
from datetime import date
from typing import TYPE_CHECKING, Any
import warnings

from iqm.iqm_client import (
    DEFAULT_TIMEOUT_SECONDS,
    CircuitBatch,
    CircuitJob,
    CircuitMeasurementResults,
    CircuitMeasurementResultsBatch,
)
from iqm.iqm_client import (
    JobStatus as IQMJobStatus,
)
from iqm.iqm_client.errors import APITimeoutError, CircuitExecutionError
from iqm.qiskit_iqm.qiskit_to_iqm import MeasurementKey
import numpy as np
from qiskit.providers import JobStatus, JobV1
from qiskit.result import Counts, Result

from exa.common.errors.station_control_errors import StationControlError
from iqm.pulse import Circuit
from iqm.station_control.interface.models import HeraldingMode

if TYPE_CHECKING:
    from iqm.qiskit_iqm.iqm_provider import IQMBackend


class IQMJob(JobV1):
    """Implementation of Qiskit's job interface to handle circuit execution on an IQM Server.

    Args:
        backend: Backend instance initiating this job.
        job: Circuit execution job generated by IQM Server.
        kwargs: Arguments to be passed to the initializer of the parent class.

    """

    def __init__(self, backend: IQMBackend, job: CircuitJob, **kwargs):
        super().__init__(backend, job_id=str(job.job_id))
        # kwargs go into self.metadata, Qiskit doesn't currently do anything with it
        self._iqm_job = job
        self._iqm_result: None | list[tuple[str, list[str], Counts]] = None
        self.circuit_metadata: list | None = None  # Metadata that was originally associated with circuits by user

    @staticmethod
    def _iqm_format_results(
        iqm_results_batch: CircuitMeasurementResultsBatch,
        circuits: CircuitBatch,
        requested_shots: int,
        expect_exact_shots: bool,
    ) -> list[tuple[str, list[str], Counts]]:
        """Convert the measurement results for a batch of circuits into the Qiskit format.

        Args:
            iqm_results_batch: measurement results for the circuit batch
            circuits: circuits that were executed
            requested_shots: number of shots requested
            expect_exact_shots: iff True, we must get exactly as many shots as requested
        Returns:
            ``(circuit_name, measurements, counts)`` tuples, one tuple for each circuit in the batch.
            ``measurements`` are a list of bitstrings, one per shot, representing the state of the classical
            registers after the shot.

        """
        qiskit_measurement_results_batch = [
            IQMJob._iqm_format_measurement_results(measurement_results, requested_shots, expect_exact_shots)
            for measurement_results in iqm_results_batch
        ]
        return [
            (
                circuit.name if isinstance(circuit, Circuit) else f"circuit_{i}",
                qiskit_measurement_results,
                Counts(Counter(qiskit_measurement_results)),
            )
            for i, (qiskit_measurement_results, circuit) in enumerate(zip(qiskit_measurement_results_batch, circuits))
        ]

    @staticmethod
    def _iqm_format_measurement_results(
        measurement_results: CircuitMeasurementResults, requested_shots: int, expect_exact_shots: bool = True
    ) -> list[str]:
        """Convert the measurement results from a circuit into the Qiskit format.

        Args:
            measurement_results: measurement results for a single circuit
            requested_shots: number of shots requested
            expect_exact_shots: iff True, we must get exactly as many shots as requested
        Returns:
            For each shot, a bitstring representing the state of the classical registers after the
            shot, in little-endian order.

        """
        # Mapping from creg index (in the circuit) to an array with shape (shots, len(creg)) with the results.
        formatted_results: dict[int, np.ndarray] = {}
        for k, v in measurement_results.items():
            # measurement keys encode data about the classical registers in the original Qiskit circuit
            mk = MeasurementKey.from_string(k)
            res = np.array(v, dtype=int)
            shots = len(res)
            if shots == 0 and not expect_exact_shots:
                warnings.warn(
                    "Received measurement results containing zero shots. "
                    "In case you are using non-default heralding mode, this could be because of bad calibration."
                )
                res = np.array([])
            else:
                # in Qiskit each measurement is a separate single-qubit instruction. qiskit-iqm assigns unique
                # measurement key to each such instruction, so only one column is expected per measurement key.
                if res.shape[1] != 1:
                    raise ValueError(f"Measurement result {mk} has the wrong shape {res.shape}, expected (*, 1)")
                res = res[:, 0]

            if expect_exact_shots and shots != requested_shots:
                raise ValueError(f"Expected {requested_shots} shots but got {shots} for measurement result {mk}")

            # group the measurements into cregs, fill in zeros for unused bits
            creg = formatted_results.setdefault(mk.creg_idx, np.zeros((shots, mk.creg_len), dtype=int))
            creg[:, mk.clbit_idx] = res

        # TODO If the original circuit has a creg that is not used at all we won't know about it here,
        # and thus cannot include it (containing only zeros) in the result strings.

        # Number of shots is the same for all measurement keys.
        # Qiskit uses the little-endian convention in presenting the result bitstrings
        # (both between and within registers), hence the [::-1]
        return [
            " ".join("".join(map(str, res[s, :])) for _, res in sorted(formatted_results.items()))[::-1]
            for s in range(shots)
        ]

    def submit(self) -> None:
        raise NotImplementedError(
            "You should never have to submit jobs by calling this method. When running circuits through "
            "RemoteIQMBackend, the submission will happen under the hood. The job instance that you get is only for "
            "checking the progress and retrieving the results of the submitted job."
        )

    def cancel(self) -> bool:
        """Attempt to cancel the job.

        Returns:
            True if the job was cancelled successfully, False otherwise

        """
        try:
            self._iqm_job.cancel()
            return True
        except StationControlError as e:
            warnings.warn(f"Failed to cancel job: {e}")
            return False

    def result(self, *, timeout: float = DEFAULT_TIMEOUT_SECONDS, cancel_after_timeout: bool = False) -> Result:
        """Retrieve job result within defined timeout.

        Args:
            timeout: Time limit for the job to finish, in seconds.
            cancel_after_timeout: Whether client will try to cancel the job if timeout exceeded.

        Returns:
            Result if job finished within ``timeout``.

        Raises:
            APITimeoutError: Waiting for results exceeded timeout.

        """
        job = self._iqm_job
        # TODO we might want to cache these as well
        circuits, job_parameters = job.payload()

        # wait for job to finish if it hasn't already
        if job.status not in IQMJobStatus.terminal_statuses():
            status = job.wait_for_completion(timeout_secs=timeout)
            if status not in IQMJobStatus.terminal_statuses():
                msg = f"The job {job.job_id} didn't finish in {timeout} seconds."
                if cancel_after_timeout:
                    # Cancel the job if client was unable to get the results within the timeout
                    job.cancel()
                    msg += " Cancelled."
                raise APITimeoutError(msg)

        # job has reached a terminal status, which means no status update is necessary, and we can return a Result
        result_dict: dict[str, Any] = {
            "backend_name": self.backend().name,
            "backend_version": "",
            "qobj_id": "",
            "job_id": str(job.job_id),
            "success": False,
            "date": date.today().isoformat(),
            "results": [],
            # the ones below go into result._metadata
            "circuits": circuits,
            "parameters": job_parameters,
            "timeline": job.data.timeline.copy(),
        }

        # try to return results if job completed successfully
        if job.status == IQMJobStatus.COMPLETED:
            if self._iqm_result is None:
                # retrieve and cache the results
                if (result_batch := job.result()) is None:
                    raise CircuitExecutionError(
                        f"No results were available for job {job.job_id} even though it is completed."
                    )
                # IQMBackend.run() populates IQMJob.circuit_metadata, so it may be None if this IQMJob
                # was created manually from a job_id. In that case retrieve circuit metadata from
                # the job payload which we just retrieved.
                if self.circuit_metadata is None:
                    self.circuit_metadata = [c.metadata if isinstance(c, Circuit) else {} for c in circuits]

                # If no heralding, for all circuits we expect the same number of shots which is the
                # shots requested by user.
                self._iqm_result = self._iqm_format_results(
                    result_batch,
                    circuits,
                    job_parameters.shots,
                    job_parameters.heralding_mode == HeraldingMode.NONE,
                )

            # return the results
            result_dict["success"] = True  # job is successful iff it is IQMJobStatus.COMPLETED (and we got result data)
            result_dict["results"] = [
                {
                    "shots": len(measurement_results),
                    "success": True,
                    "data": {
                        "memory": measurement_results,
                        "counts": counts,
                        "metadata": self.circuit_metadata[i] if self.circuit_metadata is not None else {},
                    },
                    "header": {"name": name},
                    "calibration_set_id": job.data.compilation.calibration_set_id if job.data.compilation else None,
                }
                for i, (name, measurement_results, counts) in enumerate(self._iqm_result)
            ]
        return Result.from_dict(result_dict)

    def status(self) -> JobStatus:
        job = self._iqm_job
        # terminal statuses need not be updated, they're terminal
        status = job.status if job.status in IQMJobStatus.terminal_statuses() else job.update()

        if status == IQMJobStatus.COMPLETED:
            return JobStatus.DONE
        if status == IQMJobStatus.FAILED:
            return JobStatus.ERROR
        if status == IQMJobStatus.CANCELLED:
            return JobStatus.CANCELLED
        if status == IQMJobStatus.PROCESSING:
            return JobStatus.RUNNING
        # IQMJobStatus.WAITING
        return JobStatus.QUEUED

    def queue_position(self, refresh: bool = False) -> int | None:
        """Return the position of the job in the server queue.

        Note:
            The position is not yet implemented and this function will always
            return ``None``. The ``refresh`` argument is ignored.

        Args:
            refresh: If ``True``, re-query the server to get the latest value.
                Otherwise return the cached value.

        Returns:
            Position in the queue or ``None`` if position is unknown or not applicable.

        """
        return None

    def error_message(self) -> str | None:
        """Returns the error message if job has failed, otherwise returns None."""
        self._iqm_job.update()
        return self._iqm_job._errors or None
