import array
import pyaudio
import sys
import colorsys
import time
import alsaaudio as aa
from struct import unpack
import numpy as np
from sys import exit
import numpy as np
import unicornhat as unicorn
import math
import signal
import sys

# Gamma correction lookup table.  This is adapted from the table at:
#   https://learn.adafruit.com/led-tricks-gamma-correction/the-quick-fix
GAMMA8 = array.array('B', (
    0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
    0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  1,  1,  1,  1,
    1,  1,  1,  1,  1,  1,  1,  1,  1,  2,  2,  2,  2,  2,  2,  2,
    2,  3,  3,  3,  3,  3,  3,  3,  4,  4,  4,  4,  4,  5,  5,  5,
    5,  6,  6,  6,  6,  7,  7,  7,  7,  8,  8,  8,  9,  9,  9, 10,
   10, 10, 11, 11, 11, 12, 12, 13, 13, 13, 14, 14, 15, 15, 16, 16,
   17, 17, 18, 18, 19, 19, 20, 20, 21, 21, 22, 22, 23, 24, 24, 25,
   25, 26, 27, 27, 28, 29, 29, 30, 31, 32, 32, 33, 34, 35, 35, 36,
   37, 38, 39, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 50,
   51, 52, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 66, 67, 68,
   69, 70, 72, 73, 74, 75, 77, 78, 79, 81, 82, 83, 85, 86, 87, 89,
   90, 92, 93, 95, 96, 98, 99,101,102,104,105,107,109,110,112,114,
  115,117,119,120,122,124,126,127,129,131,133,135,137,138,140,142,
  144,146,148,150,152,154,156,158,160,162,164,167,169,171,173,175,
  177,180,182,184,186,189,191,193,196,198,200,203,205,208,210,213,
  215,218,220,223,225,228,231,233,236,239,241,244,247,249,252,255
))


SPEED = 0.5
CHUNK = 512
FORMAT = pyaudio.paInt16
CHANNELS = 1
RATE = 16000

chunk = CHUNK
sample_rate = RATE

p = pyaudio.PyAudio()

stream = p.open(
     format = FORMAT,
     channels = CHANNELS,
     rate = RATE,
     input = True,
     frames_per_buffer = CHUNK,
     )

print("""Spectrum Analyzer with Unicorn pHAT""")

unicorn.set_layout(unicorn.AUTO)
unicorn.rotation(0)
unicorn.brightness(0.4)
width,height=unicorn.get_shape()

matrix    = [0, 0, 0, 0, 0, 0, 0, 0 ]
smoothbins= [0, 0, 0, 0, 0, 0, 0, 0 ]
power     = []
weighting = [1, 1, 2, 4, 8, 8, 8, 8]
hues      = [0, 50, 60, 120, 170, 240, 280, 310]


def HSV_to_RGB(h, s, v):
    """HSV color space to RGB color space conversion. Hue (h) should be a
    degree value from 0.0 to 360.0, saturation (s) and value (v) should be a
    value from 0.0 to 1.0.  Returns a 3-tuple of gamma-corrected RGB color
    bytes (0-255).
    """
    # This is adapted from C/C++ code here:
    #  https://www.cs.rit.edu/~ncs/color/t_convert.html
    r = 0
    g = 0
    b = 0
    if s == 0.0:
        r = v
        g = v
        b = v
    else:
        h /= 60.0       # sector 0 to 5
        i = int(math.floor(h))
        f = h - i       # factorial part of h
        p = v * ( 1.0 - s )
        q = v * ( 1.0 - s * f )
        t = v * ( 1.0 - s * ( 1.0 - f ) )
        if i == 0:
            r = v
            g = t
            b = p
        elif i == 1:
            r = q
            g = v
            b = p
        elif i == 2:
            r = p
            g = v
            b = t
        elif i == 3:
            r = p
            g = q
            b = v
        elif i == 4:
            r = t
            g = p
            b = v
        else:
            r = v
            g = p
            b = q
    r = GAMMA8[int(255.0*r)]
    g = GAMMA8[int(255.0*g)]
    b = GAMMA8[int(255.0*b)]
    return (r, g, b)



def signal_handler(signal, frame):
     print("* closing sream")
     stream.stop_stream()
     stream.close()
     p.terminate()
     sys.exit(0)

def power_index(val):
     return int(2 * chunk * val / sample_rate)

def compute_fft(data, chunk, sample_rate):
     global matrix
     data = unpack("%dh" % (len(data) / 2), data)
     data = np.array(data, dtype='h')

     fourier = np.fft.rfft(data)
     fourier = np.delete(fourier, len(fourier) - 1)

     power = np.abs(fourier)
     matrix[0] = int(np.mean(power[power_index(0)    :power_index(62) :1]))
     matrix[1] = int(np.mean(power[power_index(62)  :power_index(125) :1]))
     matrix[2] = int(np.mean(power[power_index(125)  :power_index(250) :1]))
     matrix[3] = int(np.mean(power[power_index(250)  :power_index(500) :1]))
     matrix[4] = int(np.mean(power[power_index(500) :power_index(1000) :1]))
     matrix[5] = int(np.mean(power[power_index(1000) :power_index(2000) :1]))
     matrix[6] = int(np.mean(power[power_index(2000) :power_index(4000) :1]))
     matrix[7] = int(np.mean(power[power_index(2000) :power_index(8000) :1]))

     matrix = np.divide(np.multiply(matrix, weighting), 10000)
     matrix = [float(m) for m in matrix]
     matrix = np.clip(matrix,0, 4)
     return matrix

signal.signal(signal.SIGINT, signal_handler)
data = stream.read(CHUNK)


while data != '':
     matrix = compute_fft(data, chunk, sample_rate)

     for x in range(width): # scans  0 1 2 3 4 5 6 7
          avg = matrix[x]
          if avg > smoothbins[x]:
              smoothbins[x] = avg
          barheight = smoothbins[x]
          #barheight = matrix[x] # [0.0,4.0]
          #print 'Barheight {}'.format(barheight)
          for y in range(height): # scans 0 1 2 3
               Brightness = 1.0;
               if y == int(barheight):
                   Brightness = math.fmod(barheight, 1.0);
                   #print 'Half Brightness {}' .format(Brightness)
               elif y > barheight:
                   Brightness = 0
               ScaledBrightness = Brightness

               rgb = HSV_to_RGB(hues[x], 1.0, 1.0)
               #print 'RGB = {} Hue = {}' .format(rgb,hues[x])

               unicorn.set_pixel(7-x,y,int(ScaledBrightness*rgb[0]),int(ScaledBrightness*rgb[1]),int(ScaledBrightness*rgb[2]))
               unicorn.show()
          smoothbins[x] *= SPEED
     data = stream.read(CHUNK, exception_on_overflow = False)

stream = p.close