import pyaudio
import struct
import math
import time
import numpy as np

from PIL import Image, ImageDraw

import adafruit_blinka_raspberry_pi5_piomatter as piomatter

width = 192
height = 64

geometry = piomatter.Geometry(width=width, height=height, n_addr_lines=5,
                                                     rotation=piomatter.Orientation.Normal, n_planes=6, n_temporal_planes=0)

canvas = Image.new('RGB', (width, height), (0, 0, 0))
draw = ImageDraw.Draw(canvas)

framebuffer = np.asarray(canvas) + 0  # Make a mutable copy
matrix = piomatter.PioMatter(colorspace=piomatter.Colorspace.RGB888Packed,
                             pinout=piomatter.Pinout.AdafruitMatrixBonnet,
                             framebuffer=framebuffer,
                             geometry=geometry)

p = pyaudio.PyAudio()
info = p.get_host_api_info_by_index(0)
numdevices = info.get('deviceCount')

for i in range(0, numdevices):
    if (p.get_device_info_by_host_api_device_index(0, i).get('maxInputChannels')) > 0:
        print("Input Device id ", i, " - ", p.get_device_info_by_host_api_device_index(0, i).get('name'))

print(p.get_device_info_by_index(0)['defaultSampleRate'])

CHUNK = 1024
FORMAT = pyaudio.paInt16
CHANNELS = 2
RATE = 44100

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

def rms( data ):
    count = len(data)/2
    format = "%dh"%(count)
    shorts = struct.unpack( format, data )
    sum_squares = 0.0
    for sample in shorts:
        n = sample * (1.0/32768)
        sum_squares += n*n
    return math.sqrt( sum_squares / count )


run = True

max_vol = 0

while run:
    # print(rms(stream.read(CHUNK, exception_on_overflow = False)))

    print(max_vol)

    

    buffer = stream.read(CHUNK, exception_on_overflow = False)
    waveform = np.frombuffer(buffer, dtype=np.int16)

    fft_complex = np.fft.fft(waveform, n=int(CHUNK*2))

    max_val = math.sqrt(max(v.real * v.real + v.imag * v.imag for v in fft_complex))

    if (max_val > max_vol):
        max_vol = max_val
    
    # factor out the scale multiply.
    

    draw.rectangle((0,0,width,height), fill=0x000000)

    for i in range(192):

        scale_value = (height / max_vol) * (1 + (i/100))

        freq = i * 5

        if (i < 2):
            scale_value = (height / max_vol) * 0.9

        if (i < 192):
           
            freq = i
    
        target = int(freq)

        # print(target)
        

        v = fft_complex[target]
        
        dist = math.sqrt(v.real * v.real + v.imag * v.imag)

        mapped_dist = dist * scale_value

        if (max_vol < 500000):
            mapped_dist = 1
                
        draw.rectangle((i, height-mapped_dist, i, height), fill=0x880000)

    """
    for i,v in enumerate(fft_complex):
        dist = math.sqrt(v.real * v.real + v.imag * v.imag)
        mapped_dist = dist * scale_value
        
        draw.rectangle((i, 0, i, mapped_dist), fill=0x880000)
    """

    run=True

    max_vol = max_vol * 0.99
    
    framebuffer[:] = np.asarray(canvas)
    matrix.show()

    time.sleep(0.01)