"""Calibration preserves the runtime INT4 contract and invalidates stale assets."""
import json
from pathlib import Path
import sys
import tempfile
import unittest

import numpy as np
import torch
from safetensors.torch import save_file

sys.path.insert(0,str(Path(__file__).resolve().parents[1]/'tools'))
from calibrate_w4 import choose_scales,safe_rows
from quantize_model import AssetWriter, source_linear_names


class CalibrationTests(unittest.TestCase):
    def test_weighted_error_and_zero_rows(self):
        torch.manual_seed(42)
        w = torch.randn(65,128)
        w[:,0] *= 7
        w[-1].zero_()
        h = torch.linspace(.1,2,128)
        s, old, new = choose_scales(w,h,rows=32)
        baseline_q, baseline_s = AssetWriter._quantize_rows(w)
        ref = ((w-baseline_q*baseline_s[:,None]).square()*h).sum(1)
        got = ((w-torch.round(w/s[:,None]).clamp(-8,7)*s[:,None]).square()*h).sum(1)
        self.assertTrue(torch.all(got <= ref+1e-5))
        self.assertLess(new,old)
        self.assertTrue(torch.isfinite(s).all() and (s>0).all())
        self.assertEqual(float(got[-1]),0.)

    def test_mapping(self):
        self.assertEqual(source_linear_names('model.layers.3.attention.o_proj'),
                         ['model.layers.3.attention.dense.weight'])
        self.assertEqual(source_linear_names('model.layers.2.attention.o_proj'),
                         ['model.layers.2.attention.o_proj.weight'])
        self.assertEqual(source_linear_names('model.layers.1.mlp.experts.7.gate_up'),
                         ['model.layers.1.mlp.experts.7.gate_proj.weight',
                          'model.layers.1.mlp.experts.7.up_proj.weight'])

    def test_accumulator_constraint(self):
        q=torch.full((1,2048),7.)
        self.assertFalse(bool(safe_rows(q,1)[0]))
        self.assertTrue(bool(safe_rows(q,4)[0]))
        with self.assertRaises(RuntimeError):
            AssetWriter._safe_splits(q.to(torch.int8),required_parts=1)
        self.assertEqual(AssetWriter._safe_splits(q.to(torch.int8),required_parts=4),4)
        torch.manual_seed(14)
        w=torch.randn(8,2048)
        s,_,_=choose_scales(w,torch.ones(2048),k_splits=1)
        self.assertTrue(safe_rows(torch.round(w/s[:,None]).clamp(-8,7),1).all())

    def test_packing_padding_correction_and_cache(self):
        with tempfile.TemporaryDirectory(prefix='ling3-calibration-') as directory:
            root=Path(directory); path=root/'scales.safetensors'
            name='model.layers.0.attention.qkvfgb'
            names=source_linear_names(name)
            torch.manual_seed(5)
            tensors=[torch.randn(16,32) for _ in names]
            scales={n:torch.full((16,),.2+i*.01) for i,n in enumerate(names)}
            save_file(scales,str(path))
            writer=AssetWriter(root/'assets-output',path)
            writer.linear(name,tensors,pad_n=True)
            folder=writer.assets/'linear'/name
            data=(folder/'weight-int4-kn.bin').read_bytes()
            packed=np.frombuffer(data,dtype=np.uint8)
            raw=np.stack((packed&15,packed>>4),axis=1).reshape(32,128).T.astype(np.int8)
            raw[raw>=8]-=16
            expected=torch.round(torch.cat(tensors)/torch.cat(list(scales.values()))[:,None]).clamp(-8,7)
            np.testing.assert_array_equal(raw[:96],expected.numpy())
            self.assertFalse(raw[96:].any())
            corr=np.fromfile(folder/'correction-i32.bin',dtype=np.int32)
            np.testing.assert_array_equal(corr,8*raw.astype(np.int32).sum(1))
            old_sha=json.loads((folder/'spec.json').read_text())['calibration_sha256']
            save_file({n:s*1.3 for n,s in scales.items()},str(path))
            writer=AssetWriter(root/'assets-output',path)
            writer.linear(name,tensors,pad_n=True)
            self.assertNotEqual(old_sha,json.loads((folder/'spec.json').read_text())['calibration_sha256'])
            self.assertNotEqual(data,(folder/'weight-int4-kn.bin').read_bytes())
            # A non-calibrated export must not silently reuse calibrated bytes either.
            writer=AssetWriter(root/'assets-output')
            writer.linear(name,tensors,pad_n=True)
            self.assertNotIn('calibration_sha256',json.loads((folder/'spec.json').read_text()))


if __name__=='__main__':
    unittest.main()