from __future__ import annotations
from typing import (
    TYPE_CHECKING,
    Optional,
    List,
    Dict,
    Literal,
    Any,
)


import matplotlib.pyplot as plt

from ionq.schemas.job import (
    JobRequest,
    MultiCircuitWorkload,
    WorkloadInput,
    JobType,
    CircuitWorkload,
    GateSet,
)

from .base import Result, Job, Workload

if TYPE_CHECKING:
    from .backend import Backend


class CircuitResult(Result):
    def __init__(
        self, probabilities: Dict[str, float], counts: Optional[Dict[str, int]] = None
    ):
        self.counts = counts
        self.probabilities = probabilities

    def top_candidates(self, n: Optional[int] = None) -> List[str]:
        if n is None:
            n = len(self.probabilities)
        return sorted(self.probabilities, key=self.probabilities.get, reverse=True)[:n]  # type: ignore

    def plot_results(self, ax=None):
        if ax is None:
            fig, ax = plt.subplots(figsize=(10, 6))

        ax.bar(list(self.probabilities.keys()), list(self.probabilities.values()))
        ax.set_ylabel("Probability")
        ax.set_xlabel("State")
        ax.set_title("Circuit Result")

        if ax is None:
            plt.show()


class CircuitJob(Job[CircuitResult]):

    def _get_counts(self, probabilities: Dict[str, float]) -> Optional[Dict[str, int]]:
        if self.details.shots is None:
            return None
        return {k: int(v * self.details.shots) for k, v in probabilities.items()}

    def results(self) -> List[CircuitResult]:
        if self.id is None:
            raise ValueError("Circuit Job not yet run")

        results = self.backend.results(self.id)

        if self.details.type == JobType.multi_circuit:
            all_probabilities = self.backend.get_job_results(results.get("probabilities", {}).get("url"))
            return [
                CircuitResult(
                    probabilities=probabilities,
                    counts=self._get_counts(probabilities),
                )
                for probabilities in all_probabilities.values()
            ]
        else:
            if "histogram" in results:
                counts = self.backend.get_job_results(results["histogram"]["url"])
                probabilities = {
                    k: v / self.details.shots for k, v in counts.items()
                } if self.details.shots else {}
                return [
                    CircuitResult(
                        probabilities=probabilities,
                        counts=counts,
                    )
                ]
            probabilities = self.backend.get_job_results(results.get("probabilities", {}).get("url"))
            return [
                CircuitResult(
                    probabilities=probabilities,
                    counts=self._get_counts(probabilities),
                )
            ]


class Circuit(Workload[CircuitJob]):

    gateset: Literal["qis", "native"]

    def __init__(
        self,
        circuits: List[Dict[str, Any]],
        gateset: Literal["qis", "native"] = "qis",
        qubits: Optional[int] = None,
        name: Optional[str] = None,
    ):
        self.circuits = circuits
        self.gateset = gateset
        self.qubits = qubits
        self.name = name

    def to_workload_input(
        self,
        params: Optional[List[float]] = None,
    ) -> WorkloadInput:
        if len(self.circuits) > 1:
            return WorkloadInput(
                type=JobType.multi_circuit,
                input=MultiCircuitWorkload(
                    gateset=GateSet(self.gateset),
                    qubits=self.qubits,
                    circuits=self.circuits,
                ),
            )
        return WorkloadInput(
            type=JobType.circuit,
            input=CircuitWorkload(
                gateset=GateSet(self.gateset),
                qubits=self.qubits,
                circuit=self.circuits[0]["circuit"],
            ),
        )

    def create_job(self, id: str, request: JobRequest, backend: Backend) -> CircuitJob:
        return CircuitJob(id, request, backend)
