#!/usr/bin/env python3
"""
Convert E4B model weights to MarkBase format.

Reads the MLX-community 4-bit quantized Gemma 4 E4B model and re-quantizes
the vision and audio tower weights into the MarkBase 4-bit packed format.
"""

import os
import sys
import shutil
import struct
import json
import numpy as np
import safetensors

SRC_DIR = os.path.expanduser(
    "~/.cache/huggingface/hub/models--mlx-community--gemma-4-e4b-it-4bit/"
    "snapshots/deb1db712068b1c9f83fb1c97f08c1204b9459a1"
)
DST_DIR = "/Users/accusys/MarkBase12B/models/E4B-MarkBase"
GROUP_SIZE = 64

SRC_MODEL = os.path.join(SRC_DIR, "model.safetensors")
DST_MODEL = os.path.join(DST_DIR, "model.safetensors")
DST_INDEX = os.path.join(DST_DIR, "model.safetensors.index.json")


def bf16_bytes_to_f32(data: bytes) -> np.ndarray:
    """Convert BF16 bytes to float32 numpy array."""
    raw = np.frombuffer(data, dtype=np.uint16)
    f32 = (raw.astype(np.uint32) << 16).view(np.float32)
    return f32


def quantize_weight_4bit(weight_f32: np.ndarray, group_size: int = GROUP_SIZE):
    """
    Group-wise 4-bit affine quantization.

    weight_f32: shape [outDim, inDim], float32
    Returns (packed_u32, scales_f32, biases_f32)
      packed_u32: uint32 [outDim, inDim/8]
      scales_f32: float32 [outDim, inDim/64]
      biases_f32: float32 [outDim, inDim/64]
    """
    out_dim, in_dim = weight_f32.shape
    assert in_dim % group_size == 0, f"in_dim {in_dim} not divisible by group_size {group_size}"
    assert in_dim % 8 == 0

    # Reshape to [outDim, inDim/64, 64]
    num_groups = in_dim // group_size
    groups = weight_f32.reshape(out_dim, num_groups, group_size)

    # Per-group min/max
    group_min = groups.min(axis=-1)  # [outDim, num_groups]
    group_max = groups.max(axis=-1)  # [outDim, num_groups]

    scale = (group_max - group_min) / 15.0
    bias = group_min

    # Handle zero-range groups (scale == 0)
    zero_scale_mask = scale == 0.0
    scale[zero_scale_mask] = 1.0  # avoid division by zero

    # Quantize
    quantized = np.round((groups - bias[:, :, np.newaxis]) / scale[:, :, np.newaxis])
    quantized = np.clip(quantized, 0, 15).astype(np.uint32)

    # Pack 8 values into one U32
    # pack order: q[0] at bits 0-3, q[1] at bits 4-7, ..., q[7] at bits 28-31
    packed_shape = (out_dim, in_dim // 8)
    packed = np.zeros(packed_shape, dtype=np.uint32)

    # Reshape quantized to [outDim, inDim/8, 8]
    quantized_reshaped = quantized.reshape(out_dim, in_dim // 8, 8)
    for j in range(8):
        packed |= (quantized_reshaped[:, :, j] << (j * 4))

    return packed, scale.astype(np.float32), bias.astype(np.float32)


def is_linear_weight(name: str) -> bool:
    """Check if a tensor name corresponds to a linear weight that needs quantization."""
    if name.startswith("language_model."):
        return False
    if name.startswith("embed_vision.") or name.startswith("embed_audio."):
        return False
    if name.endswith(".linear.weight"):
        return True
    if name == "vision_tower.patch_embedder.input_proj.weight":
        return True
    if name == "audio_tower.output_proj.weight":
        return True
    return False


def is_norm_weight(name: str) -> bool:
    """Check if a tensor name corresponds to a norm/weight that should stay as F32."""
    if name.startswith("language_model."):
        return False
    if name.startswith("embed_vision.") or name.startswith("embed_audio."):
        return False
    if "norm" in name.lower() and name.endswith(".weight"):
        return True
    if name == "vision_tower.patch_embedder.position_embedding_table":
        return True
    if "per_dim_scale" in name:
        return True
    if "relative_k_proj.weight" in name:
        return True
    if "depthwise_conv1d.weight" in name:
        return True
    if "conv_norm.weight" in name:
        return True
    if "pre_layer_norm.weight" in name:
        return True
    if "post_layer_norm.weight" in name:
        return True
    if "norm_out.weight" in name:
        return True
    if "norm_pre_attn.weight" in name:
        return True
    if "norm_post_attn.weight" in name:
        return True
    if name.startswith("audio_tower.subsample_conv"):
        return True
    if name.startswith("audio_tower.subsample_conv_projection"):
        return True
    if name.endswith("output_proj.bias"):
        return True
    return False


def process_tensors():
    """Read source tensors, convert, and write output."""
    os.makedirs(DST_DIR, exist_ok=True)

    print("Reading source safetensor...")
    with open(SRC_MODEL, "rb") as f:
        src_data = f.read()
    tensors = safetensors.deserialize(src_data)
    print(f"  Found {len(tensors)} tensors")

    output_tensor_dict = {}
    text_copied = 0
    embed_copied = 0
    vision_quantized = 0
    vision_f32 = 0
    audio_quantized = 0
    audio_f32 = 0
    skipped = 0

    # Rust safetensors backend only accepts lowercase dtype names for writing.
    DTYPE_MAP = {
        "U32": "uint32",
        "BF16": "bfloat16",
        "F16": "float16",
        "F32": "float32",
        "I32": "int32",
        "I64": "int64",
    }

    def add_tensor(name, shape, dtype, data):
        if isinstance(data, bytearray):
            data = bytes(data)
        elif not isinstance(data, bytes):
            data = data.tobytes()
        py_dtype = DTYPE_MAP.get(dtype, dtype)
        output_tensor_dict[name] = {
            "shape": list(shape) if isinstance(shape, (list, tuple)) else shape,
            "dtype": py_dtype,
            "data": data,
        }

    for name, info in tensors:
        dtype_tag = info["dtype"]
        shape = info["shape"]
        data = info["data"]

        # --- Text model: copy as-is ---
        if name.startswith("language_model."):
            add_tensor(name, shape, dtype_tag, data)
            text_copied += 1
            continue

        # --- Embed vision/audio projections: already quantized, copy as-is ---
        if name.startswith("embed_vision.") or name.startswith("embed_audio."):
            add_tensor(name, shape, dtype_tag, data)
            embed_copied += 1
            continue

        # --- Vision/Audio tower input_min/max/output_min/output_max: skip ---
        if ("input_min" in name or "input_max" in name or
            "output_min" in name or "output_max" in name):
            skipped += 1
            continue

        # --- Quantize linear weights ---
        if is_linear_weight(name):
            print(f"  Quantizing: {name}  shape={shape}")
            weight_f32 = bf16_bytes_to_f32(data).reshape(shape)
            packed, scales, biases = quantize_weight_4bit(weight_f32)

            out_dim, in_dim = shape

            # Strip `.linear` if present; otherwise use name directly.
            # e.g. "q_proj.linear.weight" → "q_proj.weight" / ".scales" / ".biases"
            #       "input_proj.weight"   → "input_proj.weight" / ".scales" / ".biases"
            if ".linear.weight" in name:
                weight_name = name.replace(".linear.weight", ".weight")
                scales_name = name.replace(".linear.weight", ".scales")
                biases_name = name.replace(".linear.weight", ".biases")
            else:
                weight_name = name
                scales_name = name.replace(".weight", ".scales")
                biases_name = name.replace(".weight", ".biases")

            add_tensor(weight_name, [out_dim, in_dim // 8], "U32", packed)
            add_tensor(scales_name, [out_dim, in_dim // GROUP_SIZE], "F32", scales)
            add_tensor(biases_name, [out_dim, in_dim // GROUP_SIZE], "F32", biases)

            if name.startswith("vision_tower."):
                vision_quantized += 1
            else:
                audio_quantized += 1
            continue

        # --- Convert BF16 norms and special tensors to float32 ---
        if is_norm_weight(name):
            f32_array = bf16_bytes_to_f32(data).reshape(shape)
            add_tensor(name, shape, "F32", f32_array)
            if name.startswith("vision_tower."):
                vision_f32 += 1
            else:
                audio_f32 += 1
            continue

        # --- Catch-all: convert any remaining BF16 to F32 ---
        if dtype_tag == "BF16":
            print(f"  Converting BF16->F32: {name}  shape={shape}")
            f32_array = bf16_bytes_to_f32(data).reshape(shape)
            add_tensor(name, shape, "F32", f32_array)
            if name.startswith("vision_tower."):
                vision_f32 += 1
            else:
                audio_f32 += 1
            continue

        # Should not reach here for any known tensor
        print(f"  WARNING: unhandled tensor: {name}  dtype={dtype_tag}  shape={shape}")
        add_tensor(name, shape, dtype_tag, data)

    print()
    print(f"Text tensors copied: {text_copied}")
    print(f"Embed tensors copied: {embed_copied}")
    print(f"Vision quantized: {vision_quantized}")
    print(f"Vision f32: {vision_f32}")
    print(f"Audio quantized: {audio_quantized}")
    print(f"Audio f32: {audio_f32}")
    print(f"Input/output minmax skipped: {skipped}")
    print(f"Total output tensors: {len(output_tensor_dict)}")

    print("\nSerializing output safetensor...")
    serialized = safetensors.serialize(output_tensor_dict)
    with open(DST_MODEL, "wb") as f:
        f.write(serialized)

    # Write index file
    total_bytes = os.path.getsize(DST_MODEL)
    weight_map = {name: "model.safetensors" for name in output_tensor_dict}
    index_data = {
        "metadata": {"total_size": total_bytes},
        "weight_map": weight_map,
    }
    with open(DST_INDEX, "w") as f:
        json.dump(index_data, f, indent=2)

    print(f"Output model: {DST_MODEL} ({total_bytes} bytes)")
    print(f"Output index: {DST_INDEX}")


def copy_metadata_files():
    """Copy non-weight files from source to destination."""
    os.makedirs(DST_DIR, exist_ok=True)
    metadata_files = [
        "config.json",
        "generation_config.json",
        "tokenizer.json",
        "tokenizer_config.json",
        "processor_config.json",
        "chat_template.jinja",
    ]
    for fn in metadata_files:
        src = os.path.join(SRC_DIR, fn)
        dst = os.path.join(DST_DIR, fn)
        if os.path.exists(src):
            shutil.copy2(src, dst)
            print(f"  Copied {fn}")
        else:
            print(f"  WARNING: {fn} not found in source")


def verify_output():
    """Verify the output safetensor can be read back."""
    print("\nVerifying output...")
    try:
        with open(DST_MODEL, "rb") as f:
            data = f.read()
        tensors = safetensors.deserialize(data)
        keys = [t[0] for t in tensors]
        print(f"  Output contains {len(keys)} tensors")

        # Build a lookup by key
        tensor_map = {t[0]: t[1] for t in tensors}

        # Verify a text tensor
        sample = "language_model.model.layers.0.input_layernorm.weight"
        if sample in tensor_map:
            info = tensor_map[sample]
            print(f"  {sample}: dtype={info['dtype']} shape={info['shape']}")

        # Verify a vision quantized tensor
        sample = "vision_tower.encoder.layers.0.self_attn.q_proj.weight"
        if sample in tensor_map:
            info = tensor_map[sample]
            print(f"  {sample}: dtype={info['dtype']} shape={info['shape']}")
        sample = "vision_tower.encoder.layers.0.self_attn.q_proj.scales"
        if sample in tensor_map:
            info = tensor_map[sample]
            print(f"  {sample}: dtype={info['dtype']} shape={info['shape']}")

        # Verify a vision norm tensor
        sample = "vision_tower.encoder.layers.0.input_layernorm.weight"
        if sample in tensor_map:
            info = tensor_map[sample]
            print(f"  {sample}: dtype={info['dtype']} shape={info['shape']}")

        # Verify audio quantized
        sample = "audio_tower.layers.0.self_attn.q_proj.weight"
        if sample in tensor_map:
            info = tensor_map[sample]
            print(f"  {sample}: dtype={info['dtype']} shape={info['shape']}")

        # Verify embed_vision
        sample = "embed_vision.embedding_projection.weight"
        if sample in tensor_map:
            info = tensor_map[sample]
            print(f"  {sample}: dtype={info['dtype']} shape={info['shape']}")

        print("  All checks passed!")
    except Exception as e:
        print(f"  ERROR during verification: {e}")
        sys.exit(1)


def main():
    print("=== E4B to MarkBase Converter ===\n")
    print(f"Source: {SRC_MODEL}")
    print(f"Destination: {DST_DIR}\n")

    copy_metadata_files()
    process_tensors()
    verify_output()

    print("\nDone! Model converted successfully.")


if __name__ == "__main__":
    main()