"""
PicoCalc-Inspired Wokwi Template
MicroPython code for Raspberry Pi Pico with ST7789 display and keypad
"""

from machine import Pin, SPI
import time

# Display Configuration (ST7789 - 240x240)
class ST7789:
    def __init__(self, spi, cs, dc, rst, width=240, height=240):
        self.spi = spi
        self.cs = cs
        self.dc = dc
        self.rst = rst
        self.width = width
        self.height = height
        
        # Initialize pins
        self.cs.init(Pin.OUT, value=1)
        self.dc.init(Pin.OUT, value=0)
        self.rst.init(Pin.OUT, value=1)
        
        self.init_display()
    
    def write_cmd(self, cmd):
        self.cs(0)
        self.dc(0)
        self.spi.write(bytearray([cmd]))
        self.cs(1)
    
    def write_data(self, data):
        self.cs(0)
        self.dc(1)
        if isinstance(data, int):
            self.spi.write(bytearray([data]))
        else:
            self.spi.write(data)
        self.cs(1)
    
    def init_display(self):
        print("ST7789: Starting initialization...")
        
        # Hardware reset
        print("ST7789: Hardware reset...")
        self.rst(1)
        time.sleep_ms(50)
        self.rst(0)
        time.sleep_ms(50)
        self.rst(1)
        time.sleep_ms(150)
        
        # Software reset
        print("ST7789: Software reset...")
        self.write_cmd(0x01)
        time.sleep_ms(150)
        
        # Sleep out
        print("ST7789: Sleep out...")
        self.write_cmd(0x11)
        time.sleep_ms(255)
        
        # Color mode - 16-bit color
        print("ST7789: Setting color mode...")
        self.write_cmd(0x3A)
        self.write_data(0x55)
        time.sleep_ms(10)
        
        # Memory data access control
        print("ST7789: Setting rotation...")
        self.write_cmd(0x36)
        self.write_data(0x00)
        
        # Porch setting
        self.write_cmd(0xB2)
        self.write_data(bytearray([0x0C, 0x0C, 0x00, 0x33, 0x33]))
        
        # Gate control
        self.write_cmd(0xB7)
        self.write_data(0x35)
        
        # VCOM setting
        self.write_cmd(0xBB)
        self.write_data(0x19)
        
        # LCM control
        self.write_cmd(0xC0)
        self.write_data(0x2C)
        
        # VDV and VRH command enable
        self.write_cmd(0xC2)
        self.write_data(0x01)
        
        # VRH set
        self.write_cmd(0xC3)
        self.write_data(0x12)
        
        # VDV set
        self.write_cmd(0xC4)
        self.write_data(0x20)
        
        # Frame rate control
        self.write_cmd(0xC6)
        self.write_data(0x0F)
        
        # Power control
        self.write_cmd(0xD0)
        self.write_data(bytearray([0xA4, 0xA1]))
        
        # Positive voltage gamma control
        self.write_cmd(0xE0)
        self.write_data(bytearray([0xD0, 0x04, 0x0D, 0x11, 0x13, 0x2B, 0x3F, 0x54, 0x4C, 0x18, 0x0D, 0x0B, 0x1F, 0x23]))
        
        # Negative voltage gamma control
        self.write_cmd(0xE1)
        self.write_data(bytearray([0xD0, 0x04, 0x0C, 0x11, 0x13, 0x2C, 0x3F, 0x44, 0x51, 0x2F, 0x1F, 0x1F, 0x20, 0x23]))
        
        # Inversion on
        print("ST7789: Setting inversion...")
        self.write_cmd(0x21)
        
        # Normal display mode on
        print("ST7789: Normal display mode...")
        self.write_cmd(0x13)
        
        # Display on
        print("ST7789: Display on...")
        self.write_cmd(0x29)
        time.sleep_ms(100)
        
        print("ST7789: Initialization complete!")
    
    def set_window(self, x0, y0, x1, y1):
        """Set drawing window"""
        self.write_cmd(0x2A)  # Column address set
        self.write_data(bytearray([x0 >> 8, x0 & 0xFF, x1 >> 8, x1 & 0xFF]))
        
        self.write_cmd(0x2B)  # Row address set
        self.write_data(bytearray([y0 >> 8, y0 & 0xFF, y1 >> 8, y1 & 0xFF]))
        
        self.write_cmd(0x2C)  # Memory write
    
    def fill(self, color):
        """Fill screen with color (RGB565)"""
        print(f"Filling screen with color: 0x{color:04X}")
        self.set_window(0, 0, self.width - 1, self.height - 1)
        
        color_high = color >> 8
        color_low = color & 0xFF
        chunk = bytearray(512)
        for i in range(0, 512, 2):
            chunk[i] = color_high
            chunk[i+1] = color_low
        
        self.cs(0)
        self.dc(1)
        pixels = self.width * self.height
        full_chunks = pixels // 256
        remainder = (pixels % 256) * 2
        
        for _ in range(full_chunks):
            self.spi.write(chunk)
        if remainder > 0:
            self.spi.write(chunk[:remainder])
        self.cs(1)
        print("Fill complete")
    
    def fill_rect(self, x, y, w, h, color):
        """Fill a rectangle with color"""
        if x >= self.width or y >= self.height:
            return
        if x + w > self.width:
            w = self.width - x
        if y + h > self.height:
            h = self.height - y
        
        self.set_window(x, y, x + w - 1, y + h - 1)
        
        color_high = color >> 8
        color_low = color & 0xFF
        chunk = bytearray(200)
        for i in range(0, 200, 2):
            chunk[i] = color_high
            chunk[i+1] = color_low
        
        self.cs(0)
        self.dc(1)
        pixels = w * h
        full_chunks = pixels // 100
        remainder = (pixels % 100) * 2
        
        for _ in range(full_chunks):
            self.spi.write(chunk)
        if remainder > 0:
            self.spi.write(chunk[:remainder])
        self.cs(1)
    
    def text_large(self, string, x, y, color, bg_color=None):
        """Draw large text (simple 7-segment style digits)"""
        # Clear background if specified
        if bg_color is not None:
            self.fill_rect(x, y, len(string) * 30, 50, bg_color)
        
        char_x = x
        for char in string:
            self.draw_char_large(char, char_x, y, color)
            char_x += 30
    
    def draw_char_large(self, char, x, y, color):
        """Draw a single large character (7-segment style)"""
        # Simple 7-segment display patterns
        segments = {
            '0': [1,1,1,1,1,1,0],
            '1': [0,1,1,0,0,0,0],
            '2': [1,1,0,1,1,0,1],
            '3': [1,1,1,1,0,0,1],
            '4': [0,1,1,0,0,1,1],
            '5': [1,0,1,1,0,1,1],
            '6': [1,0,1,1,1,1,1],
            '7': [1,1,1,0,0,0,0],
            '8': [1,1,1,1,1,1,1],
            '9': [1,1,1,1,0,1,1],
            '-': [0,0,0,0,0,0,1],
            '+': [0,1,0,0,0,1,1],
            '*': [0,0,0,0,0,0,0],
            '/': [0,0,0,0,0,0,0],
            '=': [0,0,0,0,0,0,1],
            'E': [1,0,0,1,1,1,1],
            'r': [0,0,0,0,1,0,1],
            'C': [1,0,0,1,1,1,0],
            '.': [0,0,0,0,0,0,0],
        }
        
        if char not in segments:
            return
        
        seg = segments[char]
        
        # Draw segments (simplified rectangles)
        if seg[0]:  # top
            self.fill_rect(x+2, y, 20, 4, color)
        if seg[1]:  # top-right
            self.fill_rect(x+20, y+2, 4, 20, color)
        if seg[2]:  # bottom-right
            self.fill_rect(x+20, y+24, 4, 20, color)
        if seg[3]:  # bottom
            self.fill_rect(x+2, y+42, 20, 4, color)
        if seg[4]:  # bottom-left
            self.fill_rect(x, y+24, 4, 20, color)
        if seg[5]:  # top-left
            self.fill_rect(x, y+2, 4, 20, color)
        if seg[6]:  # middle
            self.fill_rect(x+2, y+21, 20, 4, color)
        
        # Special case for decimal point
        if char == '.':
            self.fill_rect(x+10, y+40, 4, 4, color)


# Keypad Configuration (4x4 matrix)
class Keypad:
    def __init__(self, row_pins, col_pins):
        self.rows = [Pin(p, Pin.OUT) for p in row_pins]
        self.cols = [Pin(p, Pin.IN, Pin.PULL_DOWN) for p in col_pins]
        
        # Key mapping for calculator (Wokwi membrane keypad layout)
        # Wokwi keypad is: 1,2,3,A / 4,5,6,B / 7,8,9,C / *,0,#,D
        self.keys = [
            ['1', '2', '3', '/'],  # Row 1: 1,2,3,A -> map A to /
            ['4', '5', '6', '*'],  # Row 2: 4,5,6,B -> map B to *
            ['7', '8', '9', '-'],  # Row 3: 7,8,9,C -> map C to -
            ['C', '0', '=', '+']   # Row 4: *,0,#,D -> map * to C, # to =, D to +
        ]
    
    def scan(self):
        """Scan keypad and return pressed key"""
        for i, row in enumerate(self.rows):
            row.value(1)
            time.sleep_ms(1)
            
            for j, col in enumerate(self.cols):
                if col.value():
                    row.value(0)
                    time.sleep_ms(200)  # Debounce
                    return self.keys[i][j]
            
            row.value(0)
        
        return None


# Calculator Logic
class Calculator:
    def __init__(self):
        self.display_value = "0"
        self.operation = None
        self.first_operand = None
        self.clear_display = False
    
    def process_key(self, key):
        """Process calculator key press"""
        if key is None:
            return
        
        if key == 'C':
            self.display_value = "0"
            self.operation = None
            self.first_operand = None
            self.clear_display = False
        
        elif key in '0123456789':
            if self.clear_display or self.display_value == "0":
                self.display_value = key
                self.clear_display = False
            else:
                self.display_value += key
        
        elif key in '+-*/':
            self.first_operand = float(self.display_value)
            self.operation = key
            self.clear_display = True
        
        elif key == '=':
            if self.operation and self.first_operand is not None:
                second_operand = float(self.display_value)
                
                if self.operation == '+':
                    result = self.first_operand + second_operand
                elif self.operation == '-':
                    result = self.first_operand - second_operand
                elif self.operation == '*':
                    result = self.first_operand * second_operand
                elif self.operation == '/':
                    result = self.first_operand / second_operand if second_operand != 0 else "Error"
                
                if isinstance(result, float):
                    # Check if result is a whole number
                    if result == int(result):
                        self.display_value = str(int(result))
                    else:
                        self.display_value = str(round(result, 4))
                else:
                    self.display_value = result
                
                self.operation = None
                self.first_operand = None
                self.clear_display = True


# Main Program
def main():
    print("=" * 50)
    print("PicoCalc Simulator Starting...")
    print("=" * 50)
    
    # Initialize SPI for display
    print("Initializing SPI...")
    spi = SPI(0, baudrate=40000000, sck=Pin(18), mosi=Pin(19))
    
    # Initialize display (ST7789)
    cs = Pin(17, Pin.OUT)
    dc = Pin(21, Pin.OUT)
    rst = Pin(20, Pin.OUT)
    
    print("Creating display object...")
    display = ST7789(spi, cs, dc, rst, width=240, height=240)
    
    # Initialize keypad (rows: GP2-5, cols: GP6-9)
    print("Initializing keypad...")
    keypad = Keypad([2, 3, 4, 5], [6, 7, 8, 9])
    
    # Initialize calculator
    print("Initializing calculator...")
    calc = Calculator()
    
    # Test pattern sequence
    print("\n" + "=" * 50)
    print("DISPLAY TEST SEQUENCE")
    print("=" * 50)
    
    # Test 1: Red screen
    print("\nTest 1: RED screen (2 seconds)...")
    display.fill(0xF800)  # Red
    time.sleep(2)
    
    # Test 2: Green screen
    print("Test 2: GREEN screen (2 seconds)...")
    display.fill(0x07E0)  # Green
    time.sleep(2)
    
    # Test 3: Blue screen
    print("Test 3: BLUE screen (2 seconds)...")
    display.fill(0x001F)  # Blue
    time.sleep(2)
    
    # Test 4: White screen
    print("Test 4: WHITE screen (2 seconds)...")
    display.fill(0xFFFF)  # White
    time.sleep(2)
    
    # Draw calculator interface
    print("\n" + "=" * 50)
    print("Drawing calculator interface...")
    print("=" * 50)
    
    print("Setting background...")
    display.fill(0x2104)  # Dark blue-gray background
    time.sleep_ms(500)
    
    # Display area (white background)
    print("Drawing display area...")
    display.fill_rect(10, 10, 220, 60, 0xFFFF)
    time.sleep_ms(500)
    
    # Initial display value
    print("Drawing initial value '0'...")
    display.text_large(calc.display_value, 20, 20, 0x0000, 0xFFFF)
    
    print("\n" + "=" * 50)
    print("Calculator ready!")
    print("=" * 50)
    print("\nKeypad layout:")
    print("  1-9: Number keys")
    print("  A: / (divide)")
    print("  B: * (multiply)")
    print("  C: - (minus)")
    print("  D: + (plus)")
    print("  *: C (clear)")
    print("  #: = (equals)")
    print("=" * 50 + "\n")
    
    last_display = calc.display_value
    
    # Main loop
    while True:
        key = keypad.scan()
        
        if key:
            print(f">>> Key pressed: {key}")
            calc.process_key(key)
            print(f">>> Display: {calc.display_value}")
            
            # Update display only if value changed
            if calc.display_value != last_display:
                # Clear display area
                display.fill_rect(10, 10, 220, 60, 0xFFFF)
                
                # Truncate if too long
                display_text = calc.display_value[:7]
                
                # Draw new value
                display.text_large(display_text, 20, 20, 0x0000, 0xFFFF)
                
                last_display = calc.display_value
        
        time.sleep_ms(50)


if __name__ == "__main__":
    main()