from serial import Serial
import numpy as np
import atexit, time, math, psutil
from PIL import Image, ImageDraw, ImageFont, BdfFontFile

# somewhat from https://github.com/Encampeded/led_matrix.py

def send(command_id: int, parameters: list):
    serial_port.write(bytearray([0x32, 0xAC, command_id]) + bytearray(parameters))

#Sends in columns. Slower, but supports different brightness levels
def csend(m, brightness):
    send(0x00,[brightness])
    for i in range(9):
        send(0x07, [i, *m[i]])
    send(0x08, [])

#Sends whole image at once. Faster, but 1 bit display
def qsend(m, brightness):
    send(0x06, np.packbits(m.swapaxes(1,0) > 127,bitorder='little').tobytes())
    send(0x00,[brightness])

#calculate brightness from ambient light sensor, with smoothing
def calcAmbientBrightness(brightness):
    with open("/sys/bus/iio/devices/iio:device0/in_illuminance_raw") as f:
        currentBrightness = min(float(f.read())**.75 + 3,255) #curve based on personal preference
        if brightness == 0: #initialize
            return currentBrightness
        else:
            return brightness * .8 + currentBrightness * .2 #smoothing of brightness value over time

#calculate brightness based on screen
def calcScreenBrightness():
    with open("/sys/class/backlight/amdgpu_bl1/actual_brightness") as f:
        return float(f.read())/2000 + 2 #curve based on personal preference (update calcBrightnessMeta to match maximum)

#use ambient light sensor if it's brighter and the screen is at max brightness
def calcBrightnessMeta(brightness):
    screenBrightness = calcScreenBrightness()
    if screenBrightness > 34:
        brightnessWithAmbient = calcAmbientBrightness(brightness)
        return max(screenBrightness, brightnessWithAmbient)
    else:
        return screenBrightness

def fillLEDInLine(m, x, y, value, i):
    if i == math.floor(value):
        m[x,y] = 255*(value%1)
    else:
        m[x,y] = 255

def batteryIndicator(m, charge_full):
    with open("/sys/class/power_supply/BAT1/charge_now") as f: #more precise than reading the percentage
        value = 82*float(f.read())/charge_full
    num_LEDs = math.ceil(value)
    for i in range(0,min(num_LEDs,34)):
        fillLEDInLine(m,0,i,value,i)
    for i in range(34,min(num_LEDs,42)):
        fillLEDInLine(m,i-33,33,value,i)
    for i in range(42,min(num_LEDs,75)):
        fillLEDInLine(m,8,74-i,value,i)
    for i in range(75,min(num_LEDs,82)):
        fillLEDInLine(m,82-i,0,value,i)

def cpuUsageIndicator(m):
    arr = np.reshape(psutil.cpu_percent(percpu=True),(4,4)) * 2.55 #shape (currently 4x4) must mach number of CPU threads
    for index, value in np.ndenumerate(arr):
        m[index[0]+1,index[1]+29] = value

def gpuUsageIndicator(m):
    with open("/sys/class/drm/card1/device/gpu_busy_percent") as f:
        value = 12*float(f.read())/100
    num_LEDs = math.ceil(value)
    for x in range(3):
        for i in range(4*x,min(num_LEDs,4+4*x)):
            fillLEDInLine(m,x+5,32+4*x-i,value,i)

def ramIndicator(m):
    value = 14 * psutil.virtual_memory().percent/100
    num_LEDs = math.ceil(value)
    for y in range(2):
        for i in range(7*y,min(num_LEDs,7+7*y)):
            fillLEDInLine(m,1+i-7*y,28-y,value,i)

def swapIndicator(m):
    value = 14 * psutil.swap_memory().percent/100
    num_LEDs = math.ceil(value)
    for y in range(2):
        for i in range(7*y,min(num_LEDs,7+7*y)):
            fillLEDInLine(m,1+i-7*y,26-y,value,i)

def clock(m, font):
    img = Image.new("L",(34,9))
    draw = ImageDraw.Draw(img)
    t = time.strftime("%I:%M:%S", time.localtime())[0:7] #Change %I to %H for 24-hour time. Last seconds digit is truncated for space.
    draw.text((1,2), t, fill=255, font=font)
    m += np.flip(np.asarray(img),axis=0)

#Initialization
serial_port = Serial("/dev/ttyACM0", 115200)
atexit.register(serial_port.close)
brightness = 0
with open("/sys/class/power_supply/BAT1/charge_full") as f:
    charge_full = float(f.read())
with open("tinyfont.bdf", "rb") as fp:
        font = BdfFontFile.BdfFontFile(fp).to_imagefont()

#Refresh
while True:
    with open("/sys/class/backlight/amdgpu_bl1/device/enabled") as f:
        if f.read() == "disabled\n": #when display turns off for power saving
            send(0x03,[True]) #sleep LED matrix
            time.sleep(.25)
            continue

    brightness = calcBrightnessMeta(brightness)

    m = np.zeros((9,34),dtype=np.uint8) #image to be displayed

    cpuUsageIndicator(m)
    gpuUsageIndicator(m)
    ramIndicator(m)
    swapIndicator(m)
    clock(m, font)
    batteryIndicator(m, charge_full)

    csend(m, int(brightness)) #send image to display, in practice the speed of the script is limited by this step