#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""patch_infix2.py — S3/S5 in 维修正: gguf-py t.shape=(in,out) 逻辑形, in=shape[0] (numpy data=(out,in))"""
CH_OLD = '''            if hn.ndim == 1 and hn.shape[0] == shape[1]:
                H = torch.diag(torch.from_numpy(hn).to(dev) + 1e-8)
            elif hn.shape == (shape[1], shape[1]):
                H = torch.from_numpy(hn).to(dev)
            if H is not None and shape[1] > 4200:
                Hb = torch.zeros_like(H)
                for i in range(0, shape[1], 1024):
                    Hb[i:i + 1024, i:i + 1024] = H[i:i + 1024, i:i + 1024]
                H = Hb
            if H is not None:
                H = H / max(1e-9, float(torch.trace(H))) + 1e-6 * torch.eye(shape[1], device=dev)'''
CH_NEW = '''            if hn.ndim == 1 and hn.shape[0] == shape[0]:
                H = torch.diag(torch.from_numpy(hn).to(dev) + 1e-8)
            elif hn.shape == (shape[0], shape[0]):
                H = torch.from_numpy(hn).to(dev)
            if H is not None and shape[0] > 4200:
                Hb = torch.zeros_like(H)
                for i in range(0, shape[0], 1024):
                    Hb[i:i + 1024, i:i + 1024] = H[i:i + 1024, i:i + 1024]
                H = Hb
            if H is not None:
                H = H / max(1e-9, float(torch.trace(H))) + 1e-6 * torch.eye(shape[0], device=dev)'''
CH_EYE_OLD = '''        if H is None:
            H = torch.eye(shape[1], device=dev)'''
CH_EYE_NEW = '''        if H is None:
            H = torch.eye(shape[0], device=dev)'''
RF_LD_OLD = '''        H = load_H(imx[hk], shape[1], dev)'''
RF_LD_NEW = '''        H = load_H(imx[hk], shape[0], dev)'''
RF_EYE_OLD = '''            H = torch.eye(shape[1], device=dev)
            P("S5 %s: H missing -> identity" % base)'''
RF_EYE_NEW = '''            H = torch.eye(shape[0], device=dev)
            P("S5 %s: H missing -> identity" % base)'''

def main():
    for path, olds, news in (
        ("/tmp/k8b/r13_chain.py", [CH_OLD, CH_EYE_OLD], [CH_NEW, CH_EYE_NEW]),
        ("/tmp/k8b/r13_refine.py", [RF_LD_OLD, RF_EYE_OLD], [RF_LD_NEW, RF_EYE_NEW]),
    ):
        src = open(path).read()
        for o, n in zip(olds, news):
            cnt = src.count(o)
            assert cnt == 1, "anchor not unique in %s: count=%d" % (path, cnt)
            src = src.replace(o, n, 1)
        compile(src, path, "exec")
        open(path, "w").write(src)
        print("PATCHED", path)
    print("PATCH_OK infix2 shape[0] in-dim fix")
main()