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

stop_bits = 1.5             # Standard RTTY stop length


# ==================================================
# ITA2 / BAUDOT CODES
# ==================================================

LETTERS = {
    '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
}


FIGURES = {
    '1': 0b10111,
    '2': 0b10011,
    '3': 0b00001,
    '4': 0b01010,
    '5': 0b10000,
    '6': 0b10101,
    '7': 0b00111,
    '8': 0b00110,
    '9': 0b11000,
    '0': 0b10110,

    '-': 0b00011,
    '/': 0b11100,
    ':': 0b10010,
    '?': 0b11001,
    '.': 0b11101,
    ',': 0b01100,
    '(': 0b11110,
    ')': 0b01111,
    '$': 0b10010
}


# Special RTTY codes

FIGURES_SHIFT = 0b1101       # FIGS
LETTERS_SHIFT = 0b11111      # LTRS


# ==================================================
# CREATE RTTY BIT STREAM
# ==================================================

text = "HELLO WORLD 123"

bits = []


current_mode = "LETTERS"


for char in text.upper():

    # ----------------------------------------------
    # Space
    # ----------------------------------------------

    if char == ' ':

        code = LETTERS[' ']

        for i in range(5):
            bits.append((code >> i) & 1)

        continue


    # ----------------------------------------------
    # Letter
    # ----------------------------------------------

    if char in LETTERS:

        if current_mode != "LETTERS":

            # Send LETTERS shift
            code = LETTERS_SHIFT

            for i in range(5):
                bits.append((code >> i) & 1)

            current_mode = "LETTERS"


        code = LETTERS[char]


    # ----------------------------------------------
    # Number / Figure
    # ----------------------------------------------

    elif char in FIGURES:

        if current_mode != "FIGURES":

            # Send FIGURES shift
            code = FIGURES_SHIFT

            for i in range(5):
                bits.append((code >> i) & 1)

            current_mode = "FIGURES"


        code = FIGURES[char]


    else:

        # Ignore unsupported characters
        continue


    # ----------------------------------------------
    # Add 5 data bits
    # LSB first
    # ----------------------------------------------

    for i in range(5):

        bits.append((code >> i) & 1)


# ==================================================
# ADD RTTY START / DATA / STOP FRAMING
# ==================================================

framed_bits = []


for i in range(0, len(bits), 5):

    data = bits[i:i + 5]

    # ----------------------------------------------
    # START BIT
    #
    # RTTY idle = MARK
    # START = SPACE
    # ----------------------------------------------

    framed_bits.append(0)


    # ----------------------------------------------
    # DATA BITS
    # ----------------------------------------------

    framed_bits.extend(data)


    # ----------------------------------------------
    # STOP
    #
    # MARK for 1.5 bit times
    # ----------------------------------------------

    framed_bits.append(1)


# ==================================================
# FSK MODULATOR
# ==================================================

phase = 0.0

signal = []
time = []

sample_count = 0


# Fractional sample accumulator
sample_remainder = 0.0


for bit_index, bit in enumerate(framed_bits):

    # ----------------------------------------------
    # Select MARK or SPACE
    # ----------------------------------------------

    if bit == 1:
        frequency = mark_freq
    else:
        frequency = space_freq


    # ----------------------------------------------
    # Determine symbol duration
    # ----------------------------------------------

    if bit_index == 0:

        symbol_length = 1.0

    else:

        symbol_length = 1.0


    samples_float = Fs / baud * symbol_length

    samples = int(samples_float)

    sample_remainder += samples_float - samples


    if sample_remainder >= 1.0:

        samples += 1
        sample_remainder -= 1.0


    # ----------------------------------------------
    # Generate samples
    # ----------------------------------------------

    for _ in range(samples):

        phase += 2 * math.pi * frequency / 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, 6))

plt.plot(time, signal)

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

plt.title(
    "RTTY AFSK - 45.45 Baud - "
    "2125 Hz MARK / 2295 Hz SPACE"
)

plt.grid()

plt.tight_layout()

plt.show()


# ==================================================
# INFORMATION
# ==================================================

print()
print("RTTY TRANSMISSION")
print("------------------")
print("Text:          ", text)
print("Baud rate:     ", baud)
print("MARK:          ", mark_freq, "Hz")
print("SPACE:         ", space_freq, "Hz")
print("Shift:         ", space_freq - mark_freq, "Hz")
print("Sample rate:   ", Fs, "Hz")
print("Samples:       ", len(signal))
print("Duration:      ", len(signal) / Fs, "seconds")