# ssd1306.py
# MicroPython SSD1306 OLED driver (I2C)
# Compatible with 128x64 and 128x32 displays

from machine import I2C, Pin
import framebuf
import time

# SSD1306 commands
SET_CONTRAST = 0x81
DISPLAY_ALL_ON_RESUME = 0xA4
DISPLAY_ALL_ON = 0xA5
NORMAL_DISPLAY = 0xA6
INVERT_DISPLAY = 0xA7
DISPLAY_OFF = 0xAE
DISPLAY_ON = 0xAF
SET_DISPLAY_OFFSET = 0xD3
SET_COM_PINS = 0xDA
SET_VCOM_DETECT = 0xDB
SET_DISPLAY_CLOCK_DIV = 0xD5
SET_PRECHARGE = 0xD9
SET_MULTIPLEX = 0xA8
SET_LOW_COLUMN = 0x00
SET_HIGH_COLUMN = 0x10
SET_START_LINE = 0x40
MEMORY_MODE = 0x20
COLUMN_ADDR = 0x21
PAGE_ADDR = 0x22
COM_SCAN_INC = 0xC0
COM_SCAN_DEC = 0xC8
SEG_REMAP = 0xA0
CHARGE_PUMP = 0x8D
EXTERNAL_VCC = 0x1
SWITCH_CAP_VCC = 0x2

class SSD1306_I2C(framebuf.FrameBuffer):
    def __init__(self, width, height, i2c, addr=0x3c):
        self.width = width
        self.height = height
        self.i2c = i2c
        self.addr = addr
        self.pages = self.height // 8
        self.buffer = bytearray(self.pages * self.width)
        super().__init__(self.buffer, self.width, self.height, framebuf.MONO_VLSB)
        self.init_display()

    def write_cmd(self, cmd):
        self.i2c.writeto(self.addr, bytearray([0x00, cmd]))

    def init_display(self):
        for cmd in (
            DISPLAY_OFF,
            SET_DISPLAY_CLOCK_DIV, 0x80,
            SET_MULTIPLEX, self.height - 1,
            SET_DISPLAY_OFFSET, 0x00,
            SET_START_LINE | 0x00,
            CHARGE_PUMP, 0x14,
            MEMORY_MODE, 0x00,
            SEG_REMAP | 0x1,
            COM_SCAN_DEC,
            SET_COM_PINS, 0x12 if self.height == 64 else 0x02,
            SET_CONTRAST, 0xCF,
            SET_PRECHARGE, 0xF1,
            SET_VCOM_DETECT, 0x40,
            DISPLAY_ALL_ON_RESUME,
            NORMAL_DISPLAY,
            DISPLAY_ON
        ):
            self.write_cmd(cmd)
        self.fill(0)
        self.show()

    def show(self):
        for page in range(self.pages):
            self.write_cmd(0xB0 + page)
            self.write_cmd(SET_LOW_COLUMN)
            self.write_cmd(SET_HIGH_COLUMN)
            start = self.width * page
            end = start + self.width
            self.i2c.writeto(self.addr, b'\x40' + self.buffer[start:end])

    def poweroff(self):
        self.write_cmd(DISPLAY_OFF)

    def poweron(self):
        self.write_cmd(DISPLAY_ON)

    def invert(self, invert):
        self.write_cmd(INVERT_DISPLAY if invert else NORMAL_DISPLAY)