import torch
from safetensors.torch import save_file

weights = {}

# T Flip-Flop (Toggle)
# Inputs: T (toggle), Q_prev (previous state)
# Outputs: Q, Qn
#
# When T=1: Q toggles (Q = NOT Q_prev)
# When T=0: Q holds (Q = Q_prev)
#
# Q = T XOR Q_prev = (T AND NOT_Q_prev) OR (NOT_T AND Q_prev)
# Qn = NOT Q
#
# XOR is not linearly separable, requires 2 layers:
# XOR(a,b) = AND(OR(a,b), NAND(a,b))

# Layer 1: OR and NAND for Q, and for Qn
# OR(T, Q_prev)
weights['or.weight'] = torch.tensor([[1.0, 1.0]], dtype=torch.float32)
weights['or.bias'] = torch.tensor([-1.0], dtype=torch.float32)

# NAND(T, Q_prev)
weights['nand.weight'] = torch.tensor([[-1.0, -1.0]], dtype=torch.float32)
weights['nand.bias'] = torch.tensor([1.0], dtype=torch.float32)

# For Qn = NOT(T XOR Q_prev) = XNOR(T, Q_prev) = (NOT_T AND NOT_Q_prev) OR (T AND Q_prev)
# XNOR = AND(OR(NOT_T, NOT_Q_prev), NAND(NOT_T, NOT_Q_prev))
#      = AND(NAND(T, Q_prev), OR(T, Q_prev))... wait, that's not right
# Let me think: XNOR(a,b) = NOT XOR(a,b) = NOR(a,b) OR AND(a,b)
# Actually: XNOR = (a AND b) OR (NOT_a AND NOT_b)
# Which in threshold: we need NOR and AND, then OR them
# NOR(T, Q_prev): fires when both are 0
weights['nor.weight'] = torch.tensor([[-1.0, -1.0]], dtype=torch.float32)
weights['nor.bias'] = torch.tensor([0.0], dtype=torch.float32)

# AND(T, Q_prev): fires when both are 1
weights['and.weight'] = torch.tensor([[1.0, 1.0]], dtype=torch.float32)
weights['and.bias'] = torch.tensor([-2.0], dtype=torch.float32)

# Layer 2: Combine
# Q = AND(OR, NAND) = XOR
weights['q.weight'] = torch.tensor([[1.0, 1.0]], dtype=torch.float32)
weights['q.bias'] = torch.tensor([-2.0], dtype=torch.float32)

# Qn = OR(NOR, AND) = XNOR
weights['qn.weight'] = torch.tensor([[1.0, 1.0]], dtype=torch.float32)
weights['qn.bias'] = torch.tensor([-1.0], dtype=torch.float32)

save_file(weights, 'model.safetensors')

def t_flipflop(t, q_prev):
    inp = torch.tensor([float(t), float(q_prev)])

    # Layer 1
    or_out = int((inp @ weights['or.weight'].T + weights['or.bias'] >= 0).item())
    nand_out = int((inp @ weights['nand.weight'].T + weights['nand.bias'] >= 0).item())
    nor_out = int((inp @ weights['nor.weight'].T + weights['nor.bias'] >= 0).item())
    and_out = int((inp @ weights['and.weight'].T + weights['and.bias'] >= 0).item())

    # Layer 2
    l1_q = torch.tensor([float(or_out), float(nand_out)])
    q = int((l1_q @ weights['q.weight'].T + weights['q.bias'] >= 0).item())

    l1_qn = torch.tensor([float(nor_out), float(and_out)])
    qn = int((l1_qn @ weights['qn.weight'].T + weights['qn.bias'] >= 0).item())

    return q, qn

def reference_t_ff(t, q_prev):
    if t == 1:
        return 1 - q_prev, q_prev
    else:
        return q_prev, 1 - q_prev

print("Verifying T Flip-Flop...")
errors = 0
for t in range(2):
    for q_prev in range(2):
        result = t_flipflop(t, q_prev)
        expected = reference_t_ff(t, q_prev)
        if result != expected:
            errors += 1
            print(f"ERROR: T={t}, Q_prev={q_prev} -> {result}, expected {expected}")

if errors == 0:
    print("All 4 test cases passed!")
else:
    print(f"FAILED: {errors} errors")

print("\nTruth Table:")
print("T Q_prev | Q Qn | Mode")
print("-" * 25)
for t in range(2):
    for q_prev in range(2):
        q, qn = t_flipflop(t, q_prev)
        mode = "Toggle" if t == 1 else "Hold"
        print(f"{t}   {q_prev}   | {q}  {qn} | {mode}")

mag = sum(t.abs().sum().item() for t in weights.values())
print(f"\nMagnitude: {mag:.0f}")
print(f"Parameters: {sum(t.numel() for t in weights.values())}")
print(f"Neurons: {len([k for k in weights.keys() if 'weight' in k])}")