import math
import random
import matplotlib.pyplot as plt


# --------------------------------------------------
# FSK parameters
# --------------------------------------------------

Fs = 100000         # Sample rate: 100 kHz

Rb = 100            # Bit rate: 100 bits/sec

f0 = 1000           # Frequency for bit 0
f1 = 2000           # Frequency for bit 1

bits = [1, 0, 1, 1, 0, 0, 1, 0]


# --------------------------------------------------
# Calculate samples per bit
# --------------------------------------------------

samples_per_bit = int(Fs / Rb)


# --------------------------------------------------
# FSK modulator
# --------------------------------------------------

phase = 0

time = []
signal = []
bit_signal = []


for bit in bits:

    # Select frequency based on the bit
    if bit == 0:
        f = f0
    else:
        f = f1

    # Generate samples for this bit
    for i in range(samples_per_bit):

        # Advance phase
        phase += 2 * math.pi * f / Fs

        # Keep phase between 0 and 2*pi
        phase %= 2 * math.pi

        # Generate carrier sample
        sample = math.cos(phase)

        # Store results
        time.append(len(time) / Fs)
        signal.append(sample)
        bit_signal.append(bit)


# --------------------------------------------------
# Noisy channel
# --------------------------------------------------

noise_amplitude = 0.3

received_signal = []

for sample in signal:

    noise = random.gauss(0, noise_amplitude)

    received_sample = sample + noise

    received_signal.append(received_sample)


# --------------------------------------------------
# Plot everything
# --------------------------------------------------

fig, (ax1, ax2, ax3) = plt.subplots(
    3, 1,
    figsize=(12, 8),
    sharex=True
)


# Original bits
ax1.step(time, bit_signal, where="post")

ax1.set_ylabel("Bit")
ax1.set_title("Original Data")
ax1.grid()


# Clean FSK signal
ax2.plot(time, signal)

ax2.set_ylabel("Amplitude")
ax2.set_title("Transmitted FSK")
ax2.grid()


# Noisy received signal
ax3.plot(time, received_signal)

ax3.set_xlabel("Time (seconds)")
ax3.set_ylabel("Amplitude")
ax3.set_title("Received FSK + Noise")
ax3.grid()


plt.tight_layout()
plt.show()