import math
import matplotlib.pyplot as plt


# --------------------------------------------------
# RTTY parameters
# --------------------------------------------------

Fs = 100000             # Sample rate: 100 kHz

baud = 45.45            # RTTY baud rate

mark_freq = 2125        # MARK frequency
space_freq = 2295       # SPACE frequency

samples_per_symbol = Fs / baud


# --------------------------------------------------
# ITA2 / Baudot character table
# --------------------------------------------------

baudot = {
    'A': 0b00011,
    'B': 0b11001,
    'C': 0b01110,
    'D': 0b01001,
    'E': 0b00001,
    'F': 0b01101,
    'G': 0b11010,
    'H': 0b10100,
    'I': 0b00110,
    'J': 0b01011,
    'K': 0b01111,
    'L': 0b10010,
    'M': 0b11100,
    'N': 0b01100,
    'O': 0b11000,
    'P': 0b10110,
    'Q': 0b10111,
    'R': 0b01010,
    'S': 0b00101,
    'T': 0b10000,
    'U': 0b00111,
    'V': 0b11110,
    'W': 0b10011,
    'X': 0b11101,
    'Y': 0b10101,
    'Z': 0b10001,
    ' ': 0b00100
}


# --------------------------------------------------
# Text to Baudot bits
# --------------------------------------------------

text = "HELLO WORLD"

bits = []

for char in text:

    code = baudot[char]

    # RTTY sends least-significant bit first
    for i in range(5):
        bits.append((code >> i) & 1)


# --------------------------------------------------
# FSK / AFSK generator
# --------------------------------------------------

phase = 0

signal = []
time = []

sample_count = 0


for bit in bits:

    # Select MARK or SPACE frequency
    if bit == 0:
        f = space_freq
    else:
        f = mark_freq

    # Generate one RTTY symbol
    for i in range(round(samples_per_symbol)):

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

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

        sample = math.cos(phase)

        signal.append(sample)
        time.append(sample_count / Fs)

        sample_count += 1


# --------------------------------------------------
# Plot
# --------------------------------------------------

plt.figure(figsize=(14, 5))

plt.plot(time, signal)

plt.xlabel("Time (seconds)")
plt.ylabel("Amplitude")

plt.title("RTTY AFSK - 2125 / 2295 Hz")

plt.grid()

plt.show()