#!/usr/bin/env python3
"""Adapt an APEX generate_config.sh output for Granite-4.0-H-Tiny (granitemoehybrid).

Hybrid: 36 Mamba-2 layers (ssm_in/out/conv1d) + 4 attention layers (plain MHA) +
MoE (64 routed + shared expert) on every layer. The stock generator emits
ffn_*_exps/shexp, attn_q/k/v/output, and ssm_out — but MISSES the Mamba-2
input/conv projections. We add ssm_in + ssm_conv1d at the mixer precision, and
drop tensors Granite doesn't have. 1-D SSM state (ssm_a/d/dt/norm) stays F32.

--ssm-type pins the Mamba-2 mixer projections (ssm_in/out/conv1d) higher (e.g.
Q8_0) to test whether protecting the recurrence helps (the hand-roll variant).
"""
import argparse, re
DROP = {"attn_gate","attn_qkv","shortconv.in_proj","shortconv.out_proj","ssm_alpha","ssm_beta"}
ADD  = ["ssm_in","ssm_conv1d"]   # ssm_out already emitted by the generator
SSM_ALL = ["ssm_in","ssm_out","ssm_conv1d"]

def main():
    ap=argparse.ArgumentParser(); ap.add_argument("infile"); ap.add_argument("outfile")
    ap.add_argument("--ssm-type",default=None,help="override Mamba-2 mixer tensors (e.g. Q8_0)")
    a=ap.parse_args()
    out=[]
    for l in (x.strip() for x in open(a.infile) if x.strip()):
        m=re.match(r"blk\.\d+\.([\w.]+)=",l); tail=m.group(1) if m else ""
        if tail in DROP: continue
        # optionally bump ssm_out (already present) to the mixer override type
        mo_out=re.match(r"(blk\.\d+\.ssm_out)=(\S+)",l)
        if mo_out and a.ssm_type:
            out.append(f"{mo_out.group(1)}={a.ssm_type}"); continue
        out.append(l)
        mo=re.match(r"blk\.(\d+)\.attn_output=(\S+)",l)
        if mo:
            n,t=mo.group(1),mo.group(2)
            st=a.ssm_type or t
            for name in ADD: out.append(f"blk.{n}.{name}={st}")
    open(a.outfile,"w").write("\n".join(out)+"\n")
    print(f"wrote {a.outfile} ({len(out)} lines)")
if __name__=="__main__": main()