#!/usr/bin/env python3

# Convert .tn.model files from Keras 2 to Keras 3 format for libfdeep-dev 0.18.2.
#
# Differences handled:
#   1. InputLayer config: "batch_input_shape" -> "batch_shape"
#   2. inbound_nodes: [[["layer", n, t, {}]]] -> [{"args": [{"config": {"keras_history": ["layer", n, t]}}]}]

import json
import sys


def convert_inbound_nodes(nodes_k2):
    """Convert Keras 2 inbound_nodes list to Keras 3 format expected by fdeep 0.18.2."""
    nodes_k3 = []
    for conn_list in nodes_k2:
        args = [{"config": {"keras_history": conn[:3]}} for conn in conn_list]
        nodes_k3.append({"args": args})
    return nodes_k3


def convert_model(data):
    arch = data.get("architecture")
    if not isinstance(arch, dict):
        return data

    config = arch.get("config", {})
    layers = config.get("layers", [])

    for layer in layers:
        layer_config = layer.get("config", {})

        # Fix 1: batch_input_shape -> batch_shape
        if "batch_input_shape" in layer_config:
            layer_config["batch_shape"] = layer_config.pop("batch_input_shape")

        # Fix 2: inbound_nodes Keras 2 -> Keras 3
        inbound = layer.get("inbound_nodes")
        if inbound:
            # Detect Keras 2 format: outer list elements are arrays (not dicts)
            if isinstance(inbound[0], list):
                layer["inbound_nodes"] = convert_inbound_nodes(inbound)

    return data


def main():
    for path in sys.argv[1:]:
        with open(path) as f:
            data = json.load(f)
        data = convert_model(data)
        with open(path, "w") as f:
            json.dump(data, f, separators=(",", ":"))
        print(f"Converted: {path}", file=sys.stderr)


if __name__ == "__main__":
    main()
