#!/usr/bin/python3

import argparse
import os
import sys
from warnings import warn

import numpy as np
from scipy import signal
from scipy.io import wavfile

argument_parser = argparse.ArgumentParser(description='Generate a digital FIR filter from a frequency response specification. Basically a wrapper around scipy.signal.firwin2.')
argument_parser.add_argument('--frequency-response-spec-file', help='path to the filter frequency response specification file in REW export format; phase is ignored', required=True)
argument_parser.add_argument('--sample-rate-hz', help='sample rate to generate the FIR filter for in Hz', type=float, default=48000)
argument_parser.add_argument('--taps', help='number of FIR filter taps', type=int, default=127)
argument_parser.add_argument('--antisymmetric', help='generate an antisymmetric filter', action='store_true')
argument_parser.add_argument('--dc-gain-db', help='gain at DC in dB; required if missing from input', type=float)
argument_parser.add_argument('--nyquist-gain-db', help='gain at the Nyquist frequency (sample-rate/2) in dB; required if missing from input', type=float)
argument_parser.add_argument('--impulse-response-wav-file', help='WAV file to write the filter impulse response to')
argument_parser.add_argument('--print-sox-fir', help='Output a SoX fir filter command line', action='store_true')
argument_parser.add_argument('--sox-fir-coefs-file', help='Write to SoX fir filter coefs file')
args = argument_parser.parse_args()

if (args.sample_rate_hz <= 0):
	raise RuntimeError('invalid sample rate')

nyquist_frequency_hz = args.sample_rate_hz / 2

frequency_response_spec = np.loadtxt(
	fname=args.frequency_response_spec_file,
	usecols=(0, 1),  # Ignore phase
	dtype=[('frequency_hz', None), ('gain_db', None)],
	comments='*'
)
frequency_response_spec.sort(kind='mergesort', order='frequency_hz')

if frequency_response_spec['frequency_hz'][0] < 0:
	raise RuntimeError('lowest frequency in specification is less than zero')
while frequency_response_spec['frequency_hz'][-1] > nyquist_frequency_hz:
	warn('dropping ' + str(frequency_response_spec['gain_db'][-1]) + ' dB at ' + str(frequency_response_spec['frequency_hz'][-1]) + ' Hz because it is above the Nyquist frequency (' + str(nyquist_frequency_hz) + ' Hz)')
	frequency_response_spec = frequency_response_spec[:-1]

if args.dc_gain_db is not None:
	if frequency_response_spec['frequency_hz'][0] == 0:
		warn('input specification already contains a gain at DC, ignoring --dc-gain-db')
	else:
		frequency_response_spec = np.concatenate((np.array([(0, args.dc_gain_db)], dtype=frequency_response_spec.dtype), frequency_response_spec))

if args.nyquist_gain_db is not None:
	if frequency_response_spec['frequency_hz'][-1] == nyquist_frequency_hz:
		warn('input specification already contains a gain at Nyquist, ignoring --nyquist-gain-db')
	else:
		frequency_response_spec = np.concatenate((frequency_response_spec, np.array([(nyquist_frequency_hz, args.nyquist_gain_db)], dtype=frequency_response_spec.dtype)))

spec_frequencies_hz = frequency_response_spec['frequency_hz']
spec_gains_factor = np.power(10, frequency_response_spec['gain_db'] / 20)  # Convert from decibels

impulse_response = signal.firwin2(
	numtaps=args.taps,
	freq=spec_frequencies_hz,
	gain=spec_gains_factor,
	nyq=nyquist_frequency_hz,
	antisymmetric=args.antisymmetric
)
if args.impulse_response_wav_file is not None:
	wavfile.write(args.impulse_response_wav_file, args.sample_rate_hz, impulse_response)
if args.print_sox_fir:
	sys.stdout.write('fir ')
	np.savetxt(sys.stdout, impulse_response, newline=' ')
	sys.stdout.write(os.linesep)
if args.sox_fir_coefs_file is not None:
	np.savetxt(args.sox_fir_coefs_file, impulse_response, header=('Generated by: ' + ' '.join(sys.argv[:]) + os.linesep + 'Sample rate: ' + str(args.sample_rate_hz) + ' Hz'))
