#!/usr/bin/python

# ----------------------------------------------------------------------------
#
#  Copyright (C) 2008-2022 Fons Adriaensen <fons@linuxaudio.org>
#    
#  This program is free software; you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  the Free Software Foundation; either version 3 of the License, or
#  (at your option) any later version.
#
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
#  You should have received a copy of the GNU General Public License
#  along with this program.  If not, see <http:#www.gnu.org/licenses/>.
#
# ----------------------------------------------------------------------------

import sys
import numpy as np
import matplotlib.pyplot as plt
from scipy.signal import convolve # much faster than numpy.convolve
from zita_jacktools.jacksignal import JackSignal
sys.path.append ('..')
from utils.logsweep import genlogsweep


# ------- Frequency response using log sweep, deconvolution and fft ---------
#
# This method can be used to obtain the impulse response or frequency
# response in noisy conditions, e.g. for measuring a room or a speaker.
# This works because a sweep has much more energy than a single sample
# as used in freqresp1.py. To see this work you could repeat this test
# with the pink noise output of japa connected to jnoisemeter. The noise
# level is -20 dB, the same as the sweep used here, so the S/N ratio is
# as bad as 0 dB. Still the measured impulse will be quite clean.
#
# An important thing to understand is that the impulse or frequency
# response you get with this method will be that of the system being
# tested in series with a bandpass filter. The shape of that filter
# is determined by the sweep parameters, and the filter will be a
# linear-phase one. That means there will some response even in
# negative time - before any output should appear at all. To get
# the correct frequency response, this must be included in the
# input to the Fourier transform. In this test we use 50 ms.
#
#
# Make sure jnoisemeter is running, using
# input 1, and select the A, C or ITU filter.


# Create a JackSignal object and connect.
#
J = JackSignal("JackSignal")
if J.get_state() < 0:
    print ("Failed to create JackSignal -- is Jack running ?")
    exit(1)
J.create_output (0, "out")
J.create_input (0, "in")
J.silence()
J.connect_output (0, "jnoisemeter:in_1")
J.connect_input (0, "jnoisemeter:out")

# Get Jack info.
#
name, fsamp, period = J.get_jack_info ()

# Sweep parameters.
#
amp = 0.1
fmin = 10
fmax = 22e3
tfade1 = 0.25
tsweep = 8.00
tfade2 = 0.05

# FFT size
#
LFFT = 0x4000
NBIN = LFFT // 2 + 1

# How many samples before impulse to keep.
# Should be <= LFFT / 2.
#
OFFS = fsamp // 20  # 50 ms

# Generate sweep.
#
Afwd, Ainv, p = genlogsweep (amp, fmin, fmax, fsamp, tfade1, tsweep, tfade2)

# Create input buffer.
#
Ainp = np.zeros ([Afwd.shape [0] + LFFT,], dtype = np.float32)

# Connect buffer to ports.
#
J.set_output_data (0, Afwd)
J.set_input_data (0, Ainp, nskip = period)

# Run the test
#
print ("Measuring....")
J.process()
J.wait()
del J

# Deconvolve sweep and keep LFFT samples.
#
print ("Convolving....")
k = Ainv.shape [0] - 1 - OFFS
T = convolve (Ainp, Ainv, mode = 'full')
T = T [k:k + LFFT]
F = np.fft.rfft (T, LFFT)

# Display impulse and magnitude response.
#
print ("Display...")

# Time axis in samples or seconds.
Tax = np.linspace (0, LFFT, LFFT, endpoint = False) - OFFS
#Tax /= fsamp

fig = plt.figure (figsize=(8,6), facecolor='white')
ax1 = fig.add_axes ([0.07, 0.05, 0.90, 0.44])
ax1.set_ylim (-1.5, 1.5)
ax1.plot (Tax, T, color='b', lw=1)
ax1.grid ()

# Frequency axis in Hz.
Fax = np.linspace (0, fsamp / 2, NBIN, endpoint = True)

ax2 = fig.add_axes ([0.07, 0.53, 0.90, 0.44])
ax2.set_xlim (20, 20e3)
ax2.set_ylim (-60, 15)
ax2.set_xscale ('log')
ax2.plot (Fax, 20 * np.log10 (np.abs (F) + 1e-10), color='b', lw=1)
ax2.grid ()
plt.show()
