# SPDX-License-Identifier: MIT
"""YS-EN-0006 · check Equation 1 and the parameter counts on the networks themselves.

  python check_rf.py

Feeds a 400-hour input through an (untrained) 8-layer dilated stack and finds which input
hours the gradient of its last position touches: the receptive field.  Then counts the
parameters of the convolution stack and of each model used in the note.
"""
from __future__ import annotations

import torch

import fc

for layers in (4, 8):
    m = fc.make_model("tcn", 1, False, n_layers=layers, head="report")
    x = torch.randn(1, 1, 400, requires_grad=True)
    m.net(x)[0, :, -1].sum().backward()
    hit = (x.grad[0, 0] != 0).nonzero().flatten()
    print(f"{layers} layers: last position depends on {int(hit.max() - hit.min() + 1)} input hours "
          f"(Equation 1 with k = 2: {2 ** layers})")

m = fc.make_model("tcn", 1, False, head="report")
conv = sum(p.numel() for n, p in m.named_parameters() if n.startswith("net."))
print(f"conv stack, 1 input channel, 32 filters: {conv} parameters")
for name, args in [("TCN, 2023 design", ("tcn", 1, False, "report")),
                   ("TCN, 2023 design + calendar", ("tcn", 2, True, "report")),
                   ("TCN + calendar (direct head)", ("tcn", 2, True, "direct")),
                   ("GRU + calendar", ("gru", 2, True, "direct")),
                   ("Linear + calendar", ("linear", 2, True, "direct"))]:
    kind, ch, cal, head = args
    print(f"{name}: {fc.n_params(fc.make_model(kind, ch, cal, head=head))} parameters")