import argparse
import os
import pickle
import time
from pathlib import Path

import numpy as np
import onnx
import onnxruntime as ort
from onnx import numpy_helper

from build_fpga import BuildError, SYNTHESIS_STAGES, synthesize_fpga_kernel
from extraction import load_splice_function
from quantization import quantize_model
from utils import Configuration, load_config

ONNX_DESIRED_OPSET = 21


def get_info(data_dir):
    """Print some basic info on nodes and weights of model"""
    model_path = f"{data_dir}model.onnx"
    onnx_model = onnx.load(model_path)
    onnx.checker.check_model(onnx_model)

    for node in onnx_model.graph.node:
        print("name=%r type=%r input=%r output=%r" % (
        node.name, node.op_type, node.input, node.output))

    print("=============weights===============")

    weights = onnx_model.graph.initializer

    for w in weights:
        if w.name == 'onnx::MatMul_1204':
            print(w)
            print("======================")
            print(onnx.numpy_helper.to_array(w))


def get_zero_point(layer_name: str, all_layers):
    """Return the matching quantization zero-point tensor name, if present."""
    if f"{layer_name}_zero_point" in all_layers:
        return f"{layer_name}_zero_point"

    if layer_name.endswith("_quantized"):
        return get_zero_point(layer_name.rstrip("_quantized"), all_layers)

    return None


def get_scale(layer_name: str, all_layers):
    """Return the matching quantization scale tensor name, if present."""
    if f"{layer_name}_scale" in all_layers:
        return f"{layer_name}_scale"

    if layer_name.endswith("_quantized"):
        return get_scale(layer_name.rstrip("_quantized"), all_layers)

    return None


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="FAT - FPGA Acceleration of Transformers")
    parser.add_argument("data_directory", type=str, help="Path to the data directory")
    parser.add_argument("-q", "--quantize", action="store_true", help="Quantize model.onnx in data dir")
    parser.add_argument("-v", "--view", action="store_true", help="View the ONNX Model structure and weights")
    parser.add_argument("-m", "--model", type=str, default="model.onnx", help="ONNX model filename in data_directory, defaults to model.onnx")
    parser.add_argument("-s", "--split", action="store_true", help="Performing splitting (slicing) of model.")
    parser.add_argument("-i", "--inference", action="store_true", help="Run inference on FPGA")
    parser.add_argument("-e", "--evaluate", action="store_true", help="Run inference on FPGA and compare with ONNX")
    parser.add_argument("-o", "--onnx", action="store_true", help="Run inference on ONNX only")
    parser.add_argument(
        "--synthesis",
        nargs="?",
        const="both",
        choices=SYNTHESIS_STAGES,
        default=None,
        help=(
            "Generate and synthesize the FPGA kernel. This fills in the HLS templates based on model.yaml, "
            "runs the required Tcl scripts, and places the final .bit and .hwh files in the data directory. "
            "Use '--synthesis vitis' or '--synthesis vivado' to run a single stage; if no stage is specified, "
            "both stages run sequentially."
        ),
    )
    parser.add_argument("--build-jobs", type=int, default=None, help="Parallel job count for Vivado; defaults to FAT_VIVADO_JOBS or 8")
    parser.add_argument("--build-dry-run", action="store_true", help="Print the resolved FPGA build workflow without launching vendor tools")
    parser.add_argument("--random-inputs", action="store_true", help="Use randomly generated inputs instead of input.pkl")
    parser.add_argument("--random-count", type=int, default=1, help="Number of random input samples to generate when --random-inputs is set")
    parser.add_argument("--random-seed", type=int, default=0, help="Seed for reproducible random input generation")

    args = parser.parse_args()

    data_dir = args.data_directory
    os.environ["MODEL"] = args.model

    if data_dir[-1] != "/":
        data_dir = data_dir + "/"

    # Synthesis can run on its own, or as a first step before another action.
    synthesis_stage = args.synthesis

    if synthesis_stage or args.build_dry_run:
        synthesis_stage = synthesis_stage or "both"
        try:
            bit_path, hwh_path = synthesize_fpga_kernel(
                data_dir,
                stage=synthesis_stage,
                jobs=args.build_jobs,
                dry_run=args.build_dry_run,
            )
        except BuildError as exc:
            print(f"FPGA build failed: {exc}")
            raise SystemExit(1) from exc

        if args.build_dry_run:
            print(f"Dry run complete. Planned artifacts: {bit_path}, {hwh_path}")
        elif bit_path is not None and hwh_path is not None:
            print(f"FPGA build complete. Artifacts copied to {bit_path.parent}")
        else:
            print(f"FPGA synthesis stage '{synthesis_stage}' completed.")

        if not any([args.quantize, args.view, args.split, args.inference, args.evaluate, args.onnx]):
            raise SystemExit(0)

    # Keep the source model on the opset expected by ONNX extraction/runtime.
    model = onnx.load(f'{data_dir + os.environ["MODEL"]}')
    opset = model.opset_import[0].version
    if opset < ONNX_DESIRED_OPSET:
        print(f"Model opset version is {opset}, converting to {ONNX_DESIRED_OPSET}")
        model = onnx.version_converter.convert_version(model, ONNX_DESIRED_OPSET)
        onnx.save(model, f'{data_dir + os.environ["MODEL"]}')
        print(f"Model converted and saved to {data_dir + os.environ['MODEL']}")

    if args.view:
        model = onnx.load(f'{data_dir + os.environ["MODEL"]}')
        init = {t.name: numpy_helper.to_array(t) for t in model.graph.initializer}
        print("=================")
        for name, data in init.items():
            if name.endswith("_zero_point") or name.endswith("_scale"):
                continue
            zp = get_zero_point(name, init.keys())
            scale = get_scale(name, init.keys())

            info = ""
            if zp:
                info += f"zero: {init[zp]}, "
            if scale:
                info += f"scale: {init[scale]}"
            print(f"{name}: shape: {data.shape}, {info}")

        exit(0)

    if args.quantize:
        print(f"Quantizing model.onnx in {data_dir}")
        quantize_model(data_dir)

    if args.onnx:
        from inference import run_inference, run_inference_original
        print("Running original ONNX inference")
        original_outputs = run_inference_original(data_dir, random_inputs=args.random_inputs, random_count=args.random_count, random_seed=args.random_seed)

        with open(f"{data_dir}output.pkl", 'wb') as f:
            pickle.dump(original_outputs, f)

        exit(0)

    if args.split:
        # The split points are model-specific, so each data directory supplies a
        # small splice.py with a splice(data_dir, layer_count) function.
        configuration: Configuration = load_config(data_dir)
        layer_count = configuration.kernel.encoder_layer_count
        splice_fn = load_splice_function(data_dir)
        splice_fn(data_dir, layer_count)

    if args.inference:
        from inference import load_input_data, run_inference, run_inference_original, run_inference_enccla
        configuration: Configuration = load_config(data_dir)
        layer_count = configuration.kernel.encoder_layer_count
        print(f"Running inference with {layer_count} layers offloaded to FPGA")

        outputs = run_inference_enccla(data_dir, layer_count, evaluate=False, random_inputs=args.random_inputs, random_count=args.random_count, random_seed=args.random_seed)
        with open(f"{data_dir}output.pkl", 'wb') as f:
            pickle.dump(outputs, f)

        exit(0)

    if args.evaluate:
        from inference import load_input_data, run_inference, run_inference_original, run_inference_enccla, run_inference_generic
        configuration: Configuration = load_config(data_dir)
        layer_count = configuration.kernel.encoder_layer_count
        print(f"Running inference with {layer_count} layers offloaded to FPGA")

        shared_input_data = load_input_data(
            data_dir,
            configuration,
            random_inputs=args.random_inputs,
            random_count=args.random_count,
            random_seed=args.random_seed,
        )

        outputs = run_inference_generic(data_dir, layer_count, random_inputs=args.random_inputs, random_count=args.random_count, random_seed=args.random_seed, input_data=shared_input_data)
        ref = run_inference_original(data_dir, random_inputs=args.random_inputs, random_count=args.random_count, random_seed=args.random_seed, input_data=shared_input_data)

        for i, (out, r) in enumerate(zip(outputs, ref)):
            if not np.allclose(out, r, atol=1e-2):
                print(f"Output {i} does not match reference!")
                print("Output:", out)
                print("Reference:", r)

        outputs_flat = np.asarray(outputs, dtype=np.float32).ravel()
        ref_flat = np.asarray(ref, dtype=np.float32).ravel()

        if outputs_flat.shape != ref_flat.shape:
            raise ValueError(f"Shape mismatch: outputs={outputs_flat.shape}, ref={ref_flat.shape}")

        rmse = np.sqrt(np.mean((outputs_flat - ref_flat) ** 2))
        print(f"Global RMSE: {rmse:.6f}")

        exit(0)
