import torch
class FaultTolerance:
    """
    - 冗余副本存入尾数低 7 位 (Bit 0~6)
    - 指数偶校验位存入指数最高位 (Bit 7)
    - 尾数偶校验位存入尾数倒数第 8 位 (Bit 7) [紧邻低 7 位]
    - 恢复时指数最高位强制设为 0
    - 都错时逐位修复：相同保留，不同设为 1
    """

    @staticmethod
    def _calculate_parity_vectorized(x: torch.Tensor, num_bits: int = 7) -> torch.Tensor:
        """计算 x 的低 num_bits 位的偶校验（返回 0 或 1）"""
        parity = torch.zeros_like(x)
        for i in range(num_bits):
            parity ^= (x >> i) & 1
        return parity

    @staticmethod
    def protect_model(model: torch.nn.Module) -> None:
        """为模型所有 float32 参数添加容错保护"""
        count = 0
        for name, param in model.named_parameters():
            if param.dtype == torch.float32:
                param.data = FaultTolerance._add_protection_to_tensor(param.data)
                count += 1
        print(f"✅ 容错保护 (7-bit) 已应用到 {count} 个 float32 参数")

    @staticmethod
    def recover_model(model: torch.nn.Module) -> None:
        """从受保护状态恢复模型所有 float32 参数"""
        count = 0
        for param in model.parameters():
            if param.dtype == torch.float32:
                param.data = FaultTolerance._recover_tensor(param.data)
                count += 1
        print(f"✅ {count} 个 float32 参数已恢复")

    @staticmethod
    def _add_protection_to_tensor(tensor: torch.Tensor) -> torch.Tensor:
        """将 float32 张量转换为带容错保护的表示"""
        if tensor.dtype != torch.float32:
            return tensor

        original_shape = tensor.shape
        flat = tensor.view(-1)
        bits = flat.view(torch.int32)

        sign = (bits >> 31) & 0x1
        exponent = (bits >> 23) & 0xFF
        mantissa = bits & 0x7FFFFF

        # 1. 提取指数低 7 位 (bit0~6)
        exp_low7 = exponent & 0x7F  # 0b01111111

        # 2. 计算偶校验
        # 校验对象：指数低 7 位
        exp_parity = FaultTolerance._calculate_parity_vectorized(exp_low7, num_bits=7)
        # 校验对象：冗余副本 (此时还没存，但数据相同，所以校验值也相同)
        # 我们直接计算 exp_low7 的校验值作为两个位置的校验码
        red_parity = exp_parity.clone()

        # 3. 写入冗余：覆盖 mantissa 低 7 位 (Bit 0~6)
        # 清除低 7 位，然后写入
        new_mantissa = (mantissa & 0xFFFFFF80) | (exp_low7 & 0x7F)

        # 4. 写入尾数校验位：尾数倒数第 8 位 (Bit 7)
        # 清除 Bit 7，然后写入
        new_mantissa = (new_mantissa & 0xFFFFFF7F) | (red_parity << 7)

        # 5. 构造新指数：
        #   Bit 7 (最高位) = exp_parity
        #   Bit 6~0       = 原始低 7 位
        new_exponent = (exp_parity << 7) | exp_low7

        protected_bits = (sign << 31) | (new_exponent << 23) | new_mantissa

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

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

        original_shape = tensor.shape
        flat = tensor.view(-1)
        bits = flat.view(torch.int32)

        sign = (bits >> 31) & 0x1
        exponent = (bits >> 23) & 0xFF
        mantissa = bits & 0x7FFFFF

        # --- 读取数据 ---

        # 从 mantissa 低 7 位读取冗余副本 (Bit 0~6)
        redundant_low7 = mantissa & 0x7F
        
        # 从 mantissa 倒数第 8 位读取冗余校验位 (Bit 7)
        red_parity_stored = (mantissa >> 7) & 1

        # 从 exponent 中提取：
        stored_exp_low7 = exponent & 0x7F          # 指数低 7 位 (可能已损坏)
        exp_parity_stored = (exponent >> 7) & 1    # 指数最高位存储的校验位

        # --- 校验逻辑 ---

        # 重新计算期望的校验值
        # 1. 基于存储的指数低 7 位计算
        expected_exp_parity = FaultTolerance._calculate_parity_vectorized(stored_exp_low7, num_bits=7)
        # 2. 基于存储的冗余副本计算
        expected_red_parity = FaultTolerance._calculate_parity_vectorized(redundant_low7, num_bits=7)

        # 比对校验位
        exp_ok = (exp_parity_stored == expected_exp_parity)
        red_ok = (red_parity_stored == expected_red_parity)

        corrected_low7 = stored_exp_low7.clone()
        
        # 情况 1: 两者都对 (或至少有一个对且数据一致)，优先取原始值 (stored_exp_low7)
        # 情况 2: 原始错，冗余对 → 使用冗余
        mask2 = (~exp_ok) & red_ok
        if mask2.any():
            corrected_low7[mask2] = redundant_low7[mask2]

        # 情况 3: 两者都错 → 逐位修复
        mask3 = (~exp_ok) & (~red_ok)
        if mask3.any():
            exp_vals = stored_exp_low7[mask3]
            red_vals = redundant_low7[mask3]
            corrected_vals = torch.zeros_like(exp_vals)

            for bit in range(7):
                e_bit = (exp_vals >> bit) & 1
                r_bit = (red_vals >> bit) & 1
                # 相同则保留该值，不同则设为 1
                agreed_bit = torch.where(e_bit == r_bit, e_bit, torch.tensor(1, dtype=exp_vals.dtype, device=exp_vals.device))
                corrected_vals |= (agreed_bit << bit)

            corrected_low7[mask3] = corrected_vals

        # --- 重组指数 ---
        # 恢复指数最高位为 0 (Bit 7 = 0)
        # 指数结构：0 (Bit 7) | corrected_low7 (Bit 6~0)
        recovered_exponent = (0 << 7) | (corrected_low7 & 0x7F)

        # --- 清理尾数 ---
        # 清除用于存储冗余的低 7 位 (Bit 0~6) 和校验位 (Bit 7)
        # 掩码：清除低 8 位 (0xFF = 11111111)
        clean_mantissa = mantissa & 0xFFFFFF00

        recovered_bits = (sign << 31) | (recovered_exponent << 23) | clean_mantissa

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