"""Depth recurrence without new parameters: rebuild a GGUF whose layer sequence revisits a block of layers.
e.g. --order 0-13,10-13,14-23 runs layers 10..13 twice (the second pass reuses the same weights; it gets its own
KV/conv state because llama.cpp sees it as a separate layer). Per-layer metadata arrays are remapped.
usage: loop_layers.py in.gguf out.gguf --order 0-13,10-13,14-23"""
import argparse, re, numpy as np
import gguf
ap = argparse.ArgumentParser()
ap.add_argument("inp"); ap.add_argument("out"); ap.add_argument("--order", default="0-23")
ap.add_argument("--window", type=int, default=None, help="override the sliding-attention window")
ap.add_argument("--global-swa", action="store_true", help="turn every sliding-attention layer into full attention")
a = ap.parse_args()
order = []
for part in a.order.split(","):
    lo, hi = map(int, part.split("-")) if "-" in part else (int(part), int(part))
    order += list(range(lo, hi + 1))
r = gguf.GGUFReader(a.inp)
arch = r.fields["general.architecture"].contents()
w = gguf.GGUFWriter(a.out, arch)
per_layer = {f"{arch}.attention.head_count_kv", f"{arch}.attention.sliding_window_pattern"}
for f in r.fields.values():
    if f.name == "general.architecture" or f.name.startswith("GGUF."):
        continue
    vt = f.types[0]
    st = f.types[-1] if vt == gguf.GGUFValueType.ARRAY else None
    v = f.contents()
    if f.name == f"{arch}.block_count":
        v = len(order)
    elif f.name in per_layer:
        v = [v[i] for i in order]
        if a.global_swa and f.name.endswith("sliding_window_pattern"):
            v = [False] * len(v)
    elif f.name == f"{arch}.attention.sliding_window" and a.window:
        v = a.window
    w.add_key_value(f.name, v, vt, sub_type=st)
w.add_string("general.layer_order", a.order)
src = {t.name: t for t in r.tensors}
out = []
for t in r.tensors:
    if not t.name.startswith("blk."):
        out.append((t.name, t))
for new_i, old_i in enumerate(order):
    for name, t in src.items():
        m = re.match(rf"blk\.{old_i}\.(.+)", name)
        if m:
            out.append((f"blk.{new_i}.{m.group(1)}", t))
for name, t in out:
    w.add_tensor_info(name, t.data.shape, t.data.dtype, t.data.nbytes, t.tensor_type)
w.write_header_to_file(); w.write_kv_data_to_file(); w.write_ti_data_to_file()
for name, t in out:
    w.write_tensor_data(t.data, tensor_endianess=r.endianess)
w.close()
print("wrote", a.out, len(order), "layers:", order)