# Programmer: Seonwoo Kang
# Purpose: Maker faire exhibit (2024.11 Orlando)
# Date: 2024.9
# 
import pygame
import serial
import time
from serial.tools import list_ports
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_community.llms import Ollama

# Initialize Pygame
pygame.init()

# Set up display
width, height = 800, 600
window = pygame.display.set_mode((width, height))
pygame.display.set_caption("Serial Port to Language Model")

# Set up fonts
font = pygame.font.Font(None, 48)
small_font = pygame.font.Font(None, 36)

# Set up colors
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)

# Initialize serial port
def find_com_port():
    ports = list_ports.comports()
    for port in ports:
        print(f"Found port: {port.device}")
    return ports[0].device if ports else None

com_port = find_com_port()
if com_port:
    ser = serial.Serial(port=com_port, baudrate=9600, timeout=1)
else:
    raise Exception("No COM port found")

# Function to read data from serial port
def read_from_serial():
    try:
        line = ser.readline().decode('utf-8').strip()
        return line
    except serial.SerialException:
        return "Error: Serial port not accessible"

# Define the prompt template and model
prompt = ChatPromptTemplate.from_template("tell me a short joke about {topic}")
model = Ollama(model="llama2")
output_parser = StrOutputParser()

chain = prompt | model | output_parser

# Main loop
running = True
morse_string = ""
generated_joke = ""
while running:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False

    # Read data from serial port
    serial_data = read_from_serial()
    if len(serial_data) > 1:
        morse_string = serial_data
        result = chain.invoke({"topic": morse_string})
        generated_joke = result

    # Clear screen
    window.fill(BLACK)

    # Render serial data
    serial_text = font.render("Morse message: " + morse_string, True, WHITE)
    window.blit(serial_text, (20, 20))

    # Render generated joke
    joke_lines = generated_joke.split('\n')
    y_offset = 100
    for line in joke_lines:
        joke_text = small_font.render(line, True, WHITE)
        window.blit(joke_text, (20, y_offset))
        y_offset += small_font.get_height() + 5  # Add some space between lines

    # Update display
    pygame.display.flip()

    # Delay to control the loop speed
    time.sleep(0.1)

# Close serial port
ser.close()

# Quit Pygame
pygame.quit()