import M5
import time
import random
import gc

================= 硬件与内存锁定 =================

M5.begin()
lcd = M5.Lcd
imu = M5.Imu
gc.collect()
gc.threshold(1024)

try:
lcd.setBrightness(255)
except: pass

WIDTH = lcd.width() # 240
HEIGHT = lcd.height() # 135

SCALE = 5
W = WIDTH // SCALE # 48
H = HEIGHT // SCALE # 27
TOTAL = W * H # 1296

grid = bytearray(TOTAL) # bit0=流体, bit7=本帧已移动锁

OFFSETS = (-1, 1, -W, W, -W-1, -W+1, W-1, W+1)

COLOR_BG = 0x0000
COLOR_FLUID = 0xFFE0 # 💡 已修改为黄色 (RGB565)
COLOR_PAUSE = 0xF800

VISCOSITY_THRESHOLD = 3
SIDE_SLIP_PROB = 70 # 侧滑概率 (0-100)，控制安息角
GRAV_DX, GRAV_DY = 0, 1

idle_frames = 0
MAX_IDLE_FRAMES = 5
is_paused = False
last_space_state = False
frame_counter = 0 # 用于每帧改变随机种子

def draw_cell(x, y, color):
lcd.fillRect(x * SCALE, y * SCALE, SCALE, SCALE, color)

def update_gravity():
global GRAV_DX, GRAV_DY
try:
acc = imu.getAccel()
gyro = imu.getGyro()
ax, ay = int(-acc[0] * 10), int(acc[1] * 10)
gz = int(gyro[2])
if abs(ax) < 2: ax = 0
if abs(ay) < 2: ay = 0
if abs(ay) >= abs(ax):
GRAV_DX, GRAV_DY = 0, (1 if ay > 0 else -1)
else:
GRAV_DX, GRAV_DY = (1 if ax > 0 else -1), 0
if abs(gz) > 15:
if GRAV_DY != 0: GRAV_DX = 1 if gz > 0 else -1
else: GRAV_DY = 1 if gz > 0 else -1
except: pass

def check_keyboard():
global is_paused, last_space_state, idle_frames
current_space = False
try:
keys = M5.Keyboard.getKeys()
for k in keys:
if k == 32 or (isinstance(k, dict) and k.get('key') == 32):
current_space = True; break
except:
try: current_space = M5.BtnA.isPressed()
except: pass
if current_space and not last_space_state:
is_paused = not is_paused
idle_frames = 0
last_space_state = current_space

================= 自然堆积求解器 =================

def solve_fluid():
global frame_counter
has_movement = False
frame_counter += 1

# 奇偶帧交替扫描方向，打破逐行刷新伪影
if frame_counter & 1:
    y_range = range(H-1, -1, -1)
else:
    y_range = range(H-2, -1, -2)

if GRAV_DX > 0: x_range = range(W-1, -1, -1)
elif GRAV_DX < 0: x_range = range(0, W, 1)
else: x_range = range(0, W, 1)

for y in y_range:
    row_off = y * W
    for x in x_range:
        idx = row_off + x
        val = grid[idx]
        if (val & 0x01) == 0 or (val & 0x80) != 0:
            continue
        
        moved = False
        nx, ny = x + GRAV_DX, y + GRAV_DY
        
        # 重力方向移动
        if 0 <= nx < W and 0 <= ny < H:
            t_idx = ny * W + nx
            if (grid[t_idx] & 0x01) == 0:
                grid[idx] = 0
                grid[t_idx] = 0x81
                draw_cell(nx, ny, COLOR_FLUID)
                draw_cell(x, y, COLOR_BG)
                moved = True
                has_movement = True
        
        # 概率性粘性侧滑 → 形成三角形安息角
        if not moved:
            neighbor_count = 0
            for off in OFFSETS:
                n_idx = idx + off
                if 0 <= n_idx < TOTAL and (grid[n_idx] & 0x01) == 1:
                    neighbor_count += 1
            
            if neighbor_count >= VISCOSITY_THRESHOLD:
                if random.randint(0, 99) < SIDE_SLIP_PROB:
                    if GRAV_DY != 0:
                        left_ok = (x > 0 and (grid[idx-1] & 0x01) == 0)
                        right_ok = (x < W-1 and (grid[idx+1] & 0x01) == 0)
                    else:
                        left_ok = (y > 0 and (grid[idx-W] & 0x01) == 0)
                        right_ok = (y < H-1 and (grid[idx+W] & 0x01) == 0)
                    
                    # 伪随机偏好替代确定性(x^y)&1
                    prefer_left = ((x * 7 + y * 13 + frame_counter) & 3) < 2
                    
                    target = -1
                    if prefer_left and left_ok:
                        target = idx-1 if GRAV_DY != 0 else idx-W
                    elif right_ok:
                        target = idx+1 if GRAV_DY != 0 else idx+W
                    elif left_ok:
                        target = idx-1 if GRAV_DY != 0 else idx-W
                    
                    if target >= 0:
                        grid[idx] = 0
                        grid[target] = 0x81
                        tx, ty = target % W, target // W
                        draw_cell(tx, ty, COLOR_FLUID)
                        draw_cell(x, y, COLOR_BG)
                        moved = True
                        has_movement = True

# 清除移动锁
for i in range(TOTAL):
    if grid[i] & 0x80:
        grid[i] &= 0x01

return has_movement

================= 初始化 =================

lcd.clear(COLOR_BG)
water_count = 0
attempts = 0
INIT_WATER = TOTAL // 3

while water_count < INIT_WATER and attempts < INIT_WATER * 8:
idx = random.randint(0, TOTAL - 1)
attempts += 1
if grid[idx] == 0:
grid[idx] = 1
water_count += 1
draw_cell(idx % W, idx // W, COLOR_FLUID)

lcd.show()
print(f"[OK] Yellow fluid: {W}x{H}, slip_prob={SIDE_SLIP_PROB}%")

================= 主循环 =================

while True:
check_keyboard()
if is_paused:
draw_cell(0, 0, COLOR_PAUSE)
lcd.show()
time.sleep(0.05)
continue

update_gravity()
has_movement = solve_fluid()

if has_movement:
    lcd.show()
    idle_frames = 0
else:
    idle_frames += 1
    if idle_frames > MAX_IDLE_FRAMES:
        time.sleep(0.1)
time.sleep(0.016)