#!env python3

# Generate M-sequence,
# Interpolate it such that bit rate of m-sequence is fraction of Fs,
# Shift it to some Intermediate Frequency in order to work in less noisy range,
# Fade In and Out with some arbitrary chosen window.

from scipy.fft import fft
from scipy.io.wavfile import write
from scipy.signal import convolve
from scipy.signal import max_len_seq
from scipy.signal.windows import hamming
from scipy.signal.windows import hann
import matplotlib.pyplot as plt
import numpy as np
import os
import textwrap

###############################################################################

def spectrum(s, fs_=48000):
    ftest = fft(s, norm='forward')
    xf = np.linspace(0.0, fs_ / 2.0, ftest.shape[-1] // 2)
    return xf, 20 * np.log10(np.abs(ftest[0:ftest.shape[-1] // 2]))

###############################################################################

fs = 48000  # SPS

f0 = 8000  # Intermediate frequency.
bw = 14e3
n_filt = 1024
nbits = 13
N = np.power(2, nbits)

###############################################################################

impulse = max_len_seq(nbits)[0].astype(np.float64) - 0.5

# Bandpass filter by windowed and shifted sinc.
idx_bpf = np.arange(-n_filt / 2, n_filt / 2)
h_bpf = np.sinc(idx_bpf * bw / fs) * np.sin(2 * np.pi / fs * f0 * np.arange(1, n_filt + 1)) * hann(n_filt)

# Circular convolution
impulse = convolve(impulse, h_bpf, 'full')
impulse[:n_filt - 1] += impulse[N - 1:]
impulse = impulse[:N]
impulse = impulse / np.max(np.abs(impulse)) / 2

# Fade In n Out
win = hamming(256)
nwin_half = win.shape[0] // 2
win_raise = win[:nwin_half]
win_fall = win[nwin_half:]
impulse[:nwin_half] *= win_raise
impulse[-nwin_half:] *= win_fall

###############################################################################

# Plot spectrum and auto-correlation
if False:
    # noinspection PyUnreachableCode
    plt.subplot(211)
    plt.plot(*spectrum(h_bpf))
    plt.xlim([0, 20e3])
    plt.grid(True)

    plt.subplot(212)
    plt.plot(convolve(impulse, np.flip(impulse)))
    plt.grid(True)
    plt.xlim([8137, 8245])

    plt.show()

###############################################################################

os.chdir(os.path.join(
    os.path.dirname(os.path.abspath(__file__)), '..'))

###############################################################################

# Make probe signal stereo, but occupy only left channel.
zeros_to_right = np.zeros(impulse.shape[0], dtype=np.float32)
output = np.stack((impulse, zeros_to_right), axis=1)

if False:
    wav_path = 'src/base/processing/Impulse.wav'
    write(wav_path, fs, output.astype(np.float32))

###############################################################################
# Generate header

cpp_path = 'src/base/processing/Impulse.cpp'
with open(cpp_path, 'w', encoding = 'utf-8') as f:
    f.write(textwrap.dedent('''
    // Copyright (c) Signal Estimator authors
    // Licensed under MIT

    // Generated by generate_impulse.py. DO NOT EDIT!

    #include "processing/Impulse.hpp"

    const std::array<float, 8192> impulse = {
    ''').lstrip())
    for i in range(output.shape[0]):
        f.write(f"    (float){output[i, 0]},\n")
    f.write(textwrap.dedent('''
    };
    '''.lstrip()))
