# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.

from time import monotonic
from typing import Any, Callable, cast, Dict, List, Optional, Tuple, Union, Literal
from .._fs import read_file, list_directory, resolve
from .._http import fetch_github
from .._native import QasmError, Output, run_qasm_program  # type: ignore
from .._types import (
    BitFlipNoise,
    DepolarizingNoise,
    PauliNoise,
    PhaseFlipNoise,
    ShotResult,
    StateDump,
)
from .._interpreter import (
    get_interpreter,
    ipython_helper,
    python_args_to_interpreter_args,
)
from .._native import StateDumpData, NoiseConfig
from .. import telemetry_events
from ._ipython import display_or_print


def run(
    source: Union[str, Callable],
    shots: int = 1024,
    *args: Any,
    on_result: Optional[Callable[[ShotResult], None]] = None,
    save_events: bool = False,
    noise: Optional[
        Union[
            Tuple[float, float, float],
            PauliNoise,
            BitFlipNoise,
            PhaseFlipNoise,
            DepolarizingNoise,
            NoiseConfig,
        ]
    ] = None,
    qubit_loss: Optional[float] = None,
    as_bitstring: bool = False,
    type: Optional[Literal["sparse", "clifford"]] = None,
    num_qubits: Optional[int] = None,
    **kwargs: Any,
) -> List[Any]:
    """
    Runs the given OpenQASM program for the given number of shots.
    Either a full program or a callable with arguments must be provided.
    Each shot uses an independent instance of the simulator.

    :param source: An OpenQASM program. Alternatively, a callable can be provided,
        which must be an already imported global callable.
    :type source: str or Callable
    :param shots: The number of shots to run. Defaults to ``1024``.
    :type shots: int
    :param *args: The arguments to pass to the callable, if one is provided.
    :param on_result: A callback function that will be called with each result.
        Only supported when a callable is provided; raises :class:`~qdk.openqasm.QasmError`
        if used otherwise.
    :type on_result: Callable
    :param save_events: If true, the output of each shot will be saved. If false, they will be printed.
        Only supported when a callable is provided; raises :class:`~qdk.openqasm.QasmError`
        if used otherwise.
    :type save_events: bool
    :param noise: The noise to use in simulation.
    :type noise: Union[Tuple[float, float, float], PauliNoise, BitFlipNoise, PhaseFlipNoise, DepolarizingNoise, NoiseConfig]
    :param qubit_loss: The probability of qubit loss in simulation.
    :type qubit_loss: float
    :param as_bitstring: If true, the result registers will be converted to bitstrings.
    :type as_bitstring: bool
    :param type: The type of simulator to use. If not specified, the default sparse state vector simulation will be used.
    :param num_qubits: The number of qubits to use for the simulation type "clifford".
        If not specified, the Clifford simulator assumes a default of 1000 qubits.
    :param **kwargs: Additional keyword arguments for compiling the source program. Common options:

        - ``name`` (str): The name of the circuit. This is used as the entry point for the program.
        - ``target_profile`` (TargetProfile): The target profile to use for code generation.
        - ``search_path`` (str): The optional search path for resolving file references.
        - ``output_semantics`` (OutputSemantics): The output semantics for the compilation.
        - ``seed`` (int): The seed to use for the random number generator.
    :return: A list of results or runtime errors. If ``save_events`` is true, a list of ``ShotResult`` values is returned.
    :rtype: List[Any]
    :raises QasmError: If there is an error generating, parsing, or analyzing the OpenQASM source.
    :raises QSharpError: If there is an error interpreting the input.
    :raises ValueError: If the number of shots is less than 1.
    :raises QasmError: If ``on_result`` or ``save_events`` are used when running OpenQASM programs.
    """

    ipython_helper()

    if shots < 1:
        raise ValueError("The number of shots must be greater than 0.")

    telemetry_events.on_run_qasm(
        shots, noise=noise is not None, qubit_loss=qubit_loss is not None
    )
    start_time = monotonic()

    results: List[ShotResult] = []

    def on_save_events(output: Output) -> None:
        # Append the output to the last shot's output list
        results[-1]["events"].append(output)
        if output.is_matrix():
            results[-1]["matrices"].append(output)
        elif output.is_state_dump():
            dump_data = cast(StateDumpData, output.state_dump())
            results[-1]["dumps"].append(StateDump(dump_data))
        elif output.is_message():
            results[-1]["messages"].append(str(output))

    callable = None
    source_str: Optional[str] = None
    if isinstance(source, Callable) and hasattr(source, "__global_callable"):
        args = python_args_to_interpreter_args(args)
        callable = source.__global_callable
    elif isinstance(source, str):
        source_str = source

    if type is not None and type == "clifford":
        if noise is not None and not isinstance(noise, NoiseConfig):
            raise ValueError(
                "only `NoiseConfig` is supported when using noise with the clifford simulator."
            )

    noise_config = None
    if isinstance(noise, NoiseConfig):
        noise_config = noise
        noise = None

    if callable:
        for _ in range(shots):
            results.append(
                {
                    "result": None,
                    "events": [],
                    "matrices": [],
                    "dumps": [],
                    "messages": [],
                }
            )
            run_results = get_interpreter().run(
                source_str,
                on_save_events if save_events else display_or_print,
                noise_config,
                noise,
                qubit_loss=qubit_loss,
                callable=callable,
                args=args,
                seed=kwargs.get("seed"),
                sim_type=type,
                num_qubits=num_qubits,
            )
            results[-1]["result"] = run_results

            if on_result:
                on_result(results[-1])

        if not save_events:
            # If we are not saving events, we can just return the results
            # as a list of results.
            results = [result["result"] for result in results]
    else:
        # running the QASM program in isolation means we can't use the
        # interpreter to run the program, so we can't cache the compilation
        # results. This means we need to compile the program for each
        # shot, or we push the shots into the QASM program and compile it once.
        #
        # This breaks the output streaming and event saving.
        if on_result or save_events:
            raise QasmError(
                "The `on_result` and `save_events` parameters are not supported when running QASM programs."
            )

        if source_str is None:
            raise QasmError(
                "source must be a string or a callable with __global_callable attribute"
            )

        # remove any entries from kwargs with a None key or None value
        kwargs = {k: v for k, v in kwargs.items() if k is not None and v is not None}

        if "search_path" not in kwargs:
            kwargs["search_path"] = "."

        kwargs["shots"] = shots
        if type is not None:
            kwargs["sim_type"] = type
        if num_qubits is not None:
            kwargs["num_qubits"] = num_qubits

        results = run_qasm_program(
            source_str,
            display_or_print,
            noise_config,
            noise,
            qubit_loss,
            read_file,
            list_directory,
            resolve,
            fetch_github,
            **kwargs,
        )

    durationMs = (monotonic() - start_time) * 1000
    telemetry_events.on_run_qasm_end(durationMs, shots)

    if as_bitstring:
        from ._utils import as_bitstring as convert_to_bitstring

        results = convert_to_bitstring(results)

    return results
