# Major library imports
import atexit
import pyaudio
from numpy import zeros, short, fromstring, array
from numpy.fft import fft

import colorsys
import time
import blup.frame
import blup.output

#NUM_SAMPLES = 512
NUM_SAMPLES = 36
#NUM_SAMPLES = 18
#NUM_SAMPLES = 16
#SAMPLING_RATE = 11025
#SAMPLING_RATE = 11024 / 2
SAMPLING_RATE = 1100

_stream = None

def read_fft():
    global _stream
    pa = None

    def cleanup_audio():
        if _stream:
            _stream.stop_stream()
            _stream.close()
        pa.terminate()

    if _stream is None:
        pa = pyaudio.PyAudio()
        _stream = pa.open(format=pyaudio.paInt16, channels=1,
                           rate=SAMPLING_RATE,
                           input=True, frames_per_buffer=NUM_SAMPLES)
        atexit.register(cleanup_audio)

    audio_data  = fromstring(_stream.read(NUM_SAMPLES), dtype=short)
    normalized_data = audio_data / 32768.0

    return fft(normalized_data)[1:1+NUM_SAMPLES/2]

def flatten_fft(scale = 1.0):
    """
    Produces a nicer graph, I'm not sure if this is correct
    """
    for i, v in enumerate(read_fft()):
        yield scale * (i * v) / NUM_SAMPLES


def makeColor(x):
	hue = (7-x) / 7.0 / 3.0 -0.1
	if hue < 0:
		hue += 1
	rgb = map(lambda x: int(round(x*7)), colorsys.hsv_to_rgb(hue, 1, 1))
	return rgb
def makeSpectrumFrame(spectrum):
	f = blup.frame.Frame(blup.frame.FrameDimension(18,8,8,3))
	for x, y in enumerate(spectrum):
		if y > 7:
			y = 7
		y = int(y)
		for i in range(y):
			#val = (7,7,7)
			val = makeColor(i)
			f.setPixel(x, 7-i, val)
	return f

def draw():
    '''Draw 3 different colour graphs'''
    global NUM_SAMPLES
    audio = array(list(flatten_fft(scale = 80)))
    freqs = len(audio)
    bass, mid, treble = triple(audio)

    colours = (0.5, 1.0, 0.5), (1, 1, 0), (1, 0.2, 0.5)

    fill(0, 0, 1)
    rect(0, 0, WIDTH, 400)
    translate(50, 200)

    for spectrum, col in zip((bass, mid, treble), colours):
        fill(col)
        for i, s in enumerate(spectrum):
            rect(i, 0, 1, -abs(s))
        else:
            translate(i, 0)

    audio = array(list(flatten_fft(scale = 80)))


out = blup.output.getOutput('colorfulshell')
while True:
	audio = array(list(flatten_fft(scale=80)))
	#audio = map(abs, audio)
	#print len(audio)
	frame = makeSpectrumFrame([ int((x/18.0)*8) for x in range(18)])
	frame = makeSpectrumFrame(audio)
	#print map(int,audio)
	out.sendFrame(frame)
	#time.sleep(0.1)