# ============================================================================
# GGUF Pipeline: Q4_K_M → float32 merge → Q4_0
#
# Disk-efficient sequential download:
#   Peak disk ≈ 3 × model_size (not 4×)
#
# Flow:
#   0) Download A (Q4_K_M), convert → Q4_0 via llama-quantize, delete A
#   1) Download B, C (Q4_K_M)
#   2) Cross-quant merge in float32: A(Q4_0) + (B−C)(Q4_K_M) → write Q4_0
#   3) Delete B, C
#   4) (optional) llama-quantize Q4_0 → Q4_K_M (single clean quantization disabled by default)
#
# Only ONE quantization step touch the final data:
#   - merge result → Q4_0 (intermediate)
# ============================================================================

import subprocess, sys, os, gc, time, shutil
from math import gcd
import numpy as np

WORK_DIR = "/content/gguf_workspace"
os.makedirs(WORK_DIR, exist_ok=True)

def install():
    for pkg in ["huggingface_hub[hf_xfer]", "gguf"]:
        subprocess.check_call(
            [sys.executable, "-m", "pip", "install", "-q", pkg],
            stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
        )
    print("[OK] Dependencies installed.")

install()

from huggingface_hub import hf_hub_download
from gguf import GGUFReader
from gguf.constants import GGMLQuantizationType, GGML_QUANT_SIZES
from gguf import quants as gguf_quants

# ============================================================================
# PATCH IQ block sizes
# ============================================================================

_IQ = {
    "IQ1_S":(256,50),"IQ1_M":(256,56),"IQ2_XXS":(256,66),"IQ2_XS":(256,74),
    "IQ2_S":(256,82),"IQ3_XXS":(256,98),"IQ3_S":(256,110),
    "IQ4_NL":(32,18),"IQ4_XS":(256,136),
}
for _n, _g in _IQ.items():
    _qt = getattr(GGMLQuantizationType, _n, None)
    if _qt is not None and _qt not in GGML_QUANT_SIZES:
        GGML_QUANT_SIZES[_qt] = _g

# ============================================================================
# PROBE quantize support
# ============================================================================

_QUANTIZE_SUPPORTED = set()

def _probe_quantize():
    for qt in GGMLQuantizationType:
        if qt in (GGMLQuantizationType.F32, GGMLQuantizationType.F16,
                  GGMLQuantizationType.BF16):
            _QUANTIZE_SUPPORTED.add(qt)
        elif qt in GGML_QUANT_SIZES:
            be, _ = GGML_QUANT_SIZES[qt]
            try:
                r = gguf_quants.quantize(np.zeros(be, np.float32), qt)
                if r is not None and len(r.ravel().tobytes()) > 0:
                    _QUANTIZE_SUPPORTED.add(qt)
            except Exception:
                pass

_probe_quantize()
print(f"[INFO] quantize() supported: {sorted([q.name for q in _QUANTIZE_SUPPORTED])}")

# ============================================================================
# CONFIGURATION
# ============================================================================

MODEL_A = {
    "repo": "HauhauCS/Qwen3.5-35B-A3B-Uncensored-HauhauCS-Aggressive",
    "file": "Qwen3.5-35B-A3B-Uncensored-HauhauCS-Aggressive-Q4_K_M.gguf",
}
MODEL_B = {
    "repo": "samuelcardillo/Qwen3.5-35B-A3B-Claude-4.6-Opus-Reasoning-Distilled-GGUF",
    "file": "Qwen3.5-35B-A3B-Claude-4.6-Opus-Reasoning-Distilled-Q4_K_M.gguf",
}
MODEL_C = {
    "repo": "unsloth/Qwen3.5-35B-A3B-GGUF",
    "file": "Qwen3.5-35B-A3B-Q4_K_M.gguf",
}

MERGE_WEIGHT = 1.0
MERGED_NAME  = "Qwen3.5-35B-Q4_0-merged.gguf"
MERGED_PATH  = os.path.join(WORK_DIR, MERGED_NAME)

MERGE_CHUNK_ELEMS = 2 * 1024 * 1024

# --- Final output (set None to "Q4_K_M" if you want to quantize final model) ---
FINAL_QTYPE = None

# llama.cpp Vulkan build
LLAMA_URL = "https://github.com/ggml-org/llama.cpp/releases/download/b8468/llama-b8468-bin-ubuntu-vulkan-x64.tar.gz"
LLAMA_TAR = "llama-b8468-bin-ubuntu-vulkan-x64.tar.gz"

# ============================================================================
# CONSTANTS
# ============================================================================

FLOAT_QTYPES = {
    GGMLQuantizationType.F32,
    GGMLQuantizationType.F16,
    GGMLQuantizationType.BF16,
}

FINETUNE_SUFFIXES = (
    "attn_q.weight","attn_k.weight","attn_v.weight","attn_output.weight",
    "ssm_in.weight","ssm_out.weight","ssm_a.weight",
    "ssm_d.weight","ssm_dt.weight",
    "ffn_gate.weight","ffn_up.weight","ffn_down.weight",
    "ffn_gate_exps.weight","ffn_up_exps.weight","ffn_down_exps.weight",
    "ffn_gate_inp.weight","ffn_gate_shexp.weight",
    "ffn_up_shexp.weight","ffn_down_shexp.weight",
    "attn_gate.weight",
)

PROTECTED_SUBSTRINGS = (
    "token_embd","output.weight",
    "_norm",".bias",
    "rope_freqs","pos_embd",
    "v.blk.","v.patch_embd","v.position_embd",
    "v.post_ln","v.pre_ln","v.output_norm","v.proj",
    "mm.","mtp.",
)

# ============================================================================
# QUANT GEOMETRY
# ============================================================================

def quant_block_elements(qt):
    if qt in FLOAT_QTYPES: return 1
    if qt in GGML_QUANT_SIZES: return GGML_QUANT_SIZES[qt][0]
    return 1

def quant_block_bytes(qt):
    if qt == GGMLQuantizationType.F32: return 4
    if qt == GGMLQuantizationType.F16: return 2
    if qt == GGMLQuantizationType.BF16: return 2
    if qt in GGML_QUANT_SIZES: return GGML_QUANT_SIZES[qt][1]
    return 1

def elems_to_raw(n_elems, qt):
    be = quant_block_elements(qt)
    bb = quant_block_bytes(qt)
    return (n_elems // be) * bb

# ============================================================================
# UTILITY
# ============================================================================

def disk_free_gb(p=WORK_DIR):
    st = os.statvfs(p); return (st.f_bavail * st.f_frsize) / 1073741824

def file_size_gb(p):
    return os.path.getsize(p) / 1073741824

def fmt_bytes(n):
    for u in ("B","KB","MB","GB"):
        if n < 1024: return f"{n:.1f} {u}"
        n /= 1024
    return f"{n:.1f} TB"

def ram_used_gb():
    try:
        with open("/proc/self/status") as f:
            for line in f:
                if line.startswith("VmRSS:"):
                    return int(line.split()[1]) / 1048576
    except Exception:
        pass
    return -1

def safe_remove(path, protect=None):
    try:
        a = os.path.abspath(path)
        if protect and a == os.path.abspath(protect): return
        if os.path.isfile(a):
            sz = file_size_gb(a)
            os.remove(a)
            print(f"[CLEANUP] Removed {os.path.basename(a)} ({sz:.2f} GB)")
    except OSError:
        pass

def flush_and_close(f, label="file"):
    f.flush(); os.fsync(f.fileno()); f.close()
    print(f"[FLUSH] {label}: flush + fsync done.")

def align_down(v, a):
    if a <= 0: return v
    return (v // a) * a

def check_gpu():
    print("\n[GPU] Checking GPU...")
    try:
        result = subprocess.run(
            ["nvidia-smi", "--query-gpu=name,memory.total,driver_version",
             "--format=csv,noheader"],
            capture_output=True, text=True, timeout=10,
        )
        if result.returncode == 0:
            for line in result.stdout.strip().split('\n'):
                print(f"[GPU] {line.strip()}")
            return True
    except Exception:
        pass
    print("[GPU] nvidia-smi not available.")
    return False

# ============================================================================
# GGUF STRUCTURE
# ============================================================================

def parse_gguf_structure(path):
    reader = GGUFReader(path)
    tensors = {}
    for t in reader.tensors:
        tensors[t.name] = {
            "offset":     int(t.data_offset),
            "size":       int(t.n_bytes),
            "qtype":      t.tensor_type,
            "qtype_name": t.tensor_type.name,
            "shape":      tuple(t.shape),
            "n_elements": int(t.n_elements),
        }
    n = len(tensors)
    del reader; gc.collect()
    return tensors, n

# ============================================================================
# RAW FILE I/O
# ============================================================================

def read_chunk(f, tinfo, offset, size):
    actual = min(size, tinfo["size"] - offset)
    if actual <= 0: return b""
    f.seek(tinfo["offset"] + offset)
    return f.read(actual)

def write_chunk(f, tinfo, offset, data):
    f.seek(tinfo["offset"] + offset)
    f.write(data)

# ============================================================================
# DEQUANTIZE / QUANTIZE
# ============================================================================

def bytes_to_f32(raw, qt):
    if qt == GGMLQuantizationType.F32:
        return np.frombuffer(raw, dtype=np.float32).copy()
    if qt == GGMLQuantizationType.F16:
        return np.frombuffer(raw, dtype=np.float16).astype(np.float32)
    if qt == GGMLQuantizationType.BF16:
        u16 = np.frombuffer(raw, dtype=np.uint16).copy()
        return (u16.astype(np.uint32) << 16).view(np.float32).copy()
    u8 = np.frombuffer(raw, dtype=np.uint8).copy()
    return gguf_quants.dequantize(u8, qt).ravel().astype(np.float32)

def f32_to_bytes_aligned(f32, qt, raw_chunk_size):
    f32 = np.nan_to_num(f32, nan=0.0, posinf=0.0, neginf=0.0)
    if qt == GGMLQuantizationType.F32:
        r = f32.astype(np.float32).tobytes()
    elif qt == GGMLQuantizationType.F16:
        r = f32.astype(np.float16).tobytes()
    elif qt == GGMLQuantizationType.BF16:
        u32 = f32.astype(np.float32).view(np.uint32)
        r = (u32 >> 16).astype(np.uint16).tobytes()
    elif qt in _QUANTIZE_SUPPORTED:
        be = quant_block_elements(qt)
        n = align_down(len(f32), be)
        if n == 0: return None
        try:
            r = gguf_quants.quantize(f32[:n], qt).ravel().tobytes()
        except Exception:
            return None
    else:
        return None
    return r if len(r) == raw_chunk_size else None

# ============================================================================
# DOWNLOAD
# ============================================================================

def download(spec, label):
    print(f"[DL] {label}: {spec['repo']} / {spec['file']}")
    path = hf_hub_download(repo_id=spec["repo"], filename=spec["file"],
                           local_dir=WORK_DIR)
    print(f"[DL] {label}: {file_size_gb(path):.2f} GB -> {path}")
    return path

# ============================================================================
# LLAMA-QUANTIZE
# ============================================================================

LLAMA_BIN_DIR = os.path.join(WORK_DIR, "llama-bin")

def get_llama_quantize():
    for root, dirs, files in os.walk(LLAMA_BIN_DIR):
        for f in files:
            if f == "llama-quantize":
                p = os.path.join(root, f)
                if os.access(p, os.X_OK):
                    print(f"[BIN] Already available: {p}")
                    return p

    print(f"[BIN] Downloading Vulkan build...")
    t0 = time.time()
    tar_path = os.path.join(WORK_DIR, LLAMA_TAR)

    downloaded = False
    for cmd in [
        ["wget", "-q", "--show-progress", "-O", tar_path, LLAMA_URL],
        ["curl", "-sL", "-o", tar_path, LLAMA_URL],
    ]:
        try:
            subprocess.check_call(cmd, timeout=300)
            if os.path.isfile(tar_path) and os.path.getsize(tar_path) > 10240:
                downloaded = True; break
        except Exception:
            continue

    if not downloaded:
        print("[BIN] ERROR: Download failed."); return None

    os.makedirs(LLAMA_BIN_DIR, exist_ok=True)
    for strip in ["--strip-components=1", None]:
        try:
            cmd = ["tar", "xzf", tar_path, "-C", LLAMA_BIN_DIR]
            if strip: cmd.append(strip)
            subprocess.check_call(cmd); break
        except subprocess.CalledProcessError:
            continue

    if os.path.isfile(tar_path): os.remove(tar_path)

    for root, dirs, files in os.walk(LLAMA_BIN_DIR):
        for f in files:
            if f == "llama-quantize":
                binary = os.path.join(root, f)
                os.chmod(binary, 0o755)
                print(f"[BIN] Ready in {time.time()-t0:.1f}s")
                return binary

    print("[BIN] ERROR: llama-quantize not found."); return None

def run_llama_quantize(input_path, output_path, quant_type):
    binary = get_llama_quantize()
    if binary is None: return None

    bin_dir = os.path.dirname(binary)
    env = os.environ.copy()
    ld = env.get("LD_LIBRARY_PATH", "")
    env["LD_LIBRARY_PATH"] = f"{bin_dir}:{ld}" if ld else bin_dir

    cmd = [binary, "--allow-requantize", input_path, output_path, quant_type]
    print(f"[QUANT] {' '.join(cmd)}")
    t0 = time.time()

    try:
        proc = subprocess.Popen(cmd, stdout=subprocess.PIPE,
                                stderr=subprocess.STDOUT, text=True, env=env, bufsize=1)
        for line in proc.stdout:
            line = line.rstrip()
            if line: print(f"  {line}")
        proc.wait()
        if proc.returncode != 0:
            print(f"[QUANT] ERROR: exit {proc.returncode}"); return None
    except Exception as e:
        print(f"[QUANT] ERROR: {e}"); return None

    if not os.path.isfile(output_path):
        print("[QUANT] ERROR: Output not created."); return None

    elapsed = time.time() - t0
    print(f"[QUANT] Done in {elapsed:.0f}s ({elapsed/60:.1f} min)")
    return output_path

# ============================================================================
# STAGE 1: CROSS-QUANT MERGE (Q4_K_M sources → float32 → Q4_0 output)
# ============================================================================

def is_finetune_weight(name):
    return any(name.endswith(s) for s in FINETUNE_SUFFIXES)

def is_protected(name):
    return any(p in name for p in PROTECTED_SUBSTRINGS)

def classify_tensor_merge(name, ti_out, ti_b, ti_c):
    if is_protected(name): return "protect", "structural/norm/embd"
    if not is_finetune_weight(name): return "protect", "not fine-tunable"
    if ti_b is None or ti_c is None: return "skip", "missing in B/C"
    if ti_out["n_elements"] != ti_b["n_elements"]:
        return "skip", f"elem mismatch {ti_out['n_elements']}��{ti_b['n_elements']}"
    if ti_b["n_elements"] != ti_c["n_elements"]:
        return "skip", f"elem mismatch B≠C"
    return "merge", "fine-tunable"

def merge_tensor_cross(f_out, ti_out, f_b, ti_b, f_c, ti_c, weight,
                        qt_out, qt_b, qt_c):
    total_elems = ti_out["n_elements"]

    be_out = quant_block_elements(qt_out)
    bb_out = quant_block_bytes(qt_out)
    be_b = quant_block_elements(qt_b)
    bb_b = quant_block_bytes(qt_b)
    be_c = quant_block_elements(qt_c)
    bb_c = quant_block_bytes(qt_c)

    lcm_bc = be_b * be_c // gcd(be_b, be_c)
    lcm_all = lcm_bc * be_out // gcd(lcm_bc, be_out)

    chunk_elems = align_down(MERGE_CHUNK_ELEMS, lcm_all)
    if chunk_elems <= 0: chunk_elems = lcm_all

    changed = 0; failed = 0; elem_off = 0

    while elem_off < total_elems:
        n = min(chunk_elems, total_elems - elem_off)
        n = align_down(n, lcm_all)
        if n <= 0: break

        out_boff = elems_to_raw(elem_off, qt_out)
        out_bsz  = elems_to_raw(n, qt_out)
        b_boff   = elems_to_raw(elem_off, qt_b)
        b_bsz    = elems_to_raw(n, qt_b)
        c_boff   = elems_to_raw(elem_off, qt_c)
        c_bsz    = elems_to_raw(n, qt_c)

        ra = read_chunk(f_out, ti_out, out_boff, out_bsz)
        rb = read_chunk(f_b, ti_b, b_boff, b_bsz)
        rc = read_chunk(f_c, ti_c, c_boff, c_bsz)

        if len(ra) != out_bsz or len(rb) != b_bsz or len(rc) != c_bsz:
            failed += 1; elem_off += n; continue

        try:
            fa = bytes_to_f32(ra, qt_out)
            fb = bytes_to_f32(rb, qt_b)
            fc = bytes_to_f32(rc, qt_c)
        except Exception:
            failed += 1; del ra, rb, rc; elem_off += n; continue

        k = min(len(fa), len(fb), len(fc))
        k = align_down(k, be_out)
        if k == 0:
            failed += 1; del ra, rb, rc, fa, fb, fc; elem_off += n; continue

        # *** MERGE IN FLOAT32 — NO PRECISION LOSS ***
        merged = fa[:k] + (fb[:k] - fc[:k]) * np.float32(weight)
        del fa, fb, fc

        new_raw = f32_to_bytes_aligned(merged, qt_out, out_bsz)
        del merged

        if new_raw is not None:
            old_arr = np.frombuffer(ra, dtype=np.uint8)
            new_arr = np.frombuffer(new_raw, dtype=np.uint8)
            changed += int(np.sum(old_arr != new_arr))
            write_chunk(f_out, ti_out, out_boff, new_raw)
            del old_arr, new_arr
        else:
            failed += 1

        del ra, rb, rc
        if new_raw is not None: del new_raw
        elem_off += n

    return changed, failed

def run_merge(q4_0_path, path_b, path_c, weight):
    t0 = time.time()

    print("[MERGE] Parsing structures...")
    info_out, n_out = parse_gguf_structure(q4_0_path)
    info_b, n_b = parse_gguf_structure(path_b)
    info_c, n_c = parse_gguf_structure(path_c)

    print(f"[MERGE] Output (Q4_0):  {n_out} tensors ({file_size_gb(q4_0_path):.2f} GB)")
    print(f"[MERGE] B (Q4_K_M):     {n_b} tensors ({file_size_gb(path_b):.2f} GB)")
    print(f"[MERGE] C (Q4_K_M):     {n_c} tensors ({file_size_gb(path_c):.2f} GB)")

    plan = {}
    for name in sorted(info_out.keys()):
        ti_out = info_out[name]
        ti_b = info_b.get(name)
        ti_c = info_c.get(name)
        action, reason = classify_tensor_merge(name, ti_out, ti_b, ti_c)
        plan[name] = (action, reason)

    n_merge = sum(1 for v in plan.values() if v[0] == "merge")
    n_protect = sum(1 for v in plan.values() if v[0] == "protect")
    n_skip = sum(1 for v in plan.values() if v[0] == "skip")

    print(f"\n[MERGE] MERGE={n_merge}  PROTECT={n_protect}  SKIP={n_skip}")

    if n_skip > 0:
        for name, (action, reason) in sorted(plan.items()):
            if action == "skip":
                print(f"  SKIP: {name:<48} {reason}")

    if n_merge == 0:
        print("[MERGE] Nothing to merge."); return {}

    f_out = open(q4_0_path, "r+b")
    f_b   = open(path_b, "rb")
    f_c   = open(path_c, "rb")

    print(f"\n[MERGE] output = A + (B - C) * {weight}")
    print(f"[MERGE] All arithmetic in float32")
    print(f"{'#':>4} {'Tensor':<50} {'Out':<6} {'Src':<8} "
          f"{'Size':>9}  {'Action':<10} Detail")
    print("-" * 115)

    stats = {"merged":0,"protected":0,"skipped":0,
             "total_changed":0,"identical":0,"failed":0}

    for idx, name in enumerate(sorted(info_out.keys())):
        ti_out = info_out[name]
        qt_out_name = ti_out["qtype_name"]
        sz = ti_out["size"]
        action, reason = plan[name]

        if action != "merge":
            stats["protected" if action == "protect" else "skipped"] += 1
            print(f"{idx+1:>4} {name:<50} {qt_out_name:<6} {'--':<8} "
                  f"{fmt_bytes(sz):>9}  {action.upper():<10} {reason}")
            continue

        ti_b = info_b[name]; ti_c = info_c[name]
        qt_out = ti_out["qtype"]; qt_b = ti_b["qtype"]; qt_c = ti_c["qtype"]

        changed, failed = merge_tensor_cross(
            f_out, ti_out, f_b, ti_b, f_c, ti_c, weight, qt_out, qt_b, qt_c)

        stats["merged"] += 1; stats["total_changed"] += changed; stats["failed"] += failed

        src_qt = ti_b["qtype_name"]
        if changed > 0:
            f_s = f" ({failed} fail)" if failed else ""
            print(f"{idx+1:>4} {name:<50} {qt_out_name:<6} {src_qt:<8} "
                  f"{fmt_bytes(sz):>9}  {'MERGED':<10} {changed:,}B{f_s}")
        else:
            stats["identical"] += 1
            print(f"{idx+1:>4} {name:<50} {qt_out_name:<6} {src_qt:<8} "
                  f"{fmt_bytes(sz):>9}  {'MERGED':<10} (identical)")

        if (idx + 1) % 50 == 0: gc.collect()

    flush_and_close(f_out, "output"); f_b.close(); f_c.close()
    elapsed = time.time() - t0

    print(f"\n{'='*75}")
    print(f"MERGE SUMMARY: {stats['merged']} merged, {stats['protected']} protected, "
          f"{stats['identical']} identical, {elapsed:.1f}s")
    if stats['failed']: print(f"  Failed chunks: {stats['failed']}")
    print(f"Output: {file_size_gb(q4_0_path):.2f} GB")
    print(f"{'='*75}")

    return stats

# ============================================================================
# VERIFY
# ============================================================================

def verify_output(output_path):
    print(f"\n[VERIFY] {os.path.basename(output_path)}...")
    try:
        info_out, n_out = parse_gguf_structure(output_path)
        print(f"[VERIFY] {n_out} tensors parsed.")
    except Exception as e:
        print(f"[VERIFY] ERROR: {e}"); return False

    f = open(output_path, "rb")
    ok = warn = 0
    for name in sorted(info_out.keys())[:30]:
        ti = info_out[name]; cs = min(4096, ti["size"])
        raw = read_chunk(f, ti, 0, cs)
        arr = np.frombuffer(raw, dtype=np.uint8)
        if arr.sum() == 0 and cs > 64:
            warn += 1; print(f"  WARNING: '{name}' all zeros!")
        else: ok += 1
        del raw, arr
    f.close()
    print(f"[VERIFY] {ok} OK" + (f", {warn} WARNINGS" if warn else ""))
    return warn == 0

# ============================================================================
# MAIN PIPELINE
# ============================================================================

def main():
    print("=" * 75)
    print("GGUF PIPELINE: Q4_K_M → float32 merge → Q4_0 → Q4_K_M")
    print()
    print(f"  Model A:      {MODEL_A['repo']} / {MODEL_A['file']}")
    print(f"  Model B:      {MODEL_B['repo']} / {MODEL_B['file']}")
    print(f"  Model C:      {MODEL_C['repo']} / {MODEL_C['file']}")
    print(f"  Merge weight: {MERGE_WEIGHT}")
    print(f"  Final quant:  {FINAL_QTYPE if FINAL_QTYPE else 'None (keep Q4_0)'}")
    print(f"  Disk free:    {disk_free_gb():.1f} GB")
    print(f"  RAM used:     {ram_used_gb():.2f} GB")
    print()
    print("  Data flow (minimal loss):")
    print("    Q4_K_M(A) ──llama-quantize──→ Q4_0(A')  [1 requant]")
    print("    Q4_K_M(B,C) ──dequant──→ float32 ──merge──→ float32  [lossless]")
    print("    float32 ──quantize──→ Q4_0(merged)  [1 quantization]")
    print("=" * 75)

    check_gpu()

    # ================================================================
    # STAGE 0: Download A, convert to Q4_0, free disk
    # ================================================================
    print("\n" + "=" * 75)
    print("STAGE 0: Download A (Q4_K_M) → convert Q4_0")
    print("=" * 75)

    path_a = download(MODEL_A, "A (base)")
    a_gb = file_size_gb(path_a)
    print(f"[FS] A: {a_gb:.2f} GB  |  Disk free: {disk_free_gb():.1f} GB")

    print(f"\n[CONV] Converting A: Q4_K_M → Q4_0 via llama-quantize...")
    result = run_llama_quantize(path_a, MERGED_PATH, "Q4_0")
    if result is None:
        sys.exit("[ERROR] Q4_0 conversion failed.")

    merged_gb = file_size_gb(MERGED_PATH)
    print(f"[CONV] {a_gb:.2f} GB → {merged_gb:.2f} GB")

    safe_remove(path_a, protect=MERGED_PATH)
    print(f"[FS] Disk free: {disk_free_gb():.1f} GB")

    # ================================================================
    # STAGE 1: Download B, C → cross-quant merge → delete B, C
    # ================================================================
    print("\n" + "=" * 75)
    print("STAGE 1: Download B, C → merge into Q4_0")
    print("  Dequant Q4_K_M → float32 → merge → requant Q4_0")
    print("=" * 75)

    path_b = download(MODEL_B, "B (fine-tuned)")
    path_c = download(MODEL_C, "C (reference)")
    print(f"[FS] Disk free: {disk_free_gb():.1f} GB")

    run_merge(MERGED_PATH, path_b, path_c, MERGE_WEIGHT)
    gc.collect()

    verify_output(MERGED_PATH)

    print("\n[CLEANUP] Removing B, C...")
    safe_remove(path_c, protect=MERGED_PATH)
    safe_remove(path_b, protect=MERGED_PATH)
    gc.collect()
    print(f"[FS] Disk free: {disk_free_gb():.1f} GB")

    # ================================================================
    # STAGE 2: Final quantization Q4_0 → Q4_K_M
    # ================================================================
    final_path = MERGED_PATH

    if FINAL_QTYPE:
        print("\n" + "=" * 75)
        print(f"STAGE 2: QUANTIZE Q4_0 → {FINAL_QTYPE}")
        print("=" * 75)

        base = os.path.splitext(os.path.basename(MERGED_PATH))[0]
        for sfx in ["-Q4_0","-Q4_K_M","-merged"]:
            if base.endswith(sfx): base = base[:-len(sfx)]

        output_name = f"{base}-{FINAL_QTYPE}.gguf"
        output_path = os.path.join(WORK_DIR, output_name)

        input_gb = file_size_gb(MERGED_PATH)
        print(f"[QUANT] Input:  {input_gb:.2f} GB  |  Disk free: {disk_free_gb():.1f} GB")

        result = run_llama_quantize(MERGED_PATH, output_path, FINAL_QTYPE)

        if result and os.path.isfile(output_path):
            final_path = output_path
            output_gb = file_size_gb(output_path)
            print(f"[QUANT] {input_gb:.2f} GB → {output_gb:.2f} GB")

            safe_remove(MERGED_PATH, protect=final_path)
            verify_output(final_path)
        else:
            print("[QUANT] Failed. Keeping Q4_0.")
    else:
        print("\n[INFO] Final quantize disabled. Keeping Q4_0.")

    # ================================================================
    # DONE
    # ================================================================
    print("\n" + "=" * 75)
    print("PIPELINE COMPLETE")
    print("=" * 75)
    print(f"\n[FS] {WORK_DIR}:")
    for fn in sorted(os.listdir(WORK_DIR)):
        fp = os.path.join(WORK_DIR, fn)
        if os.path.isfile(fp):
            print(f"  {fn}  ({file_size_gb(fp):.2f} GB)")
    print(f"\n  Final:     {final_path}")
    print(f"  Size:      {file_size_gb(final_path):.2f} GB")
    print(f"  Disk free: {disk_free_gb():.1f} GB")
    print(f"  RAM used:  {ram_used_gb():.2f} GB")

main()