#!/usr/bin/env python3
"""Convert interp_lb_profile.bin to a C header for compile-time embedding."""
import struct, sys

if len(sys.argv) < 3:
    print(f"Usage: {sys.argv[0]} <profile.bin> <output.h>")
    sys.exit(1)

with open(sys.argv[1], 'rb') as f:
    magic, version, nprocs, num_heavy = struct.unpack('IIii', f.read(16))
    threshold = struct.unpack('d', f.read(8))[0]
    times = list(struct.unpack(f'{nprocs}d', f.read(nprocs * 8)))
    heavy = list(struct.unpack(f'{num_heavy}i', f.read(num_heavy * 4)))

# For each heavy rank, compute split: left half -> lighter neighbor, right half -> heavy rank
# (or vice versa depending on which neighbor is lighter)
splits = []
for hr in heavy:
    prev_t = times[hr - 1] if hr > 0 else 1e30
    next_t = times[hr + 1] if hr < nprocs - 1 else 1e30
    if prev_t <= next_t:
        splits.append((hr, hr - 1, hr))  # (block_id, r_left, r_right)
    else:
        splits.append((hr, hr, hr + 1))

# Also remap the displaced neighbor blocks
remaps = {}
for hr, r_l, r_r in splits:
    if r_l != hr:
        # We took r_l's slot, so remap block r_l to its other neighbor
        displaced = r_l
        if displaced > 0 and displaced - 1 not in [s[0] for s in splits]:
            remaps[displaced] = displaced - 1
        elif displaced < nprocs - 1:
            remaps[displaced] = displaced + 1
    else:
        displaced = r_r
        if displaced < nprocs - 1 and displaced + 1 not in [s[0] for s in splits]:
            remaps[displaced] = displaced + 1
        elif displaced > 0:
            remaps[displaced] = displaced - 1

with open(sys.argv[2], 'w') as out:
    out.write("/* Auto-generated from interp_lb_profile.bin — do not edit */\n")
    out.write("#ifndef INTERP_LB_PROFILE_DATA_H\n")
    out.write("#define INTERP_LB_PROFILE_DATA_H\n\n")
    out.write(f"#define INTERP_LB_NPROCS {nprocs}\n")
    out.write(f"#define INTERP_LB_NUM_HEAVY {num_heavy}\n\n")
    out.write(f"static const int interp_lb_heavy_blocks[{num_heavy}] = {{")
    out.write(", ".join(str(h) for h in heavy))
    out.write("};\n\n")
    out.write("/* Split table: {block_id, r_left, r_right} */\n")
    out.write(f"static const int interp_lb_splits[{num_heavy}][3] = {{\n")
    for bid, rl, rr in splits:
        out.write(f"    {{{bid}, {rl}, {rr}}},\n")
    out.write("};\n\n")
    out.write("/* Rank remap for displaced neighbor blocks */\n")
    out.write(f"static const int interp_lb_num_remaps = {len(remaps)};\n")
    out.write(f"static const int interp_lb_remaps[][2] = {{\n")
    for src, dst in sorted(remaps.items()):
        out.write(f"    {{{src}, {dst}}},\n")
    if not remaps:
        out.write("    {-1, -1},\n")
    out.write("};\n\n")
    out.write("#endif /* INTERP_LB_PROFILE_DATA_H */\n")

print(f"Generated {sys.argv[2]}:")
print(f"  {num_heavy} heavy blocks to split: {heavy}")
for bid, rl, rr in splits:
    print(f"    block {bid}: split -> rank {rl} (left), rank {rr} (right)")
for src, dst in sorted(remaps.items()):
    print(f"    block {src}: remap -> rank {dst}")