import os
import sys
import time
import spidev
import numpy as np
import wiringpi as wpi
from smbus2 import SMBus
from PIL import Image, ImageDraw, ImageFont

# Screen resolution
LCD_WIDTH  = 320
LCD_HEIGHT = 240
MENU_W     = 80

# Pin definitions
LCD_RST_PIN = 3  # BCM 22
LCD_DC_PIN  = 6  # BCM 25
LCD_BL_PIN  = 5  # BCM 24
TP_RST_PIN  = 2  # BCM 27


class LCD_2Inch:
    def __init__(self):
        wpi.wiringPiSetup()
        wpi.pinMode(LCD_DC_PIN, wpi.OUTPUT)
        wpi.pinMode(LCD_RST_PIN, wpi.OUTPUT)
        wpi.pinMode(LCD_BL_PIN, wpi.OUTPUT)
        
        wpi.softPwmCreate(LCD_BL_PIN, 0, 100)
        wpi.softPwmWrite(LCD_BL_PIN, 100)     # Set backlight brightness to 100%

        self.spi = spidev.SpiDev()
        self.spi.open(0, 0)
        self.spi.max_speed_hz = 16000000 
        self.spi.mode = 0b00
        
        self.lcd_init()

    def write_cmd(self, cmd):
        wpi.digitalWrite(LCD_DC_PIN, wpi.LOW)
        self.spi.writebytes([cmd])

    def write_data(self, value):
        wpi.digitalWrite(LCD_DC_PIN, wpi.HIGH)
        if isinstance(value, int):
            self.spi.writebytes([value])
        else:
            self.spi.writebytes(value)

    def reset(self):
        wpi.digitalWrite(LCD_RST_PIN, wpi.HIGH)
        time.sleep(0.01)
        wpi.digitalWrite(LCD_RST_PIN, wpi.LOW)
        time.sleep(0.01)
        wpi.digitalWrite(LCD_RST_PIN, wpi.HIGH)
        time.sleep(0.1)

    def lcd_init(self):
        self.reset()
        self.write_cmd(0x36); self.write_data(0x00)  # Memory data access control
        self.write_cmd(0x3A); self.write_data(0x05)  # Interface pixel format (16-bit)
        self.write_cmd(0x21)   # Display inversion on
        self.write_cmd(0xB2); self.write_data([0x0C, 0x0C, 0x00, 0x33, 0x33])
        self.write_cmd(0xB7); self.write_data(0x35)
        self.write_cmd(0xBB); self.write_data(0x1F)
        self.write_cmd(0xC0); self.write_data(0x2C)
        self.write_cmd(0xC2); self.write_data(0x01)
        self.write_cmd(0xC3); self.write_data(0x12)
        self.write_cmd(0xC4); self.write_data(0x20)
        self.write_cmd(0xC6); self.write_data(0x0F)
        self.write_cmd(0xD0); self.write_data([0xA4, 0xA1])
        self.write_cmd(0xE0); self.write_data([0xD0,0x08,0x11,0x08,0x0C,0x15,0x39,0x33,0x50,0x36,0x13,0x14,0x29,0x2D])
        self.write_cmd(0xE1); self.write_data([0xD0,0x08,0x10,0x08,0x06,0x06,0x39,0x44,0x51,0x0B,0x16,0x14,0x2F,0x31])
        self.write_cmd(0x11)
        time.sleep(0.12)
        self.write_cmd(0x29)

    def display(self, img):
        if img.size != (LCD_WIDTH, LCD_HEIGHT):
            img = img.resize((LCD_WIDTH, LCD_HEIGHT))
            
        self.write_cmd(0x36); self.write_data(0x70)    # Set landscape mode
        self.write_cmd(0x2A); self.write_data([0, 0, (LCD_WIDTH-1)>>8, (LCD_WIDTH-1)&0xFF])
        self.write_cmd(0x2B); self.write_data([0, 0, (LCD_HEIGHT-1)>>8, (LCD_HEIGHT-1)&0xFF])
        self.write_cmd(0x2C)
        
        wpi.digitalWrite(LCD_DC_PIN, wpi.HIGH)
        
         # Convert RGB888 to RGB565 format
        img_arr = np.asarray(img)
        r = (img_arr[..., 0] & 0xF8).astype(np.uint16)
        g = (img_arr[..., 1] & 0xFC).astype(np.uint16)
        b = (img_arr[..., 2] & 0xF8).astype(np.uint16)
        rgb565 = (r << 8) | (g << 3) | (b >> 3)
        
        high_byte = (rgb565 >> 8).astype(np.uint8)
        low_byte = (rgb565 & 0xFF).astype(np.uint8)
        b_data = np.stack((high_byte, low_byte), axis=-1).flatten().tobytes()
        
        for i in range(0, len(b_data), 4096):
            self.spi.writebytes(list(b_data[i:i+4096]))

    def draw_point_fast(self, x, y, color_rgb):
        if x >= LCD_WIDTH-2 or y >= LCD_HEIGHT-2: return
        r, g, b = color_rgb
        hi = (r & 0xF8) | (g >> 5)
        lo = ((g << 3) & 0xE0) | (b >> 3)
        self.write_cmd(0x2A); self.write_data([x>>8, x&0xFF, (x+2)>>8, (x+2)&0xFF])
        self.write_cmd(0x2B); self.write_data([y>>8, y&0xFF, (y+2)>>8, (y+2)&0xFF])
        self.write_cmd(0x2C)
        wpi.digitalWrite(LCD_DC_PIN, wpi.HIGH)
        self.spi.writebytes([hi, lo] * 9)

class Touch_CST816:
    def __init__(self):
        wpi.pinMode(TP_RST_PIN, wpi.OUTPUT)
        wpi.digitalWrite(TP_RST_PIN, wpi.HIGH); time.sleep(0.05)
        wpi.digitalWrite(TP_RST_PIN, wpi.LOW); time.sleep(0.02)
        wpi.digitalWrite(TP_RST_PIN, wpi.HIGH); time.sleep(0.3)
        try:
            self.bus = SMBus(1)
            self.bus.read_byte(0x15)
        except:
            print("Touch I2C Fail"); sys.exit(1)

    def get_coord(self):
        try:
            data = self.bus.read_i2c_block_data(0x15, 0x01, 6)
            if data[1] > 0:
                y = ((data[2] & 0x0F) << 8) | data[3]
                x = ((data[4] & 0x0F) << 8) | data[5]
                return x, y
        except: pass
        return None

def main():
    lcd = LCD_2Inch()
    touch = Touch_CST816()

    demo_img = Image.new("RGB", (LCD_WIDTH, LCD_HEIGHT), "WHITE")
    draw_demo = ImageDraw.Draw(demo_img)
    
    try:
        font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 20)
    except:
        font = ImageFont.load_default()

    try:
        bg = Image.open('fruits.jpg').resize((LCD_WIDTH, LCD_HEIGHT))
        demo_img.paste(bg, (0,0))
    except:
        draw_demo.rectangle([0, 0, LCD_WIDTH, LCD_HEIGHT], fill="LIGHTBLUE")
        
    lcd.display(demo_img)
    time.sleep(2)
    draw_demo.rectangle([0, 0, LCD_WIDTH, LCD_HEIGHT], fill="WHITE")
    
    draw_demo.line([20, 80, 80, 80], fill="red", width=2)
    draw_demo.rectangle([20, 10, 70, 60], fill="WHITE", outline="BLUE") 
    draw_demo.rectangle([85, 10, 130, 60], outline="BLUE", width=3) 
    draw_demo.arc((160, 30, 210, 80),0,360,fill=(0,255,255)) 
    draw_demo.ellipse([240, 30, 290, 80], fill="WHITE", outline="ORANGE") 

    draw_demo.text((20, 120), "SEENGREAT", fill="BLUE", font=font)
    draw_demo.text((20, 200), "ab CD 3210", fill="RED", font=font)

    
    lcd.display(demo_img)
    time.sleep(3)

    # --- Initialize drawing board mode ---
    canvas = Image.new("RGB", (LCD_WIDTH, LCD_HEIGHT), "WHITE")
    draw = ImageDraw.Draw(canvas)
    draw.rectangle([0, 0, MENU_W, LCD_HEIGHT], fill="#808080")   # Menu area backgroun
    
    buttons = [
        (5, 'CLR', (0,0,255)), (44, 'RED', (255,0,0)), (83, 'GRN', (0,255,0)),
        (122, 'BLU', (0,0,255)), (161, 'BLK', (0,0,0)), (200, 'YEL', (255,255,0))
    ]
    for y, t, c in buttons:
        draw.rectangle([5, y, 75, y+34], fill=c, outline="BLACK")
        if t == 'CLR': draw.text((20, y+8), "CLR", fill="WHITE", font=font)
    
    draw.line([MENU_W, 0, MENU_W, LCD_HEIGHT], fill="BLACK", width=2)
    lcd.display(canvas)

    curr_color = (0, 0, 0)
    last_act = 0

    print("Touch Loop Started...")
    try:
        while True:
             pos = touch.get_coord()
             if pos:
                x, y = pos
                if x < MENU_W:
                    if time.time() - last_act > 0.3:
                        if 5 <= y <= 39:    # CLR button - clear drawing area
                            draw.rectangle([MENU_W+2, 0, LCD_WIDTH, LCD_HEIGHT], fill="WHITE")
                            lcd.display(canvas)
                        elif 44 <= y <= 78: curr_color = (255, 0, 0)   # RED
                        elif 83 <= y <= 117: curr_color = (0, 255, 0)   # GREEN
                        elif 122 <= y <= 156: curr_color = (0, 0, 255)   # BLUE
                        elif 161 <= y <= 195: curr_color = (0, 0, 0)   # BLACK
                        elif 200 <= y <= 234: curr_color = (255, 255, 0)   # YELLOW
                        last_act = time.time()
                else:

                    lcd.draw_point_fast(x, y, curr_color)

                    draw.point([x,y], fill=curr_color) 
             time.sleep(0.001)
    except KeyboardInterrupt:
        print("\nExit")

if __name__ == "__main__":
    main()
