#!/home/taco/venvs/visualizer/bin/python

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

# LED panel width
width = 192

# LED panel height
height = 64

# Top gradient color
color_1 = "#0000ff"

# Bottom gradient color
color_2 = "#ff0000"

# Should we mirror left and right channels?
mirror = False

# Should we force draw a single pixel line?
zero_db_line = True

# how long to hold maximum volume in seconds?
max_db_hold_time = 1

# how long between steps should we sleep?
delay = 0.01

# minimum maximum volume
minimum_max_volume = 4000000

# PyAudio config
CHUNK = 1024
FORMAT = pyaudio.paInt16
CHANNELS = 2
RATE = 44100

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

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()

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

run = True

max_vol = minimum_max_volume
target_max_vol = minimum_max_volume

left_channel = np.zeros(int(width))
hf_left_channel = np.zeros(int(width/2))
right_channel = np.zeros(int(width))
hf_right_channel = np.zeros(int(width/2))



def generate_gradient(colour1: str, colour2: str, width: int, height: int) -> Image:
    base = Image.new('RGB', (width, height), colour1)
    top = Image.new('RGB', (width, height), colour2)
    mask = Image.new('L', (width, height))
    mask_data = []
    for y in range(height):
        mask_data.extend([int(255 * (y / height))] * width)
    mask.putdata(mask_data)
    base.paste(top, (0, 0), mask)
    return base

def clamp(minimum, maximum, value):
    if (value > maximum):
        return maximum

    if (value < minimum):
        return minimum

    return value

def mathCurve(step):
    result = math.sin((step/40) + sin_offset)*4 + math.cos((step/30) + sin_offset*3)*4 + math.sin((step/35) + sin_offset*5)*4
    return math.floor(height/2 + result + 0.5)

sin_offset = 0

db_hold_time = 0

gradient = generate_gradient(color_1, color_2, width, height)

while run:

    left_channel = left_channel * 0.9
    right_channel = right_channel * 0.9

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

    waveform = np.reshape(waveform, (CHUNK, 2))

    #fft_complex_left = np.fft.fft(waveform[:, 0], n=int(CHUNK))[:width]
    #fft_complex_right = np.fft.fft(waveform[:, 1], n=int(CHUNK))[:width]

    fft_complex_left = np.fft.fft(waveform[:, 0], n=int(CHUNK))
    fft_complex_right = np.fft.fft(waveform[:, 1], n=int(CHUNK))

    max_val_left = math.sqrt(max(v.real * v.real + v.imag * v.imag for v in fft_complex_left))
    max_val_right = math.sqrt(max(v.real * v.real + v.imag * v.imag for v in fft_complex_right))

    max_val = max(max_val_left, max_val_right)

    if (max_val > target_max_vol):
        target_max_vol = max_val
        db_hold_time = time.time() + max_db_hold_time

    if (max_vol < target_max_vol):
        max_vol = max_vol + target_max_vol*0.1

    canvas.paste(gradient)

    def calcDist(step, fft, fft_hist):
        scale_value = (height / max_vol) * (1 + (step/100))

        if (step < 2):
            scale_value = (height / max_vol) * 0.9
    
        use_value = fft_hist[step]

        v = fft[step]

        dist = math.sqrt(v.real * v.real + v.imag * v.imag)

        if dist > use_value:
            fft_hist[step] = dist
            use_value = dist

        return (use_value * scale_value)
        
    # LEFT FR
    for i in range(0, int(width)):

        mapped_dist = calcDist(i, fft_complex_left, left_channel)/2

        midpoint = mathCurve(i)

        draw.rectangle((i, 0, i, clamp(1, midpoint, midpoint - mapped_dist)), fill=0x000)
    
    # RIGHT FR
    for i in range(0, int(width)):
        mapped_dist = calcDist(i, fft_complex_right, right_channel)/2

        horizontal_position = (width - 1) - i
        if (mirror):
            horizontal_position = i

        midpoint = mathCurve(horizontal_position)

        vertical_addition = 0
        if (zero_db_line):
            vertical_addition = 1

        draw.rectangle((horizontal_position, clamp(1, height, midpoint + vertical_addition + mapped_dist), horizontal_position, height), fill=0x000)

    if (max_vol > minimum_max_volume and db_hold_time < time.time()):
        max_vol = target_max_vol = max_vol * 0.99

    if (max_vol < minimum_max_volume):
        max_vol = minimum_max_volume
    
    framebuffer[:] = np.asarray(canvas)
    matrix.show()

    sin_offset = sin_offset + 0.01

    time.sleep(delay)