# Copyright 2025 RnD Center "ELVEES", JSC

import argparse
from array import array
import numpy as np
from scipy import signal

parser = argparse.ArgumentParser(
    description="Makes files with data for fir filters tests",
    formatter_class=argparse.RawTextHelpFormatter,
)
parser.add_argument("-w", "--window", type=str, help="Type of window function")
parser.add_argument(
    "-p", "--param", type=float, help="Attenuation (in dB) for parametric windows", default=0.0
)
parser.add_argument("-t", "--taps", type=int, help="Length of the filter. Should be an odd number")
parser.add_argument(
    "-w0", type=float, help="Cutoff frequency or left bound of band for bandpass or bandstop filter"
)
parser.add_argument(
    "-w1",
    type=float,
    help="Cutoff frequency or right bound of band for bandpass or bandstop filter",
)
parser.add_argument("-s", "--size", type=int, help="Size of the signal with noise")
parser.add_argument("-l", "--location", type=str, help="Location of the test_data files")
args = parser.parse_args()

# Input parameters
fs = 2
taps = args.taps
w0 = args.w0
w1 = args.w1
win_type = args.window

if args.param > 0.0:
    h_lpf = signal.firwin(taps, w0, window=(win_type, args.param), fs=fs)
    h_hpf = signal.firwin(taps, w1, window=(win_type, args.param), fs=fs, pass_zero=False)
    h_bpass = signal.firwin(taps, [w0, w1], window=(win_type, args.param), fs=fs, pass_zero=False)
    h_bstop = signal.firwin(taps, [w0, w1], window=(win_type, args.param), fs=fs)
else:
    h_lpf = signal.firwin(taps, w0, window=win_type, fs=fs)
    h_hpf = signal.firwin(taps, w1, window=win_type, fs=fs, pass_zero=False)
    h_bpass = signal.firwin(taps, [w0, w1], window=win_type, fs=fs, pass_zero=False)
    h_bstop = signal.firwin(taps, [w0, w1], window=win_type, fs=fs)

# Filter some signal
N = args.size
t = np.linspace(0, 1, N, endpoint=True)
x = np.sin(2 * np.pi * 0.75 * t) + np.sin(2 * np.pi * 8 * t)
# Add some white noise
np.random.seed(1)
xn = x + np.random.randn(N)
# Filter
y_lpf = signal.lfilter(h_lpf, 1, xn)
y_hpf = signal.lfilter(h_hpf, 1, xn)
y_bpass = signal.lfilter(h_bpass, 1, xn)
y_bstop = signal.lfilter(h_bstop, 1, xn)

# Make binary files with data for DSP tests
with open(f"{args.location}/clear_signal.bin", "wb") as write_data:
    write_data.write(array("f", x).tobytes())

with open(f"{args.location}/signal_with_noise.bin", "wb") as write_data:
    src0_array = np.zeros(N * 2)
    src0_array[1::2] = xn
    write_data.write(array("f", src0_array).tobytes())

with open(f"{args.location}/h_lowpass.bin", "wb") as write_data:
    write_data.write(array("f", h_lpf).tobytes())

with open(f"{args.location}/signal_after_lowpass.bin", "wb") as write_data:
    write_data.write(array("f", y_lpf).tobytes())

with open(f"{args.location}/h_highpass.bin", "wb") as write_data:
    write_data.write(array("f", h_hpf).tobytes())

with open(f"{args.location}/signal_after_highpass.bin", "wb") as write_data:
    write_data.write(array("f", y_hpf).tobytes())

with open(f"{args.location}/h_bandpass.bin", "wb") as write_data:
    write_data.write(array("f", h_bpass).tobytes())

with open(f"{args.location}/signal_after_bandpass.bin", "wb") as write_data:
    write_data.write(array("f", y_bpass).tobytes())

with open(f"{args.location}/h_bandstop.bin", "wb") as write_data:
    write_data.write(array("f", h_bstop).tobytes())

with open(f"{args.location}/signal_after_bandstop.bin", "wb") as write_data:
    write_data.write(array("f", y_bstop).tobytes())
