from __future__ import annotations

import logging
import time
from importlib import resources
from pathlib import Path
from typing import Any

import numpy as np
from joblib import Parallel, delayed
from mqt.bench.tket_helper import get_rebase
from mqt.predictor import Result, ml, reward, rl
from pytket import OpType
from pytket.architecture import Architecture
from pytket.extensions.qiskit import qiskit_to_tk, tk_to_qiskit
from pytket.passes import (
    FullPeepholeOptimise,
    PlacementPass,
    RoutingPass,
)
from pytket.placement import GraphPlacement
from qiskit import QuantumCircuit, transpile

logger = logging.getLogger("mqt-predictor")


def create_qiskit_result(qc: QuantumCircuit, device: dict[str, Any] | None = None) -> Result:
    """Creates a Result object for a given benchmark and device using qiskit for compilation.

    Args:
        benchmark (str): The path to the benchmark to be compiled.
        device (dict[str, Any] | None, optional): The device to be used for compilation. Defaults to None.

    Returns:
        Result: Returns a Result object containing the compiled quantum circuit.
    """
    assert device is not None
    if qc.num_qubits > device["max_qubits"]:
        return Result("qiskit_", -1, None, device["name"])
    start_time = time.time()
    try:
        transpiled_qc_qiskit = transpile(
            qc,
            basis_gates=device["native_gates"],
            coupling_map=device["cmap"],
            optimization_level=3,
            seed_transpiler=1,
        )
    except Exception as e:
        logger.warning("qiskit Transpile Error occurred for: " + device["name"] + " " + str(e))
        return Result("qiskit_" + device["name"], -1, None, device["name"])
    duration = time.time() - start_time
    return Result("qiskit_" + device["name"], duration, transpiled_qc_qiskit, device["name"])


def create_tket_result(
    qc: QuantumCircuit,
    device: dict[str, Any] | None = None,
) -> Result:
    """Creates a Result object for a given benchmark and device using tket for compilation.

    Args:
        benchmark (str): The path to the benchmark to be compiled.
        device (dict[str, Any] | None, optional): The device to be used for compilation. Defaults to None.

    Returns:
        Result: Returns a Result object containing the compiled quantum circuit.
    """
    assert device is not None
    if qc.num_qubits > device["max_qubits"]:
        return Result("tket_" + device["name"], -1, None, device["name"])
    tket_qc = qiskit_to_tk(qc)
    arch = Architecture(device["cmap"])

    native_rebase = get_rebase(device["name"].split("_")[0])
    assert native_rebase is not None

    start_time = time.time()
    try:
        native_rebase.apply(tket_qc)
        FullPeepholeOptimise(target_2qb_gate=OpType.TK2).apply(tket_qc)
        PlacementPass(GraphPlacement(arch)).apply(tket_qc)
        RoutingPass(arch).apply(tket_qc)
        native_rebase.apply(tket_qc)
        duration = time.time() - start_time
        transpiled_qc_tket = tk_to_qiskit(tket_qc)
    except Exception as e:
        logger.warning("tket Transpile Error occurred for: " + device["name"] + " " + str(e))
        return Result("tket_" + device["name"], -1, None, device["name"])

    return Result("tket_" + device["name"], duration, transpiled_qc_tket, device["name"])


def create_mqtpredictor_result(qc: QuantumCircuit, figure_of_merit: reward.figure_of_merit, filename: str) -> Result:
    dev_name = ml.helper.predict_device_for_figure_of_merit(qc, figure_of_merit)
    """ Creates a Result object for a given benchmark and figure of merit using mqt-predictor for compilation.

    Args:
        benchmark (str): The path to the benchmark to be compiled.
        figure_of_merit (reward.reward_functions): The figure of merit to be used for compilation.

    Returns:
        Result: Returns a Result object containing the compiled quantum circuit.
    """
    assert isinstance(dev_name, str)
    dev_index = rl.helper.get_device_index_of_device(dev_name)
    target_filename = filename.split("/")[-1].split(".qasm")[0] + "_" + figure_of_merit + "_" + str(dev_index)
    combined_path_filename = ml.helper.get_path_training_circuits_compiled() / (target_filename + ".qasm")
    if Path(combined_path_filename).exists():
        qc = QuantumCircuit.from_qasm_file(combined_path_filename)
        if qc:
            return Result(
                "mqt-predictor_" + figure_of_merit,
                -1,
                qc,
                dev_name,
            )
    else:
        try:
            qc_compiled = ml.qcompile(qc, figure_of_merit=figure_of_merit)
            if qc_compiled:
                assert isinstance(qc_compiled, tuple)
                return Result(
                    "mqt-predictor_" + figure_of_merit,
                    -1,
                    qc_compiled[0],
                    dev_name,
                )

        except Exception as e:
            logger.warning("mqt-predictor Transpile Error occurred for: " + filename + " " + dev_name + " " + str(e))
    return Result("mqt-predictor_" + figure_of_merit, -1, None, dev_name)


def evaluate_all_sample_circuits() -> None:
    """Evaluates all sample circuits and saves the results to a csv file."""
    res_csv = []

    results = Parallel(n_jobs=-1, verbose=3, backend="threading")(
        delayed(evaluate_sample_circuit)(str(file))
        for file in list(ml.helper.get_path_training_circuits().glob("*.qasm"))
    )
    res_csv.append(list(results[0].keys()))
    for res in results:
        res_csv.append(list(res.values()))
    np.savetxt(
        ml.helper.get_path_results(),
        res_csv,
        delimiter=",",
        fmt="%s",
    )


def evaluate_GHZ_circuits() -> None:
    """Evaluates all GHZ circuits and saves the results to a csv file."""
    res_csv = []

    path = Path(str(resources.files("mqt.predictor"))) / "ml" / "training_data" / "ghz"
    results = Parallel(n_jobs=-1, verbose=3, backend="threading")(
        delayed(evaluate_sample_circuit)(str(file)) for file in list(path.glob("*.qasm"))
    )
    res_csv.append(list(results[0].keys()))
    for res in results:
        res_csv.append(list(res.values()))
    np.savetxt(
        ml.helper.get_path_results(ghz_results=True),
        res_csv,
        delimiter=",",
        fmt="%s",
    )


def evaluate_sample_circuit(filename: str) -> dict[str, Any]:
    """Evaluates a given sample circuit and returns the results as a dictionary.

    Args:
        filename (str): The path to the sample circuit to be evaluated.

    Returns:
        dict[str, Any]: Returns a dictionary containing the results of the evaluation.
    """
    logger.info("Evaluate file: " + filename)

    results: dict[str, Any] = {
        "file_path": str(Path(filename).stem),
        "benchmark_name": str(Path(filename).stem).replace("_", " ").split(" ")[0],
        "num_qubits": str(Path(filename).stem).replace("_", " ").split(" ")[-1],
    }
    qc = QuantumCircuit.from_qasm_file(filename)
    results.update(create_mqtpredictor_result(qc, "expected_fidelity", filename=filename).get_dict())
    results.update(create_mqtpredictor_result(qc, "critical_depth", filename=filename).get_dict())

    for _i, dev in enumerate(rl.helper.get_devices()):
        results.update(create_qiskit_result(qc, dev).get_dict())
        results.update(create_tket_result(qc, dev).get_dict())

    return results
