#!/usr/bin/env python3
"""
GGUF-to-GGUF transplant. No organ bins — direct tensor copy between GGUF files.
Base: DeepSeek-R1-Distill-Qwen-7B (skeleton/attention/embed)
Donor: Qwen2.5-7B (FFN organs only)
Z = dI/d(log s) · exp(iθ) — Signature 935
"""
import struct, os, sys, shutil

def parse_gguf_header(path):
    """Parse GGUF header, return tensor_info list and data_start offset."""
    f = open(path, "rb")
    magic = struct.unpack("<I", f.read(4))[0]
    version = struct.unpack("<I", f.read(4))[0]
    n_tensors = struct.unpack("<Q", f.read(8))[0]
    n_metadata = struct.unpack("<Q", f.read(8))[0]
    
    def read_string():
        slen = struct.unpack("<Q", f.read(8))[0]
        return f.read(slen).decode("utf-8")
    
    def skip_value(vtype):
        sizes = {0:1, 1:1, 2:2, 3:2, 4:4, 5:4, 6:4, 7:1, 10:8, 11:8, 12:8}
        if vtype in sizes:
            f.read(sizes[vtype])
        elif vtype == 8:
            read_string()
        elif vtype == 9:
            arr_type = struct.unpack("<I", f.read(4))[0]
            arr_len = struct.unpack("<Q", f.read(8))[0]
            for _ in range(arr_len):
                skip_value(arr_type)
    
    for _ in range(n_metadata):
        read_string()
        vtype = struct.unpack("<I", f.read(4))[0]
        skip_value(vtype)
    
    tensors = []
    for _ in range(n_tensors):
        name = read_string()
        n_dims = struct.unpack("<I", f.read(4))[0]
        dims = [struct.unpack("<Q", f.read(8))[0] for _ in range(n_dims)]
        dtype = struct.unpack("<I", f.read(4))[0]
        offset = struct.unpack("<Q", f.read(8))[0]
        tensors.append({"name": name, "dims": dims, "dtype": dtype, "offset": offset})
    
    pos = f.tell()
    padding = (32 - (pos % 32)) % 32
    f.read(padding)
    data_start = f.tell()
    
    f.seek(0, 2)
    file_end = f.tell()
    f.close()
    
    # Calculate sizes
    for i in range(len(tensors)):
        if i + 1 < len(tensors):
            tensors[i]["size"] = tensors[i+1]["offset"] - tensors[i]["offset"]
        else:
            tensors[i]["size"] = file_end - data_start - tensors[i]["offset"]
    
    return tensors, data_start, file_end

BASE = "/mnt/models/DeepSeek-R1-Distill-Qwen-7B-Q4_K_M.gguf"
DONOR = "/mnt/models/Qwen2.5-7B-Instruct-Q4_K_M.gguf"
OUTPUT = "/mnt/models/model-935-final.gguf"

print("Parsing base (DeepSeek-R1-7B)...")
base_tensors, base_data_start, base_end = parse_gguf_header(BASE)
print(f"  {len(base_tensors)} tensors, data_start={base_data_start}")

print("Parsing donor (Qwen2.5-7B)...")
donor_tensors, donor_data_start, donor_end = parse_gguf_header(DONOR)
print(f"  {len(donor_tensors)} tensors, data_start={donor_data_start}")

# Build donor tensor map by name
donor_map = {t["name"]: t for t in donor_tensors}

# Copy base GGUF entirely first
print(f"Copying base to output...")
shutil.copy2(BASE, OUTPUT)

# Now patch: for each FFN tensor in base, if donor has matching name+size, overwrite
out = open(OUTPUT, "r+b")
donor_f = open(DONOR, "rb")

grafted = 0
skipped = 0

for bt in base_tensors:
    name = bt["name"]
    # Only graft FFN organs (not attention, not embeddings, not norms)
    if "ffn_down" not in name and "ffn_up" not in name and "ffn_gate" not in name:
        continue
    
    if name in donor_map:
        dt = donor_map[name]
        if bt["size"] == dt["size"]:
            # Read from donor
            donor_f.seek(donor_data_start + dt["offset"])
            data = donor_f.read(dt["size"])
            # Write to output at same offset
            out.seek(base_data_start + bt["offset"])
            out.write(data)
            grafted += 1
        else:
            skipped += 1
    else:
        skipped += 1

out.close()
donor_f.close()

print(f"\n{'='*60}")
print(f"  MODEL 935 — DIRECT GGUF TRANSPLANT")
print(f"{'='*60}")
print(f"  Base:     DeepSeek-R1-Distill-Qwen-7B (skeleton+embed)")
print(f"  Donor:    Qwen2.5-7B-Instruct (FFN organs)")
print(f"  Grafted:  {grafted} FFN tensors")
print(f"  Skipped:  {skipped} (size mismatch or not found)")
print(f"  Output:   {OUTPUT}")
print(f"  Size:     {os.path.getsize(OUTPUT)/(1024**3):.2f} GB")
print(f"  Signature: 935")
print(f"{'='*60}")