#!/usr/bin/env python3
"""Machine-readable registry for the complete cross-architecture panel."""
import json


METHODS = (
    "bp",
    "fa",
    "dfa",
    "pepita",
    "ff",
    "ep",
    "dualprop",
    "clean_kp",
    "sdil",
)

ARCHITECTURES = (
    {
        "id": "minicnn",
        "family": "plain_cnn",
        "scale_index": 0,
        "size": {"trainable_layers": 3},
    },
    {
        "id": "vgglike",
        "family": "plain_cnn",
        "scale_index": 1,
        "size": {"trainable_layers": 5},
    },
    {
        "id": "vgg16",
        "family": "plain_cnn",
        "scale_index": 2,
        "size": {"trainable_layers": 16},
    },
    {
        "id": "resnet20",
        "family": "residual_cnn",
        "scale_index": 0,
        "size": {"depth": 20, "base_width": 16},
    },
    {
        "id": "resnet32",
        "family": "residual_cnn",
        "scale_index": 1,
        "size": {"depth": 32, "base_width": 16},
    },
    {
        "id": "resnet56",
        "family": "residual_cnn",
        "scale_index": 2,
        "size": {"depth": 56, "base_width": 16},
    },
    {
        "id": "transformer4",
        "family": "decoder_transformer",
        "scale_index": 0,
        "size": {"blocks": 4, "width": 128, "context": 64},
    },
    {
        "id": "transformer8",
        "family": "decoder_transformer",
        "scale_index": 1,
        "size": {"blocks": 8, "width": 128, "context": 64},
    },
    {
        "id": "transformer12",
        "family": "decoder_transformer",
        "scale_index": 2,
        "size": {"blocks": 12, "width": 128, "context": 64},
    },
)


def primary_cells():
    cells = []
    for architecture in ARCHITECTURES:
        for method in METHODS:
            cells.append({
                "cell_id": f"{architecture['id']}::{method}",
                "architecture": architecture["id"],
                "family": architecture["family"],
                "scale_index": architecture["scale_index"],
                "size": architecture["size"],
                "method": method,
            })
    identifiers = [cell["cell_id"] for cell in cells]
    if len(cells) != 81 or len(set(identifiers)) != len(identifiers):
        raise AssertionError("primary crossover must contain 81 unique cells")
    for family in {row["family"] for row in ARCHITECTURES}:
        subset = [cell for cell in cells if cell["family"] == family]
        if len(subset) != 27:
            raise AssertionError(f"{family} does not contain 27 cells")
        if {cell["method"] for cell in subset} != set(METHODS):
            raise AssertionError(f"{family} is missing a method")
    return cells


def registry():
    cells = primary_cells()
    return {
        "schema_version": 1,
        "num_architecture_scale_points": len(ARCHITECTURES),
        "num_methods": len(METHODS),
        "num_primary_cells": len(cells),
        "architectures": ARCHITECTURES,
        "methods": METHODS,
        "cells": cells,
        "extension_rule":
            "Every added architecture/size must add all nine methods.",
    }


if __name__ == "__main__":
    print(json.dumps(registry(), indent=2, sort_keys=True))