import torch
from safetensors.torch import save_file

weights = {}

# D Flip-Flop (level-sensitive / D Latch behavior)
# Inputs: D, CLK, Q_prev (3 inputs)
# Outputs: Q, Qn (2 outputs)
#
# When CLK=1: Q = D (transparent)
# When CLK=0: Q = Q_prev (hold)
#
# Q = (CLK AND D) OR (NOT CLK AND Q_prev)
# Q = MUX(Q_prev, D, CLK)

# Input indices: D=0, CLK=1, Q_prev=2

# CLK AND D
weights['clk_and_d.weight'] = torch.tensor([[1.0, 1.0, 0.0]], dtype=torch.float32)
weights['clk_and_d.bias'] = torch.tensor([-2.0], dtype=torch.float32)

# NOT CLK
weights['not_clk.weight'] = torch.tensor([[0.0, -1.0, 0.0]], dtype=torch.float32)
weights['not_clk.bias'] = torch.tensor([0.0], dtype=torch.float32)

# Q_prev passthrough
weights['q_prev.weight'] = torch.tensor([[0.0, 0.0, 1.0]], dtype=torch.float32)
weights['q_prev.bias'] = torch.tensor([0.0], dtype=torch.float32)

# NOT CLK AND Q_prev
weights['hold.weight'] = torch.tensor([[1.0, 1.0]], dtype=torch.float32)
weights['hold.bias'] = torch.tensor([-2.0], dtype=torch.float32)

# Q = (CLK AND D) OR hold
weights['q.weight'] = torch.tensor([[1.0, 1.0]], dtype=torch.float32)
weights['q.bias'] = torch.tensor([-1.0], dtype=torch.float32)

# Qn = NOT Q (computed separately for clarity)
# Qn = (CLK AND NOT D) OR (NOT CLK AND NOT Q_prev)

# NOT D
weights['not_d.weight'] = torch.tensor([[-1.0, 0.0, 0.0]], dtype=torch.float32)
weights['not_d.bias'] = torch.tensor([0.0], dtype=torch.float32)

# CLK AND NOT D
weights['clk_and_notd.weight'] = torch.tensor([[1.0, 1.0]], dtype=torch.float32)
weights['clk_and_notd.bias'] = torch.tensor([-2.0], dtype=torch.float32)

# NOT Q_prev
weights['not_qprev.weight'] = torch.tensor([[0.0, 0.0, -1.0]], dtype=torch.float32)
weights['not_qprev.bias'] = torch.tensor([0.0], dtype=torch.float32)

# NOT CLK AND NOT Q_prev
weights['hold_n.weight'] = torch.tensor([[1.0, 1.0]], dtype=torch.float32)
weights['hold_n.bias'] = torch.tensor([-2.0], dtype=torch.float32)

# Qn = (CLK AND NOT D) OR hold_n
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 d_ff(d, clk, q_prev):
    if clk == 1:
        return d, 1 - d
    else:
        return q_prev, 1 - q_prev

def eval_circuit(d, clk, q_prev):
    clk_and_d = clk and d
    not_clk = 1 - clk
    hold = not_clk and q_prev
    q = clk_and_d or hold
    not_d = 1 - d
    clk_and_notd = clk and not_d
    not_qprev = 1 - q_prev
    hold_n = not_clk and not_qprev
    qn = clk_and_notd or hold_n
    return q, qn

print("Verifying D Flip-Flop...")
errors = 0
for d in range(2):
    for clk in range(2):
        for q_prev in range(2):
            expected_q, expected_qn = d_ff(d, clk, q_prev)
            calc_q, calc_qn = eval_circuit(d, clk, q_prev)
            if expected_q != calc_q or expected_qn != calc_qn:
                errors += 1
                print(f"ERROR: D={d}, CLK={clk}, Q_prev={q_prev}")

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

print("\nTruth Table:")
print("D CLK Q_prev | Q Qn | Mode")
print("-" * 32)
for d in range(2):
    for clk in range(2):
        for q_prev in range(2):
            q, qn = d_ff(d, clk, q_prev)
            mode = "Transparent" if clk else "Hold"
            print(f"{d}  {clk}    {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())}")