from math import sqrt
import matplotlib.pyplot as plt
import numpy as np
from os import listdir, chdir
from os.path import isfile, join
from pydub import AudioSegment
from scipy import polyval, polyfit
from scipy.fft import *
from scipy.io import wavfile
from termcolor import colored


path = "C:\\Users\\HP\\Downloads\\TMF\\TMF-audio\\Experiment #1\\cutted\\"
chdir(path)  #go to directory


def amplitude(file, path):

    chdir(path)  #.wav files only
    sound = AudioSegment.from_mp3(path + file)
    loudness = sound.dBFS
    return loudness

    '''sound = AudioSegment.from_mp3(path + file)

    normalized_sound = sound.apply_gain(-sound.max_dBFS)
    return normalized_sound'''

    #return AudioSegment.from_wav(path + file).max
#print("10cm1.wav".dBFS)


def duration(file, path):
    chdir(path)
    return AudioSegment.from_wav(path + file).duration_seconds * 1000  #in milliseconds


#https://stackoverflow.com/a/66892766
def freq(file, path):
    chdir(path)
    start_time = 0
    end_time = duration(file, path)

    sr, data = wavfile.read(file, path)
    if data.ndim > 1:
        data = data[:, 0]
    else:
        pass

    dataToRead = data[int(start_time * sr / 1000) : int(end_time * sr / 1000) + 1]

    N = len(dataToRead)
    yf = rfft(dataToRead)
    xf = rfftfreq(N, 1 / sr)

    idx = np.argmax(np.abs(yf))
    freq = xf[idx]

    return round(freq, 1)


my_audio = list(f for f in listdir(path) if isfile(join(path, f)))  #get files in folder

for _ in range(10):
    my_audio.append(my_audio.pop(0))

f, f2 = open("ampl.txt", "r+"), open("freq.txt", "r+")
f.truncate(0); f2.truncate(0)  #clear files

counter = 1
for audio in my_audio:
    try:
        if audio[2] in "1234567890":
            if counter != 90:
                f.write(audio[0:3] + " " + str(amplitude(audio, path)) + "\n")
                #f.write(audio[0:3] + " " + str(10**(amplitude(audio, path) / 10)) + "\n")
                f2.write(audio[0:3] + " " + str(freq(audio, path)) + "\n")
            else:
                f.write(audio[0:3] + " " + str(amplitude(audio, path)))
                #f.write(audio[0:3] + " " + str(10**(amplitude(audio, path) / 10)))
                f2.write(audio[0:3] + " " + str(freq(audio, path)))
        else:
            if counter != 90:
                f.write(audio[0:2] + " " + str(amplitude(audio, path)) + "\n")
                #f.write(audio[0:2] + " " + str(10**(amplitude(audio, path) / 10)) + "\n")
                f2.write(audio[0:2] + " " + str(freq(audio, path)) + "\n")
            else:
                f.write(audio[0:2] + " " + str(amplitude(audio, path)))
                #f.write(audio[0:2] + " " + str(10**(amplitude(audio, path) / 10)))
                f2.write(audio[0:2] + " " + str(freq(audio, path)))

        counter += 1
        print(colored((audio + " - " + str(amplitude(audio, path)) + " dBFS, " + \
            str(freq(audio, path)) + " Hz"), "green"))  #green text
    except:
        pass

f.close()
f2.close()


############################# plot #############################

m = 0.014
g = 9.81
coeff = 0.01 * m * g  #convert to Joules

heightsA, ampls = list(), list()
for line in open("ampl.txt", "r"):
  values = [float(number) for number in line.split()]
  heightsA.append(values[0] * coeff)
  ampls.append(10**(values[1] / 20))
  #ampls.append(values[1])


print(ampls)

_10cmA, _15cmA, _25cmA, _35cmA, _50cmA, _60cmA, _70cmA, _85cmA, _100cmA = \
    list(), list(), list(), list(), list(), list(), list(), list(), list()
_10cmA.append(ampls[0:10])
_15cmA.append(ampls[10:20])
_25cmA.append(ampls[20:30])
_35cmA.append(ampls[30:40])
_50cmA.append(ampls[40:50])
_60cmA.append(ampls[50:60])
_70cmA.append(ampls[60:70])
_85cmA.append(ampls[70:75])
_100cmA.append(ampls[75:])
_10cmA, _15cmA, _25cmA, _35cmA, _50cmA, _60cmA, _70cmA, _85cmA, _100cmA = \
    _10cmA[0], _15cmA[0], _25cmA[0], _35cmA[0], _50cmA[0], _60cmA[0], _70cmA[0], _85cmA[0], _100cmA[0]


'''
heightsF, frequencies = list(), list()
for line in open("freq.txt", "r"):
  values = [float(number) for number in line.split()]
  heightsF.append(values[0])
  frequencies.append(values[1])

_10cmF, _25cmF, _35cmF, _50cmF, _70cmF, _100cmF = \
    list(), list(), list(), list(), list(), list()
_10cmF.append(frequencies[0:10])
_25cmF.append(frequencies[10:20])
_35cmF.append(frequencies[20:30])
_50cmF.append(frequencies[30:40])
_70cmF.append(frequencies[40:50])
_100cmF.append(frequencies[50:])
_10cmF, _25cmF, _35cmF, _50cmF, _70cmF, _100cmF = \
    _10cmF[0], _25cmF[0], _35cmF[0], _50cmF[0], _70cmF[0], _100cmF[0]
'''


def average(list):
    return sum(list) / len(list)


def dispersion(list):
  aver = average(list)
  sum = 0
  for item in list:
    a = (aver - item)**2
    sum += a
  disp = sum / len(list)
  return sqrt(disp) / sqrt(len(list))


h1 = 10*coeff
h2 = 25*coeff
h3 = 35*coeff
h4 = 50*coeff
h5 = 70*coeff
h6 = 100*coeff
h7 = 15*coeff
h8 = 60*coeff
h9 = 85*coeff


'''plot1 = plt.scatter(heightsF, frequencies, c = "darkgreen", s = 2)
plt.show()'''

experiment_size = 1
average_size = 25
average_color = "red"
errorbar_color = "darkgreen"
errorbar_width = 2

#plot2 = plt.scatter(heightsA, ampls, c = "darkblue", s = experiment_size)

plt.scatter(h1, average(_10cmA), c = average_color, s = average_size)
plt.errorbar(h1, average(_10cmA), xerr = 0, \
    yerr = dispersion(_10cmA), c = errorbar_color, linewidth = errorbar_width)

plt.scatter(h2, average(_25cmA), c = average_color, s = average_size)
plt.errorbar(h2, average(_25cmA), xerr = 0, \
    yerr = dispersion(_25cmA), c = errorbar_color, linewidth = errorbar_width)

plt.scatter(h3, average(_35cmA), c = average_color, s = average_size)
plt.errorbar(h3, average(_35cmA), xerr = 0, \
    yerr = dispersion(_35cmA), c = errorbar_color, linewidth = errorbar_width)

plt.scatter(h4, average(_50cmA), c = average_color, s = average_size)
plt.errorbar(h4, average(_50cmA), xerr = 0, \
    yerr = dispersion(_50cmA), c = errorbar_color, linewidth = errorbar_width)

plt.scatter(h5, average(_70cmA), c = average_color, s = average_size)
plt.errorbar(h5, average(_70cmA), xerr = 0, \
    yerr = dispersion(_70cmA), c = errorbar_color, linewidth = errorbar_width)

plt.scatter(h6, average(_100cmA), c = average_color, s = average_size)
plt.errorbar(h6, average(_100cmA), xerr = 0, \
    yerr = dispersion(_100cmA), c = errorbar_color, linewidth = errorbar_width)



plt.scatter(h7, average(_15cmA), c = average_color, s = average_size)
plt.errorbar(h7, average(_15cmA), xerr = 0, \
    yerr = dispersion(_15cmA), c = errorbar_color, linewidth = errorbar_width)

plt.scatter(h8, average(_60cmA), c = average_color, s = average_size)
plt.errorbar(h8, average(_60cmA), xerr = 0, \
    yerr = dispersion(_60cmA), c = errorbar_color, linewidth = errorbar_width)

#plt.scatter(h9, average(_85cmA), c = average_color, s = average_size)
#plt.errorbar(h9, average(_85cmA), xerr = 0, \
#    yerr = dispersion(_85cmA), c = errorbar_color, linewidth = errorbar_width)


fit_color = "blue"

#https://stackoverflow.com/a/4296777
'''x = [h1, h7, h2, h3, h4]
y = [average(_10cmA), average(_15cmA), average(_25cmA), average(_35cmA), average(_50cmA)]
a, b, c = np.polyfit(x, y, 2)
y_pred = np.polyval([a, b, c], x)
x_out = np.linspace(0, 0.07, 20)
y_pred = np.polyval([a, b, c], x_out)
plt.plot(x_out, y_pred, c = fit_color)

x = [h4, h5, h6]
y = [average(_50cmA), average(_70cmA), average(_100cmA)]
a, b = np.polyfit(x, y, 1)
y_pred = np.polyval([a, b], x)
x_out = np.linspace(0.05, 0.14, 20)
y_pred = np.polyval([a, b], x_out)
plt.plot(x_out, y_pred, c = fit_color)'''


plt.xlabel("Энергия, E [Дж]")
plt.ylabel("Амплитуда, A [отн. ед.]")

#plt.text(0.00, 30500, "$m_{шарика} = 0.014 кг$", fontsize = 15)
plt.text(0.02, 0.101, "$m_{шарика} = 0.014 кг$", fontsize = 15)

plt.grid(linestyle = "-", linewidth = 0.4, c = "black")
plt.rc("axes", axisbelow = True)

plt.tick_params(direction = "in")

plt.show()