#!/usr/bin/env python3

# RKM to WAV converter
# Aleksander Alekseev, 2023
# https://eax.me/rkm-to-wav/
# Optimized by @kimstik

import argparse

WAV_HEADER_SIZE = 44
SAMPLE_RATE = 44100
BAUD_RATE = 700
SAMPLES_PER_BIT = int(SAMPLE_RATE / BAUD_RATE)
BIT_ZERO_WAVE = [] # initialized below
BIT_ONE_WAVE = [] # initialized below

temp = []
for i in range(0, SAMPLES_PER_BIT):
    if i < SAMPLES_PER_BIT/2:
        temp += [255 - 32]
    else:
        temp += [32]

# tranfer from high to low
BIT_ZERO_WAVE = temp.copy()
temp.reverse()
# tranfer from low to high
BIT_ONE_WAVE = temp

# generated with https://fiiir.com/
FIR_WEIGHTS = [
    0.000577926080672738,
    0.000459465018502213,
    0.000314862245890514,
    0.000114058758552773,
    -0.000176232019410298,
    -0.000586368131836170,
    -0.001136970252974637,
    -0.001832200078313594,
    -0.002653914944386688,
    -0.003557457996743626,
    -0.004469734119871817,
    -0.005290040439872776,
    -0.005893881011277897,
    -0.006139718992073919,
    -0.005878331252147731,
    -0.004964157085603880,
    -0.003267801176075040,
    -0.000688684837013550,
    0.002833242904185398,
    0.007307810544085602,
    0.012687912264589127,
    0.018866545092452867,
    0.025678084518976477,
    0.032903957573255177,
    0.040282515704921813,
    0.047522557733791673,
    0.054319636536848272,
    0.060374030484608231,
    0.065409095148934615,
    0.069188648369429409,
    0.071532089395465556,
    0.072326107924878466,
    0.071532089395465556,
    0.069188648369429409,
    0.065409095148934615,
    0.060374030484608238,
    0.054319636536848279,
    0.047522557733791680,
    0.040282515704921813,
    0.032903957573255184,
    0.025678084518976480,
    0.018866545092452870,
    0.012687912264589122,
    0.007307810544085603,
    0.002833242904185400,
    -0.000688684837013550,
    -0.003267801176075043,
    -0.004964157085603879,
    -0.005878331252147732,
    -0.006139718992073923,
    -0.005893881011277897,
    -0.005290040439872781,
    -0.004469734119871819,
    -0.003557457996743629,
    -0.002653914944386688,
    -0.001832200078313594,
    -0.001136970252974637,
    -0.000586368131836170,
    -0.000176232019410298,
    0.000114058758552773,
    0.000314862245890515,
    0.000459465018502213,
    0.000577926080672738,
]
WINDOW_SIZE = len(FIR_WEIGHTS)
fir_window = [128 for i in range(0,WINDOW_SIZE)]
last_data = BIT_ZERO_WAVE

# without the filter Mikrosha has difficulties interpreting the data
def fir_low_pass_filter(arr):
    global fir_window, FIR_WEIGHTS
    res = []
    for i in range(0, len(arr)):
        fir_window = [arr[i]] + fir_window[0:WINDOW_SIZE-1]
        filtered = int(sum([fir_window[j]*FIR_WEIGHTS[j]
                            for j in range(0,WINDOW_SIZE)]))
        res += [ filtered ]
    return bytes(res)

W = ( fir_low_pass_filter(BIT_ZERO_WAVE + BIT_ZERO_WAVE)[-SAMPLES_PER_BIT:],
      fir_low_pass_filter(BIT_ZERO_WAVE + BIT_ONE_WAVE )[-SAMPLES_PER_BIT:],
      fir_low_pass_filter(BIT_ONE_WAVE  + BIT_ZERO_WAVE)[-SAMPLES_PER_BIT:],
      fir_low_pass_filter(BIT_ONE_WAVE  + BIT_ONE_WAVE )[-SAMPLES_PER_BIT:])

# writes one encoded byte to the output WAV file
# returns the amount of bytes written
def wav_write_byte(f, b):
    global last_data
    for i in range(0,8):
        last_data = BIT_ONE_WAVE if (b >> (7-i)) & 1 else BIT_ZERO_WAVE
#        f.write(bytes(fir_low_pass_filter(data)))
        bit = (b >> (7-i)) & 1
        pair = (wav_write_byte.last<<1) | bit
        f.write(W[pair])
        wav_write_byte.last = bit
    return SAMPLES_PER_BIT*8

wav_write_byte.last = 0

def debug(msg):
    if args.debug:
        print("DEBUG: " + msg)

parser = argparse.ArgumentParser(description='Convert RKM to WAV')
parser.add_argument(
    '-i', '--input', metavar='F', type=str, required = True,
    help='input .rkm file')
parser.add_argument(
    '-o', '--output', metavar='F', type=str, required = True,
    help='output .wav file')
parser.add_argument(
    '-d', '--debug', action="store_true",
    help='enable debug output')
args = parser.parse_args()

debug("input: {}".format(args.input))
debug("output: {}".format(args.output))

with open(args.input, 'rb') as fin, open(args.output, 'wb') as fout:
    addr_start = int.from_bytes(fin.read(2), 'big')
    addr_end = int.from_bytes(fin.read(2), 'big')
    data_length = addr_end - addr_start + 1
    fin.seek(4 + data_length)
    checksum = int.from_bytes(fin.read(2), 'big')
    fin.seek(4)

    debug("start address: 0x{:04X}".format(addr_start))
    debug("end address: 0x{:04X}".format(addr_end))
    debug("checksum: 0x{:04X}".format(checksum))

    # reserve space for the WAV header
    fout.write(bytes([0 for i in range(0,WAV_HEADER_SIZE)]))
    file_size = WAV_HEADER_SIZE;

    # preamble. this sounds like an ~3 second long 700 Hz tone
    for i in range(0,256):
        file_size += wav_write_byte(fout, 0x00)

    # sync byte
    file_size += wav_write_byte(fout, 0xE6)

    # start address, big-endian
    for i in range(0,2):
        file_size += wav_write_byte(fout, (addr_start >> (1-i)*8) & 0xFF)

    # end address, big-endian
    for i in range(0,2):
        file_size += wav_write_byte(fout, (addr_end >> (1-i)*8) & 0xFF)

    # write the data
    for i in range(0, data_length):
        b = int.from_bytes(fin.read(1), 'big')
        file_size += wav_write_byte(fout, b)

    # checksum, big-endian
    for i in range(0,2):
        file_size += wav_write_byte(fout, (checksum >> (1-i)*8) & 0xFF)

    # update the FIR filter state as if last_data was encoded without caching
    fir_low_pass_filter(last_data)
    # some intermediate values in the end, for the FIR filter
    fout.write(bytes(fir_low_pass_filter([128 for i in range(0,WINDOW_SIZE)])))
    file_size += WINDOW_SIZE

    # write the WAV header
    fout.seek(0)
    fout.write(bytes([ord(c) for c in "RIFF"]))
    temp = file_size - 8;
    fout.write(bytes([(temp >> 8*i) & 0xFF for i in range(0,4)]))
    fout.write(bytes([ord(c) for c in "WAVEfmt "]))
    temp = 16 # size of the rest of the WAV header, for PCM
    fout.write(bytes([(temp >> 8*i) & 0xFF for i in range(0,4)]))
    temp = 1 # no compression, for PCM
    fout.write(bytes([(temp >> 8*i) & 0xFF for i in range(0,2)]))
    temp = 1 # number of channels
    fout.write(bytes([(temp >> 8*i) & 0xFF for i in range(0,2)]))
    temp = SAMPLE_RATE # sample rate
    fout.write(bytes([(temp >> 8*i) & 0xFF for i in range(0,4)]))
    temp = SAMPLE_RATE*1 # bytes per one second
    fout.write(bytes([(temp >> 8*i) & 0xFF for i in range(0,4)]))
    temp = 1 # bytes per sample, total by channels
    fout.write(bytes([(temp >> 8*i) & 0xFF for i in range(0,2)]))
    temp = 8 # _bits_ per sample per one channel
    fout.write(bytes([(temp >> 8*i) & 0xFF for i in range(0,2)]))
    fout.write(bytes([ord(c) for c in "data"]))
    temp = file_size - WAV_HEADER_SIZE;
    fout.write(bytes([(temp >> 8*i) & 0xFF for i in range(0,4)]))

debug("done")