import torch
import torch.nn as nn
import torchao

def bitnet_b158_quantize(tensor):
    """
    Konvertiert ein Gewicht-Tensor in das ternäre 1.58-Bit Format {-1, 0, +1}
    inklusive der notwendigen Per-Channel Skalierung.
    """
    # 1. Berechne den durchschnittlichen Absolutwert pro Kanal (Zeile)
    scale = tensor.abs().mean(dim=-1, keepdim=True).clamp(min=1e-5)
    
    # 2. Skaliere den Tensor und runde auf die nächste ganze Zahl
    quantized = torch.round(tensor / scale)
    
    # 3. Zwinge die Werte strikt in den Bereich von -1 bis +1
    quantized = torch.clamp(quantized, min=-1.0, max=1.0)
    
    return quantized, scale

class BitLinear158(nn.Module):
    """
    Ein Ersatz für nn.Linear, der die 1.58-Bit Ternary-Inferenz ausführt.
    """
    def __init__(self, in_features, out_features, bias=False):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.register_buffer("weight_158", torch.zeros((out_features, in_features), dtype=torch.int8))
        self.register_buffer("scale", torch.zeros((out_features, 1), dtype=torch.bfloat16))
        
    @torch.no_grad()
    def from_float(self, float_layer):
        # Transformiere die originalen Gewichte
        q_w, scale = bitnet_b158_quantize(float_layer.weight.data)
        self.weight_158.copy_(q_w.to(torch.int8))
        self.scale.copy_(scale.to(torch.bfloat16))
        return self

    def forward(self, x):
        # Die Magie: x wird mit den Integer-Gewichten (-1, 0, 1) verarbeitet
        # Auf der Hardwarebene entspricht dies reinen Additionen/Subtraktionen
        out = nn.functional.linear(x, self.weight_158.to(x.dtype))
        return out * self.scale