import struct
import torch
import numpy as np

class HGSFormat:
    MAGIC = b'HGS1'

    @staticmethod
    def quantize_tensor(tensor: torch.Tensor, bits=8):
        tensor_cpu = tensor.detach().cpu()
        min_val = tensor_cpu.min()
        max_val = tensor_cpu.max()
        scale = (max_val - min_val) / (2**bits - 1)
        if scale == 0:
            scale = 1e-8
        q_tensor = ((tensor_cpu - min_val) / scale).round().clamp(0, 2**bits - 1).to(torch.uint8)
        return q_tensor.numpy(), float(min_val), float(scale)

    @staticmethod
    def dequantize_tensor(q_data, min_val, scale):
        return torch.tensor(q_data, dtype=torch.float32) * scale + min_val

    @staticmethod
    def save_hgs(model_state_dict, filepath):
        layers = list(model_state_dict.items())
        with open(filepath, 'wb') as f:
            f.write(HGSFormat.MAGIC)
            f.write(struct.pack('<I', len(layers)))

            index_pos = f.tell()
            index_data = []
            for name, tensor in layers:
                name_bytes = name.encode('utf-8')
                q_data, min_val, scale = HGSFormat.quantize_tensor(tensor, bits=8)
                index_data.append((name_bytes, tensor.shape, tensor.numel(), min_val, scale, 0, q_data.nbytes))
                f.write(struct.pack('<I', len(name_bytes)))
                f.write(name_bytes)
                f.write(struct.pack('<I', len(tensor.shape)))
                for dim in tensor.shape:
                    f.write(struct.pack('<Q', dim))
                f.write(struct.pack('<f', min_val))
                f.write(struct.pack('<f', scale))
                f.write(struct.pack('<Q', tensor.numel()))
                f.write(struct.pack('<Q', q_data.nbytes))
                f.write(struct.pack('<Q', 0))

            offsets = []
            for _, _, _, _, _, _, size_bytes in index_data:
                offsets.append(f.tell())
                q_data = index_data.pop(0)[6]

    @staticmethod
    def save_hgs(model_state_dict, filepath):
        layers = list(model_state_dict.items())
        with open(filepath, 'wb') as f:
            f.write(HGSFormat.MAGIC)
            f.write(struct.pack('<I', len(layers)))

            index_pos = f.tell()
            index_data = []
            quantized_datas = []
            for name, tensor in layers:
                name_bytes = name.encode('utf-8')
                q_data, min_val, scale = HGSFormat.quantize_tensor(tensor, bits=8)
                index_data.append({
                    'name_bytes': name_bytes,
                    'shape': tensor.shape,
                    'numel': tensor.numel(),
                    'min_val': min_val,
                    'scale': scale,
                    'q_data_len': q_data.nbytes,
                    'offset': 0
                })
                quantized_datas.append(q_data)

            for entry in index_data:
                f.write(struct.pack('<I', len(entry['name_bytes'])))
                f.write(entry['name_bytes'])
                f.write(struct.pack('<I', len(entry['shape'])))
                for dim in entry['shape']:
                    f.write(struct.pack('<Q', dim))
                f.write(struct.pack('<f', entry['min_val']))
                f.write(struct.pack('<f', entry['scale']))
                f.write(struct.pack('<Q', entry['numel']))
                f.write(struct.pack('<Q', entry['q_data_len']))
                f.write(struct.pack('<Q', 0))

            for i, q_data in enumerate(quantized_datas):
                index_data[i]['offset'] = f.tell()
                f.write(q_data.tobytes())
            f.seek(index_pos)
            for entry in index_data:
                f.seek(f.tell() + 4 + len(entry['name_bytes']) + 4 + 8 * len(entry['shape']) + 4 + 4 + 8 + 8)
                f.write(struct.pack('<Q', entry['offset']))

    @staticmethod
    def load_hgs(filepath, device='cpu'):
        state_dict = {}
        with open(filepath, 'rb') as f:
            magic = f.read(4)
            if magic != HGSFormat.MAGIC:
                raise RuntimeError("Неверный формат файла")

            num_layers = struct.unpack('<I', f.read(4))[0]

            index = []
            for _ in range(num_layers):
                name_len = struct.unpack('<I', f.read(4))[0]
                name = f.read(name_len).decode('utf-8')
                shape_len = struct.unpack('<I', f.read(4))[0]
                shape = tuple(struct.unpack('<Q', f.read(8))[0] for _ in range(shape_len))
                min_val = struct.unpack('<f', f.read(4))[0]
                scale = struct.unpack('<f', f.read(4))[0]
                numel = struct.unpack('<Q', f.read(8))[0]
                q_data_len = struct.unpack('<Q', f.read(8))[0]
                offset = struct.unpack('<Q', f.read(8))[0]
                index.append({
                    'name': name,
                    'shape': shape,
                    'min_val': min_val,
                    'scale': scale,
                    'numel': numel,
                    'q_data_len': q_data_len,
                    'offset': offset
                })

            for entry in index:
                f.seek(entry['offset'])
                q_data_bytes = f.read(entry['q_data_len'])
                q_data = np.frombuffer(q_data_bytes, dtype=np.uint8)
                tensor = HGSFormat.dequantize_tensor(q_data, entry['min_val'], entry['scale'])
                tensor = tensor.reshape(entry['shape']).to(device)
                state_dict[entry['name']] = tensor

        return state_dict