model_group_split.py

"""Split a Keras .h5 model into **groups** of consecutive layers for split computing.

Three operating modes:

  --analyze      Show model architecture, valid split points, and bottleneck
                 ranking (no files written).
  --split-at N[,N,...]  Manually specify layer indices at which to split.
  --auto N       Automatically find the N best bottleneck split points.

Each group is exported as .h5, .tflite and .h (C array for microcontrollers),
mirroring the per-layer export of model_split.py.
"""

import argparse
import os
import sys
from functools import reduce
from operator import mul

# Use the Keras 2 API (tf-keras): the .h5 models were saved with Keras 2 and rely
# on layer.input_shape and the DepthwiseConv2D 'groups' arg, both gone in Keras 3.
os.environ["TF_USE_LEGACY_KERAS"] = "1"
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"  # suppress TF info/warning noise

import numpy as np
import tensorflow as tf
from tensorflow.keras import Input
from tensorflow.keras import layers
from tensorflow.keras.models import Model

MODELS_DIR = "models"

# ─── Helpers ────────────────────────────────────────────────────────────────────


def init_folders(root_folder: str) -> None:
    """Create the output directory tree for group exports."""
    os.makedirs(f"{root_folder}/groups/h5/", exist_ok=True)
    os.makedirs(f"{root_folder}/groups/tflite/", exist_ok=True)
    os.makedirs(f"{root_folder}/groups/h/", exist_ok=True)


def to_tflite(keras_model: Model, save: bool, save_dir: str, name: str) -> bytes:
    """Convert a Keras model to TFLite format, optionally saving to disk."""
    converter = tf.lite.TFLiteConverter.from_keras_model(keras_model)
    tflite_model = converter.convert()
    if save:
        with open(f"{save_dir}/{name}.tflite", "wb") as f:
            f.write(tflite_model)
    return tflite_model


def load_h5(name: str, dir_path: str) -> Model:
    """Load a Keras .h5 model from *dir_path*/*name*.h5."""
    return tf.keras.models.load_model(f"{dir_path}/{name}.h5")


def tensor_size(shape) -> int:
    """Number of elements in a tensor shape, excluding the batch dimension.

    Handles the case where *shape* is a list of shapes (multi-output layer) by
    summing the element counts of each output.
    """
    if isinstance(shape, list):
        return sum(tensor_size(s) for s in shape)
    dims = shape[1:]  # drop batch dimension
    if not dims:
        return 0
    return reduce(mul, dims, 1)


# ─── Computation-graph analysis ─────────────────────────────────────────────────


def build_layer_graph(model) -> dict:
    """Return a mapping ``layer_idx → {source layer indices}`` by inspecting
    the Keras computation graph tensors.

    The mapping tells us, for every layer in *model*, which other layers feed
    into it.
    """
    # Step 1: map each output tensor (by ``id``) to its producing layer index.
    tensor_to_layer: dict[int, int] = {}
    for idx, layer in enumerate(model.layers):
        output = layer.output
        if isinstance(output, list):
            for t in output:
                tensor_to_layer[id(t)] = idx
        else:
            tensor_to_layer[id(output)] = idx

    # Step 2: for each layer, look at its input tensor(s) and resolve which
    # layer produced them.
    inbound: dict[int, set[int]] = {}
    for idx, layer in enumerate(model.layers):
        inbound[idx] = set()
        if isinstance(layer, layers.InputLayer):
            continue  # no inbound connections
        inp = layer.input
        if isinstance(inp, list):
            for t in inp:
                src = tensor_to_layer.get(id(t))
                if src is not None:
                    inbound[idx].add(src)
        else:
            src = tensor_to_layer.get(id(inp))
            if src is not None:
                inbound[idx].add(src)

    return inbound


def find_valid_split_points(model) -> list[int]:
    """Return the sorted list of layer indices where splitting is valid.

    A split at index *i* partitions the model into:

    * **Part A** – layers ``[0 .. i]`` (runs on device 1, transmits layer *i*'s
      output tensor).
    * **Part B** – layers ``[i+1 .. N-1]`` (runs on device 2).

    The split is **valid** if and only if every layer in Part B receives all its
    inputs either from other Part B layers or from layer *i* itself (the single
    tensor transmitted across devices).  In other words, no skip / residual
    connection may cross the split boundary except through the split-point
    layer.
    """
    inbound = build_layer_graph(model)
    n = len(model.layers)
    valid: list[int] = []

    for split_idx in range(1, n - 1):
        if isinstance(model.layers[split_idx], layers.InputLayer):
            continue

        is_valid = True
        for b_idx in range(split_idx + 1, n):
            for src_idx in inbound[b_idx]:
                if src_idx < split_idx:  # source is before the split-point layer
                    is_valid = False
                    break
            if not is_valid:
                break

        if is_valid:
            valid.append(split_idx)

    return valid


def compute_bottleneck_scores(model, valid_points: list[int]):
    """Rank valid split points by ascending tensor-transfer size.

    Returns a list of ``(layer_idx, tensor_size, layer_name, output_shape)``
    tuples sorted from smallest (best bottleneck) to largest.
    """
    results = []
    for idx in valid_points:
        layer = model.layers[idx]
        out_shape = layer.output_shape
        size = tensor_size(out_shape)
        results.append((idx, size, layer.name, out_shape))
    results.sort(key=lambda x: x[1])
    return results


# ─── Analyse mode ───────────────────────────────────────────────────────────────


def analyze_model(model) -> None:
    """Pretty-print the model architecture with bottleneck analysis."""
    n_layers = len(model.layers)

    print(f"\n{'=' * 90}")
    print("  Model Analysis")
    print(f"{'=' * 90}")
    print(f"  Total layers: {n_layers}")
    print(f"  Input shape:  {model.input_shape}")
    print(f"  Output shape: {model.output_shape}")
    print()

    # ── All layers ──────────────────────────────────────────────────────────
    header = (f"{'Idx':>4}  {'Type':<25}  {'Name':<30}  "
              f"{'Output Shape':<28}  {'Tensor Size':>12}")
    print(header)
    print(f"{'─' * 4}  {'─' * 25}  {'─' * 30}  {'─' * 28}  {'─' * 12}")
    for i, layer in enumerate(model.layers):
        out_shape = layer.output_shape
        size = tensor_size(out_shape)
        print(f"{i:4d}  {layer.__class__.__name__:<25}  {layer.name:<30}  "
              f"{str(out_shape):<28}  {size:>12,}")

    # ── Valid split points ──────────────────────────────────────────────────
    valid = find_valid_split_points(model)
    if not valid:
        print(f"\n⚠  No valid split points found "
              "(the model may have skip connections spanning the entire network).")
        return

    ranked = compute_bottleneck_scores(model, valid)

    print(f"\n{'=' * 90}")
    print(f"  Valid Split Points  ({len(valid)} found)")
    print(f"{'=' * 90}")

    # In layer order
    print(f"\n  In layer order:")
    print(f"  {'Idx':>4}  {'Type':<25}  {'Name':<30}  {'Tensor Size':>12}")
    print(f"  {'─' * 4}  {'─' * 25}  {'─' * 30}  {'─' * 12}")
    for idx in valid:
        layer = model.layers[idx]
        size = tensor_size(layer.output_shape)
        print(f"  {idx:4d}  {layer.__class__.__name__:<25}  "
              f"{layer.name:<30}  {size:>12,}")

    # Ranked by bottleneck quality
    print(f"\n  Ranked by bottleneck (smallest tensor → best for split computing):")
    print(f"  {'Rank':>4}  {'Idx':>4}  {'Type':<25}  {'Name':<30}  {'Tensor Size':>12}")
    print(f"  {'─' * 4}  {'─' * 4}  {'─' * 25}  {'─' * 30}  {'─' * 12}")
    for rank, (idx, size, name, _) in enumerate(ranked, 1):
        layer = model.layers[idx]
        marker = "  ★" if rank <= 3 else ""
        print(f"  {rank:4d}  {idx:4d}  {layer.__class__.__name__:<25}  "
              f"{name:<30}  {size:>12,}{marker}")

    print()


# ─── Sub-model extraction ──────────────────────────────────────────────────────


def extract_group_submodel(model, start_idx: int, end_idx: int) -> Model:
    """Extract layers ``[start_idx .. end_idx]`` as a standalone Keras model.

    **First group** (starting from the model input): reuses the original
    computation graph directly, so internal skip connections are preserved
    automatically.

    **Subsequent groups**: rebuilds the sub-graph with a fresh ``Input``
    tensor whose shape matches the previous group's output.  A ``tensor_map``
    translates original-graph tensors to new-graph tensors, correctly handling
    internal skip / residual connections within the group.
    """
    # Identify the first "real" (non-InputLayer) layer index
    first_real = 1 if isinstance(model.layers[0], layers.InputLayer) else 0

    if start_idx <= first_real:
        # First group – reference the original graph (preserves topology).
        return Model(inputs=model.input, outputs=model.layers[end_idx].output)

    # ── Subsequent groups: rebuild the sub-graph ────────────────────────────

    # 1. Snapshot every original tensor reference *before* we call any layer a
    #    second time (which would create extra inbound nodes).
    original_inputs: dict[int, object] = {}
    original_outputs: dict[int, object] = {}
    for idx in range(start_idx, end_idx + 1):
        original_inputs[idx] = model.layers[idx].input
        original_outputs[idx] = model.layers[idx].output

    # 2. Create a new Input that matches the previous group's output shape.
    prev_output = model.layers[start_idx - 1].output
    prev_shape = model.layers[start_idx - 1].output_shape

    if isinstance(prev_shape, list):
        new_input = [Input(shape=s[1:]) for s in prev_shape]
        tensor_map: dict[int, object] = {}
        for orig_t, new_t in zip(prev_output, new_input):
            tensor_map[id(orig_t)] = new_t
    else:
        new_input = Input(shape=prev_shape[1:])
        tensor_map = {id(prev_output): new_input}

    # 3. Re-execute each layer with mapped tensors.
    for idx in range(start_idx, end_idx + 1):
        layer = model.layers[idx]
        orig_inp = original_inputs[idx]

        if isinstance(orig_inp, list):
            mapped = [tensor_map[id(t)] for t in orig_inp]
        else:
            mapped = tensor_map[id(orig_inp)]

        new_out = layer(mapped)
        tensor_map[id(original_outputs[idx])] = new_out

    final_output = tensor_map[id(original_outputs[end_idx])]
    return Model(inputs=new_input, outputs=final_output)


# ─── Validation helpers ─────────────────────────────────────────────────────────


def validate_split_points(model, split_points: list[int]):
    """Check that every requested split point is within range and graph-valid.

    Returns ``(ok: bool, message: str)``.
    """
    valid_set = set(find_valid_split_points(model))
    n = len(model.layers)

    for sp in split_points:
        if sp < 1 or sp >= n - 1:
            return False, f"Split point {sp} is out of range (valid: 1 .. {n - 2})."
        if sp not in valid_set:
            nearest = min(valid_set, key=lambda v: abs(v - sp)) if valid_set else None
            msg = (f"Split point {sp} is not valid "
                   "(skip connections cross this boundary).")
            if nearest is not None:
                msg += f" Nearest valid point: {nearest}."
            return False, msg

    if sorted(set(split_points)) != sorted(split_points):
        return False, "Split points must be unique and in ascending order."

    return True, "OK"


def find_auto_split_points(model, n_splits: int) -> list[int]:
    """Return the *n_splits* valid split points with the smallest tensor size."""
    valid = find_valid_split_points(model)
    if not valid:
        return []

    ranked = compute_bottleneck_scores(model, valid)

    if n_splits >= len(ranked):
        return sorted(r[0] for r in ranked)

    return sorted(r[0] for r in ranked[:n_splits])


# ─── Verification ───────────────────────────────────────────────────────────────


def verify_split(model, submodels: dict, boundaries: list) -> bool:
    """Run a random input through original and split pipeline, compare outputs.

    Returns ``True`` if the outputs match within tolerance.
    """
    print("\nVerifying split correctness …")

    # Random input
    input_shape = model.input_shape
    if isinstance(input_shape, list):
        test_input = [np.random.randn(1, *s[1:]).astype(np.float32)
                      for s in input_shape]
    else:
        test_input = np.random.randn(1, *input_shape[1:]).astype(np.float32)

    # Original model
    original_output = model.predict(test_input, verbose=0)

    # Sequential pass through groups
    x = test_input
    for group_name in sorted(submodels.keys()):
        x = submodels[group_name].predict(x, verbose=0)
    split_output = x

    # Compare
    if isinstance(original_output, list):
        diffs = [np.max(np.abs(o - s))
                 for o, s in zip(original_output, split_output)]
        max_diff = max(diffs)
    else:
        max_diff = float(np.max(np.abs(original_output - split_output)))

    atol = 1e-4
    if max_diff < atol:
        print(f"  ✓ Verification PASSED (max diff = {max_diff:.2e}, "
              f"tolerance = {atol:.0e})")
        return True
    else:
        print(f"  ✗ Verification FAILED (max diff = {max_diff:.2e}, "
              f"tolerance = {atol:.0e})")
        print("    This may indicate skip connections crossing group boundaries.")
        return False


# ─── Full export pipeline ───────────────────────────────────────────────────────


def perform_group_split(model, split_points: list[int],
                        model_name: str, main_folder: str) -> dict:
    """Split *model* at *split_points* and export each group in h5/tflite/h.

    Returns a ``{group_name: keras.Model}`` dict.
    """
    init_folders(main_folder)

    n = len(model.layers)
    first_idx = 1 if isinstance(model.layers[0], layers.InputLayer) else 0

    # ── Compute group boundaries ────────────────────────────────────────────
    boundaries: list[tuple[int, int]] = []
    prev = first_idx
    for sp in sorted(split_points):
        boundaries.append((prev, sp))
        prev = sp + 1
    boundaries.append((prev, n - 1))

    print(f"\nSplitting model into {len(boundaries)} groups:")
    for i, (start, end) in enumerate(boundaries):
        out_shape = model.layers[end].output_shape
        size = tensor_size(out_shape)
        print(f"  Group {i}: layers [{start}..{end}]  "
              f"({model.layers[start].name} → {model.layers[end].name})  "
              f"│ output tensor: {size:,} elements")

    # ── Extract, convert, and save each group ───────────────────────────────
    submodels: dict[str, Model] = {}
    macro_def = "#define LOAD_GROUP() "
    h_dir = f"{main_folder}/groups/h"

    with open(f"{h_dir}/groups.h", "w") as super_header:
        for group_idx, (start, end) in enumerate(boundaries):
            gname = f"group_{group_idx}"
            print(f"\n─── Group {group_idx}  (layers {start}..{end}) "
                  f"{'─' * 40}")

            # Keras sub-model
            print("  Extracting sub-model …")
            sub = extract_group_submodel(model, start, end)
            sub.summary(print_fn=lambda x: print(f"    {x}"))
            submodels[gname] = sub

            # .h5
            h5_path = f"{main_folder}/groups/h5/{gname}.h5"
            sub.save(h5_path)
            print(f"  ✓ {h5_path}")

            # .tflite
            print("  Converting to TFLite …")
            tflite_bytes = to_tflite(
                sub, save=True,
                save_dir=f"{main_folder}/groups/tflite",
                name=gname,
            )
            print(f"  ✓ {main_folder}/groups/tflite/{gname}.tflite")

            # .h (C header)
            model_array = ", ".join(str(b) for b in tflite_bytes)
            h_path = f"{h_dir}/{gname}.h"
            with open(h_path, "w") as hf:
                hf.write("#pragma once\n\n")
                hf.write("#include <cstdint>\n\n")
                hf.write(f"const uint8_t {gname}[] = {{\n")
                hf.write(model_array)
                hf.write("\n};\n")
            print(f"  ✓ {h_path}")

            super_header.write(f'#include "{gname}.h"\n')
            macro_def += (
                f'if(group_name.equals("{gname}")) '
                f"model = tflite::GetModel({gname});\\\n"
            )

        super_header.write(macro_def)

    # ── Verify correctness ──────────────────────────────────────────────────
    verify_split(model, submodels, boundaries)

    print(f"\n{'=' * 70}")
    print(f"  Done! {len(boundaries)} groups saved to {main_folder}/groups/")
    print(f"{'=' * 70}\n")
    return submodels


# ─── CLI ────────────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description=(
            f"Split a Keras .h5 model (from '{MODELS_DIR}/') into groups of "
            f"consecutive layers for split computing."
        ),
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""\
examples:
  %(prog)s FOMO_224.h5 --analyze
      Show model architecture and bottleneck analysis.

  %(prog)s FOMO_224.h5 --split-at 10
      Split into 2 groups: [0..10] and [11..end].

  %(prog)s FOMO_224.h5 --split-at 10,20
      Split into 3 groups: [0..10], [11..20], and [21..end].

  %(prog)s FOMO_224.h5 --auto 1
      Auto-find the best bottleneck and split there (→ 2 groups).

  %(prog)s FOMO_224.h5 --auto 2
      Auto-find 2 best bottlenecks (→ 3 groups).
""",
    )
    parser.add_argument(
        "model_file",
        help=f"Name of the .h5 model file in '{MODELS_DIR}/' (e.g. FOMO_224.h5)",
    )

    mode = parser.add_mutually_exclusive_group(required=True)
    mode.add_argument(
        "--analyze",
        action="store_true",
        help="Show architecture, valid split points, and bottleneck ranking.",
    )
    mode.add_argument(
        "--split-at",
        type=str,
        metavar="N[,N,...]",
        help="Comma-separated layer indices at which to split (e.g. --split-at 10 or --split-at 10,20).",
    )
    mode.add_argument(
        "--auto",
        type=int,
        metavar="N",
        help="Automatically find the N best bottleneck split points.",
    )

    args = parser.parse_args()

    # ── Resolve model path ──────────────────────────────────────────────────
    model_name = os.path.splitext(args.model_file)[0]
    model_path = f"{MODELS_DIR}/{args.model_file}"
    if not os.path.isfile(model_path):
        parser.error(f"Model file not found: {model_path}")

    print(f"Loading model: {model_path}")
    model = load_h5(name=model_name, dir_path=MODELS_DIR)

    # ── Dispatch ────────────────────────────────────────────────────────────
    if args.analyze:
        analyze_model(model)

    elif args.split_at is not None:
        try:
            split_points = sorted(int(x) for x in args.split_at.split(","))
        except ValueError:
            parser.error(f"--split-at values must be integers separated by commas, got: '{args.split_at}'")
        ok, msg = validate_split_points(model, split_points)
        if not ok:
            print(f"\n✗ {msg}")
            print("  Run with --analyze to see valid split points.")
            sys.exit(1)

        main_folder = f"{MODELS_DIR}/{model_name}"
        perform_group_split(model, split_points, model_name, main_folder)

    elif args.auto is not None:
        if args.auto < 1:
            parser.error("--auto N must be ≥ 1")

        split_points = find_auto_split_points(model, args.auto)
        if not split_points:
            print("\n✗ No valid split points found.")
            sys.exit(1)

        if len(split_points) < args.auto:
            print(f"\n⚠ Only {len(split_points)} valid split point(s) found "
                  f"(requested {args.auto}).")

        print(f"\nAuto-selected split points: {split_points}")
        for sp in split_points:
            layer = model.layers[sp]
            size = tensor_size(layer.output_shape)
            print(f"  Layer {sp} ({layer.name}): tensor_size = {size:,}")

        main_folder = f"{MODELS_DIR}/{model_name}"
        perform_group_split(model, split_points, model_name, main_folder)