from machine import Pin, I2C
from time import sleep
from rpn import *

from ssd1306 import SSD1306_I2C
import framebuf, sys
import utime

# import ulab as np


# Functions
## Hardware


def read_key(col_pins: list, row_pins: list):
    for cp in col_pins:
        cp.value(0)  # Set column pin to LOW
        for i, row_pin in enumerate(row_pins):
            if not row_pin.value():  # If row pin reads LOW
                return (i, col_pins.index(cp))
                cp.value(1)  # Set column pin back to HIGH

        cp.value(1)
    return (None, None)


# Setup
sleep(0.1)  # Wait for USB to become ready
STACK_DEPTH = 20
STACK_HISTORY_DEPTH = 5
stack = [None] * STACK_DEPTH
stack_history = [[None]*STACK_DEPTH]*STACK_HISTORY_DEPTH

stack_commands = stack.copy()
stack_history_commands = stack_history.copy()

row_pins = [Pin(pin) for pin in [8, 9, 10, 11, 12]]
col_pins = [Pin(pin) for pin in [13, 14, 15, 16, 17]]

layer1 = [
    [None, "Swap", "Duplicate", "Undo", "Backspace"],
    [7, 8, 9, "1/x", "+-"],
    [4, 5, 6, "*", "/"],
    [1, 2, 3, "+", "-"],
    ["SH", 0, ".", "E", "Enter"],
]

for pin in row_pins:
    pin.init(Pin.IN, Pin.PULL_UP)

for pin in col_pins:
    pin.init(Pin.OUT)


# Set up display

pix_res_x = 128 
pix_res_y = 64 
disp_col = pix_res_x//8
disp_row = pix_res_y//8
font_size = 8

def init_i2c(scl_pin, sda_pin):
    # Initialize I2C device
    i2c_dev = I2C(1, scl=Pin(scl_pin), sda=Pin(sda_pin), freq=200000)
    i2c_addr = [hex(ii) for ii in i2c_dev.scan()]
    
    if not i2c_addr:
        print('No I2C Display Found')
        sys.exit()
    else:
        print("I2C Address      : {}".format(i2c_addr[0]))
        print("I2C Configuration: {}".format(i2c_dev))
    
    return i2c_dev

i2c_dev = init_i2c(scl_pin=27, sda_pin=26)
oled = SSD1306_I2C(pix_res_x, pix_res_y, i2c_dev)

# Write side numbers

# for i in range(disp_row-1):
#     oled.text(f"{disp_row-1-i}:", 0, i*font_size)
# oled.text(">", 0, (disp_row-1)*font_size)
# oled.show()

max_cmd_length = len("x1+x2")+1

# Funtions to update "current" line
def clear_current(oled, text):
    oled.fill_rect((disp_col - len(text))*font_size, (disp_row-1)*font_size, len(text)*font_size, font_size, 0)
def write_current(oled, text):
    if len(text) > disp_col:
        text = text[-disp_col:]
    oled.text(text, (disp_col - len(text))*font_size, (disp_row-1)*font_size)

def format_number(number):
    if int(number) == number:
        number = int(number)
    if len(str(number)) > disp_col - max_cmd_length:
        return f"{number:.{disp_col-2-1-2-4}e}"
    else:
        return f"{number}"

def display_stack(oled, stack):
    for i, v in enumerate(stack[:disp_row]):
        oled.fill_rect(max_cmd_length*font_size, (disp_row-2-i)*font_size, (disp_col-max_cmd_length)*font_size, font_size, 0)
        if v is not None:
            num_s = format_number(v)
            oled.text(f"{num_s}", (disp_col - len(str(num_s)))*font_size, (disp_row-2-i)*font_size)

def display_commands(oled, stack_commands):
    for i, v in enumerate(stack_commands[:disp_row]):
        oled.fill_rect(0, (disp_row-2-i)*font_size, max_cmd_length*font_size, font_size, 0 )
        if v is not None:
            oled.text(v, 0, (disp_row-2-i)*font_size)
current = ""
# Run
print("Start")
while True:
    key = read_key(col_pins, row_pins)
    if key != (None, None):
        sleep(0.2)  # Debounce
        value = layer1[key[0]][key[1]]

        if value in [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]:
            clear_current(oled, current)
            current += str(value)
            write_current(oled, current)
            oled.show()
        elif (value == ".") and ("." not in current) and ("E" not in current):
                clear_current(oled, current)
                current += str(value)
                write_current(oled, current)
                oled.show()
        elif (value == "E") and ("E" not in current):
                clear_current(oled, current)
                current += str(value)
                write_current(oled, current)
                oled.show()
  
        elif (value == "Backspace"):
            if (current != ""):
                clear_current(oled, current)
                current = current[:-1]
                write_current(oled, current)
                oled.show()
            else:
                stack_history = update_stack_history(stack, stack_history)
                stack_history_commands = update_stack_history(stack_commands, stack_history_commands)
                stack, _ = pop(stack)
                stack_commands, _ = pop(stack_commands)

                display_stack(oled, stack)
                display_commands(oled, stack_commands)
                oled.show()

        elif value == "Enter":
            if (current != "") and (current[-1] != "E"):
                stack_history = update_stack_history(stack, stack_history)
                stack_history_commands = update_stack_history(stack_commands, stack_history_commands)
                stack = push(stack, float(current))
                stack_commands = push(stack_commands, ">")

                clear_current(oled, current)
                current = ""
                display_stack(oled, stack)
                display_commands(oled, stack_commands)
                oled.show()
        elif (value == "Swap") and (current == ""):
            stack_history = update_stack_history(stack, stack_history)
            stack_history_commands = update_stack_history(stack_commands, stack_history_commands)
            stack = swap(stack)
            stack_commands = swap(stack_commands)

            display_stack(oled, stack)
            display_commands(oled, stack_commands)
            oled.show()
        elif (value == "Duplicate") and (current == ""):
            stack_history = update_stack_history(stack, stack_history)
            stack_history_commands = update_stack_history(stack_commands, stack_history_commands)
            stack = duplicate(stack)
            stack_commands = push(stack_commands, "\"")

            display_stack(oled, stack)
            display_commands(oled, stack_commands)
            oled.show()
        elif (value == "Undo"):
            stack, stack_history = undo(stack_history)
            stack_commands, stack_history_commands = undo(stack_history_commands)

            clear_current(oled, current)
            current = ""
            display_stack(oled, stack)
            display_commands(oled, stack_commands)
            oled.show()

        elif value in functions_x.keys():
            if current != "":
                stack_history = update_stack_history(stack, stack_history)
                stack_history_commands = update_stack_history(stack_commands, stack_history_commands)
                stack = push(stack, float(current))
                stack_commands = push(stack_commands, ">")

                clear_current(oled, current)
                current = ""
            if (stack[0] is not None):
                stack_history = update_stack_history(stack, stack_history)
                stack_history_commands = update_stack_history(stack_commands, stack_history_commands)
                stack = apply_function_x(stack, functions_x[value])
                
                stack_commands, _ = pop(stack_commands)
                stack_commands = push(stack_commands, f"{value}")
                
                display_stack(oled, stack)
                display_commands(oled, stack_commands)
                oled.show()
        elif value in functions_xy.keys():
            if current != "":
                stack_history = update_stack_history(stack, stack_history)
                stack_history_commands = update_stack_history(stack_commands, stack_history_commands)
                stack = push(stack, float(current))
                stack_commands = push(stack_commands, ">")

                clear_current(oled, current)
                current = ""
                
            if (stack[0] is not None) and (stack[1] is not None):
                stack_history = update_stack_history(stack, stack_history)
                stack_history_commands = update_stack_history(stack_commands, stack_history_commands)
                stack = apply_function_xy(stack, functions_xy[value])
                stack_commands, _ = pop(stack_commands)
                stack_commands, _ = pop(stack_commands)
                stack_commands = push(stack_commands, f"x1{value}x2")
            if current == "":
                display_stack(oled, stack)
                display_commands(oled, stack_commands)
                oled.show()
        elif value == "SH":
            print(stack_commands)