import torch

from norse.torch.module.snn import SNNCell
from norse.torch.functional.iaf import (
    iaf_feed_forward_step,
    IAFFeedForwardState,
    IAFParameters,
)


class IAFCell(SNNCell):
    def __init__(self, p: IAFParameters = IAFParameters(), dt: float = 0.001):
        r"""Feedforward step of an integrate-and-fire neuron, computing a single step

        .. math::
            \dot{v} = v

        together with the jump condition

        .. math::
            z = \Theta(v - v_{\text{th}})

        and transition equation

        .. math::
            v = (1-z) v + z v_{\text{reset}}

        Parameters:
            p (IAFParameters): parameters of a leaky integrate and fire neuron
            dt (float): Integration timestep to use (unused, but added for compatibility)
        """
        super().__init__(iaf_feed_forward_step, self.initial_state, p, dt)

    def initial_state(self, input_tensor: torch.Tensor) -> IAFFeedForwardState:
        state = IAFFeedForwardState(
            v=torch.full(
                input_tensor.shape,
                torch.as_tensor(self.p.v_reset.detach()),
                device=input_tensor.device,
                dtype=torch.float32,
            ),
        )
        state.v.requires_grad = True
        return state
