# ==============================================================================
# COPYRIGHT (C) 2026 KONSTANTIN VLADIMIROVICH GRABKO. ALL RIGHTS RESERVED.
# PATENT PENDING | CMS MANHATTAN JIRACK TECHNOLOGY
# ==============================================================================
#
# grow_to_7b.py
# -------------
# Grow a trained JiRackNative checkpoint into a larger model by:
#   (A) widening the FFN of every existing layer to a target intermediate size
#       (new neurons: input side random, output side zero  -> alive, not dead)
#   (B) inserting NEW transformer layers initialized as IDENTITY:
#         out_proj.weight = 0 , ffn_w2.weight = 0   -> block computes x -> x
#         q/k/v/ffn_w1/ffn_w3 = small random         -> gradient flows, learns
#         norm1/norm2 = 1
#       New layers are interleaved among the trained ones (not appended in one
#       block), which trains more stably.
#
# Why not duplicate layers? Copied layers share identical gradients and stay
# locked together -> no real added capacity. Identity-init new layers do add
# capacity while preserving the function at initialization (no loss spike).
#
# After running:
#   1) set NUM_LAYERS = <--layers> and INTERMEDIATE_SIZE = <--intermediate>
#      in JiRackNative_3b.py
#   2) load the new checkpoint and continue training (normal LR; a short warmup
#      helps the new layers settle)
#
# Example (3B -> ~7B, 36 layers, intermediate 16384):
#   python grow_to_7b.py --in jirack_3b.safetensors --out jirack_7b.safetensors \
#          --layers 36 --intermediate 16384
# ==============================================================================

import argparse
import glob
import math
import os
import torch
from safetensors.torch import load_file, save_file

INIT_STD = 0.02
torch.manual_seed(0)

# suffix -> how to initialize a NEW layer's tensor
#   "rand" small gaussian | "zero" zeros | "one" ones
NEW_LAYER_INIT = {
    "q_proj.weight":  "rand",
    "k_proj.weight":  "rand",
    "v_proj.weight":  "rand",
    "out_proj.weight": "zero",   # -> attention residual = 0 at init
    "ffn_w1.weight":  "rand",
    "ffn_w3.weight":  "rand",
    "ffn_w2.weight":  "zero",    # -> ffn residual = 0 at init
    "norm1.weight":   "one",
    "norm2.weight":   "one",
}


def load_state(path):
    if os.path.isdir(path):
        sd = {}
        for f in sorted(glob.glob(os.path.join(path, "*.safetensors"))):
            sd.update(load_file(f))
        if not sd:
            raise FileNotFoundError(f"No .safetensors in {path}")
        return sd
    return load_file(path)


def detect_layers(sd):
    idx = set()
    for k in sd:
        if k.startswith("blocks."):
            idx.add(int(k.split(".")[1]))
    return sorted(idx)


def widen_ffn(block, target_I):
    """In-place widen one block's FFN tensors to target intermediate size."""
    w1 = block["ffn_w1.weight"]
    cur_I, H = w1.shape
    if target_I <= cur_I:
        return block
    add = target_I - cur_I
    dt, dev = w1.dtype, w1.device
    for name in ("ffn_w1.weight", "ffn_w3.weight"):        # [I, H]: new rows random
        w = block[name]
        new_rows = torch.randn(add, H, dtype=dt, device=dev) * INIT_STD
        block[name] = torch.cat([w, new_rows], dim=0)
    w2 = block["ffn_w2.weight"]                            # [H, I]: new cols zero
    new_cols = torch.zeros(w2.shape[0], add, dtype=dt, device=dev)
    block["ffn_w2.weight"] = torch.cat([w2, new_cols], dim=1)
    return block


def make_identity_block(template):
    """Build a fresh identity-initialized block from {suffix: tensor} template."""
    out = {}
    for suffix, ref in template.items():
        rule = NEW_LAYER_INIT[suffix]
        if rule == "zero":
            out[suffix] = torch.zeros_like(ref)
        elif rule == "one":
            out[suffix] = torch.ones_like(ref)
        else:  # rand
            out[suffix] = torch.randn_like(ref) * INIT_STD
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--in", dest="inp", required=True)
    ap.add_argument("--out", dest="out", required=True)
    ap.add_argument("--layers", type=int, required=True, help="target NUM_LAYERS")
    ap.add_argument("--intermediate", type=int, default=0,
                    help="target INTERMEDIATE_SIZE (0 = keep current)")
    args = ap.parse_args()

    print(f"Loading {args.inp} ...")
    sd = load_state(args.inp)
    old_idx = detect_layers(sd)
    old_L = len(old_idx)
    if args.layers < old_L:
        raise ValueError(f"--layers ({args.layers}) < existing layers ({old_L})")

    # --- pull each existing block into a dict of {suffix: tensor} ---
    def block_of(j):
        pfx = f"blocks.{j}."
        return {k[len(pfx):]: v for k, v in sd.items() if k.startswith(pfx)}

    orig_blocks = [block_of(j) for j in old_idx]

    # --- (A) widen FFN of every existing block ---
    tgt_I = args.intermediate or orig_blocks[0]["ffn_w1.weight"].shape[0]
    for b in orig_blocks:
        widen_ffn(b, tgt_I)
    template = orig_blocks[0]  # correct shapes AFTER widening, for identity blocks

    # --- (B) decide where trained layers sit; fill gaps with identity ---
    L_new = args.layers
    step = L_new / old_L
    positions = [int(math.floor(i * step)) for i in range(old_L)]  # strictly increasing
    pos_to_orig = {p: i for i, p in enumerate(positions)}

    new_sd = {}
    # copy non-layer tensors unchanged (token_emb, ln_f, lm_head, buffers)
    for k, v in sd.items():
        if not k.startswith("blocks."):
            new_sd[k] = v

    n_identity = 0
    for n in range(L_new):
        if n in pos_to_orig:
            block = orig_blocks[pos_to_orig[n]]
        else:
            block = make_identity_block(template)
            n_identity += 1
        for suffix, tensor in block.items():
            new_sd[f"blocks.{n}.{suffix}"] = tensor

    # --- report ---
    total = sum(t.numel() for t in new_sd.values())
    print(f"Existing layers : {old_L}  -> new layers: {L_new} "
          f"({n_identity} identity-initialized)")
    print(f"Intermediate    : {tgt_I}")
    print(f"Trained layers placed at indices: {positions}")
    print(f"Total parameters: {total/1e9:.3f} B")

    print(f"Saving {args.out} ...")
    save_file(new_sd, args.out)
    print("Done. New layers start as identity (no loss spike) and learn from ~step 2.")
    print("Remember to update NUM_LAYERS and INTERMEDIATE_SIZE in the model file.")


if __name__ == "__main__":
    main()