#!/usr/bin/env python

# 8 band Audio equalizer using wav file as the source.
# Source 2013 - https://www.raspberrypi.org/forums/viewtopic.php?p=314087
# Source 2012 - https://learn.adafruit.com/piccolo

# 8 band Audio equaliser from wav file
 
import alsaaudio as aa
#import smbus
from adafruit_servokit import ServoKit
from time import sleep
from struct import unpack
import numpy as np
import wave

# Set channels to the number of servo channels on your kit.
# 8 for FeatherWing, 16 for Shield/HAT/Bonnet.
kit = ServoKit(channels=16)

def ResetAllServos():
	for i in range(len(kit.servo)):
		kit.servo[i].angle = 0

def Set_Column(row, col):
	# Since this code is designed for 8x8 matrix, there are 8 rows and 8 columns.
	# Servo angle could be 0-160 degrees, so multiply row by 20 to get angle.
	my_angle = row*20
	print("row {}".format(row))
	print("ang {}".format(my_angle))
	kit.servo[col].angle = my_angle

# Initialize position of all servos
ResetAllServos()
matrix    = [0,0,0,0,0,0,0,0]
power     = []
weighting = [2,2,8,8,16,32,64,64] # Change these according to taste

# Set up audio
wavfile = wave.open('/home/pi/heartache.wav','r')
sample_rate = wavfile.getframerate()
no_channels = wavfile.getnchannels()
chunk       = 4096 # Use a multiple of 8
output = aa.PCM(aa.PCM_PLAYBACK, aa.PCM_NORMAL)
output.setchannels(no_channels)
output.setrate(sample_rate)
output.setformat(aa.PCM_FORMAT_S16_LE)
output.setperiodsize(chunk)

# Return power array index corresponding to a particular frequency
def piff(val):
	return int(2*chunk*val/sample_rate)
	
def calculate_levels(data, chunk, sample_rate):
	global matrix
	# Convert raw data (ASCII string) to numpy array
	data = unpack("%dh"%(len(data)/2),data)
	data = np.array(data, dtype='h')
	# Apply FFT - real data
	fourier=np.fft.rfft(data)
	# Remove last element in array to make it the same size as chunk
	fourier=np.delete(fourier,len(fourier)-1)
	# Find average 'amplitude' for specific frequency ranges in Hz
	power = np.abs(fourier)	
	matrix[0]= int(np.mean(power[piff(0)    :piff(156):1]))
	matrix[1]= int(np.mean(power[piff(156)  :piff(313):1]))
	matrix[2]= int(np.mean(power[piff(313)  :piff(625):1]))
	matrix[3]= int(np.mean(power[piff(625)  :piff(1250):1]))
	matrix[4]= int(np.mean(power[piff(1250) :piff(2500):1]))
	matrix[5]= int(np.mean(power[piff(2500) :piff(5000):1]))
	matrix[6]= int(np.mean(power[piff(5000) :piff(10000):1]))
	matrix[7]= int(np.mean(power[piff(10000):piff(20000):1]))
	# Tidy up column values for the LED matrix
	matrix=np.divide(np.multiply(matrix,weighting),1000000)
	# Set floor at 0 and ceiling at 8 for LED matrix
	matrix=matrix.clip(0,8) 
	return matrix

print ("Processing.....")

data = wavfile.readframes(chunk)
while data!='':
	output.write(data)	
	matrix=calculate_levels(data, chunk, sample_rate)
	# np.rint() changes 7.999 to 8. and astype(int) changes it to just 8
	matrix=np.rint(matrix).astype(int)
	#print (matrix)
	for i in range (0,8):
		# Took out complicated math. Matrix should have values 0-8 as for loops through rows 0-8.
		#Set_Column((1<<matrix[i])-1,0xFF^(1<<i))
		Set_Column(matrix[i],i)
	data = wavfile.readframes(chunk)
	# LED refresh speed was WAY too fast for servos, so added sleep.
	sleep(0.02)
	#ResetAllServos()