"""
Inference code for mod5-verified threshold network.

This network computes MOD-5 (Hamming weight mod 5) on 8-bit binary inputs.
"""

import torch
import torch.nn as nn
from safetensors.torch import load_file


def heaviside(x):
    """Heaviside step function: 1 if x >= 0, else 0."""
    return (x >= 0).float()


class Mod5Network(nn.Module):
    """
    Verified threshold network for MOD-5 computation.

    Architecture: 8 -> 9 -> 4 -> 5
    - Layer 1: Thermometer encoding (9 neurons detect HW >= k)
    - Layer 2: MOD-5 detection using (1,1,1,1,-4) weight pattern
    - Output: 5-class classification
    """

    def __init__(self):
        super().__init__()
        self.layer1 = nn.Linear(8, 9)
        self.layer2 = nn.Linear(9, 4)
        self.output = nn.Linear(4, 5)

    def forward(self, x):
        """Forward pass with Heaviside activation."""
        x = x.float()
        x = heaviside(self.layer1(x))
        x = heaviside(self.layer2(x))
        x = self.output(x)
        return x

    def predict(self, x):
        """Get predicted class (0, 1, 2, 3, or 4)."""
        return self.forward(x).argmax(dim=-1)

    @classmethod
    def from_safetensors(cls, path):
        """Load model from safetensors file."""
        model = cls()
        weights = load_file(path)

        model.layer1.weight.data = weights['layer1.weight']
        model.layer1.bias.data = weights['layer1.bias']
        model.layer2.weight.data = weights['layer2.weight']
        model.layer2.bias.data = weights['layer2.bias']
        model.output.weight.data = weights['output.weight']
        model.output.bias.data = weights['output.bias']

        return model


def mod5_reference(x):
    """Reference implementation: Hamming weight mod 5."""
    return (x.sum(dim=-1) % 5).long()


def verify(model, verbose=True):
    """Verify model on all 256 inputs."""
    inputs = torch.zeros(256, 8)
    for i in range(256):
        for j in range(8):
            inputs[i, j] = (i >> j) & 1

    targets = mod5_reference(inputs)
    predictions = model.predict(inputs)

    correct = (predictions == targets).sum().item()

    if verbose:
        print(f"Verification: {correct}/256 ({100*correct/256:.1f}%)")

        if correct < 256:
            errors = (predictions != targets).nonzero(as_tuple=True)[0]
            print(f"Errors at indices: {errors[:10].tolist()}")

    return correct == 256


def demo():
    """Demonstration of MOD-5 computation."""
    print("Loading mod5-verified model...")
    model = Mod5Network.from_safetensors('model.safetensors')

    print("\nVerifying on all 256 inputs...")
    verify(model)

    print("\nExample predictions:")
    test_cases = [
        [0, 0, 0, 0, 0, 0, 0, 0],
        [1, 0, 0, 0, 0, 0, 0, 0],
        [1, 1, 0, 0, 0, 0, 0, 0],
        [1, 1, 1, 0, 0, 0, 0, 0],
        [1, 1, 1, 1, 0, 0, 0, 0],
        [1, 1, 1, 1, 1, 0, 0, 0],
        [1, 1, 1, 1, 1, 1, 0, 0],
        [1, 1, 1, 1, 1, 1, 1, 0],
        [1, 1, 1, 1, 1, 1, 1, 1],
    ]

    for bits in test_cases:
        x = torch.tensor([bits], dtype=torch.float32)
        hw = sum(bits)
        pred = model.predict(x).item()
        expected = hw % 5
        status = "OK" if pred == expected else "ERROR"
        print(f"  {bits} -> HW={hw}, pred={pred}, expected={expected} [{status}]")


if __name__ == '__main__':
    demo()