import pyaudio
import numpy as np
# from matplotlib import pyplot as plt
from spectro_gen import plot_stft, generate_spec_frec


def record_spectrogram(n_sec, plot=False, playback=False):
    SAMPLE_RATE = 22050
    N_CHANNELS = 2
    N_SEC = n_sec
    CHUNKSIZE = int(SAMPLE_RATE * N_SEC / N_CHANNELS)  # fixed chunk size
    # show_record_prompt()
    input('Press [Enter] to start recording sample... ')
    p_inp = pyaudio.PyAudio()
    stream = p_inp.open(
        format=pyaudio.paFloat32,
        channels=N_CHANNELS,
        rate=SAMPLE_RATE,
        input=True,
        frames_per_buffer=CHUNKSIZE)
    data = stream.read(CHUNKSIZE)
    numpydata = np.frombuffer(data, dtype=np.float32)
    multi_channel = np.abs(np.reshape(numpydata, (-1, 2))).mean(axis=1)
    one_channel = np.asarray([multi_channel, -1 * multi_channel]).T.reshape(-1)
    mean_channel_data = one_channel.tobytes()
    stream.stop_stream()
    stream.close()
    p_inp.terminate()
    if plot:
        plot_stft(one_channel, SAMPLE_RATE)
    if playback:
        p_oup = pyaudio.PyAudio()
        stream = p_oup.open(
            format=pyaudio.paFloat32,
            channels=2,
            rate=SAMPLE_RATE,
            output=True)
        stream.write(mean_channel_data)
        stream.close()
        p_oup.terminate()
    ims, _ = generate_spec_frec(one_channel, SAMPLE_RATE)
    return ims