"""Read safetensors BF16 directly; never instantiate a model or use pickle."""
import json
import math
import mmap
import re
import struct

import numpy as np

DTYPES = {'BF16': ('<u2', 2), 'F16': ('<f2', 2), 'F32': ('<f4', 4), 'F64': ('<f8', 8)}


def decode(raw, dtype):
    if dtype == 'BF16':
        return (np.asarray(raw).astype(np.uint32) << 16).view(np.float32)
    return np.asarray(raw, dtype=np.float32)


class Tensor:
    def __init__(self, file, base, name, meta):
        self.name, self.shape, self.dtype = name, tuple(meta['shape']), meta['dtype']
        self.numel = math.prod(self.shape)
        self.bytes = meta['data_offsets'][1] - meta['data_offsets'][0]
        code, itemsize = DTYPES[self.dtype]
        if self.bytes != self.numel * itemsize:
            raise ValueError(f'Invalid byte length: {name}')
        self.raw = np.memmap(file, mode='r', dtype=code,
                             offset=base + meta['data_offsets'][0], shape=(self.numel,))
        self.role = re.sub(r'\.(?:layers|blocks)\.\d+', '.layers.*', name)
        self.role = re.sub(r'\.experts\.\d+', '.experts.*', self.role)
        m = re.search(r'\.(?:layers|blocks)\.(\d+)', name)
        self.layer = int(m[1]) if m else -1

    def read(self, selector=slice(None)):
        return decode(self.raw[selector], self.dtype)

    def chunked(self, size=262144):
        for start in range(0, self.numel, size):
            yield self.read(slice(start, min(start + size, self.numel)))

    def evict(self):
        if hasattr(self.raw._mmap, 'madvise'):
            self.raw._mmap.madvise(mmap.MADV_DONTNEED)


def load_tensors(directory):
    tensors = {}
    for file in sorted(directory.glob('*.safetensors')):
        with file.open('rb') as f:
            size = struct.unpack('<Q', f.read(8))[0]
            if size > 100_000_000:
                raise ValueError('Oversized safetensors header')
            header = json.loads(f.read(size))
        for name, meta in header.items():
            if name == '__metadata__':
                continue
            if name in tensors:
                raise ValueError(f'Duplicate tensor name: {name}')
            tensors[name] = Tensor(file, 8 + size, name, meta)
    if not tensors:
        raise ValueError('No safetensors found')
    return tensors


def metrics(x, y):
    """x is member; y is representative. Float64 accumulation."""
    x, y = x.astype(np.float64), y.astype(np.float64)
    xx, yy, xy = float(x @ x), float(y @ y), float(x @ y)
    delta = float((x - y) @ (x - y))
    return finish(xx, yy, xy, delta, float(x.sum()), float(y.sum()), x.size)


def finish(xx, yy, xy, delta, sx, sy, n):
    tiny = 1e-30
    vx, vy = max(0., xx - sx * sx / n), max(0., yy - sy * sy / n)
    return {'cosine': xy / np.sqrt(xx * yy) if xx * yy > tiny else None,
            'pearson': (xy - sx * sy / n) / np.sqrt(vx * vy) if vx * vy > tiny else None,
            'error_member_to_rep': np.sqrt(delta / max(xx, tiny)),
            'error_rep_to_member': np.sqrt(delta / max(yy, tiny)),
            'symmetric_error': np.sqrt(2 * delta / max(xx + yy, tiny)),
            'scaled_error': np.sqrt(max(0., xx - xy * xy / max(yy, tiny)) / max(xx, tiny)),
            'scale': xy / max(yy, tiny)}