#!/usr/bin/env python3

import argparse
import json
from pathlib import Path


def build_latency_config(core_count, crossbar_count, crossbar_size, layout):
    return {
        "chip_config": {
            "core_config": {
                "rob_size": 1,
                "period": 0.5,
                "matrix_config": {
                    "xbar_array_count": crossbar_count,
                    "period": 0.5,
                    "pipeline_mode": True,
                    "dac_resolution": 4,
                    "dac_count": crossbar_size,
                    "xbar_size": [crossbar_size, crossbar_size],
                    "cell_precision": 2,
                    "xbar_latency": 10,
                    "sample_hold_latency_cycle": 1,
                    "adc_resolution": 8,
                    "adc_latency_cycle": 4,
                    "adc_count": max(128, crossbar_size // 8),
                    "shift_adder_latency_cycle": 1,
                    "output_buffer_latency_cycle": 1,
                },
                "vector_width": 64,
                "vector_latency_cycle": 2,
                "local_memory_config": {
                    "data_width": 256,
                    "period": 0.5,
                    "write_latency_cycle": 1,
                    "read_latency_cycle": 1,
                },
                "global_memory_switch_id": -10,
            },
            "global_memory_config": {
                "data_width": 256,
                "period": 0.5,
                "write_latency_cycle": 1,
                "read_latency_cycle": 1,
            },
            "network_config": {
                "bus_topology": "mesh",
                "bus_width": 256,
                "layout": [layout[0], layout[1]],
                "net_config_file_path": f"network_mesh_{core_count}.json",
            },
            "core_cnt": core_count,
            "global_memory_switch_id": -10,
        },
        "sim_config": {
            "sim_mode": 1,
            "sim_time": 1,
            "report_verbose_level": 0,
        },
    }


def build_network(core_count, layout):
    rows, cols = layout
    assert rows * cols == core_count

    global_memory_latency = 6
    global_memory_energy = 24
    per_hop_latency = 1
    base_latency = 2
    per_hop_energy = 3
    base_energy = 12

    latency = {"-10": {}}
    energy = {"-10": {}}

    for dst in range(core_count):
        latency["-10"][str(dst)] = global_memory_latency
        energy["-10"][str(dst)] = global_memory_energy

    for src in range(core_count):
        src_row, src_col = divmod(src, cols)
        latency[str(src)] = {"-10": global_memory_latency}
        energy[str(src)] = {"-10": global_memory_energy}
        for dst in range(core_count):
            if src == dst:
                continue
            dst_row, dst_col = divmod(dst, cols)
            hops = abs(src_row - dst_row) + abs(src_col - dst_col)
            latency[str(src)][str(dst)] = base_latency + per_hop_latency * hops
            energy[str(src)][str(dst)] = base_energy + per_hop_energy * hops

    return {"latency": latency, "energy": energy}


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--output-dir", required=True)
    parser.add_argument("--core-count", type=int, required=True)
    parser.add_argument("--crossbar-count", type=int, required=True)
    parser.add_argument("--crossbar-size", type=int, required=True)
    parser.add_argument("--mesh-rows", type=int, required=True)
    parser.add_argument("--mesh-cols", type=int, required=True)
    args = parser.parse_args()

    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    layout = (args.mesh_rows, args.mesh_cols)
    latency_config = build_latency_config(
        args.core_count, args.crossbar_count, args.crossbar_size, layout
    )
    network = build_network(args.core_count, layout)

    with open(output_dir / "latency_config.json", "w", encoding="utf-8") as f:
        json.dump(latency_config, f, indent=2)
        f.write("\n")

    network_path = output_dir / f"network_mesh_{args.core_count}.json"
    with open(network_path, "w", encoding="utf-8") as f:
        json.dump(network, f, separators=(",", ":"))
        f.write("\n")


if __name__ == "__main__":
    main()