import importlib.util
import os
import sys
from pathlib import Path

import onnx.utils

def load_splice_function(data_dir: str):
    """Load the per-data-directory model slicing function."""
    splice_path = Path(data_dir) / "splice.py"
    if not splice_path.exists():
        raise FileNotFoundError(f"No splice.py found in {splice_path.parent}")

    # Each experiment can define different ONNX split points. Loading the
    # module under a data-directory-specific name avoids collisions when
    # multiple experiments are used in one Python process.
    module_name = f"fat_splice_{splice_path.parent.name}"
    spec = importlib.util.spec_from_file_location(module_name, splice_path)
    if spec is None or spec.loader is None:
        raise ImportError(f"Could not load splice module from {splice_path}")

    module = importlib.util.module_from_spec(spec)
    sys.modules[module_name] = module
    spec.loader.exec_module(module)

    splice = getattr(module, "splice", None)
    if splice is None or not callable(splice):
        raise AttributeError(f"{splice_path} must define a callable splice(data_dir, count)")

    return splice


def delete_submodels(data_dir):
    """ Delete any existing submodels in data_dir"""
    subparts = ["a", "b", "c"]

    for x in subparts:
        file_path = f"{data_dir}model_{x}.onnx"

        if os.path.exists(file_path):
            os.remove(file_path)


def split_into_2(data_dir, input_names, output_names):
    """ Split model in two parts. 
        input_names and output_names refer to the exact names of input/output tensors of the first part. 
        Model will be split into two parts and saved as a and b in data_dir.
        Will delete any existing submodels data_dir.
    """
   
    delete_submodels(data_dir)

    model_path = f"{data_dir}model.onnx"

    # Extract the first path
    onnx.utils.extract_model(model_path, f"{data_dir}model_a.onnx", input_names, output_names)
    
    model = onnx.load(model_path)
    # figure out the inputs and outputs of original model
    input = model.graph.input
    input_names_2 = [x.name for x in input if x.name not in input_names]
    input_names_2 = input_names_2 + output_names
    output = model.graph.output
    output_names = [x.name for x in output]

    onnx.utils.extract_model(model_path, f"{data_dir}model_b.onnx", input_names_2, output_names)
    return

def split_into_3(data_dir, first, middle, last):
    """ Split model in three parts. 
        input_names and output_names refer to the exact names of input/output tensors of the first part. 
        Model will be split into three parts and saved as a, b and c in data_dir.
        Will delete any existing submodels data_dir.
    """
    
    delete_submodels(data_dir)

    model_path = f"{data_dir}model.onnx"

    model = onnx.load(model_path)


    # Extract the first path
    onnx.utils.extract_model(model_path, f"{data_dir}model_a.onnx", first[0], first[1])

    # Extract the middle path
    onnx.utils.extract_model(model_path, f"{data_dir}model_b.onnx", middle[0], middle[1])

    # Extract the last path
    onnx.utils.extract_model(model_path, f"{data_dir}model_c.onnx", last[0], last[1])

   
    
    return

import numpy as np
from onnx import numpy_helper


EXPECTED_KEYS = [
    "W_QKV", "B_QKV", "W_O", "B_O",
    "ff_w1", "ff_b1", "ff_w2", "ff_b2",
    "norm1_gamma", "norm1_beta",
    "norm2_gamma", "norm2_beta",
]


def find_node_by_name(model, node_name: str):
    for node in model.graph.node:
        if node.name == node_name:
            return node
    raise ValueError(f"Node '{node_name}' not found.")


def find_node_by_name_contains(model, pattern: str):
    matches = [node for node in model.graph.node if pattern in node.name]
    if not matches:
        raise ValueError(f"No node name contains '{pattern}'.")
    if len(matches) > 1:
        raise ValueError(
            f"Multiple node names contain '{pattern}': {[n.name for n in matches]}"
        )
    return matches[0]


def find_initializer_by_contains(init_dict, pattern: str):
    matches = [name for name in init_dict if pattern in name]
    if not matches:
        raise ValueError(f"No initializer contains '{pattern}'.")
    if len(matches) > 1:
        raise ValueError(
            f"Multiple initializers contain '{pattern}': {matches}"
        )
    return matches[0]


def get_input_from_node_by_match(node, input_contains: str):
    for inp in node.input:
        if input_contains in inp:
            return inp
    raise ValueError(
        f"No input containing '{input_contains}' found for node '{node.name}'. "
        f"Inputs were: {list(node.input)}"
    )


def build_node_input_consumers(model):
    consumers = {}
    for node in model.graph.node:
        for inp in node.input:
            consumers.setdefault(inp, []).append(node)
    return consumers


def find_weight_matmul_nodes(model, init_dict):
    consumers_by_output = build_node_input_consumers(model)
    result = []

    for node in model.graph.node:
        if node.op_type != "MatMul":
            continue

        consumers = []
        for out in node.output:
            consumers.extend(consumers_by_output.get(out, []))

        has_initializer_consumer = any(
            any(inp in init_dict for inp in consumer.input)
            for consumer in consumers
        )
        if has_initializer_consumer:
            result.append(node)

    return result


def get_initializer_input_from_node(node, init_dict):
    init_inputs = [inp for inp in node.input if inp in init_dict]
    if not init_inputs:
        raise ValueError(
            f"No initializer input found for node '{node.name}'. "
            f"Inputs were: {list(node.input)}"
        )
    if len(init_inputs) > 1:
        raise ValueError(
            f"Multiple initializer inputs found for node '{node.name}': {init_inputs}"
        )
    return init_inputs[0]


def resolve_layer_weight_matmul(model, init_dict, layer_idx: int, position_within_layer: int):
    weight_matmuls = find_weight_matmul_nodes(model, init_dict)

    if len(weight_matmuls) < 2:
        raise ValueError("Could not find enough weight-bearing MatMul nodes.")

    encoder_matmuls = weight_matmuls[1:-1]
    start = layer_idx * 3
    layer_nodes = encoder_matmuls[start:start + 3]

    if len(layer_nodes) != 3:
        raise ValueError(
            f"Could not resolve 3 weight MatMul nodes for layer {layer_idx}. "
            f"Got {[node.name for node in layer_nodes]}"
        )

    try:
        node = layer_nodes[position_within_layer]
    except IndexError as exc:
        raise ValueError(
            f"Invalid position_within_layer={position_within_layer} for layer {layer_idx}"
        ) from exc

    return get_initializer_input_from_node(node, init_dict)


def resolve_tensor_name(model, init_dict, spec: dict, layer_idx: int):
    """Resolve one parameter spec from model.yaml to an ONNX initializer name."""
    spec_type = spec["type"]

    if spec_type == "initializer":
        return spec["name"].format(layer_idx=layer_idx)

    if spec_type == "initializer_contains":
        pattern = spec["name_contains"].format(layer_idx=layer_idx)
        return find_initializer_by_contains(init_dict, pattern)

    if spec_type == "node_input_contains":
        if "node_name" in spec:
            node = find_node_by_name(
                model,
                spec["node_name"].format(layer_idx=layer_idx)
            )
        elif "node_name_contains" in spec:
            node = find_node_by_name_contains(
                model,
                spec["node_name_contains"].format(layer_idx=layer_idx)
            )
        else:
            raise ValueError(
                "node_input_contains spec requires 'node_name' or 'node_name_contains'."
            )

        return get_input_from_node_by_match(node, spec["input_contains"])

    if spec_type == "layer_weight_matmul":
        return resolve_layer_weight_matmul(
            model,
            init_dict,
            layer_idx=layer_idx,
            position_within_layer=spec["position_within_layer"],
        )

    raise ValueError(f"Unsupported parameter spec type: '{spec_type}'")


def extract_parameters(model, layer_count, config, store_type=np.float32):
    """Extract and flatten encoder parameters in the order expected by HLS."""
    init = {t.name: numpy_helper.to_array(t) for t in model.graph.initializer}
    param_specs = config.get("parameters", config)

    missing = [k for k in EXPECTED_KEYS if k not in param_specs]
    if missing:
        raise ValueError(f"Config is missing parameter definitions for: {missing}")

    per_layer = []

    for layer_idx in range(layer_count):
        layer_params = {}

        for key in EXPECTED_KEYS:
            tensor_name = resolve_tensor_name(model, init, param_specs[key], layer_idx)

            if tensor_name not in init:
                raise KeyError(
                    f"Resolved tensor '{tensor_name}' for key '{key}' "
                    f"(layer {layer_idx}) is not in model initializers."
                )

            layer_params[key] = init[tensor_name].astype(store_type)

        per_layer.append(layer_params)

    def cat_flat(key):
        return np.concatenate(
            [per_layer[layer_idx][key].reshape(-1) for layer_idx in range(layer_count)],
            axis=0
        )

    return {key: cat_flat(key) for key in EXPECTED_KEYS}


import numpy as np


def compare_parameter_dicts(old_params, new_params, atol=1e-6, rtol=1e-6, verbose=True):
    """
    Compare two parameter dictionaries.

    Returns True if everything matches, otherwise False.
    """

    old_keys = set(old_params.keys())
    new_keys = set(new_params.keys())

    if old_keys != new_keys:
        print("Key mismatch detected.")
        print("Only in old:", sorted(old_keys - new_keys))
        print("Only in new:", sorted(new_keys - old_keys))
        return False

    all_ok = True

    for key in sorted(old_keys):
        a = old_params[key]
        b = new_params[key]

        same_shape = a.shape == b.shape
        same_dtype = a.dtype == b.dtype
        same_values = np.allclose(a, b, atol=atol, rtol=rtol)

        if verbose:
            print(f"[{key}]")
            print(f"  old shape={a.shape}, dtype={a.dtype}")
            print(f"  new shape={b.shape}, dtype={b.dtype}")
            print(f"  shape_match={same_shape}, dtype_match={same_dtype}, value_match={same_values}")

        if not (same_shape and same_dtype and same_values):
            all_ok = False

            if same_shape:
                diff = np.abs(a - b)
                if verbose:
                    print("  max_abs_diff:", diff.max())
                    print("  mean_abs_diff:", diff.mean())

            if verbose:
                print("  old values:")
                print(a)
                print("  new values:")
                print(b)

        if verbose:
            print()

    if verbose:
        if all_ok:
            print("SUCCESS: parameters match.")
        else:
            print("MISMATCH: parameters differ.")

    return all_ok
