import argparse
import json
import numpy as np
import subprocess
from pathlib import Path
from colorama import Style, Fore
from onnx_utils import gen_random_inputs, save_inputs_to_files, onnx_io, write_inputs_to_memory_bin, _ONNX_TO_NP
from raptor import compile_with_raptor
from gen_network_runner import gen_network_runner


def compile_onnx_network(network_onnx_path, raptor_path, raptor_dir, runner_dir):
    subprocess.run([raptor_path, network_onnx_path, "--EmitONNXIR"], check=True)
    subprocess.run([raptor_path, network_onnx_path], check=True)
    parent = network_onnx_path.parent
    stem = network_onnx_path.stem
    so_path = parent / f"{stem}.so"
    mlir_path = parent / f"{stem}.onnx.mlir"
    tmp_path = parent / f"{stem}.tmp"
    moved_so = runner_dir / so_path.name
    moved_mlir = raptor_dir / mlir_path.name
    so_path.rename(moved_so)
    mlir_path.rename(moved_mlir)
    tmp_path.unlink(missing_ok=True)
    return moved_so, moved_mlir


def build_onnx_runner(source_dir, build_dir):
    subprocess.run(["cmake", source_dir], cwd=build_dir, check=True)
    subprocess.run(["cmake", "--build", ".", "-j"], cwd=build_dir, check=True)
    return build_dir / "runner"


def build_dump_ranges(config_path, outputs_descriptor):
    with open(config_path) as f:
        output_addresses = json.load(f)["outputs_addresses"]
    ranges = []
    for addr, (_, _, dtype_code, shape) in zip(output_addresses, outputs_descriptor):
        byte_size = int(np.prod(shape)) * np.dtype(_ONNX_TO_NP[dtype_code]).itemsize
        ranges.append(f"{addr},{byte_size}")
    return ",".join(ranges)


def run_pim_simulator(simulator_dir, pim_dir, output_bin_path, dump_ranges):
    subprocess.run(
        ["cargo", "run", "--release", "--package", "pim-simulator", "--bin", "pim-simulator", "--",
         "-f", str(pim_dir), "-o", str(output_bin_path), "-d", dump_ranges],
        cwd=simulator_dir, check=True
    )


def parse_pim_simulator_outputs(output_bin_path, outputs_descriptor):
    raw = output_bin_path.read_bytes()
    arrays = []
    offset = 0
    for _, _, dtype_code, shape in outputs_descriptor:
        dtype = np.dtype(_ONNX_TO_NP[dtype_code])
        count = int(np.prod(shape))
        array = np.frombuffer(raw, dtype=dtype, count=count, offset=offset).reshape(shape)
        offset += count * dtype.itemsize
        arrays.append(array)
    return arrays


def validate_outputs(sim_arrays, runner_out_dir, outputs_descriptor, threshold=1e-3):
    all_passed = True
    for sim_array, (oi, name, _, shape) in zip(sim_arrays, outputs_descriptor):
        csv_name = f"output{oi}_{name}.csv"
        runner_array = np.loadtxt(runner_out_dir / csv_name, delimiter=',', dtype=np.float32).reshape(shape)
        max_diff = float(np.max(np.abs(sim_array.astype(np.float64) - runner_array.astype(np.float64))))
        passed = max_diff <= threshold
        status = Fore.GREEN + "[PASS]" if passed else Fore.RED + "[FAIL]"
        print(f"  {name}: max diff = {max_diff:.6e} {status}" + Style.RESET_ALL)
        if not passed:
            all_passed = False
    return all_passed


def validate_network(network_onnx_path, raptor_path, onnx_include_dir,
                     simulator_dir, crossbar_size=64, crossbar_count=8, threshold=1e-3):
    network_onnx_path = Path(network_onnx_path).resolve()
    raptor_path = Path(raptor_path).resolve()
    onnx_include_dir = Path(onnx_include_dir).resolve()
    simulator_dir = Path(simulator_dir).resolve()

    workspace_dir = network_onnx_path.parent
    raptor_dir = workspace_dir / "raptor"
    runner_dir = workspace_dir / "runner"
    runner_build_dir = runner_dir / "build"
    Path.mkdir(raptor_dir, exist_ok=True)
    Path.mkdir(runner_build_dir, parents=True, exist_ok=True)

    print(Style.BRIGHT + "\nCompiling the onnx network:" + Style.RESET_ALL)
    network_so_path, network_mlir_path = compile_onnx_network(network_onnx_path, raptor_path, raptor_dir, runner_dir)

    print(Style.BRIGHT + "\nGenerating and building the runner:" + Style.RESET_ALL)
    gen_network_runner(network_onnx_path, network_so_path, onnx_include_dir, out=runner_dir / "runner.c")
    runner_path = build_onnx_runner(runner_dir, runner_build_dir)

    print(Style.BRIGHT + "\nGenerating random inputs:" + Style.RESET_ALL)
    inputs_descriptor, outputs_descriptor = onnx_io(network_onnx_path)
    inputs_list, _inputs_dict = gen_random_inputs(inputs_descriptor)
    flags, _files = save_inputs_to_files(network_onnx_path, inputs_list, out_dir=workspace_dir / "inputs")

    print(Style.BRIGHT + "\nRunning inference with the runner:" + Style.RESET_ALL)
    out_dir = workspace_dir / "outputs"
    Path.mkdir(out_dir, exist_ok=True)
    run_cmd = [runner_path, *flags]
    run_cmd += ["--save-csv-dir", f"{out_dir}"]
    subprocess.run(run_cmd, cwd=runner_build_dir, check=True)

    print(Style.BRIGHT + "\nCompiling for PIM with Raptor:" + Style.RESET_ALL)
    compile_with_raptor(network_mlir_path, raptor_path, crossbar_size, crossbar_count)

    print(Style.BRIGHT + "\nRunning PIM simulation:" + Style.RESET_ALL)
    pim_dir = raptor_dir / "pim"
    write_inputs_to_memory_bin(pim_dir / "memory.bin", pim_dir / "config.json", inputs_list)
    simulation_dir = workspace_dir / "simulation"
    Path.mkdir(simulation_dir, exist_ok=True)
    dump_ranges = build_dump_ranges(pim_dir / "config.json", outputs_descriptor)
    output_bin_path = simulation_dir / "out.bin"
    run_pim_simulator(simulator_dir, pim_dir, output_bin_path, dump_ranges)

    print(Style.BRIGHT + "\nValidating the results:" + Style.RESET_ALL)
    sim_arrays = parse_pim_simulator_outputs(output_bin_path, outputs_descriptor)
    return validate_outputs(sim_arrays, out_dir, outputs_descriptor, threshold)


if __name__ == '__main__':
    ap = argparse.ArgumentParser()
    ap.add_argument("--network-onnx", required=True)
    ap.add_argument("--raptor-path", required=True)
    ap.add_argument("--onnx-include-dir", required=True)
    a = ap.parse_args()

    simulator_dir = Path(__file__).parent.resolve() / ".." / "backend-simulators" / "pim" / "pim-simulator"

    passed = validate_network(
        a.network_onnx, a.raptor_path, a.onnx_include_dir, simulator_dir
    )
    raise SystemExit(0 if passed else 1)