#!/usr/bin/env python3

import sys
import numpy as np
from numpy import fft
import matplotlib.pyplot as plt
from scipy.io import wavfile

def audfft(file):
  sr, data = wavfile.read(file)
  N = len(data)
  S=sr # Taking a 1s sample
  data = data[int(N/2-S/2):int(N/2-S/2)+S]
  return sr, abs(fft.rfft(data))

# This averages data using avg function. They are binned into bins of size 1,1,2,4,8,16,...
# If the number of samples is >2**N, it'll linearely split it into 2*N smaller bins
# Returns an array for the X and an array for the Y coordinates
def logbin(data, avg, N=1):
  N=2**N
  X=[]
  Y=[]
  i=0
  o=0
  n=len(data)
  while True:
    if o >= n:
      break
    i = i * 2
    if i == 0:
      i = 1
    if i < N or N < 2:
      v = avg(data[o:min(o+i,n)])
      X.append(o)
      Y.append(v)
      X.append(o+i)
      Y.append(v)
      o += i
    else:
      j = i//N
      for k in range(N):
        if o > n:
          break
        v = avg(data[o:min(o+j,n)])
        X.append(o)
        Y.append(v)
        X.append(o+j)
        Y.append(v)
        o += j
  return np.array(X), np.array(Y)

# Log bins will be split linearely into 2**s smaller segments, of which
s=3

sr1, A = audfft(sys.argv[1])
sr2, B = audfft(sys.argv[2])
assert(sr1 == sr2)
X1,Al = logbin(A, np.average, s)
X2,Bl = logbin(B, np.average, s)
diff = Bl-Al


# Output audacity EQ filter curve
f=0.75
MX=[]
MY=[]
ymin=float("inf")
ymax=float("-inf")
for i in range(len(X1)//2):
  MX.append((X1[i*2+0]+X1[i*2+1])/2)
  v = diff[i*2]
  MY.append(v)
  if v < ymin:
    ymin = v
  if v > ymax:
    ymax = v
res  = 'FilterCurve:'
for i in range(len(MX)):
  res += f'f{i}="{MX[i]}" '
res += 'FilterLength="8191" InterpolateLin="0" InterpolationMethod="B-spline" '
for i in range(len(MY)):
  v=(MY[i]-ymin)/(ymax-ymin)
  x = -v*2
  #x=1-1/(1-v*f)
  res += f'v{i}="{x}" '
print(res)


# Plot stuff

#plt.figure(figsize = (100, 100))
plt.xlim(10, 20000)
plt.xscale('log')
#plt.yscale('log')
plt.xlabel('Freq (Hz)')
plt.ylabel('FFT Amplitude |X(freq)|')
plt.plot(X1, diff)
plt.show()

plt.xscale('log')
plt.xlabel('Freq (Hz)')
plt.ylabel('FFT Amplitude |X(freq)|')
plt.plot(A)
plt.plot(B)
plt.plot(X1,Al)
plt.plot(X2,Bl)
plt.show()