"""Week 07 lecture replay. Requires PyTorch; runs on CPU, without data downloads.
Run: python3 scalar-rnn.py
Change the three weights or inputs below, predict the effect, then rerun.
"""
import torch


def worked_example():
    x = [1.0, 0.5, -0.5]
    y = 0.6
    w_x, w_h, w_o = [
        torch.tensor(v, dtype=torch.float64, requires_grad=True)
        for v in [0.5, 0.8, 1.0]
    ]
    h = torch.tensor(0.0, dtype=torch.float64)
    states, preactivations = [], []
    for x_t in x:
        z = w_x * x_t + w_h * h
        h = torch.tanh(z)
        z.retain_grad()
        h.retain_grad()
        preactivations.append(z)
        states.append(h)
    prediction = w_o * h
    loss = 0.5 * (prediction - y) ** 2
    loss.backward()
    print(' t       z_t       h_t    dL/dh_t    dL/dz_t')
    for t, (z, h) in enumerate(zip(preactivations, states)):
        print(f'{t:2d} {z.item():9.6f} {h.item():9.6f}'
              f' {h.grad.item():10.6f} {z.grad.item():10.6f}')
    print(f'prediction={prediction.item():.6f}, loss={loss.item():.6f}')
    print('dL/d(w_x,w_h,w_o):', [round(w.grad.item(), 6)
                                       for w in (w_x, w_h, w_o)])
    return preactivations, states, (w_x, w_h, w_o), loss


def state_sensitivity(w_h=1.2, input_value=0.0, steps=30):
    """Differentiate the final state with respect to the initial state.

    This is a state sensitivity, not a training-loss or weight gradient.
    w_x=1, tanh activation, constant inputs. No optimizer is involved.
    """
    initial = torch.tensor(0.0, dtype=torch.float64, requires_grad=True)
    h = initial
    for _ in range(steps):
        h = torch.tanh(input_value + w_h * h)
    sensitivity, = torch.autograd.grad(h, initial)
    return h.item(), sensitivity.item()


if __name__ == '__main__':
    worked_example()
    print('\nSame recurrent weight, different input trajectory (30 steps):')
    for input_value in [0.0, 1.0]:
        state, derivative = state_sensitivity(input_value=input_value)
        print(f'x_t={input_value:.1f}: final h={state:.6f}, '
              f'dh_final/dh_initial={derivative:.3e}')