#!/usr/bin/env python3
"""Activation-diagonal weighted W4 scale search; unchanged per-row INT4 layout.

Calibration and held-out fixtures must be different. This is not GPTQ/AWQ:
no full Hessian compensation, channel rescaling, or optimizer training is used.
"""
import argparse
import hashlib
import json
import time
from pathlib import Path

import torch
from safetensors.torch import save_file
from transformers import AutoTokenizer
from quant_loss_reference import load_model, evaluate


def safe_rows(codes, splits):
    """Exact worst-case signed INT4 products, matching SplitAligned(K, splits, 32)."""
    blocks = codes.shape[1] // 32
    if codes.shape[1] % 32 or not 1 <= splits <= blocks:
        raise ValueError('invalid aligned K partition')
    size, extra = divmod(blocks, splits)
    begin = 0
    safe = torch.ones(codes.shape[0], dtype=torch.bool, device=codes.device)
    for part in range(splits):
        end = begin + (size + (part < extra)) * 32
        q = codes[:,begin:end].float()
        positive = q.clamp_min(0).sum(1)
        negative = (-q).clamp_min(0).sum(1)
        safe &= (7*positive+8*negative <= 32767) & (8*positive+7*negative <= 32768)
        begin = end
    return safe


def choose_scales(weight, diagonal, rows=256, k_splits=None):
    """Search clipping and least-squares scales; retain baseline if no improvement."""
    selected = []
    before, after = 0.0, 0.0
    for offset in range(0, weight.shape[0], rows):
        w = weight[offset:offset+rows].float()
        base = torch.maximum(w.amax(1).clamp_min(0)/7, (-w.amin(1)).clamp_min(0)/8)
        base = torch.where(base > 0, base, torch.ones_like(base))
        def error(scale):
            q = torch.round(w/scale[:, None]).clamp(-8, 7)
            return ((w-q*scale[:, None]).square()*diagonal).sum(1), q
        best_error, baseline_codes = error(base)
        if k_splits is not None and not safe_rows(baseline_codes,k_splits).all():
            raise ValueError('baseline scale is unsafe for required K partitions')
        before += best_error.double().sum().item()
        best = base.clone()
        # Include the old solution exactly. Search never worsens this local objective.
        for fraction in (1.0, .95, .9, .85, .8, .75, .7, .65, .6, .55, .5):
            scale = base * fraction
            for iteration in range(3):
                loss, q = error(scale)
                improve = loss < best_error
                if k_splits is not None:
                    improve &= safe_rows(q,k_splits)
                best = torch.where(improve, scale, best)
                best_error = torch.where(improve, loss, best_error)
                if iteration < 2:
                    scale = ((w*q*diagonal).sum(1) /
                             (q.square()*diagonal).sum(1).clamp_min(1e-30)).clamp_min(1e-20)
        after += best_error.double().sum().item()
        selected.append(best)
    return torch.cat(selected), before, after


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument('--source', type=Path, required=True)
    p.add_argument('--corpus', type=Path, required=True)
    p.add_argument('--suite', type=Path, required=True)
    p.add_argument('--output', type=Path, required=True)
    p.add_argument('--max-tokens', type=int, default=512)
    p.add_argument('--reuse-activations', action='store_true')
    p.add_argument('--baseline-manifest', type=Path, help='Preserve baseline INT16-safe K partitions')
    a = p.parse_args()
    torch.set_num_threads(4)
    torch.backends.cuda.matmul.allow_tf32 = False
    a.output.mkdir(parents=True, exist_ok=True)
    corpus = json.loads(a.corpus.read_text())
    suite = json.loads(a.suite.read_text())
    # Disallow accidental reuse of held-out questions or answers in calibration.
    for c in corpus:
        for e in suite['cases']:
            assert c['prompt'] not in e['prompt_text'] and c['answer'] not in e['target_text']
    model = load_model(a.source)
    modules = {name: m for name, m in model.named_modules() if isinstance(m, torch.nn.Linear)}
    partitions = {}
    if a.baseline_manifest:
        from quantize_model import source_linear_names
        for item in json.loads(a.baseline_manifest.read_text())['tensors']:
            if item['role'] == 'linear_weight':
                for n in source_linear_names(item['name'].removesuffix('.weight')):
                    partitions[n.removesuffix('.weight')] = item['flags']
        assert set(partitions) == set(modules)
    stats_path = a.output/'activations.pt'
    if a.reuse_activations:
        stats = torch.load(stats_path, map_location='cuda', weights_only=True)
        assert stats['corpus_sha256'] == hashlib.sha256(a.corpus.read_bytes()).hexdigest()
    else:
        sums, counts = {}, {}
        def hook(name):
            def accumulate(module, inputs):
                x = inputs[0].detach().reshape(-1, inputs[0].shape[-1]).float()
                if not x.shape[0]:
                    return
                value = x.square().sum(0)
                if name in sums:
                    sums[name].add_(value)
                else:
                    sums[name] = value
                counts[name] = counts.get(name, 0) + x.shape[0]
            return accumulate
        hooks = [m.register_forward_pre_hook(hook(n)) for n,m in modules.items()]
        tokenizer = AutoTokenizer.from_pretrained(a.source, trust_remote_code=True, local_files_only=True)
        lengths = []
        with torch.inference_mode():
            for i, c in enumerate(corpus):
                for thinking in (False, True):
                    prompt = tokenizer.apply_chat_template([{'role':'user','content':c['prompt']}],
                        tokenize=False, add_generation_prompt=True, enable_thinking=thinking)
                    target = (c['reasoning']+'</think>' if thinking else '') + c['answer'] + tokenizer.eos_token
                    ids = tokenizer.encode(prompt+target, add_special_tokens=False)[:a.max_tokens]
                    output = model(input_ids=torch.tensor([ids],device='cuda'), use_cache=False)
                    del output
                    lengths.append(len(ids))
                print(json.dumps({'calibration_cases':i+1, 'total_tokens':sum(lengths),
                                  'matrices_observed':len(sums)}), flush=True)
        for h in hooks:
            h.remove()
        stats = {'sum_squares':{k:v.cpu() for k,v in sums.items()},'counts':counts,
                 'lengths':lengths, 'corpus_sha256':hashlib.sha256(a.corpus.read_bytes()).hexdigest()}
        torch.save(stats, stats_path)
        del sums
    scales, records = {}, []
    start = time.monotonic()
    with torch.inference_mode():
        for i,(name,m) in enumerate(modules.items()):
            count = stats['counts'].get(name,0)
            diag = stats['sum_squares'].get(name)
            if diag is None or count == 0:
                # Unobserved experts retain baseline weights, not a guessed calibration.
                w = m.weight.float()
                scale = torch.maximum(w.amax(1).clamp_min(0)/7,(-w.amin(1)).clamp_min(0)/8)
                scale = torch.where(scale>0,scale,torch.ones_like(scale))
                old = new = None
            else:
                diag = diag.to('cuda') / count
                mean = diag.mean().clamp_min(1e-20)
                # Regularize rarely visited experts more strongly towards uniform weights.
                strength = .05 + .95 * max(0., 1.-count/64.)
                diag = (1-strength)*diag/mean + strength
                scale, old, new = choose_scales(m.weight,diag,k_splits=partitions.get(name))
            scales[name+'.weight'] = scale.cpu().contiguous()
            records.append({'name':name,'observed_tokens':count,'weighted_error_before':old,
                            'weighted_error_after':new})
            for row in range(0,m.weight.shape[0],1024):
                w = m.weight[row:row+1024].float()
                s = scale[row:row+1024,None]
                m.weight[row:row+1024].copy_((torch.round(w/s).clamp(-8,7)*s).to(m.weight.dtype))
            if (i+1)%200 == 0:
                print(json.dumps({'calibrated_matrices':i+1,'seconds':time.monotonic()-start}),flush=True)
    save_file(scales,str(a.output/'scales.safetensors'))
    report = {'method':'activation_diagonal_weighted_scale_search_v1','corpus_sha256':stats['corpus_sha256'],
              'calibration_sequences':len(stats['lengths']),'calibration_tokens':sum(stats['lengths']),
              'matrices':records,'seconds':time.monotonic()-start,'max_tokens':a.max_tokens,
              'baseline_manifest_sha256':hashlib.sha256(a.baseline_manifest.read_bytes()).hexdigest() if a.baseline_manifest else None,
              'preserve_baseline_k_partitions':bool(a.baseline_manifest)}
    (a.output/'calibration.json').write_text(json.dumps(report,indent=2)+'\n')
    evaluate(model,suite['cases'],a.output,'calibrated_bf16')
    print('PASS: calibrated scales and held-out logits exported',flush=True)


if __name__ == '__main__':
    main()