import torch
class ZMORP:
    """
    超显存优化版：GPU 加速的 Float16 容错保护
    """

    @staticmethod
    def _calculate_parity_vectorized(x: torch.Tensor, num_bits=5) -> torch.Tensor:
        """向量化计算偶校验（针对5位指数）"""
        parity = torch.zeros_like(x)
        for i in range(num_bits):
            parity ^= ((x >> i) & 1)
        return parity

    @staticmethod
    def protect_model(model: nn.Module) -> None:
        """为模型所有 float16 参数添加容错保护"""
        count = 0
        for name, param in model.named_parameters():
            if param.dtype == torch.float16:
                param.data = ZMORP._add_protection_to_tensor(param.data)
                count += 1
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
        print(f"✅ 容错保护已应用到 {count} 个 float16 参数（超显存优化版）")

    @staticmethod
    def recover_model(model: nn.Module) -> None:
        """恢复模型 float16 参数"""
        count = 0
        for param in model.parameters():
            if param.dtype == torch.float16:
                param.data = ZMORP._recover_tensor(param.data)
                count += 1
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
        print(f"✅ {count} 个 float16 参数已恢复（超显存优化版）")

    @staticmethod
    def _add_protection_to_tensor(tensor: torch.Tensor) -> torch.Tensor:
        if tensor.dtype != torch.float16:
            return tensor

        original_shape = tensor.shape
        flat_tensor = tensor.view(-1)
        bits = flat_tensor.view(torch.int16)  # float16 → int16

        sign = (bits >> 15) & 0x1
        exponent = (bits >> 10) & 0x1F        # 5 bits
        mantissa = bits & 0x3FF               # 10 bits

        redundant_exp = exponent
        exp_parity = ZMORP._calculate_parity_vectorized(exponent, num_bits=5)
        red_parity = ZMORP._calculate_parity_vectorized(redundant_exp, num_bits=5)

        protection_bits = (red_parity << 6) | (redundant_exp << 1) | exp_parity  

        mantissa_high = mantissa & 0x380
        new_mantissa = mantissa_high | (protection_bits & 0x7F) 

        protected_bits = (sign << 15) | (exponent << 10) | new_mantissa

        del bits, sign, exponent, mantissa, redundant_exp, exp_parity, red_parity, protection_bits

        return protected_bits.view(torch.float16).view(original_shape)

    @staticmethod
    def _recover_tensor(tensor: torch.Tensor) -> torch.Tensor:
        """从受保护的 float16 张量中恢复原始值"""
        if tensor.dtype != torch.float16:
            return tensor

        original_shape = tensor.shape
        flat_tensor = tensor.view(-1)
        bits = flat_tensor.view(torch.int16)

        sign = (bits >> 15) & 0x1
        exponent = (bits >> 10) & 0x1F
        mantissa = bits & 0x3FF

        # 从 mantissa 中提取保护信息
        protected_part = mantissa  # 全部10位都用于保护
        redundant_exp = (protected_part >> 1) & 0x1F      # 5 bits
        exp_parity = protected_part & 0x1                 # LSB
        red_parity = (protected_part >> 6) & 0x1          # MSB of mantissa

        expected_exp_parity = ZMORP._calculate_parity_vectorized(exponent, num_bits=5)
        expected_red_parity = ZMORP._calculate_parity_vectorized(redundant_exp, num_bits=5)

        exp_ok = (exp_parity == expected_exp_parity)
        red_ok = (red_parity == expected_red_parity)

        corrected_exp = exponent.clone()

        # 情况2: 指数错，冗余对 → 用冗余
        mask2 = (~exp_ok) & red_ok
        if mask2.any():
            corrected_exp[mask2] = redundant_exp[mask2]

        # 情况3: 都错 → 逐位投票（5位）
        mask3 = (~exp_ok) & (~red_ok)
        if mask3.any():
            exp_vals = exponent[mask3]
            red_vals = redundant_exp[mask3]
            corrected_vals = torch.zeros_like(exp_vals)

            for bit in range(5):  # 只有5位指数
                exp_bit = (exp_vals >> bit) & 1
                red_bit = (red_vals >> bit) & 1
                agreed_bit = torch.where(exp_bit == red_bit, exp_bit,
                                         torch.tensor(0, dtype=exp_vals.dtype, device=exp_vals.device))
                corrected_vals += (agreed_bit << bit)

            corrected_exp[mask3] = corrected_vals

        # 恢复时，清除保护位（尾数设为0，或保留原始？这里清零）
        clean_mantissa = mantissa & 0x380  # 或者你想保留原始尾数？但已被覆盖

        recovered_bits = (sign << 15) | (corrected_exp << 10) | clean_mantissa

        del bits, sign, exponent, mantissa, redundant_exp, exp_parity, red_parity, \
            expected_exp_parity, expected_red_parity, exp_ok, red_ok, mask2, mask3, \
            corrected_exp, clean_mantissa

        return recovered_bits.view(torch.float16).view(original_shape)