"""Activation-aware calibrated FP8 (W8A8 E4M3) for Image 2.1 Turbo; no evaluation prompts used.

Per layer: five SmoothQuant exponents x {three weight-clipping factors with round-to-nearest,
GPTQ error-feedback rounding with full-token second moments}. Every candidate is scored with the
actual runtime path (Triton per-token activation quantization + native SM120 FP8 GEMM, FP32
accumulation) on 512 fit rows; the selected candidate is reported on 512 disjoint diagnostic rows.
Layers above 5% diagnostic output NRMSE would stay BF16.
"""
import json, math, shutil, statistics, sys
from pathlib import Path
import torch
from safetensors.torch import load_file, save_file
from diffusers import QwenImage21Transformer2DModel
from common import ROOT, BASE, BASE_REVISION, CALIBRATION, eligible_modules
from fp8_runtime import CalibratedFP8Linear, replace_module, quantize_activation
from hessian_store import hessian
from gptq import gptq_fp8

OUT_ROOT = ROOT / 'fp8' / 'release'
OUT = OUT_ROOT / 'transformer'
CACHE = ROOT / 'fp8' / 'calibrated-layers'
NOTICE = 'Built with Qwen. ProCreations modified Qwen-Image-2.1-Turbo transformer weights through calibrated FP8 quantization, 2026-10-09.'


def split(x):
    g = torch.arange(len(x) // 4, device=x.device)
    return x[g * 4 + g % 4].contiguous(), x[g * 4 + (g + 1) % 4].contiguous()


def fp8_forward(x, q, scales, smooth):
    a, asc = quantize_activation(x, smooth)
    return torch._scaled_mm(a, q.T, scale_a=asc, scale_b=scales.T, out_dtype=torch.bfloat16, use_fast_accum=False).float()


def nmse(y, ref):
    return ((y - ref).square().mean() / ref.square().mean().clamp_min(1e-12)).item()


def bf16_weight(name):
    from quantize_nvfp4 import weight
    return weight(name)


@torch.inference_mode()
def calibrate_layer(name, layer_weight):
    x = load_file(str(CALIBRATION / (name + '.safetensors')))['activations'].to('cuda')
    fit, check = split(x)
    H, tokens = hessian(name)
    w = layer_weight.float()
    amax = fit.float().abs().amax(0).clamp_min(1e-5)
    wmax = w.abs().amax(0).clamp_min(1e-5)
    ref = torch.nn.functional.linear(fit, layer_weight).float()
    candidates, best = [], None
    for alpha in [0., .25, .5, .75, 1.]:
        smooth = torch.ones_like(amax) if alpha == 0 else (amax.pow(alpha) / wmax.pow(1 - alpha)).clamp(.01, 100.)
        smooth = smooth / smooth.log().mean().exp()
        ws = w * smooth[None, :]
        options = []
        for clip in [1., .995, .99]:
            scales = (ws.abs().amax(1, keepdim=True) * clip / 448).clamp_min(1e-12)
            options.append(({'rounding': 'nearest', 'weight_amax_clip_factor': clip}, scales,
                            (ws / scales).clamp(-448, 448).to(torch.float8_e4m3fn)))
        scales = (ws.abs().amax(1, keepdim=True) / 448).clamp_min(1e-12)
        Hz = H / (smooth[:, None] * smooth[None, :])
        options.append(({'rounding': 'gptq', 'weight_amax_clip_factor': 1.}, scales, gptq_fp8(ws, Hz, scales)))
        for item, scales, q in options:
            err = nmse(fp8_forward(fit, q, scales, smooth), ref)
            item = {'alpha': alpha, **item, 'fit_nmse': err}
            candidates.append(item)
            if best is None or err < best[0]:
                best = (err, q.clone(), scales.clone(), smooth.clone(), item)
    _, q, scales, smooth, item = best
    ref_check = torch.nn.functional.linear(check, layer_weight).float()
    diag = nmse(fp8_forward(check, q, scales, smooth), ref_check)
    # Round-to-nearest reference at the same smoothing, for reporting the GPTQ contribution.
    rtn = min((c for c in candidates if c['alpha'] == item['alpha'] and c['rounding'] == 'nearest'), key=lambda c: c['fit_nmse'])
    result = {**item, 'diagnostic_nrmse': math.sqrt(diag), 'fit_nrmse': math.sqrt(item['fit_nmse']),
              'nearest_fit_nrmse_same_alpha': math.sqrt(rtn['fit_nmse']), 'shape': list(q.shape),
              'sample_rows': len(x), 'fit_rows': len(fit), 'diagnostic_rows': len(check),
              'hessian_tokens': tokens, 'candidates': candidates}
    return result, {'weight': q, 'weight_scale': scales.float(), 'smooth': smooth.float()}


def compute_shard(index, count):
    names = json.loads((CALIBRATION / 'manifest.json').read_text())['modules']
    for i, name in enumerate(names):
        record, saved = CACHE / (name + '.json'), CACHE / (name + '.safetensors')
        if i % count != index or (record.exists() and saved.exists()):
            continue
        result, state = calibrate_layer(name, bf16_weight(name))
        save_file({k: v.cpu().contiguous() for k, v in state.items()}, str(saved))
        record.write_text(json.dumps(result, indent=2))
        print(json.dumps({'event': 'layer', 'shard': index, 'done': i + 1, 'layer': name, 'nrmse': result['diagnostic_nrmse'],
                          'alpha': result['alpha'], 'rounding': result['rounding']}), flush=True)


@torch.inference_mode()
def main():
    CACHE.mkdir(parents=True, exist_ok=True); OUT.mkdir(parents=True, exist_ok=True)
    if len(sys.argv) == 3 and sys.argv[1] == '--shard':
        return compute_shard(*map(int, sys.argv[2].split('/')))
    model = QwenImage21Transformer2DModel.from_pretrained(BASE, subfolder='transformer', torch_dtype=torch.bfloat16).to('cuda').eval()
    modules = eligible_modules(model)
    results, quantized, held = {}, {}, {}
    for i, (name, layer) in enumerate(modules.items()):
        record, saved = CACHE / (name + '.json'), CACHE / (name + '.safetensors')
        if record.exists() and saved.exists():
            result = json.loads(record.read_text()); state = {k: v.cuda() for k, v in load_file(str(saved)).items()}
        else:
            result, state = calibrate_layer(name, layer.weight)
            save_file({k: v.cpu().contiguous() for k, v in state.items()}, str(saved))
            record.write_text(json.dumps(result, indent=2))
        preserve = not math.isfinite(result['diagnostic_nrmse']) or result['diagnostic_nrmse'] > .05
        result['kept_bf16'] = preserve
        results[name] = result
        if preserve:
            held[name] = result
        else:
            quantized[name] = {'shape': result['shape'], 'alpha': result['alpha'], 'rounding': result['rounding'],
                               'weight_amax_clip_factor': result['weight_amax_clip_factor']}
            replace_module(model, name, CalibratedFP8Linear(state['weight'], state['weight_scale'], state['smooth']))
        print(json.dumps({'event': 'quantize', 'done': i + 1, 'total': len(modules), 'layer': name,
                          'nrmse': result['diagnostic_nrmse'], 'alpha': result['alpha'], 'rounding': result['rounding'],
                          'nearest_fit': result['nearest_fit_nrmse_same_alpha'], 'fit': result['fit_nrmse']}), flush=True)
    model.save_config(OUT)
    config = {'format': 'procreations-calibrated-fp8-v1', 'base_model': 'Qwen/Qwen-Image-2.1-Turbo',
              'base_revision': BASE_REVISION, 'weight_dtype': 'float8_e4m3fn', 'activation_dtype': 'float8_e4m3fn',
              'weight_scaling': 'per-output-channel FP32', 'activation_scaling': 'dynamic per-token FP32',
              'weight_rounding': 'per layer: round-to-nearest with clipping search, or GPTQ error feedback (full-token second moments)',
              'accumulation': 'FP32, use_fast_accum=False', 'output_dtype': 'bfloat16',
              'kernel': 'torch._scaled_mm native CUTLASS SM120, Triton fused smoothing and activation quantization',
              'quantized_modules': quantized, 'kept_bf16_modules': list(held), 'modified_notice': NOTICE}
    (OUT / 'quantization_config.json').write_text(json.dumps(config, indent=2))
    report = OUT_ROOT / 'reports'; report.mkdir(parents=True, exist_ok=True)
    (report / 'calibration_search.json').write_text(json.dumps(results, indent=2))
    shutil.copy(CALIBRATION / 'manifest.json', report / 'calibration_manifest.json')
    state = model.state_dict(); shards, shard, size = [], {}, 0
    for n, v in state.items():
        nb = v.numel() * v.element_size()
        if shard and size + nb > 3_500_000_000:
            shards.append(shard); shard, size = {}, 0
        shard[n] = v.detach().cpu().contiguous(); size += nb
    if shard:
        shards.append(shard)
    for old in OUT.glob('model-*.safetensors'):
        old.unlink()
    for i, s in enumerate(shards):
        save_file(s, str(OUT / f'model-{i + 1:05d}-of-{len(shards):05d}.safetensors'), metadata={'format': 'pt', 'notice': NOTICE})
    total = sum(v.numel() * v.element_size() for v in state.values())
    vals = list(results.values())
    summary = {'quantized_modules': len(quantized), 'bf16_outliers': list(held), 'transformer_bytes': total,
               'mean_diagnostic_nrmse': statistics.mean(v['diagnostic_nrmse'] for v in vals),
               'max_diagnostic_nrmse': max(v['diagnostic_nrmse'] for v in vals),
               'gptq_selected': sum(v['rounding'] == 'gptq' for v in vals),
               'mean_fit_nrmse': statistics.mean(v['fit_nrmse'] for v in vals),
               'mean_nearest_fit_nrmse_same_alpha': statistics.mean(v['nearest_fit_nrmse_same_alpha'] for v in vals)}
    (report / 'quantization_summary.json').write_text(json.dumps(summary, indent=2))
    print('QUANTIZATION_COMPLETE', json.dumps(summary), flush=True)


if __name__ == '__main__':
    main()