from __future__ import print_function, division

import thinkdsp
import thinkplot
import thinkstats2

import numpy as np
import scipy.fftpack

import warnings
warnings.filterwarnings('ignore')

import dct

%matplotlib qt5

# Saxaphone recording for sample
wave = thinkdsp.read_wave('100475__iluppai__saxophone-weep.wav')
wave.make_audio()

# A short segment
segment = wave.segment(start=1.2, duration=0.5)
segment.normalize()
segment.make_audio()

# A DCT of that segment
seg_dct = segment.make_dct()
seg_dct.plot(high=4000)
thinkplot.config(xlabel='Frequency (Hz)', ylabel='DCT')

#This function takes a DCT and sets all elements below threshold to 0
def compress(dct, thresh=1):
    count = 0
    for i, amp in enumerate(dct.amps):
        if abs(amp) < thresh:
            dct.hs[i] = 0
            count += 1

    n = len(dct.amps)
    print(count, n, 100 * count / n, sep='\t')

# Applied to a segment
seg_dct = segment.make_dct()
compress(seg_dct, thresh=10)
seg_dct.plot(high=4000)

# And the results sound more or less the same
seg2 = seg_dct.make_wave()
seg2.make_audio()

# To compress a longer segment, a DCT spectrogram is created
def make_dct_spectrogram(wave, seg_length):

    window = np.hamming(seg_length)
    i, j = 0, seg_length
    step = seg_length / 2

    # map from time to Spectrum
    spec_map = {}

    while j < len(wave.ys):
        segment = wave.slice(i, j)
        segment.window(window)

        # the nominal time for this segment is the midpoint
        t = (segment.start + segment.end) / 2
        spec_map[t] = segment.make_dct()

        i += step
        j += step

    return thinkdsp.Spectrogram(spec_map, seg_length)

    Applying `compress` to each segment:
    spectro = make_dct_spectrogram(wave, seg_length=1024)
for t, dct in sorted(spectro.spec_map.items()):
    compress(dct, thresh=0.2)

    # In most segments, the compression is 75-80%.

    #Convert the spectrogram back to wave to hear what it sounds like
    wave2 = spectro.make_wave()
    wave2.make_audio()

    # the original
    wave.make_audio()
