#%%
import numpy as np
import matplotlib.pyplot as plt

# encode a message using poisson rate coding
def encode_message(symbol: str,
                   signal_rate: float,
                   duration: int,
                   **kwargs):
  codebook = kwargs.get("codebook",None)
  noise_rate = kwargs.get("noise_rate", 0)

  if codebook is None: raise TypeError("Invalid code book.")


  rate = signal_rate + noise_rate

  # generate spikes more than the required size.
  isi = np.random.poisson(lam=rate, size= 2 * rate * duration)
  spikes = np.cumsum(isi) # convert intervals to a train
  spikes = spikes * (2*duration/spikes[-1]) # normalize

  # add leak spikes

  # cull spikes
  spikes = spikes[np.where(spikes<=duration)]
  return spikes

# generate spike train with signal rate of 2 and noise rate of 2.
spikes = encode_message("A", 2, duration=10, codebook=['A','B'],noise_rate=2)


#%% [markdown]

# the following code plots signal train alognside the noise spikes.

# generate signal spikes
spikes = encode_message("A", 2, duration=10, codebook=['A','B'])

# generate noise spikes
noise_spikes = encode_message("", 2, duration=10, codebook=[])

plt.eventplot(spikes, color='white')
plt.eventplot(noise_spikes, color='red')
plt.show()

#%%