#!/usr/bin/env python3
"""Debug E4B inference using MLX to compare with Swift Metal."""

import mlx.core as mx
import numpy as np
import json
from pathlib import Path
import struct

MODEL_PATH = Path("/Users/accusys/MarkBase12B/models/E4B-MarkBase")

def load_config():
    with open(MODEL_PATH / "config.json") as f:
        return json.load(f)

def bf16_to_f32_bytes(bf16_data):
    """Convert raw BF16 bytes to float32 numpy array."""
    bf16_uint = np.frombuffer(bf16_data, dtype=np.uint16)
    f32_uint = bf16_uint.astype(np.uint32) << 16
    return f32_uint.view(np.float32)

def load_safetensors_raw():
    """Load safetensors file and return raw data + metadata."""
    model_file = MODEL_PATH / "model.safetensors"
    with open(model_file, 'rb') as f:
        header_size = struct.unpack('<Q', f.read(8))[0]
        header_json = f.read(header_size)
        metadata = json.loads(header_json)
        data_start = 8 + header_size
        raw_data = f.read()
    return metadata, data_start, raw_data

def load_tensor_raw(name, metadata, data_start, raw_data):
    """Load tensor from raw safetensors data."""
    info = metadata[name]
    dtype = info['dtype']
    shape = info['shape']
    offsets = info['data_offsets']
    
    start = offsets[0]
    end = offsets[1]
    tensor_bytes = raw_data[start:end]
    
    if dtype == 'BF16':
        arr = bf16_to_f32_bytes(tensor_bytes)
    elif dtype == 'U32':
        arr = np.frombuffer(tensor_bytes, dtype=np.uint32)
    elif dtype == 'F32':
        arr = np.frombuffer(tensor_bytes, dtype=np.float32)
    else:
        raise ValueError(f"Unknown dtype: {dtype}")
    
    # Reshape to proper dimensions
    return arr.reshape(shape)

def debug_inference():
    print("=" * 50)
    print("MLX/NumPy E4B Debug Inference")
    print("=" * 50)
    
    config = load_config()
    text_config = config["text_config"]
    hidden_size = text_config['hidden_size']
    n_heads = text_config['num_attention_heads']
    n_kv_heads = text_config['num_key_value_heads']
    head_dim = text_config['head_dim']
    
    print(f"\nConfig: hidden={hidden_size}, head_dim={head_dim}")
    print(f"  n_heads={n_heads}, n_kv_heads={n_kv_heads}")
    
    # Load safetensors
    metadata, data_start, raw_data = load_safetensors_raw()
    
    def load(name):
        return load_tensor_raw(name, metadata, data_start, raw_data)
    
    print("\n--- Loading weights ---")
    embed_weight = load("language_model.model.embed_tokens.weight")
    embed_scales = load("language_model.model.embed_tokens.scales")
    embed_biases = load("language_model.model.embed_tokens.biases")
    
    print(f"embed_weight: shape={embed_weight.shape}, dtype={embed_weight.dtype}")
    print(f"embed_scales: shape={embed_scales.shape}, dtype={embed_scales.dtype}")
    
    # Dequantize BOS embedding
    bos_token = 2
    vocab_size, packed_dim = embed_weight.shape
    num_groups = hidden_size // 64
    
    print(f"\n--- Dequantizing BOS (token={bos_token}) ---")
    w_row = embed_weight[bos_token]
    s_row = embed_scales[bos_token]
    b_row = embed_biases[bos_token]
    
    embedding = np.zeros(hidden_size, dtype=np.float32)
    for g in range(num_groups):
        scale = float(s_row[g])
        bias = float(b_row[g])
        for j in range(64):
            packed_idx = g * 8 + j // 8
            shift = (j % 8) * 4
            qval = int((w_row[packed_idx] >> shift) & 0xF)
            embedding[g * 64 + j] = qval * scale + bias
    
    print(f"Embedding: min={embedding.min():.6f}, max={embedding.max():.6f}")
    print(f"  first 20: {embedding[:20].tolist()}")
    
    # Scale
    embed_scale = np.sqrt(float(hidden_size))
    embedding_scaled = embedding * embed_scale
    print(f"\nScaled (x{embed_scale:.2f}): min={embedding_scaled.min():.6f}, max={embedding_scaled.max():.6f}")
    print(f"  first 20: {embedding_scaled[:20].tolist()}")
    
    # Layer 0
    print("\n--- Layer 0 ---")
    prefix = "language_model.model.layers.0"
    
    input_norm_w = load(f"{prefix}.input_layernorm.weight")
    q_norm_w = load(f"{prefix}.self_attn.q_norm.weight")
    k_norm_w = load(f"{prefix}.self_attn.k_norm.weight")
    
    q_proj_w = load(f"{prefix}.self_attn.q_proj.weight")
    q_proj_s = load(f"{prefix}.self_attn.q_proj.scales")
    q_proj_b = load(f"{prefix}.self_attn.q_proj.biases")
    
    k_proj_w = load(f"{prefix}.self_attn.k_proj.weight")
    k_proj_s = load(f"{prefix}.self_attn.k_proj.scales")
    k_proj_b = load(f"{prefix}.self_attn.k_proj.biases")
    
    v_proj_w = load(f"{prefix}.self_attn.v_proj.weight")
    v_proj_s = load(f"{prefix}.self_attn.v_proj.scales")
    v_proj_b = load(f"{prefix}.self_attn.v_proj.biases")
    
    o_proj_w = load(f"{prefix}.self_attn.o_proj.weight")
    o_proj_s = load(f"{prefix}.self_attn.o_proj.scales")
    o_proj_b = load(f"{prefix}.self_attn.o_proj.biases")
    
    print(f"input_norm: {input_norm_w.shape}, q_norm: {q_norm_w.shape}")
    print(f"q_proj: w={q_proj_w.shape}, s={q_proj_s.shape}")
    
    # Forward
    def dequantize_matmul(x, w, s, b):
        out_dim, packed_dim = w.shape
        num_groups = packed_dim * 8 // 64
        result = np.zeros(out_dim, dtype=np.float32)
        for out_idx in range(out_dim):
            sum_val = 0.0
            for g in range(num_groups):
                scale = float(s[out_idx, g])
                bias_q = float(b[out_idx, g])
                for j in range(64):
                    packed_idx = g * 8 + j // 8
                    shift = (j % 8) * 4
                    qval = int((w[out_idx, packed_idx] >> shift) & 0xF)
                    dq = qval * scale + bias_q
                    sum_val += dq * x[g * 64 + j]
            result[out_idx] = sum_val
        return result
    
    hidden = embedding_scaled
    eps = 1e-6
    
    # Input norm
    ss = np.sum(hidden * hidden) / hidden_size
    rms = np.sqrt(ss + eps)
    hidden_normed = hidden / rms * input_norm_w
    print(f"\nInput norm: rms={rms:.6f}")
    print(f"  first 5: {hidden_normed[:5].tolist()}")
    
    # Q, K, V
    q = dequantize_matmul(hidden_normed, q_proj_w, q_proj_s, q_proj_b)
    k = dequantize_matmul(hidden_normed, k_proj_w, k_proj_s, k_proj_b)
    v = dequantize_matmul(hidden_normed, v_proj_w, v_proj_s, v_proj_b)
    
    print(f"\nQ: min={q.min():.6f}, max={q.max():.6f}")
    print(f"  first 5: {q[:5].tolist()}")
    print(f"\nK: min={k.min():.6f}, max={k.max():.6f}")
    print(f"  first 5: {k[:5].tolist()}")
    print(f"\nV: min={v.min():.6f}, max={v.max():.6f}")
    print(f"  first 5: {v[:5].tolist()}")
    
    # Q, K norms
    q_reshaped = q.reshape(n_heads, head_dim)
    q_normed = np.zeros_like(q_reshaped)
    for h in range(n_heads):
        h_q = q_reshaped[h]
        ss = np.sum(h_q * h_q) / head_dim
        rms = np.sqrt(ss + eps)
        q_normed[h] = h_q / rms * q_norm_w
    q_normed = q_normed.flatten()
    
    k_reshaped = k.reshape(n_kv_heads, head_dim)
    k_normed = np.zeros_like(k_reshaped)
    for h in range(n_kv_heads):
        h_k = k_reshaped[h]
        ss = np.sum(h_k * h_k) / head_dim
        rms = np.sqrt(ss + eps)
        k_normed[h] = h_k / rms * k_norm_w
    k_normed = k_normed.flatten()
    
    print(f"\nQ normed: min={q_normed.min():.6f}, max={q_normed.max():.6f}")
    print(f"  first 5: {q_normed[:5].tolist()}")
    print(f"\nK normed: min={k_normed.min():.6f}, max={k_normed.max():.6f}")
    print(f"  first 5: {k_normed[:5].tolist()}")
    
    # Attention (pos 0)
    attn_scale = 1.0 / np.sqrt(float(head_dim))
    attn_out = np.zeros(n_heads * head_dim, dtype=np.float32)
    scores = []
    for h in range(n_heads):
        kv_head = h % n_kv_heads
        h_q = q_normed[h * head_dim : (h+1) * head_dim]
        h_k = k_normed[kv_head * head_dim : (kv_head+1) * head_dim]
        h_v = v[kv_head * head_dim : (kv_head+1) * head_dim]
        score = float(np.sum(h_q * h_k) * attn_scale)
        scores.append(score)
        attn_out[h * head_dim : (h+1) * head_dim] = h_v
    
    print(f"\nAttention scores: {scores}")
    print(f"Attn out: min={attn_out.min():.6f}, max={attn_out.max():.6f}")
    print(f"  first 5: {attn_out[:5].tolist()}")
    
    # O proj
    o_out = dequantize_matmul(attn_out, o_proj_w, o_proj_s, o_proj_b)
    print(f"\nO out: min={o_out.min():.6f}, max={o_out.max():.6f}")
    print(f"  first 5: {o_out[:5].tolist()}")
    
    # Residual
    hidden_after = hidden + o_out
    print(f"\nAfter attn: min={hidden_after.min():.6f}, max={hidden_after.max():.6f}")
    print(f"  first 5: {hidden_after[:5].tolist()}")
    
    print("\n" + "=" * 50)
    print("KEY VALUES TO COMPARE:")
    print("=" * 50)
    print(f"\nEmbedding scaled: {embedding_scaled[:5].tolist()}")
    print(f"After input norm: {hidden_normed[:5].tolist()}")
    print(f"Q projection: {q[:5].tolist()}")
    print(f"Q normed: {q_normed[:5].tolist()}")
    print(f"Attention out: {attn_out[:5].tolist()}")
    print(f"O projection: {o_out[:5].tolist()}")
    print(f"Hidden after attn: {hidden_after[:5].tolist()}")

if __name__ == "__main__":
    debug_inference()