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

from pathlib import Path
import random
from typing import Callable, Literal, List, Optional, Tuple, TypeAlias, Union
import pyqir
from ._native import (
    QirInstructionId,
    QirInstruction,
    run_clifford,
    run_parallel_shots,
    run_cpu_full_state,
    NoiseConfig,
    GpuContext,
    try_create_gpu_adapter,
)
from pyqir import (
    Function,
    FunctionType,
    Type,
    qubit_type,
    Linkage,
)
from ._qsharp import QirInputData, Result
from typing import TYPE_CHECKING

if TYPE_CHECKING:  # This is in the pyi file only
    from ._native import GpuShotResults


class AggregateGatesPass(pyqir.QirModuleVisitor):
    def __init__(self):
        super().__init__()
        self.gates: List[QirInstruction | Tuple] = []
        self.required_num_qubits = None
        self.required_num_results = None

    def _get_value_as_string(self, value: pyqir.Value) -> str:
        value = pyqir.extract_byte_string(value)
        if value is None:
            return ""
        value = value.decode("utf-8")
        return value

    def run(self, mod: pyqir.Module) -> Tuple[List[QirInstruction | Tuple], int, int]:
        errors = mod.verify()
        if errors is not None:
            raise ValueError(f"Module verification failed: {errors}")

        # verify that the module is base profile
        func = next(filter(pyqir.is_entry_point, mod.functions))
        self.required_num_qubits = pyqir.required_num_qubits(func)
        self.required_num_results = pyqir.required_num_results(func)

        super().run(mod)
        return (self.gates, self.required_num_qubits, self.required_num_results)

    def _on_block(self, block):
        if (
            block.terminator
            and block.terminator.opcode == pyqir.Opcode.BR
            and len(block.terminator.operands) > 1
        ):
            raise ValueError(
                "simulation of programs with branching control flow is not supported"
            )
        super()._on_block(block)

    def _on_call_instr(self, call: pyqir.Call) -> None:
        callee_name = call.callee.name
        if callee_name == "__quantum__qis__ccx__body":
            self.gates.append(
                (
                    QirInstructionId.CCX,
                    pyqir.qubit_id(call.args[0]),
                    pyqir.qubit_id(call.args[1]),
                    pyqir.qubit_id(call.args[2]),
                )
            )
        elif callee_name == "__quantum__qis__cx__body":
            self.gates.append(
                (
                    QirInstructionId.CX,
                    pyqir.qubit_id(call.args[0]),
                    pyqir.qubit_id(call.args[1]),
                )
            )
        elif callee_name == "__quantum__qis__cy__body":
            self.gates.append(
                (
                    QirInstructionId.CY,
                    pyqir.qubit_id(call.args[0]),
                    pyqir.qubit_id(call.args[1]),
                )
            )
        elif callee_name == "__quantum__qis__cz__body":
            self.gates.append(
                (
                    QirInstructionId.CZ,
                    pyqir.qubit_id(call.args[0]),
                    pyqir.qubit_id(call.args[1]),
                )
            )
        elif callee_name == "__quantum__qis__swap__body":
            self.gates.append(
                (
                    QirInstructionId.SWAP,
                    pyqir.qubit_id(call.args[0]),
                    pyqir.qubit_id(call.args[1]),
                )
            )
        elif callee_name == "__quantum__qis__rx__body":
            self.gates.append(
                (
                    QirInstructionId.RX,
                    call.args[0].value,
                    pyqir.qubit_id(call.args[1]),
                )
            )
        elif callee_name == "__quantum__qis__rxx__body":
            self.gates.append(
                (
                    QirInstructionId.RXX,
                    call.args[0].value,
                    pyqir.qubit_id(call.args[1]),
                    pyqir.qubit_id(call.args[2]),
                )
            )
        elif callee_name == "__quantum__qis__ry__body":
            self.gates.append(
                (
                    QirInstructionId.RY,
                    call.args[0].value,
                    pyqir.qubit_id(call.args[1]),
                )
            )
        elif callee_name == "__quantum__qis__ryy__body":
            self.gates.append(
                (
                    QirInstructionId.RYY,
                    call.args[0].value,
                    pyqir.qubit_id(call.args[1]),
                    pyqir.qubit_id(call.args[2]),
                )
            )
        elif callee_name == "__quantum__qis__rz__body":
            self.gates.append(
                (
                    QirInstructionId.RZ,
                    call.args[0].value,
                    pyqir.qubit_id(call.args[1]),
                )
            )
        elif callee_name == "__quantum__qis__rzz__body":
            self.gates.append(
                (
                    QirInstructionId.RZZ,
                    call.args[0].value,
                    pyqir.qubit_id(call.args[1]),
                    pyqir.qubit_id(call.args[2]),
                )
            )
        elif callee_name == "__quantum__qis__h__body":
            self.gates.append((QirInstructionId.H, pyqir.qubit_id(call.args[0])))
        elif callee_name == "__quantum__qis__s__body":
            self.gates.append((QirInstructionId.S, pyqir.qubit_id(call.args[0])))
        elif callee_name == "__quantum__qis__s__adj":
            self.gates.append((QirInstructionId.SAdj, pyqir.qubit_id(call.args[0])))
        elif callee_name == "__quantum__qis__sx__body":
            self.gates.append((QirInstructionId.SX, pyqir.qubit_id(call.args[0])))
        elif callee_name == "__quantum__qis__t__body":
            self.gates.append((QirInstructionId.T, pyqir.qubit_id(call.args[0])))
        elif callee_name == "__quantum__qis__t__adj":
            self.gates.append((QirInstructionId.TAdj, pyqir.qubit_id(call.args[0])))
        elif callee_name == "__quantum__qis__x__body":
            self.gates.append((QirInstructionId.X, pyqir.qubit_id(call.args[0])))
        elif callee_name == "__quantum__qis__y__body":
            self.gates.append((QirInstructionId.Y, pyqir.qubit_id(call.args[0])))
        elif callee_name == "__quantum__qis__z__body":
            self.gates.append((QirInstructionId.Z, pyqir.qubit_id(call.args[0])))
        elif callee_name == "__quantum__qis__m__body":
            self.gates.append(
                (
                    QirInstructionId.M,
                    pyqir.qubit_id(call.args[0]),
                    pyqir.result_id(call.args[1]),
                )
            )
        elif callee_name == "__quantum__qis__mz__body":
            self.gates.append(
                (
                    QirInstructionId.MZ,
                    pyqir.qubit_id(call.args[0]),
                    pyqir.result_id(call.args[1]),
                )
            )
        elif callee_name == "__quantum__qis__mresetz__body":
            self.gates.append(
                (
                    QirInstructionId.MResetZ,
                    pyqir.qubit_id(call.args[0]),
                    pyqir.result_id(call.args[1]),
                )
            )
        elif callee_name == "__quantum__qis__reset__body":
            self.gates.append((QirInstructionId.RESET, pyqir.qubit_id(call.args[0])))
        elif callee_name == "__quantum__qis__move__body":
            self.gates.append(
                (
                    QirInstructionId.Move,
                    pyqir.qubit_id(call.args[0]),
                )
            )
        elif callee_name == "__quantum__rt__result_record_output":
            tag = self._get_value_as_string(call.args[1])
            self.gates.append(
                (
                    QirInstructionId.ResultRecordOutput,
                    str(pyqir.result_id(call.args[0])),
                    tag,
                )
            )
        elif callee_name == "__quantum__rt__tuple_record_output":
            tag = self._get_value_as_string(call.args[1])
            self.gates.append(
                (QirInstructionId.TupleRecordOutput, str(call.args[0].value), tag)
            )
        elif callee_name == "__quantum__rt__array_record_output":
            tag = self._get_value_as_string(call.args[1])
            self.gates.append(
                (QirInstructionId.ArrayRecordOutput, str(call.args[0].value), tag)
            )
        elif (
            callee_name == "__quantum__rt__initialize"
            or callee_name == "__quantum__rt__begin_parallel"
            or callee_name == "__quantum__rt__end_parallel"
            or callee_name == "__quantum__qis__barrier__body"
        ):
            pass
        else:
            raise ValueError(f"Unsupported call instruction: {callee_name}")


class CorrelatedNoisePass(AggregateGatesPass):
    """
    This pass replaces the QIR intrinsics that are in the provided NoiseConfig
    by correlated noise instructions that the simulator understands.
    """

    def __init__(self, noise_config: NoiseConfig):
        super().__init__()
        self.noise_intrinsics_table = noise_config.intrinsics

    def _on_call_instr(self, call: pyqir.Call) -> None:
        callee_name = call.callee.name
        if callee_name in self.noise_intrinsics_table:
            self.gates.append(
                (
                    QirInstructionId.CorrelatedNoise,
                    self.noise_intrinsics_table.get_intrinsic_id(callee_name),
                    [pyqir.qubit_id(arg) for arg in call.args],
                )
            )
        else:
            super()._on_call_instr(call)


class GpuCorrelatedNoisePass(AggregateGatesPass):
    """
    A special case of the CorrelatedNoisePass that uses data loaded
    directly from rust instead of a NoiseConfig object to detect the
    correlated noise intrinsics.
    """

    def __init__(self, noise_table: List[Tuple[int, str, int]]):
        super().__init__()
        self.noise_table = dict()
        for table_id, name, _count in noise_table:
            self.noise_table[name] = table_id

    def _on_call_instr(self, call: pyqir.Call) -> None:
        callee_name = call.callee.name
        if callee_name in self.noise_table:
            self.gates.append(
                (
                    QirInstructionId.CorrelatedNoise,
                    int(self.noise_table[callee_name]),  # Noise table ID
                    [pyqir.qubit_id(qubit) for qubit in call.args],  # qubit args
                )
            )
        else:
            super()._on_call_instr(call)


class OutputRecordingPass(pyqir.QirModuleVisitor):
    _output_str = ""
    _closers = []
    _counters = []

    def process_output(self, bitstring: str):
        return eval(
            self._output_str,
            {
                "o": [
                    Result.Zero if x == "0" else Result.One if x == "1" else Result.Loss
                    for x in bitstring
                ]
            },
        )

    def _on_function(self, function):
        if pyqir.is_entry_point(function):
            super()._on_function(function)
            while len(self._closers) > 0:
                self._output_str += self._closers.pop()
                self._counters.pop()

    def _on_rt_result_record_output(self, call, result, target):
        self._output_str += f"o[{pyqir.result_id(result)}]"
        while len(self._counters) > 0:
            self._output_str += ","
            self._counters[-1] -= 1
            if self._counters[-1] == 0:
                self._output_str += self._closers[-1]
                self._closers.pop()
                self._counters.pop()
            else:
                break

    def _on_rt_array_record_output(self, call, value, target):
        self._output_str += "["
        self._closers.append("]")
        # if len(self._counters) > 0:
        #     self._counters[-1] -= 1
        self._counters.append(value.value)

    def _on_rt_tuple_record_output(self, call, value, target):
        self._output_str += "("
        self._closers.append(")")
        # if len(self._counters) > 0:
        #     self._counters[-1] -= 1
        self._counters.append(value.value)


class DecomposeCcxPass(pyqir.QirModuleVisitor):

    h_func: Function
    t_func: Function
    tadj_func: Function
    cz_func: Function

    def __init__(self):
        super().__init__()

    def _on_module(self, module):
        void = Type.void(module.context)
        qubit_ty = qubit_type(module.context)

        # Find or create all the needed functions.
        for func in module.functions:
            match func.name:
                case "__quantum__qis__h__body":
                    self.h_func = func
                case "__quantum__qis__t__body":
                    self.t_func = func
                case "__quantum__qis__t__adj":
                    self.tadj_func = func
                case "__quantum__qis__cz__body":
                    self.cz_func = func
        if not hasattr(self, "h_func"):
            self.h_func = Function(
                FunctionType(void, [qubit_ty]),
                Linkage.EXTERNAL,
                "__quantum__qis__h__body",
                module,
            )
        if not hasattr(self, "t_func"):
            self.t_func = Function(
                FunctionType(void, [qubit_ty]),
                Linkage.EXTERNAL,
                "__quantum__qis__t__body",
                module,
            )
        if not hasattr(self, "tadj_func"):
            self.tadj_func = Function(
                FunctionType(void, [qubit_ty]),
                Linkage.EXTERNAL,
                "__quantum__qis__t__adj",
                module,
            )
        if not hasattr(self, "cz_func"):
            self.cz_func = Function(
                FunctionType(void, [qubit_ty, qubit_ty]),
                Linkage.EXTERNAL,
                "__quantum__qis__cz__body",
                module,
            )
        super()._on_module(module)

    def _on_qis_ccx(self, call, ctrl1, ctrl2, target):
        self.builder.insert_before(call)
        self.builder.call(self.h_func, [target])
        self.builder.call(self.tadj_func, [ctrl1])
        self.builder.call(self.tadj_func, [ctrl2])
        self.builder.call(self.h_func, [ctrl1])
        self.builder.call(self.cz_func, [target, ctrl1])
        self.builder.call(self.h_func, [ctrl1])
        self.builder.call(self.t_func, [ctrl1])
        self.builder.call(self.h_func, [target])
        self.builder.call(self.cz_func, [ctrl2, target])
        self.builder.call(self.h_func, [target])
        self.builder.call(self.h_func, [ctrl1])
        self.builder.call(self.cz_func, [ctrl2, ctrl1])
        self.builder.call(self.h_func, [ctrl1])
        self.builder.call(self.t_func, [target])
        self.builder.call(self.tadj_func, [ctrl1])
        self.builder.call(self.h_func, [target])
        self.builder.call(self.cz_func, [ctrl2, target])
        self.builder.call(self.h_func, [target])
        self.builder.call(self.h_func, [ctrl1])
        self.builder.call(self.cz_func, [target, ctrl1])
        self.builder.call(self.h_func, [ctrl1])
        self.builder.call(self.tadj_func, [target])
        self.builder.call(self.t_func, [ctrl1])
        self.builder.call(self.h_func, [ctrl1])
        self.builder.call(self.cz_func, [ctrl2, ctrl1])
        self.builder.call(self.h_func, [ctrl1])
        self.builder.call(self.h_func, [target])
        call.erase()


Simulator: TypeAlias = Callable[
    [List[QirInstruction], int, int, int, NoiseConfig, int], str
]


def preprocess_simulation_input(
    input: Union[QirInputData, str, bytes],
    shots: Optional[int] = 1,
    noise: Optional[NoiseConfig] = None,
    seed: Optional[int] = None,
) -> tuple[pyqir.Module, int, Optional[NoiseConfig], int]:
    if shots is None:
        shots = 1
    # If no seed specified, generate a random u32 to use
    if seed is None:
        seed = random.randint(0, 2**32 - 1)
    if isinstance(noise, tuple):
        raise ValueError(
            "Specifying Pauli noise via a tuple is not supported. Use a NoiseConfig instead."
        )

    context = pyqir.Context()
    if isinstance(input, QirInputData):
        mod = pyqir.Module.from_ir(context, str(input))
    elif isinstance(input, str):
        mod = pyqir.Module.from_ir(context, input)
    else:
        mod = pyqir.Module.from_bitcode(context, input)

    return (mod, shots, noise, seed)


def run_qir_clifford(
    input: Union[QirInputData, str, bytes],
    shots: Optional[int] = 1,
    noise: Optional[NoiseConfig] = None,
    seed: Optional[int] = None,
) -> List:
    (mod, shots, noise, seed) = preprocess_simulation_input(input, shots, noise, seed)
    if noise is None:
        (gates, num_qubits, num_results) = AggregateGatesPass().run(mod)
    else:
        (gates, num_qubits, num_results) = CorrelatedNoisePass(noise).run(mod)
    recorder = OutputRecordingPass()
    recorder.run(mod)

    return list(
        map(
            recorder.process_output,
            run_clifford(gates, num_qubits, num_results, shots, noise, seed),
        )
    )


def run_qir_cpu(
    input: Union[QirInputData, str, bytes],
    shots: Optional[int] = 1,
    noise: Optional[NoiseConfig] = None,
    seed: Optional[int] = None,
) -> List:
    (mod, shots, noise, seed) = preprocess_simulation_input(input, shots, noise, seed)
    if noise is None:
        (gates, num_qubits, num_results) = AggregateGatesPass().run(mod)
    else:
        (gates, num_qubits, num_results) = CorrelatedNoisePass(noise).run(mod)
    recorder = OutputRecordingPass()
    recorder.run(mod)

    return list(
        map(
            recorder.process_output,
            run_cpu_full_state(gates, num_qubits, num_results, shots, noise, seed),
        )
    )


def run_qir_gpu(
    input: Union[QirInputData, str, bytes],
    shots: Optional[int] = 1,
    noise: Optional[NoiseConfig] = None,
    seed: Optional[int] = None,
) -> List[str]:
    (mod, shots, noise, seed) = preprocess_simulation_input(input, shots, noise, seed)
    # Ccx is not support in the GPU simulator, decompose it
    DecomposeCcxPass().run(mod)
    if noise is None:
        (gates, num_qubits, num_results) = AggregateGatesPass().run(mod)
    else:
        (gates, num_qubits, num_results) = CorrelatedNoisePass(noise).run(mod)
    recorder = OutputRecordingPass()
    recorder.run(mod)

    return list(
        map(
            recorder.process_output,
            run_parallel_shots(gates, shots, num_qubits, num_results, noise, seed),
        )
    )


def prepare_qir_with_correlated_noise(
    input: Union[QirInputData, str, bytes],
    noise_tables: List[Tuple[int, str, int]],
) -> Tuple[List[QirInstruction], int, int]:
    # Turn the input into a QIR module
    (mod, _, _, _) = preprocess_simulation_input(input, None, None, None)

    # Ccx is not support in the GPU simulator, decompose it
    DecomposeCcxPass().run(mod)

    # Extract the gates including correlated noise instructions
    (gates, required_num_qubits, required_num_results) = GpuCorrelatedNoisePass(
        noise_tables
    ).run(mod)

    return (gates, required_num_qubits, required_num_results)


class GpuSimulator:
    """
    Represents a GPU-based QIR simulator. This is a 'full state' simulator that can simulate
    quantum programs, including non-Clifford gates, up to a limit of 27 qubits.
    """

    def __init__(self):
        self.gpu_context = GpuContext()

    def load_noise_tables(
        self,
        noise_dir: str,
    ):
        """
        Loads noise tables from the specified directory path. For each .csv file found in the directory,
        the noise table is loaded and associated with a unique identifier. The name of the file (without the .csv extension)
        is used as the label for the noise table, which should match the QIR instruction that will apply noise using this table.

        If testing various noise models, you may load new noise models at any time by calling this method again
        with a different directory path. Previously loaded noise tables will be replaced. The program currently loaded
        into the simulator (if any) will remain loaded, but any subsequent calls to `run_shots` will use the newly loaded noise tables.

        Each line of the table should be of the format: "IXYZ,1.345e-4" where IXYZ is a string of Pauli operators
        representing the error on each qubit (Z applying to the first qubit argument, Y to the second, etc.), and the second value
        is the corresponding error probability for that specific Pauli string.

        Blank lines, lines starting with #, or lines that start with the string "pauli" (i.e., a column header) are ignored.
        """
        self.tables = self.gpu_context.load_noise_tables(noise_dir)

    def set_program(self, input: Union[QirInputData, str, bytes]):
        """
        Load the QIR program into the GPU simulator, preparing it for execution. You may load and run
        multiple programs sequentially by calling this method multiple times before calling `run_shots`
        without needing to create a new simulator instance or reloading noise tables.
        """
        (self.gates, self.required_num_qubits, self.required_num_results) = (
            prepare_qir_with_correlated_noise(
                input, self.tables if not self.tables is None else []
            )
        )
        self.gpu_context.set_program(
            self.gates, self.required_num_qubits, self.required_num_results
        )

    def run_shots(self, shots: int, seed: Optional[int] = None) -> "GpuShotResults":
        """
        Run the loaded QIR program for the specified number of shots, using an optional seed for reproducibility.
        If noise is to be applied, ensure that noise has been loaded prior to running shots.
        """
        seed = seed if seed is not None else random.randint(0, 2**32 - 1)
        return self.gpu_context.run_shots(shots, seed=seed)


def run_qir(
    input: Union[QirInputData, str, bytes],
    shots: Optional[int] = 1,
    noise: Optional[NoiseConfig] = None,
    seed: Optional[int] = None,
    type: Optional[Literal["clifford", "cpu", "gpu"]] = None,
) -> List[str]:
    """
    Simulate the given QIR source.

    Args:
        input: The QIR source to simulate.
        type: The type of simulator to use.
            Use `"clifford"` if your QIR only contains Clifford gates and measurements.
            Use `"gpu"` if you have a GPU available in your system.
            Use `"cpu"` as a fallback option if you don't have a GPU in your system.
            If `None` (default), the GPU simulator will be tried first, falling back to
            CPU if a suitable GPU device could not be located.
        shots: The number of shots to run.
        noise: A noise model to use in the simulation.
        seed: A seed for reproducibility.

    Returns:
        A list of measurement results, in the order they happened during the simulation.
    """
    if type is None:
        try:
            try_create_gpu_adapter()
            type = "gpu"
        except OSError:
            type = "cpu"

    match type:
        case "clifford":
            return run_qir_clifford(input, shots, noise, seed)
        case "cpu":
            return run_qir_cpu(input, shots, noise, seed)
        case "gpu":
            return run_qir_gpu(input, shots, noise, seed)
        case _:
            raise ValueError(f"Invalid simulator type: {type}")
