from machine import Pin, I2C
from ssd1306 import SSD1306_I2C
from hx711 import HX711
import utime
import sys

# ===================== 参数区 =====================
WIDTH = 128
HEIGHT = 64

# ⚠️ Wokwi 专用标定系数（由你给的数据反推）
# 约 420 raw / kg → 1 raw ≈ 2.38 g
CALIBRATION_FACTOR = 2.382   # g / raw  （只用于 Wokwi）

# ===================== I2C =====================
def init_i2c():
    i2c = I2C(1, scl=Pin(27), sda=Pin(26), freq=400000)
    if not i2c.scan():
        print("OLED not found")
        sys.exit()
    return i2c

# ===================== HX711 扩展 =====================
class Scales(HX711):
    def __init__(self, d_out, pd_sck):
        super().__init__(d_out, pd_sck)
        self.offset = 0

    def tare(self, n=20):
        """多次平均去皮"""
        s = 0
        for _ in range(n):
            s += self.read()
            utime.sleep_ms(40)
        self.offset = s / n

    def read_raw(self, n=8):
        """
        读取去皮后的 raw
        ★ 在这里统一方向（只翻一次）
        """
        s = 0
        for _ in range(n):
            s += self.read() - self.offset
            utime.sleep_ms(15)
        return -s / n   # ← 关键：方向统一

# ===================== 等待 HX711 ready =====================
def wait_hx711(scales, timeout=5000):
    start = utime.ticks_ms()
    while True:
        try:
            scales.read()
            return
        except:
            if utime.ticks_diff(utime.ticks_ms(), start) > timeout:
                raise RuntimeError("HX711 not ready")
            utime.sleep_ms(100)

# ===================== UI =====================
def draw_static_ui(oled):
    oled.fill(0)

    # 顶部状态栏
    oled.rect(0, 0, 128, 12, 1)
    oled.text("DIGITAL SCALE", 2, 2)

    # 主显示区
    oled.rect(4, 14, 120, 36, 1)

    # 单位
    oled.text("g", 110, 42)

    oled.show()

def draw_big_number(oled, text, y):
    """2x 视觉字体"""
    x = (WIDTH - len(text) * 16) // 2
    for i, c in enumerate(text):
        px = x + i * 16
        oled.text(c, px, y)
        oled.text(c, px + 1, y)
        oled.text(c, px, y + 1)
        oled.text(c, px + 1, y + 1)

def update_display(oled, weight, stable):
    oled.fill_rect(6, 16, 116, 32, 0)

    text = "{:.1f}".format(weight)
    draw_big_number(oled, text, 22)

    # 稳定指示
    if stable:
        oled.fill_rect(96, 2, 6, 6, 1)
        oled.text("OK", 106, 2)
    else:
        oled.rect(96, 2, 6, 6, 1)
        oled.text("..", 106, 2)

    oled.show()

# ===================== 主程序 =====================
def main():
    # OLED
    oled = SSD1306_I2C(WIDTH, HEIGHT, init_i2c())

    oled.fill(0)
    oled.text("Power On", 32, 24)
    oled.show()
    utime.sleep(1)

    # HX711
    scales = Scales(d_out=14, pd_sck=15)

    oled.fill(0)
    oled.text("Apply weight", 16, 24)
    oled.show()

    wait_hx711(scales)

    oled.fill(0)
    oled.text("Taring...", 36, 24)
    oled.show()
    scales.tare()

    draw_static_ui(oled)

    last_weight = 0
    stable_count = 0

    while True:
        raw = scales.read_raw()
        weight = raw * CALIBRATION_FACTOR  # g

        # 稳定检测
        if abs(weight - last_weight) < 2:   # 2 g 阈值
            stable_count += 1
        else:
            stable_count = 0

        stable = stable_count > 5
        last_weight = weight

        update_display(oled, weight, stable)

        # 调试输出（可保留）
        print("RAW:", raw, "WEIGHT(g):", weight)

        utime.sleep_ms(200)

# =====================
if __name__ == "__main__":
    main()