"""
Direct RTN INT8 W8A8 quantization of google/diffusiongemma-26B-A4B-it.

The base on-disk stores experts FUSED: experts.gate_up_proj [E, 2I, H] and
experts.down_proj [E, H, I]. Each per-expert slice is already [out, in].
We split into per-expert Linears (RedHat's proven layout) and quantize each:

  gate_up_proj[e] -> experts.{e}.gate_proj.weight (rows 0:I) + experts.{e}.up_proj.weight (rows I:2I)
  down_proj[e]    -> experts.{e}.down_proj.weight
  dense Linears   -> quantized in place (mlp.*, self_attn.*)

  weights      -> int8, per-output-channel, symmetric  (weight i8 [out,in] + weight_scale bf16 [out,1])
  activations  -> int8, per-token, dynamic, symmetric  (runtime, nothing stored)

Ignored (copied bf16): lm_head, *embed*, *router*, *vision_tower*, *self_conditioning*.
"""
import json, glob, os, re
import torch
from safetensors import safe_open
from safetensors.torch import save_file

BASE = glob.glob('/hf/hub/models--google--diffusiongemma-26B-A4B-it/snapshots/*/')[0]
OUT = '/models/diffusiongemma-26B-A4B-it-INT8-full'
os.makedirs(OUT, exist_ok=True)

IGNORE = re.compile(r'(^|\.)lm_head|embed|router|vision_tower|self_conditioning')

def ignored(name):
    return IGNORE.search(name) is not None

def q_int8(w):
    w = w.to(torch.float32)
    scale = (w.abs().amax(dim=1, keepdim=True) / 127.0).clamp(min=1e-8)
    q = torch.clamp(torch.round(w / scale), -128, 127).to(torch.int8)
    return q, scale.to(torch.bfloat16)

shards = sorted(glob.glob(BASE + 'model-*.safetensors'))
N = len(shards)
weight_map, total_size = {}, 0

for i, shard in enumerate(shards):
    out = {}
    nq = 0
    with safe_open(shard, framework='pt') as f:
        for k in f.keys():
            t = f.get_tensor(k)
            if k.endswith('.experts.gate_up_proj') and not ignored(k):
                parent = k[:-len('.gate_up_proj')]          # ...experts
                I = t.shape[1] // 2
                for e in range(t.shape[0]):
                    we = t[e]                                # [2I, H]
                    for nm, w in (('gate_proj', we[:I]), ('up_proj', we[I:])):
                        qw, qs = q_int8(w)
                        out[f'{parent}.{e}.{nm}.weight'] = qw
                        out[f'{parent}.{e}.{nm}.weight_scale'] = qs
                        nq += 1
            elif k.endswith('.experts.down_proj') and not ignored(k):
                parent = k[:-len('.down_proj')]
                for e in range(t.shape[0]):
                    qw, qs = q_int8(t[e])                    # [H, I]
                    out[f'{parent}.{e}.down_proj.weight'] = qw
                    out[f'{parent}.{e}.down_proj.weight_scale'] = qs
                    nq += 1
            elif k.endswith('.weight') and t.ndim == 2 and not ignored(k):
                qw, qs = q_int8(t)
                out[k] = qw
                out[k[:-len('.weight')] + '.weight_scale'] = qs
                nq += 1
            else:
                out[k] = t
    fname = f'model-{i+1:05d}-of-{N:05d}.safetensors'
    save_file(out, os.path.join(OUT, fname), metadata={'format': 'pt'})
    for k, v in out.items():
        weight_map[k] = fname
        total_size += v.numel() * v.element_size()
    print(f'>> shard {i+1}/{N}: quantized {nq} weights, {len(out)} tensors -> {fname}', flush=True)

with open(os.path.join(OUT, 'model.safetensors.index.json'), 'w') as f:
    json.dump({'metadata': {'total_size': total_size}, 'weight_map': weight_map}, f, indent=2)

cfg = json.load(open(os.path.join(BASE, 'config.json')))
cfg['quantization_config'] = {
    'quant_method': 'compressed-tensors',
    'format': 'int-quantized',
    'quantization_status': 'compressed',
    'global_compression_ratio': None,
    'kv_cache_scheme': None,
    'sparsity_config': {},
    'ignore': ['lm_head', 're:.*embed.*', 're:.*router', 're:.*vision_tower.*', 're:.*self_conditioning.*'],
    'config_groups': {
        'group_0': {
            'targets': ['Linear'],
            'format': 'int-quantized',
            'output_activations': None,
            'weights': {
                'num_bits': 8, 'type': 'int', 'symmetric': True, 'strategy': 'channel',
                'dynamic': False, 'group_size': None, 'block_structure': None,
                'actorder': None, 'observer': 'memoryless_minmax', 'observer_kwargs': {},
                'scale_dtype': None, 'zp_dtype': None,
            },
            'input_activations': {
                'num_bits': 8, 'type': 'int', 'symmetric': True, 'strategy': 'token',
                'dynamic': True, 'group_size': None, 'block_structure': None,
                'actorder': None, 'observer': None, 'observer_kwargs': {},
                'scale_dtype': None, 'zp_dtype': None,
            },
        }
    },
}
json.dump(cfg, open(os.path.join(OUT, 'config.json'), 'w'), indent=2)

for aux in ['generation_config.json', 'tokenizer_config.json', 'tokenizer.json',
            'processor_config.json', 'chat_template.jinja', 'special_tokens_map.json',
            'preprocessor_config.json']:
    src = os.path.join(BASE, aux)
    if os.path.exists(src):
        open(os.path.join(OUT, aux), 'wb').write(open(src, 'rb').read())

print('>> DONE:', OUT, 'total_size_GB=%.1f' % (total_size / 1e9), flush=True)