from __future__ import annotations
import argparse
import math

def round_to(x: float, multiple: int) -> int:
    return int(multiple * round(x / multiple))

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--stem_layers", type=int, required=True)
    ap.add_argument("--hemi_layers", type=int, required=True)
    ap.add_argument("--dff_dense", type=int, required=True)
    ap.add_argument("--dff_expert", type=int, required=True)
    ap.add_argument("--topk", type=int, default=1)
    ap.add_argument("--multiple", type=int, default=256, help="round dense_dff to this multiple")
    args = ap.parse_args()

    Ls = args.stem_layers
    Lh = args.hemi_layers

    # Matches attention layer count by setting dense_layers = Ls + 2*Lh
    dense_layers = Ls + 2 * Lh

    # In our current StructuredBiHMoE implementation, MoE occurs on odd i => floor(Lh/2) MoE layers per hemisphere
    moe_per_hemi = Lh // 2
    dense_per_hemi = Lh - moe_per_hemi  # ceil(Lh/2)

    # Equivalent FFN "d_ff units" across the whole structured forward (both hemispheres + stem)
    ffn_units = (Ls * args.dff_dense) + (2 * dense_per_hemi * args.dff_dense) + (2 * moe_per_hemi * args.topk * args.dff_expert)

    dense_dff = ffn_units / dense_layers
    dense_dff_rounded = max(args.multiple, round_to(dense_dff, args.multiple))

    print("=== D_a (compute-ish match) suggestion ===")
    print(f"dense_layers = {dense_layers}   (matches attention depth: stem + 2*hemi)")
    print(f"dense_dff    ≈ {dense_dff:.2f}  -> rounded to {dense_dff_rounded} (multiple={args.multiple})")
    print()
    print("Details:")
    print(f"  stem_layers={Ls}")
    print(f"  hemi_layers={Lh} -> dense_per_hemi={dense_per_hemi} moe_per_hemi={moe_per_hemi}")
    print(f"  dff_dense={args.dff_dense} dff_expert={args.dff_expert} topk={args.topk}")
    print(f"  ffn_units(total)={ffn_units} over dense_layers={dense_layers}")

if __name__ == "__main__":
    main()