import pickle
import numpy as np
import onnxruntime as ort
import time
import onnx
from onnx import numpy_helper
from encoder import encoder_layer
import os

import json
import yaml
from extraction import split_into_3, extract_parameters, compare_parameter_dicts
from utils import load_config, load_parameter_config, Configuration
from kernel import DualEncoderKernel


def _to_numpy_input(x):
    if hasattr(x, "cpu"):
        return x.cpu().numpy()
    return np.asarray(x)


def _model_input_names(model_path):
    session = ort.InferenceSession(model_path)
    return [inp.name for inp in session.get_inputs()]


def _generate_random_inputs(configuration: Configuration, input_names, count, seed):
    rng = np.random.default_rng(seed)
    samples = []

    for _ in range(count):
        sample = rng.normal(0.0, 1.0, size=(1, configuration.model.sequence_length, 3)).astype(np.float32)

        if "padding_mask" in input_names:
            valid_length = int(rng.integers(max(1, configuration.model.sequence_length // 2), configuration.model.sequence_length + 1))
            if valid_length < configuration.model.sequence_length:
                sample[:, valid_length:, :] = -1.0

        samples.append(sample)

    return samples


def load_input_data(data_dir, configuration: Configuration, random_inputs=False, random_count=1, random_seed=0):
    """Load input.pkl or create deterministic synthetic inputs for smoke tests."""
    if not random_inputs:
        with open(f"{data_dir}input.pkl", "rb") as f:
            return pickle.load(f)

    model_name = os.environ.get("MODEL", "model.onnx")
    model_path = data_dir + model_name
    input_names = _model_input_names(model_path)
    print(f"Using {random_count} randomly generated input sample(s) with seed {random_seed}.")
    return _generate_random_inputs(configuration, input_names, random_count, random_seed)


def _compute_padding_mask(input_array):
    return np.all(input_array == -1, axis=2)


def _infer_mask_count_from_padding_mask(padding_mask):
    mask = np.asarray(padding_mask)
    if mask.ndim == 2:
        mask = mask[0]

    mask_count = 0
    while mask_count < mask.shape[0] and mask[mask_count] == False:
        mask_count += 1
    return mask_count


def _infer_mask_count_from_masked_fill(masked_fill):
    mask = np.asarray(masked_fill)
    if mask.ndim == 2:
        mask = mask[0]

    valid_mask = np.isfinite(mask) & (mask >= 0)
    return int(np.count_nonzero(valid_mask))

def run_inference_original(data_dir, random_inputs=False, random_count=1, random_seed=0, input_data=None):
    """Run inference on the original model."""

    configuration: Configuration = load_config(data_dir)
    if input_data is None:
        input_data = load_input_data(
            data_dir,
            configuration,
            random_inputs=random_inputs,
            random_count=random_count,
            random_seed=random_seed,
        )

    outputs = []
    outputs_fp16 = []
    model_name = os.environ.get("MODEL")
    session = ort.InferenceSession(data_dir + model_name)
    input_names = [inp.name for inp in session.get_inputs()]

    for i, model_input in enumerate(input_data, start=1):
        print(f"Running inference {i}/{len(input_data)}")
        input_array = _to_numpy_input(model_input).astype(np.float32, copy=False)

        if "padding_mask" in input_names:
            padding_mask = _compute_padding_mask(input_array)
            feed = {"hits": input_array, "padding_mask": padding_mask}
        else:
            feed = {input_names[0]: input_array}

        out = session.run(None, feed)
        outputs.append(out)

        out_fp16 = [o.astype(np.float16).astype(np.float32) for o in out]
        outputs_fp16.append(out_fp16)

    flat_outputs = np.asarray(outputs, dtype=np.float32).ravel()
    fp16_roundtrip = flat_outputs.astype(np.float16).astype(np.float32)

    if flat_outputs.shape != fp16_roundtrip.shape:
        raise ValueError(
            f"Shape mismatch: fp32={flat_outputs.shape}, fp16={fp16_roundtrip.shape}"
        )

    diff = flat_outputs - fp16_roundtrip
    rmse = np.sqrt(np.mean(diff ** 2))
    print(f"RMSE (float32 -> float16): {rmse:.6f}")

    return outputs_fp16

def single_numpy_encoder(model, input_data, pad_mask):

    init   = {t.name: numpy_helper.to_array(t) for t in model.graph.initializer}

    # --- 4) Build our weight dict matching ONNX naming & shapes ---
    w = {
        "n_heads":      4,
        "mha.W_qkv":    init["onnx::MatMul_1261"].astype(np.float32),
        "mha.b_qkv":    init["encoder.layers.1.self_attn.in_proj_bias"].astype(np.float32),
        "mha.W_o":      init["encoder.layers.1.self_attn.out_proj.weight"].astype(np.float32),
        "mha.b_o":      init["encoder.layers.1.self_attn.out_proj.bias"].astype(np.float32),
        "ff.W1":        init["onnx::MatMul_1315"].astype(np.float32),
        "ff.b1":        init["encoder.layers.1.linear1.bias"].astype(np.float32),
        "ff.W2":        init["onnx::MatMul_1316"].astype(np.float32),
        "ff.b2":        init["encoder.layers.1.linear2.bias"].astype(np.float32),
        "norm1.gamma":  init["encoder.layers.1.norm1.weight"].astype(np.float32),
        "norm1.beta":   init["encoder.layers.1.norm1.bias"].astype(np.float32),
        "norm2.gamma":  init["encoder.layers.1.norm2.weight"].astype(np.float32),
        "norm2.beta":   init["encoder.layers.1.norm2.bias"].astype(np.float32),
    }

   
    encoder_output = encoder_layer(input_data, pad_mask, w)
    
    return encoder_output

def load_onnx_info(config_path: str = "onnx_info.yaml"):
    """Load extraction rules from YAML or JSON."""
    if config_path.endswith((".yaml", ".yml")):
        with open(config_path, "r") as f:
            return yaml.safe_load(f)
    if config_path.endswith(".json"):
        with open(config_path, "r") as f:
            return json.load(f)
    raise ValueError("Config file must be .yaml, .yml, or .json")


def run_inference(data_dir, pipeline_latency, evaluate=True, random_inputs=False, random_count=1, random_seed=0, input_data=None):
    # model_a = onnx.load_model(data_dir +"model_a.onnx")
    model_b = onnx.load_model(data_dir + "model_b.onnx")
    # model = onnx.load_model(data_dir + "model.onnx")
    # dump_model_names(model)
    # exit(1)
    # model_c = onnx.load_model(data_dir + "model_c.onnx")
    configuration: Configuration = load_config(data_dir)

    parameter_config = load_parameter_config(data_dir)
    weights = extract_parameters(model_b, pipeline_latency, parameter_config, np.float16 if configuration.kernel.store_fp16 else np.float32)

    # weights = extract_parameters_old(model_b, pipeline_latency, np.float16 if configuration.kernel.store_fp16 else np.float32)

    # compare_parameter_dicts(weights_old, weights)
    # exit(1)

    kernel = DualEncoderKernel(data_dir, configuration)
    t0 = time.time()
    kernel.load_parameters(weights)
    t1 = time.time()
    print(f"Loaded parameters into FPGA kernel in {t1 - t0:.3f} seconds.")
    for i in range(pipeline_latency):
        kernel.push_zeros()

    if input_data is None:
        input_data = load_input_data(
            data_dir,
            configuration,
            random_inputs=random_inputs,
            random_count=random_count,
            random_seed=random_seed,
        )

    outputs = []

    # increase input data 10x to simulate more inputs
    input_data = input_data

    prev_outputs = []
    fpga_times = []
    onnx_times = []
    middle_session = ort.InferenceSession(data_dir + "model_b.onnx")

    fpga_outputs = []
    reference_onnx_outputs = []
    # Warmup + running
    for i, x in enumerate(input_data):
        # print(f"Processing input {i+1}/{len(input_data)}")
        input_array = _to_numpy_input(x).astype(np.float32, copy=False)
        padding_mask = _compute_padding_mask(input_array)
        mask_count = _infer_mask_count_from_padding_mask(padding_mask)

        # Run inference on part A
        session = ort.InferenceSession(data_dir + "model_a.onnx")
        out = session.run(None, {'hits': input_array, 'padding_mask': padding_mask})
        prev_outputs.append(out[1])

        # Run fpga
        t0 = time.perf_counter()
        fpga_out = kernel.run(out[0], mask_count)

        t1 = time.perf_counter()
        fpga_times.append(t1 - t0)
        
        if evaluate:
            # Run identical ONNX for reference
            t0 = time.perf_counter()
            middle_out = middle_session.run(None, {middle_session.get_inputs()[0].name: out[0], '/encoder/Where_output_0': out[1]})
            t1 = time.perf_counter()
            onnx_times.append(t1 - t0)
            reference_onnx_outputs.append(middle_out[0])


        # # early problem detection
        # if evaluate and i >= pipeline_latency-1:
        #     # Check if FPGA and ONNX are identical
        #     if not np.allclose(fpga_out, reference_onnx_outputs[i - (pipeline_latency-1)], atol=1e-1):
        #         print(f"Output {i} does not match ONNX reference!")
        #         print("FPGA output:", fpga_out)
        #         print("ONNX output:", reference_onnx_outputs[i - (pipeline_latency-1)])
        #         exit(1)


        # Skip this during warmup
        if i >= pipeline_latency-1:
            fpga_outputs.append(fpga_out)
            part_a_output = prev_outputs.pop(0)
            # run part c
            session = ort.InferenceSession(data_dir + "model_c.onnx")
            final_out = session.run(None, {session.get_inputs()[0].name: [fpga_out], '/encoder/Where_output_0': part_a_output})
            outputs.append(final_out)

    # cooldown
    for _ in range(pipeline_latency-1):
        t0 = time.perf_counter()
        fpga_out = kernel.push_zeros()
        t1 = time.perf_counter()
        fpga_times.append(t1 - t0)
        fpga_outputs.append(fpga_out)
        part_a_output = prev_outputs.pop(0)
        session = ort.InferenceSession(data_dir + "model_c.onnx")
        final_out = session.run(None, {session.get_inputs()[0].name: [fpga_out], '/encoder/Where_output_0': part_a_output})
        outputs.append(final_out)

    if evaluate:
        max_time = max(fpga_times)
        min_time = min(fpga_times)
        avg_time = sum(fpga_times) / len(fpga_times)
        print(f"FPGA Inference times over {len(fpga_times)} runs: max={max_time*1000:.3f} ms, min={min_time*1000:.3f} ms, avg={avg_time*1000:.3f} ms")

        
        max_time = max(onnx_times)
        min_time = min(onnx_times)
        avg_time = sum(onnx_times) / len(onnx_times)
        print(f"ONNX Inference times over {len(onnx_times)} runs: max={max_time*1000:.3f} ms, min={min_time*1000:.3f} ms, avg={avg_time*1000:.3f} ms")
        # print(onnx_times)

        # Calculate avg troughput (layers/second) (ignore warmup/cooldown)
        num_layers = len(fpga_outputs) * pipeline_latency  # Each output corresponds to pipeline_latency layers processed
        total_time = sum(fpga_times[pipeline_latency-1:])  # Ignore warmup and cooldown
        avg_throughput = num_layers / total_time if total_time > 0 else 0
        print(f"Average throughput: {avg_throughput:.2f} layers/second")

        # Calc troughput for ONNX
        num_inferences = len(reference_onnx_outputs) * pipeline_latency
        total_time = sum(onnx_times)
        onnx_throughput = num_inferences / total_time if total_time > 0 else 0
        print(f"ONNX Average throughput: {onnx_throughput:.2f} layers/second")

        # Compare outputs and calculate global RMSE (over all elements)
        # fpga_flat = np.asarray(fpga_outputs, dtype=np.float32).ravel()
        # onnx_flat = np.asarray(reference_onnx_outputs, dtype=np.float32).ravel()

        # if fpga_flat.shape != onnx_flat.shape:
        #     raise ValueError(f"Shape mismatch: fpga={fpga_flat.shape}, onnx={onnx_flat.shape}")

        # rmse = np.sqrt(np.mean((fpga_flat - onnx_flat) ** 2))
        # print(f"RMSE: {rmse:.6f}")

            # Compare outputs and calculate global RMSE (over all elements)
        fpga_flat = np.asarray(fpga_outputs, dtype=np.float32).ravel()
        onnx_flat = np.asarray(reference_onnx_outputs, dtype=np.float32).ravel()

        if fpga_flat.shape != onnx_flat.shape:
            raise ValueError(f"Shape mismatch: fpga={fpga_flat.shape}, onnx={onnx_flat.shape}")

        diff = fpga_flat - onnx_flat

        # Absolute RMSE
        rmse = np.sqrt(np.mean(diff ** 2))
        print(f"RMSE: {rmse:.6f}")

        # # L2-relative RMSE (a.k.a. relative L2 error)
        # denom = np.linalg.norm(onnx_flat, ord=2)
        # if denom == 0:
        #     rel_l2_rmse = np.nan  # or 0.0, but NaN makes the issue obvious
        #     print("Relative L2 RMSE: undefined (reference L2 norm is zero)")
        # else:
        #     rel_l2_rmse = np.linalg.norm(diff, ord=2) / denom
        #     print(f"Relative L2 RMSE: {rel_l2_rmse:.6e}")





    return outputs


def run_inference_enccla(data_dir, pipeline_latency, evaluate=True, random_inputs=False, random_count=1, random_seed=0, input_data=None):
    """Run the fixed-name split-model FPGA path kept for existing experiments."""
    model_b = onnx.load_model(data_dir + "model_b.onnx")
    configuration: Configuration = load_config(data_dir)

    parameter_config = load_parameter_config(data_dir)
    weights = extract_parameters(model_b, pipeline_latency, parameter_config, np.float16 if configuration.kernel.store_fp16 else np.float32)

    kernel = DualEncoderKernel(data_dir, configuration)
    t0 = time.time()
    kernel.load_parameters(weights)
    t1 = time.time()
    print(f"Loaded parameters into FPGA kernel in {t1 - t0:.3f} seconds.")
    for _ in range(pipeline_latency):
        kernel.push_zeros()

    if input_data is None:
        input_data = load_input_data(
            data_dir,
            configuration,
            random_inputs=random_inputs,
            random_count=random_count,
            random_seed=random_seed,
        )

    outputs = []
    carried_outputs = []
    fpga_times = []
    onnx_times = []
    fpga_outputs = []
    reference_onnx_outputs = []

    session_a = ort.InferenceSession(data_dir + "model_a.onnx")
    session_b = ort.InferenceSession(data_dir + "model_b.onnx")
    session_c = ort.InferenceSession(data_dir + "model_c.onnx")

    model_a_output_names = [out.name for out in session_a.get_outputs()]
    model_b_input_names = [inp.name for inp in session_b.get_inputs()]
    model_c_input_names = [inp.name for inp in session_c.get_inputs()]
    model_c_fpga_input_name = session_c.get_inputs()[0].name

    for i, x in enumerate(input_data):
        raw_input = _to_numpy_input(x).astype(np.float32, copy=False)
        out = session_a.run(None, {'inputs': raw_input})
        out_map = dict(zip(model_a_output_names, out))
        carried_outputs.append(out_map)

        mask_count = _infer_mask_count_from_masked_fill(out_map['masked_fill'])

        t0 = time.perf_counter()
        fpga_out = kernel.run(out_map['linear'], mask_count)
        t1 = time.perf_counter()
        fpga_times.append(t1 - t0)

        if evaluate:
            middle_feed = {name: out_map[name] for name in model_b_input_names}
            t0 = time.perf_counter()
            middle_out = session_b.run(None, middle_feed)
            t1 = time.perf_counter()
            onnx_times.append(t1 - t0)
            reference_onnx_outputs.append(middle_out[0])

        if i >= pipeline_latency - 1:
            fpga_outputs.append(fpga_out)
            carry = carried_outputs.pop(0)
            final_feed = {model_c_fpga_input_name: [fpga_out]}
            for name in model_c_input_names[1:]:
                final_feed[name] = carry[name]
            final_out = session_c.run(None, final_feed)
            outputs.append(final_out)

    for _ in range(pipeline_latency - 1):
        t0 = time.perf_counter()
        fpga_out = kernel.push_zeros()
        t1 = time.perf_counter()
        fpga_times.append(t1 - t0)
        fpga_outputs.append(fpga_out)
        carry = carried_outputs.pop(0)
        final_feed = {model_c_fpga_input_name: [fpga_out]}
        for name in model_c_input_names[1:]:
            final_feed[name] = carry[name]
        final_out = session_c.run(None, final_feed)
        outputs.append(final_out)

    if evaluate:
        max_time = max(fpga_times)
        min_time = min(fpga_times)
        avg_time = sum(fpga_times) / len(fpga_times)
        print(f"FPGA Inference times over {len(fpga_times)} runs: max={max_time*1000:.3f} ms, min={min_time*1000:.3f} ms, avg={avg_time*1000:.3f} ms")

        max_time = max(onnx_times)
        min_time = min(onnx_times)
        avg_time = sum(onnx_times) / len(onnx_times)
        print(f"ONNX Inference times over {len(onnx_times)} runs: max={max_time*1000:.3f} ms, min={min_time*1000:.3f} ms, avg={avg_time*1000:.3f} ms")

        num_layers = len(fpga_outputs) * pipeline_latency
        total_time = sum(fpga_times[pipeline_latency-1:])
        avg_throughput = num_layers / total_time if total_time > 0 else 0
        print(f"Average throughput: {avg_throughput:.2f} layers/second")

        num_inferences = len(reference_onnx_outputs) * pipeline_latency
        total_time = sum(onnx_times)
        onnx_throughput = num_inferences / total_time if total_time > 0 else 0
        print(f"ONNX Average throughput: {onnx_throughput:.2f} layers/second")

        fpga_flat = np.asarray(fpga_outputs, dtype=np.float32).ravel()
        onnx_flat = np.asarray(reference_onnx_outputs, dtype=np.float32).ravel()

        if fpga_flat.shape != onnx_flat.shape:
            raise ValueError(f"Shape mismatch: fpga={fpga_flat.shape}, onnx={onnx_flat.shape}")

        diff = fpga_flat - onnx_flat
        rmse = np.sqrt(np.mean(diff ** 2))
        print(f"RMSE: {rmse:.6f}")

    return outputs


def _build_model_a_feed(session_a, input_array):
    input_names = [inp.name for inp in session_a.get_inputs()]

    if 'padding_mask' in input_names:
        main_input_name = next(name for name in input_names if name != 'padding_mask')
        return {
            main_input_name: input_array,
            'padding_mask': _compute_padding_mask(input_array),
        }

    return {input_names[0]: input_array}


def _infer_mask_count_generic(session_a, model_a_feed, model_a_outputs, configuration: Configuration):
    input_names = {inp.name for inp in session_a.get_inputs()}

    if 'padding_mask' in input_names:
        return _infer_mask_count_from_padding_mask(model_a_feed['padding_mask'])

    if 'masked_fill' in model_a_outputs:
        return _infer_mask_count_from_masked_fill(model_a_outputs['masked_fill'])

    for name, value in model_a_outputs.items():
        lowered = name.lower()
        if 'mask' in lowered:
            arr = np.asarray(value)
            if arr.dtype == np.bool_:
                return _infer_mask_count_from_padding_mask(arr)
            if np.issubdtype(arr.dtype, np.floating):
                return _infer_mask_count_from_masked_fill(arr)

    return configuration.model.sequence_length


def run_inference_generic(data_dir, pipeline_latency, evaluate=True, random_inputs=False, random_count=1, random_seed=0, input_data=None):
    """Run split-model FPGA inference using ONNX input/output names."""
    model_b = onnx.load_model(data_dir + 'model_b.onnx')
    configuration: Configuration = load_config(data_dir)

    parameter_config = load_parameter_config(data_dir)
    weights = extract_parameters(
        model_b,
        pipeline_latency,
        parameter_config,
        np.float16 if configuration.kernel.store_fp16 else np.float32,
    )

    kernel = DualEncoderKernel(data_dir, configuration)
    t0 = time.time()
    kernel.load_parameters(weights)
    t1 = time.time()
    print(f'Loaded parameters into FPGA kernel in {t1 - t0:.3f} seconds.')
    for _ in range(pipeline_latency):
        kernel.push_zeros()

    if input_data is None:
        input_data = load_input_data(
            data_dir,
            configuration,
            random_inputs=random_inputs,
            random_count=random_count,
            random_seed=random_seed,
        )

    session_a = ort.InferenceSession(data_dir + 'model_a.onnx')
    session_b = ort.InferenceSession(data_dir + 'model_b.onnx')
    session_c = ort.InferenceSession(data_dir + 'model_c.onnx')

    model_a_output_names = [out.name for out in session_a.get_outputs()]
    model_b_input_names = [inp.name for inp in session_b.get_inputs()]
    model_b_output_name = session_b.get_outputs()[0].name
    model_c_input_names = [inp.name for inp in session_c.get_inputs()]

    carry_names = [name for name in model_c_input_names if name != model_b_output_name]

    outputs = []
    carried_outputs = []
    fpga_times = []
    onnx_times = []
    fpga_outputs = []
    reference_onnx_outputs = []

    for i, x in enumerate(input_data):
        input_array = _to_numpy_input(x).astype(np.float32, copy=False)
        model_a_feed = _build_model_a_feed(session_a, input_array)
        model_a_raw_outputs = session_a.run(None, model_a_feed)
        model_a_outputs = dict(zip(model_a_output_names, model_a_raw_outputs))
        carried_outputs.append({name: model_a_outputs[name] for name in carry_names if name in model_a_outputs})

        missing_b_inputs = [name for name in model_b_input_names if name not in model_a_outputs]
        if missing_b_inputs:
            raise KeyError(f'Model A does not provide required Model B inputs: {missing_b_inputs}')

        mask_count = _infer_mask_count_generic(session_a, model_a_feed, model_a_outputs, configuration)

        t0 = time.perf_counter()
        fpga_out = kernel.run(model_a_outputs[model_b_input_names[0]], mask_count)
        t1 = time.perf_counter()
        fpga_times.append(t1 - t0)

        if evaluate:
            middle_feed = {name: model_a_outputs[name] for name in model_b_input_names}
            t0 = time.perf_counter()
            middle_out = session_b.run(None, middle_feed)
            t1 = time.perf_counter()
            onnx_times.append(t1 - t0)
            reference_onnx_outputs.append(middle_out[0])

        if i >= pipeline_latency - 1:
            fpga_outputs.append(fpga_out)
            carry = carried_outputs.pop(0)
            final_feed = {model_b_output_name: [fpga_out]}
            missing_c_inputs = [name for name in carry_names if name not in carry]
            if missing_c_inputs:
                raise KeyError(f'Model A does not provide required Model C carry inputs: {missing_c_inputs}')
            for name in carry_names:
                final_feed[name] = carry[name]
            final_out = session_c.run(None, final_feed)
            outputs.append(final_out)

    for _ in range(pipeline_latency - 1):
        t0 = time.perf_counter()
        fpga_out = kernel.push_zeros()
        t1 = time.perf_counter()
        fpga_times.append(t1 - t0)
        fpga_outputs.append(fpga_out)
        carry = carried_outputs.pop(0)
        final_feed = {model_b_output_name: [fpga_out]}
        missing_c_inputs = [name for name in carry_names if name not in carry]
        if missing_c_inputs:
            raise KeyError(f'Model A does not provide required Model C carry inputs: {missing_c_inputs}')
        for name in carry_names:
            final_feed[name] = carry[name]
        final_out = session_c.run(None, final_feed)
        outputs.append(final_out)

    if evaluate:
        max_time = max(fpga_times)
        min_time = min(fpga_times)
        avg_time = sum(fpga_times) / len(fpga_times)
        print(f'FPGA Inference times over {len(fpga_times)} runs: max={max_time*1000:.3f} ms, min={min_time*1000:.3f} ms, avg={avg_time*1000:.3f} ms')

        max_time = max(onnx_times)
        min_time = min(onnx_times)
        avg_time = sum(onnx_times) / len(onnx_times)
        print(f'ONNX Inference times over {len(onnx_times)} runs: max={max_time*1000:.3f} ms, min={min_time*1000:.3f} ms, avg={avg_time*1000:.3f} ms')

        num_layers = len(fpga_outputs) * pipeline_latency
        total_time = sum(fpga_times[pipeline_latency-1:])
        avg_throughput = num_layers / total_time if total_time > 0 else 0
        print(f'Average throughput: {avg_throughput:.2f} layers/second')

        num_inferences = len(reference_onnx_outputs) * pipeline_latency
        total_time = sum(onnx_times)
        onnx_throughput = num_inferences / total_time if total_time > 0 else 0
        print(f'ONNX Average throughput: {onnx_throughput:.2f} layers/second')

        fpga_flat = np.asarray(fpga_outputs, dtype=np.float32).ravel()
        onnx_flat = np.asarray(reference_onnx_outputs, dtype=np.float32).ravel()

        if fpga_flat.shape != onnx_flat.shape:
            raise ValueError(f'Shape mismatch: fpga={fpga_flat.shape}, onnx={onnx_flat.shape}')

        diff = fpga_flat - onnx_flat
        rmse = np.sqrt(np.mean(diff ** 2))
        print(f'RMSE: {rmse:.6f}')

    return outputs

def run_matmul_inference(data_dir, ip):
    from fpga import run_ip_matmul
    model_a = onnx.load_model(data_dir +"model_a.onnx")
    model_b = onnx.load_model(data_dir + "model_b.onnx")


    # Load the input file with picle
    if input_data is None:
        input_data = load_input_data(
            data_dir,
            configuration,
            random_inputs=random_inputs,
            random_count=random_count,
            random_seed=random_seed,
        )

    output = []
    comparison_output = []
    for input in input_data:
        info = {}
        # Prep input
        padding_mask = (input == -1).all(dim=2)

        # Run inference on part A
        session = ort.InferenceSession(data_dir + "model_a.onnx")
        print(f"Input shape: {input.cpu().numpy().shape()}")
        out = session.run(None, {'hits': input.cpu().numpy()} )
        info["ONNX_matmul_out"] = out[0]
        #TODO Save output

        # Get weights for MatMul
        weights = model_a.graph.initializer
        for w in weights:
            if w.name == 'onnx::MatMul_1204':
                matrix = onnx.numpy_helper.to_array(w)

        #print(f"Model input shape: {input.cpu().numpy().shape}")
        #print(f"MatMul weight matrix shape: {matrix.shape}")
        #print(f"MatMul output shape: {out[0].shape}")
         
        
        if ip != None:
            # Run MatMul
            matmul_out = run_ip_matmul(ip, input.cpu().numpy(), matrix)
            matmul_out = np.expand_dims(matmul_out, axis=0)
            info["FPGA_matmul_out"] = matmul_out
        #TODO Save MatMul output
        # Run inference on part B
        session = ort.InferenceSession(data_dir + "model_b.onnx")
        final_out = session.run(None, {'/input_layer/MatMul_output_0': out[0], 'padding_mask': padding_mask.cpu().numpy()})
        info["FPGA_model_out"] = final_out
        info["padding"] = padding_mask
        comparison_output.append(info)
        #print(f"Model output shape: {final_out[0].shape}")
        output.append(final_out)



    return final_out, comparison_output

def get_matrix_input_from_node(model, node_name: str):
    """
    Given an ONNX model and a node name, find the name of the input matrix
    """
    nodes = model.graph.node
    for n in nodes:
        if n.name == node_name:
            for inp in n.input:
                if "onnx::MatMul" in inp:
                    return inp
                
    raise ValueError(f"Node {node_name} not found in model.")

def extract_parameters_old(model, layer_count, store_type=np.float32):
    """
    Extracts weights for a multi-layer encoder (currently supports 1 or 2 layers)
    and returns a SINGLE flat dictionary compatible with the HLS load_parameters() function.

    layer_count:
        Number of encoder layers to extract (1 or 2 for now).

    Output keys match kernel registers:
      W_QKV, B_QKV, W_O, B_O,
      ff_w1, ff_b1, ff_w2, ff_b2,
      norm1_gamma, norm1_beta,
      norm2_gamma, norm2_beta
    """

    init = {t.name: numpy_helper.to_array(t) for t in model.graph.initializer}


    per_layer = []

    for layer_idx in range(layer_count):
        p = f"encoder.layers.{layer_idx}."  
        node_base_name = f"/encoder/layers.{layer_idx}/"

        w = {
            "W_QKV": init[get_matrix_input_from_node(model, node_base_name + "self_attn/MatMul")].astype(store_type),  # (D_MODEL, 3*D_MODEL)
            "B_QKV": init[p + "self_attn.in_proj_bias"].astype(store_type),

            "W_O":   init[p + "self_attn.out_proj.weight"].astype(store_type),
            "B_O":   init[p + "self_attn.out_proj.bias"].astype(store_type),

            "ff_w1": init[get_matrix_input_from_node(model, node_base_name + "linear1/MatMul")].astype(store_type),
            "ff_b1": init[p + "linear1.bias"].astype(store_type),

            "ff_w2": init[get_matrix_input_from_node(model, node_base_name + "linear2/MatMul")].astype(store_type),
            "ff_b2": init[p + "linear2.bias"].astype(store_type),

            "norm1_gamma": init[p + "norm1.weight"].astype(store_type),
            "norm1_beta":  init[p + "norm1.bias"].astype(store_type),

            "norm2_gamma": init[p + "norm2.weight"].astype(store_type),
            "norm2_beta":  init[p + "norm2.bias"].astype(store_type),
        }

        per_layer.append(w)

    # ------------------------------------------------------------
    # Concatenate per-layer data into single flat arrays
    # ------------------------------------------------------------
    def cat_flat(key):
        return np.concatenate(
            [per_layer[l][key].reshape(-1) for l in range(layer_count)],
            axis=0
        )

    flat = {
        "W_QKV":        cat_flat("W_QKV"),
        "B_QKV":        cat_flat("B_QKV"),
        "W_O":          cat_flat("W_O"),
        "B_O":          cat_flat("B_O"),

        "ff_w1":        cat_flat("ff_w1"),
        "ff_b1":        cat_flat("ff_b1"),
        "ff_w2":        cat_flat("ff_w2"),
        "ff_b2":        cat_flat("ff_b2"),

        "norm1_gamma":  cat_flat("norm1_gamma"),
        "norm1_beta":   cat_flat("norm1_beta"),
        "norm2_gamma":  cat_flat("norm2_gamma"),
        "norm2_beta":   cat_flat("norm2_beta"),
    }

    return flat
