"""EOQ Loader: Load compressed EOQ model into transformers."""
import torch, torch.nn.functional as F
import json
from safetensors.torch import load_file
from transformers import AutoModelForCausalLM, AutoConfig

def dequantize_codes(codes, scales, shape, block_size=128):
    flat = codes.float().flatten()
    n = flat.numel()
    pad = (block_size - n % block_size) % block_size
    if pad > 0: flat = F.pad(flat, (0, pad))
    return (flat.view(-1, block_size) * scales.float().unsqueeze(1)).flatten()[:n].view(shape).half()

def load_eoq_model(repo_or_path, device_map="auto", trust_remote_code=True):
    """Load an EOQ-compressed model.
    
    Usage:
        from eoq_loader import load_eoq_model
        model, tokenizer = load_eoq_model("caiovicentino1/MODEL-EOQ-Q5-compressed")
    """
    from transformers import AutoTokenizer
    from huggingface_hub import hf_hub_download, snapshot_download
    import os
    
    # Download
    local_dir = snapshot_download(repo_or_path)
    
    # Load metadata
    with open(os.path.join(local_dir, "eoq_metadata.json")) as f:
        meta = json.load(f)
    
    # Load compressed weights
    index_path = os.path.join(local_dir, "model.safetensors.index.json")
    if os.path.exists(index_path):
        with open(index_path) as f:
            index = json.load(f)
        safetensor_files = set(index["weight_map"].values())
        compressed_sd = {}
        for sf in safetensor_files:
            compressed_sd.update(load_file(os.path.join(local_dir, sf)))
    else:
        compressed_sd = load_file(os.path.join(local_dir, "model.safetensors"))
    
    # Build dequantized state dict
    state_dict = {}
    for name, info in meta["tensors"].items():
        if info["quantized"]:
            codes = compressed_sd[f"{name}.codes"]
            scales = compressed_sd[f"{name}.scales"]
            shape = info["shape"]
            bs = info["block_size"]
            state_dict[name] = dequantize_codes(codes, scales, shape, bs)
        else:
            state_dict[name] = compressed_sd[name]
    
    # Load model architecture and inject weights
    config = AutoConfig.from_pretrained(local_dir, trust_remote_code=trust_remote_code)
    model = AutoModelForCausalLM.from_config(config, trust_remote_code=trust_remote_code)
    model.load_state_dict(state_dict, strict=False)
    model = model.half()
    if device_map == "auto":
        model = model.cuda()
    model.eval()
    
    tokenizer = AutoTokenizer.from_pretrained(local_dir, trust_remote_code=trust_remote_code)
    
    return model, tokenizer