import subprocess
from pathlib import Path
from colorama import Fore, Style
from subprocess_utils import run_command_with_reporter


def compile_with_raptor(network_path, raptor_onnx_path: Path, output_base: Path,
                        crossbar_size, crossbar_count, cwd=None, reporter=None):
    # Define the arguments, with the possibility to set crossbar size and count
    args = [
        network_path,
        "-o",
        output_base,
        "--maccel=PIM",
        "--EmitPimCodegen",
        # "--use-experimental-conv-impl=true",
        f"--crossbar-size={crossbar_size}",
        f"--crossbar-count={crossbar_count}",
    ]

    try:
        run_command_with_reporter(
            [str(raptor_onnx_path)] + [str(arg) for arg in args],
            cwd=cwd,
            reporter=reporter,
        )
        if reporter is None:
            print(Fore.GREEN + "Raptor execution successful" + Style.RESET_ALL)
    except subprocess.CalledProcessError:
        if reporter is None:
            print(Fore.RED + "Raptor execution failed" + Style.RESET_ALL)
        raise