from time import sleep_ms
import framebuf

# # Circuit
# VCC - VBUS	[40]
# GND - GND	[38]
# SDA - GP0	[1]
# SCL - GP1	[2]

class OledI2c:
    # MicroPython SH1106 OLED driver, I2C interfaces
    SET_CONTRAST        = 0x81
    SET_NORM_INV        = 0xa6
    SET_DISP            = 0xae
    SET_SCAN_DIR        = 0xc0
    SET_SEG_REMAP       = 0xa0

    LOW_COLUMN_ADDRESS  = 0x00
    HIGH_COLUMN_ADDRESS = 0x10
    SET_PAGE_ADDRESS    = 0xB0

    def __init__(self, i2c, width, height, res=None, addr=0x3c, external_vcc=False):
        self.i2c = i2c
        self.addr = addr
        self.res = res
        self.temp = bytearray(2)
        if res is not None:
            res.init(res.OUT, value=1)
        self.width = width
        self.height = height
        self.external_vcc = external_vcc
        self.pages = self.height // 8
        self.buffer = bytearray(self.pages * self.width)
        fb = framebuf.FrameBuffer(self.buffer, self.width, self.height, framebuf.MVLSB)
        self.framebuf = fb # set shortcuts for the methods of framebuf
        self.fill = fb.fill
        self.fill_rect = fb.fill_rect
        self.hline = fb.hline
        self.vline = fb.vline
        self.line = fb.line
        self.rect = fb.rect
        self.pixel = fb.pixel
        self.scroll = fb.scroll
        self.text = fb.text
        self.blit = fb.blit
        self.init_display()

    def init_display(self):
        self.reset()
        self.fill(0)
        self.poweron()
        self.show()

    def poweroff(self):
        self.write_cmd(self.SET_DISP | 0x00)

    def poweron(self):
        self.write_cmd(self.SET_DISP | 0x01)

    def rotate(self, flag, update=True):
        if flag:
            self.write_cmd(self.SET_SEG_REMAP | 0x01)  # mirror display vertically
            self.write_cmd(self.SET_SCAN_DIR | 0x08)  # mirror display hor.
        else:
            self.write_cmd(self.SET_SEG_REMAP | 0x00)
            self.write_cmd(self.SET_SCAN_DIR | 0x00)
        if update:
            self.show()

    def sleep(self, value):
        self.write_cmd(self.SET_DISP | (not value))

    def contrast(self, contrast):
        self.write_cmd(self.SET_CONTRAST)
        self.write_cmd(contrast)

    def invert(self, invert):
        self.write_cmd(self.SET_NORM_INV | (invert & 1))

    def show(self):
        for page in range(self.height // 8):
            self.write_cmd(self.SET_PAGE_ADDRESS | page)
            self.write_cmd(self.LOW_COLUMN_ADDRESS | 2)
            self.write_cmd(self.HIGH_COLUMN_ADDRESS | 0)
            self.write_data(self.buffer[
                self.width * page:self.width * page + self.width
            ])

    def write_cmd(self, cmd):
        self.temp[0] = 0x80  # Co=1, D/C#=0
        self.temp[1] = cmd
        self.i2c.writeto(self.addr, self.temp)

    def write_data(self, buf):
        self.i2c.writeto(self.addr, b'\x40'+buf)

    def reset(self):
        if self.res is not None:
            res(1)
            sleep_ms(1)
            res(0)
            sleep_ms(20)
            res(1)
            sleep_ms(20)

if __name__ == "__main__":
    from machine import Pin, SoftI2C
    from framebuf import FrameBuffer, MONO_HLSB

    b = bytearray(b"\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00|?\x00\x01\x86@\x80\x01\x01\x80\x80\x01\x11\x88\x80\x01\x05\xa0\x80\x00\x83\xc1\x00\x00C\xe3\x00\x00~\xfc\x00\x00L'\x00\x00\x9c\x11\x00\x00\xbf\xfd\x00\x00\xe1\x87\x00\x01\xc1\x83\x80\x02A\x82@\x02A\x82@\x02\xc1\xc2@\x02\xf6>\xc0\x01\xfc=\x80\x01\x18\x18\x80\x01\x88\x10\x80\x00\x8c!\x00\x00\x87\xf1\x00\x00\x7f\xf6\x00\x008\x1c\x00\x00\x0c \x00\x00\x03\xc0\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00")
    fb = FrameBuffer(b, 32, 32, MONO_HLSB)
    
    d = OledI2c(SoftI2C(Pin(1), Pin(0)), 128, 128)
    d.fill(0)
    d.text("ABCDEFGHIJKLMNOP", 0, 0)
    d.text("QRSUTVWXYZ012345", 0, 8)
    d.text("6789abcdefghijkl", 0, 16)
    d.text("mnopqrstuvwxyz|\\", 0, 24)
    d.text("!$%&/()=?^'*", 0, 32)
    d.text("+[]@#,;.:-<>", 0, 40)
    d.text("_{}", 0, 48)
    d.text("     Pico Pi", 0, 56)
    d.blit(fb, 96, 32)
    d.show()

    try:
        while True:
            pass
    except KeyboardInterrupt:
        d.poweroff()