#############################################
# Synthesia Basic v3.21                     #
# Expanded texture representation
#############################################
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import matplotlib.colors as colors
import matplotlib.patches as pat
import numpy as np
from scipy.io import wavfile
import wave
from pydub import AudioSegment
AudioSegment.converter = 'C:\FFmpeg\bin\ffmpeg' #local ffmpeg install
from os import walk
#defs
def chunkr(data, chunks):
    chunksize = len(data)//chunks
    tmp = [] #values in chunk
    tmp2 = [] #chunked values
    for i in range(chunks):
        for j in range(chunksize):
            tmp.append(data[(chunksize*i)+j])
        tmp2.append(np.mean(tmp))
        tmp = []
    return tmp2
def wavr(filen):
    if (switchB == 1):
        filen = 'batch/'+filen
    if ('.wav' not in filen):
        print("invalid format in batch")
        exit
    wav1 = wave.open(filen, 'r')
    if (wav1.getnchannels() == 2):
        snd = AudioSegment.from_wav(filen)
        snd = snd.set_channels(1)
        snd.export(filen, format='wav')
        print(filen+' converted to mono')
    wav1.close()
    return wavfile.read(filen)
#
switchC = 1 #0 = umodified, 1 = 'intuitive' inverted+shifted colormap
switchB = 0 #0 = manual, 1 = batch processing
switchT = 1 #0 = color only, 1 = added texture processing
#>#
queue = []
#PRC switch
if (switchB == 0):
    print('Enter name of mono .wav file:')
    queue.append(input()+'.wav')
if (switchB == 1):
    quehue = []
    print('Processing files in /batch directory...')
    for (dirpath, dirnames, filenames) in walk('batch'):
        queue.extend(filenames)
        break
for r in range(len(queue)):
    filen = queue[r]
    # Read the wav file (mono); int16 numpy array
    samplingFrequency, signalData = wavr(filen)
    print("Wavfile read")
    # HUE BEGIN
    plt.figure(1)
    mspec, freqs2, line = plt.magnitude_spectrum(signalData,Fs=samplingFrequency)
    print("Magnitude spectrum generated")
    # scale x (freqs) along Hue spectrum, and y (mags) from 0 to 1
    hueint = np.interp(freqs2, (freqs2.min(), freqs2.max()), (0, 360))
    magint = np.interp(mspec, (mspec.min(), mspec.max()), (0,1))
    print("Axes interpolated")
    # for display purposes only, average down the data set to generate 36e~ bars
    huecnk = chunkr(hueint, 360)
    magcnk = chunkr(magint, 360)
    print("Axes chunked")
    # create bar plot with new lists and generate colormap
    plt.figure(2)
    plt.axis('off')
    #INV switch; negative vmin shifts cmap leftward
    if (switchC == 0):
        specr = cm.get_cmap('hsv')
        specn = colors.Normalize(vmin=0, vmax=360)
    if (switchC == 1):
        specr = cm.get_cmap('hsv_r') 
        specn = colors.Normalize(vmin=-80, vmax=360)
    plt.bar(huecnk, magcnk, width=1, color=specr(specn(huecnk)))
    #PRC switch
    if (switchB == 0):
        plt.savefig((filen+' spectrum.png'), bbox_inches='tight', pad_inches=0)
    if (switchB == 1):
        plt.savefig(('outp/'+filen+' spectrum.png'), bbox_inches='tight', pad_inches=0)
    # get WEIGHTED mean frequency
    mhue = np.average(hueint, weights=magint)
    fig,ax = plt.subplots(1)
    plt.axis('off')
    #INV switch
    if (switchC == 0):
        inthue = mhue/360
    if (switchC == 1):
        inthue = 1-((mhue+80)/360)
    print("Mean hue: "+str(inthue*360))
    sqr = pat.Rectangle((0,0), 100, 100, color=colors.hsv_to_rgb((inthue,1,1)))
    ax.add_patch(sqr)
    #PRC switch
    if (switchB == 0):
        plt.savefig((filen+' swatch.png'), bbox_inches='tight', pad_inches=0)
    if (switchB == 1):
        plt.savefig(('outp/'+filen+' swatch.png'), bbox_inches='tight', pad_inches=0)
        quehue.append(inthue*360)
    #TEX switch
    if (switchT == 1):
        plt.figure(4)
        #sgram = [freqs,bins]
        sgram, freqs, bins, im = plt.specgram(signalData,Fs=samplingFrequency)
        print("Spectrogram generated")
        # get differentials between all bins
        difrs = []
        for f in range(len(freqs)-1):
            frame1 = sgram[f]
            frame2 = sgram[f+1]
            difrs.append(np.absolute(frame1-frame2))
        # double-reduce, scale by max possible and mult by 100
        tex0 = np.mean(difrs,axis=0)
        tex1 = np.mean(tex0)
        tex2 = (tex1/sgram.max())*100
        print("Roughness coefficient: "+str(tex2))
        # adapt difrs for display
        plt.figure(5)
        #difds = np.round(difrs,decimals=5)
        #difds = difds*1000
        plt.pcolormesh(difrs, cmap='Greys_r', vmax=10000)
        #PRC switch
        if (switchB == 0):
            plt.savefig((filen+' texture.png'), bbox_inches='tight', pad_inches=0)
        if (switchB == 1):
            plt.savefig(('outp/'+filen+' texture.png'), bbox_inches='tight', pad_inches=0)
                

#PRC switch
if (switchB == 1):
    print(str(len(queue))+' files processed into /outp')
    print('Average hue of files: '+str(np.mean(quehue)))
#if (switchB == 0):
    #plt.show()

#Note: would be optimal to assess swatch along modified HSV cmap
#compare colors.Normalize args to inthue switch