import pyb
from pyb import Pin, Timer
from oled_938 import OLED_938
import micropython
from audio import MICROPHONE
from array import array
from neopixel import NeoPixel

micropython.alloc_emergency_exception_buf(100)

# =========================================================
# Motor pins
# =========================================================
A1 = Pin('X3', Pin.OUT_PP)
A2 = Pin('X4', Pin.OUT_PP)
PWMA = Pin('X1')

B1 = Pin('X7', Pin.OUT_PP)
B2 = Pin('X8', Pin.OUT_PP)
PWMB = Pin('X2')

# PWM timer
tim = Timer(2, freq=1000)
motorA = tim.channel(1, Timer.PWM, pin=PWMA)
motorB = tim.channel(2, Timer.PWM, pin=PWMB)

# =========================================================
# OLED
# =========================================================
i2c = pyb.I2C(2, pyb.I2C.MASTER)
oled = OLED_938(
    pinout={"sda": "Y10", "scl": "Y9", "res": "Y8"},
    height=64,
    external_vcc=False,
    i2c_devid=i2c.scan()[0],
)
oled.poweron()
oled.init_display()

def oled_clear():
    try:
        oled.clear()
    except:
        try:
            oled.fill(0)
        except:
            pass

def oled_show_idle():
    oled_clear()
    oled.draw_text(0, 0, 'Challenge 4 Beat')
    oled.draw_text(0, 16, 'Load dance.txt')
    oled.draw_text(0, 32, 'Press USER')
    oled.display()

def oled_show_running():
    oled_clear()
    oled.draw_text(0, 0, 'Running')
    oled.display()

def oled_show_file_error():
    oled_clear()
    oled.draw_text(0, 0, 'File error')
    oled.draw_text(0, 16, 'dance.txt ?')
    oled.display()

# =========================================================
# NeoPixel
# IMPORTANT: change Y12 if your NeoPixel uses another pin
# =========================================================
NUM_LEDS = 8
np = NeoPixel(Pin('Y12', Pin.OUT_PP), NUM_LEDS)

def leds_off():
    for i in range(NUM_LEDS):
        np[i] = (0, 0, 0)
    np.write()

def leds_on_beat():
    for i in range(NUM_LEDS):
        np[i] = (20, 20, 20)
    np.write()

# =========================================================
# USER switch
# =========================================================
sw = pyb.Switch()

# =========================================================
# Motor control
# =========================================================
def A_forward(v):
    v = max(0, min(100, int(v)))
    A1.low()
    A2.high()
    motorA.pulse_width_percent(v)

def A_back(v):
    v = max(0, min(100, int(v)))
    A2.low()
    A1.high()
    motorA.pulse_width_percent(v)

def A_stop():
    A1.high()
    A2.high()
    motorA.pulse_width_percent(0)

def B_forward(v):
    v = max(0, min(100, int(v)))
    B2.low()
    B1.high()
    motorB.pulse_width_percent(v)

def B_back(v):
    v = max(0, min(100, int(v)))
    B1.low()
    B2.high()
    motorB.pulse_width_percent(v)

def B_stop():
    B1.high()
    B2.high()
    motorB.pulse_width_percent(0)

def stop_all():
    A_stop()
    B_stop()

def drive_AB(cmdA, cmdB):
    if cmdA > 0:
        A_forward(cmdA)
    elif cmdA < 0:
        A_back(-cmdA)
    else:
        A_stop()

    if cmdB > 0:
        B_forward(cmdB)
    elif cmdB < 0:
        B_back(-cmdB)
    else:
        B_stop()

# =========================================================
# Dance moves
# F = forward
# B = backward
# L = left
# R = right
# S = stop
# . = rest
# =========================================================
FWD_PWM  = 45
BACK_PWM = 45
TURN_PWM = 40

def do_move(ch):
    ch = ch.upper()

    if ch == 'F':
        drive_AB(FWD_PWM, FWD_PWM)

    elif ch == 'B':
        drive_AB(-BACK_PWM, -BACK_PWM)

    elif ch == 'L':
        drive_AB(-TURN_PWM, TURN_PWM)

    elif ch == 'R':
        drive_AB(TURN_PWM, -TURN_PWM)

    elif ch == 'S':
        drive_AB(0, 0)

    elif ch == '.':
        drive_AB(0, 0)

    else:
        drive_AB(0, 0)

# =========================================================
# Beat detection
# =========================================================
sample_timer = Timer(7, freq=8000)
mic = pyb.ADC(Pin('Y11'))

N = 160                    # 20 ms window
audio = MICROPHONE(sample_timer, mic, N)

M = 50                     # 50 windows = 1 second
BEAT_THRESHOLD = 2.5       # more sensitive
MIN_BEAT_PERIOD = 150      # faster successive beats allowed

e_ptr = 0
e_buf = array('L', (0 for _ in range(M)))
sum_energy = 0
last_beat_ms = pyb.millis()

# =========================================================
# Read dance file BEFORE main loop
# =========================================================
def load_dance_moves(filename='dance.txt'):
    moves = []
    f = open(filename, 'r')
    for line in f:
        for ch in line:
            if ch not in [' ', '\n', '\r', '\t']:
                moves.append(ch)
    f.close()
    return moves

# =========================================================
# Initial state
# =========================================================
leds_off()
stop_all()

oled_show_idle()
print('Challenge 4 idle')

try:
    dance_moves = load_dance_moves('dance.txt')
    print('Dance loaded:', ''.join(dance_moves))
except Exception as e:
    oled_show_file_error()
    print('Cannot load dance.txt:', e)
    raise

# =========================================================
# Wait for USER
# =========================================================
while not sw():
    pyb.delay(10)

while sw():
    pyb.delay(10)

# =========================================================
# Warm-up
# fill energy buffer before enabling beat-triggered motion
# =========================================================
oled_clear()
oled.draw_text(0, 0, 'Warm up...')
oled.display()

warm_count = 0
while warm_count < M:
    if audio.buffer_is_filled():
        E = audio.inst_energy()
        audio.reset_buffer()
        e_buf[warm_count] = E
        sum_energy += E
        warm_count += 1
    pyb.delay(1)

last_beat_ms = pyb.millis()

stop_all()
current_move = 'S'

oled_show_running()
print('Challenge 4 started')

# =========================================================
# Main loop
# =========================================================
step_index = 0
last_oled = pyb.millis()
c_ratio = 0.0

beat_flash_ms = 60
beat_led_on = False
beat_led_time = 0

try:
    while True:
        if audio.buffer_is_filled():
            E = audio.inst_energy()
            audio.reset_buffer()

            sum_energy = sum_energy - e_buf[e_ptr] + E
            e_buf[e_ptr] = E
            e_ptr = (e_ptr + 1) % M

            average_energy = sum_energy / M if sum_energy > 0 else 1
            c_ratio = E / average_energy

            now = pyb.millis()

            if (now - last_beat_ms) > MIN_BEAT_PERIOD:
                if c_ratio > BEAT_THRESHOLD:
                    leds_on_beat()
                    beat_led_on = True
                    beat_led_time = now

                    current_move = dance_moves[step_index]
                    do_move(current_move)

                    print('Beat! step', step_index, 'move', current_move, 'c=', c_ratio)

                    step_index += 1
                    if step_index >= len(dance_moves):
                        step_index = 0

                    last_beat_ms = now

        if beat_led_on and pyb.elapsed_millis(beat_led_time) > beat_flash_ms:
            leds_off()
            beat_led_on = False

except Exception as e:
    print('Error:', e)

finally:
    leds_off()
    stop_all()
    print('Motors stopped')