import pyaudio
import wave
import os

import numpy as np
import struct

class Listener:
	def __init__(self, dataTime = 1/20, agcTime = 10, input = False):
		self._dataTime = dataTime
		self._agcTime = agcTime
		self._input = input
		self.buffersize = 512
		self.p = pyaudio.PyAudio()
		
		self.left = self.right = self.fft = []
		
		self._agcMaxima = [0]
		self._agcIndex = 0
		self._agcLen = 0
		self._beatCb = None
		self._sampleRate = 0
		self._hasNewData = False
		
	def start(self, dev = None):
		self._device = dev or self.getDefaultOutputDeviceInfo()
			
		if self._device == None:
			print("no device found")
			return
			
		print("device name: {} channels: {} defaultSampleRate: {}".format(self._device["name"], self._device["channels"], self._device["defaultSampleRate"]))
			
		if self._sampleRate != self._device["defaultSampleRate"]:
			self._sampleRate = self._device["defaultSampleRate"]
			self.buffersize = int(self._sampleRate * self._dataTime)
			self.fft = self.right = self.left = np.ndarray((self.buffersize))
			
			self._agcLen = int(1 / self._dataTime * self._agcTime)
			self._agcMaxima = np.ndarray((self._agcLen))
			self._agcMaxima.fill(2**15 * 0.1)
			self._agcIndex = 0
			self._lastBeatTime = 0
			self.meanAmp = 2**15 * 0.1
		
		try:
			self._stream = self.openStream(self._device)
		except OSError:
			self._stream = None
		
		if not self._stream:
			print("stream open failed")
			return
		
		self._stream.start_stream()
		
		return self._stream.is_active()
		
	def stop(self):
		if not self._stream:# or not self._stream.is_active():
			return False
		self._stream.stop_stream()
		return True
		
	def setBeatCb(self, cb):
		self._beatCb = cb
	
	def getDefaultOutputDeviceInfo(self):
		#Set default to first in list or ask Windows
		try:
			self.p.terminate()
			self.p.__init__()
			if self._input:
				info = self.p.get_default_input_device_info()
			else:
				info = self.p.get_default_output_device_info()
		except IOError:
			info = None
		
		#Handle no devices available
		if info == None:
				print ("No device available.")
				return None

		if (self.p.get_host_api_info_by_index(info["hostApi"])["name"]).find("WASAPI") == -1:
			for i in range(0, self.p.get_device_count()):
				x = self.p.get_device_info_by_index(i)
				is_wasapi = (self.p.get_host_api_info_by_index(x["hostApi"])["name"]).find("WASAPI") != -1
				if x["name"].find(info["name"]) >= 0 and is_wasapi:
					info = x
					break

		#Handle no devices available
		if info == None:
			print ("Device doesn't support WASAPI")
			return None
			
		info["channels"] = info["maxInputChannels"] if (info["maxOutputChannels"] < info["maxInputChannels"]) else info["maxOutputChannels"]
			
		return info
		
	def openStream(self, dev):
	
		is_input = dev["maxInputChannels"] > 0
		is_wasapi = (self.p.get_host_api_info_by_index(dev["hostApi"])["name"]).find("WASAPI") != -1
		
		#print("is input: {} is wasapi: {}".format(is_input, is_wasapi))

		if not is_input and not is_wasapi:
			print ("Selection is output and does not support loopback mode.")
			return None
		if is_wasapi:
			stream = self.p.open(
				format = pyaudio.paInt16,
				channels = (dev["channels"] if dev["channels"] < 2 else 2),
				rate = int(self._sampleRate),
				input = True,
				frames_per_buffer = self.buffersize,
				input_device_index = dev["index"],
				stream_callback=self.streamCallback,
				as_loopback = False if is_input else is_wasapi)
		else:
			stream = self.p.open(
				format = pyaudio.paInt16,
				channels = (dev["channels"] if dev["channels"] < 2 else 2),
				rate = int(self._sampleRate),
				input = True,
				frames_per_buffer = self.buffersize,
				input_device_index = dev["index"],
				stream_callback=self.streamCallback)
		return stream
		
	def closeStream(self):
		if not self._stream:
			return False
		self._stream.close()
		return True
			
	def streamCallback(self, buf, frame_count, time_info, flag):
		self._buf = buf
		arr = np.array(struct.unpack("%dh" % (len(buf)/2), buf))
		
		mx = arr.max()
		self._agcIndex += 1
		if self._agcIndex >= self._agcLen:
			self._agcIndex = 0
		self._agcMaxima[self._agcIndex] = mx
			
		self.meanAmp = np.mean(np.absolute(self._agcMaxima))
		
		if self.meanAmp > 2**15 * 0.02:
			amp = 1 / self.meanAmp
		else:
			amp = 1 / (2**15 * 0.02)
		
		if self._device["channels"] >= 2:
			self.left, self.right  = arr[::2]  * amp, arr[1::2] * amp
			self.fft = self.fftCalc((self.left+self.right)/2)
		else:
			self.left = self.right = arr * amp
			self.fft = self.fftCalc(self.left)
		
		self._hasNewData = True
		
		self.agcFFT = np.fft.rfft(self._agcMaxima, self.beatnFFT) / self.beatnFFT
		
		if self._beatCb and mx * (time_info["current_time"] - self._lastBeatTime) > m * 0.5:
			self._lastBeatTime = time_info["current_time"]
			self._beatCb(self.fft)
			
		return (None, pyaudio.paContinue)
	
	def hasNewData(self):
		if not self._hasNewData:
			return False
		self._hasNewData = False
		return True
	
	def getSampleRate(self):
		return int(self._sampleRate)
	
	def getAgc(self):
		return self.meanAmp / 2**15
		
	def getVolume(self):
		if self._agcMaxima.sum() == 0:
			return 0
		return self._agcMaxima[self._agcIndex] / self.meanAmp
	
	def isActive(self):
		if not self._stream:
			return False
		return self._stream.is_active()
	
	def fftSetLimits(self, nFFT, fMin, fMax):
		self.nFFT = nFFT
		self.fftMin = int(fMin / self._sampleRate * nFFT)
		self.fftMax = int(fMax / self._sampleRate * nFFT)
		print("nFFT: {} \tfftMin: {} \tfftMax: {}".format(self.nFFT, self.fftMin, self.fftMax))

	def agcFftSetLimits(self, fMin, fMax):
		self.beatnFFT = self._agcLen
		self.beatFftMin = int(fMin * self._dataTime * self.beatnFFT)
		self.beatFftMax = int(fMax * self._dataTime * self.beatnFFT)
		print("beat nFFT: {} \tfftMin: {} \tfftMax: {}".format(self.beatnFFT, self.beatFftMin, self.beatFftMax))
		
	def fftCalc(self, data):
		return abs(np.fft.rfft(data, self.nFFT)[self.fftMin:self.fftMax]) / self.nFFT
							
	def fftGroup(self, fft, limits):
		groups = []
		for freqs in zip(limits, limits[1:]):
			a = int(freqs[0] / self._sampleRate * self.nFFT)
			b = int(freqs[1] / self._sampleRate * self.nFFT)
			#groups.append(sum(fft[a:b]) / (b-a) if (b-a) > 0 else 0)
			if b != a:
				groups.append(max(fft[a:b]))
			else:
				groups.append(fft[a])
		return groups