import numpy as np

def sigmoid(z):
    return 1 / (1 + np.exp(-z))

# ---- state and input (matches the worked example in the lecture) ----
h_prev, C_prev, x = 0.2, 0.3, 0.5

# ---- gate weights (scalar, for a single-unit illustrative LSTM cell) ----
Wfh, Wfx, bf = 0.5, 0.6, 0.1     # forget gate
Wih, Wix, bi = 0.4, 0.3, 0.0     # input gate
WCh, WCx, bC = 0.3, 0.5, 0.2     # candidate memory
Woh, Wox, bo = 0.6, 0.4, 0.1     # output gate

# ---- forward pass through one LSTM timestep, gate by gate ----
f_t = sigmoid(Wfh * h_prev + Wfx * x + bf)
i_t = sigmoid(Wih * h_prev + Wix * x + bi)
C_tilde = np.tanh(WCh * h_prev + WCx * x + bC)
C_t = f_t * C_prev + i_t * C_tilde
o_t = sigmoid(Woh * h_prev + Wox * x + bo)
h_t = o_t * np.tanh(C_t)

print(f"f_t (forget) = {f_t:.4f}")
print(f"i_t (input)  = {i_t:.4f}")
print(f"C~_t (candidate) = {C_tilde:.4f}")
print(f"C_t (cell state)  = {C_t:.4f}")
print(f"o_t (output) = {o_t:.4f}")
print(f"h_t (hidden state) = {h_t:.4f}")

expected = dict(f_t=0.6225, i_t=0.5573, C_tilde=0.4699, C_t=0.4487, o_t=0.6035, h_t=0.2540)
computed = dict(f_t=f_t, i_t=i_t, C_tilde=C_tilde, C_t=C_t, o_t=o_t, h_t=h_t)
for k in expected:
    assert abs(computed[k] - expected[k]) < 1e-3, f"{k} mismatch!"
print("\nAll values match the lecture's hand-worked example.")