from servoclass import Servo
import time
import urandom

current_pose = "pose_base"
current_state = "state_startup"
previous_state = None
animation_bool_a = False
animation_bool_b = False
last_toggle_a = time.ticks_ms()
last_toggle_b = time.ticks_ms()
new_state_flag = False

# 眼球随机看 + 身体摆动的计时器和状态
eye_last_change = time.ticks_ms()
eye_interval = 2000        # 眼球每 2 秒换一个随机目标
yaw_bool = False
yaw_last_toggle = time.ticks_ms()

# Grove 人脸追踪状态（由 main.py 更新）
grove_active = False       # Grove 是否在提供有效数据
grove_last_seen = 0        # 上次收到 Grove 数据的时间戳


servos = {
    "YAW": Servo(pin_num=6, max_speed=400, max_accel=100, min_angle=10, max_angle=170), #Base Yaw Rotation
    "ROL": Servo(pin_num=7, max_speed=600, max_accel=400, min_angle=30, max_angle=120), #Neck Roll
    "PIT": Servo(pin_num=8, max_speed=600, max_accel=400, min_angle=1, max_angle=80), #Neck Pitch
    "MOU": Servo(pin_num=19, max_speed=100000, max_accel=3500, min_angle=5, max_angle=150), #Mouth
    "EYL": Servo(pin_num=12, max_speed=200, max_accel=500, min_angle=30, max_angle=150), #Left Eyeball
    "EYR": Servo(pin_num=13, max_speed=250, max_accel=500, min_angle=30, max_angle=150), #Right Eyeball
    "LID": Servo(pin_num=14, max_speed=100000, max_accel=7000, min_angle=30, max_angle=160), #EyeLid
    "EAL": Servo(pin_num=15, max_speed=250, max_accel=200, min_angle=60, max_angle=150), #Left Ear
    "EAR": Servo(pin_num=16, max_speed=500, max_accel=200, min_angle=30, max_angle=120), #Right Ear
}

def apply_pose(pose):
    pose = pose_map.get(pose, pose_calibrate)
    for name, angle in pose.items():
        if name in servos:
            servos[name].set_target(angle)
            
def apply_state(state_name):
    state_func = state_map.get(state_name)
    if state_func:
        state_func()  # Call the function
    else:
        print(f"Unknown state: {state_name}")
        
#_________________# Poses (static servo positions used in states) #_________________#

pose_calibrate = { # Each dictionary key = servo name, value = angle
    "YAW": 90,
    "ROL": 90,
    "PIT": 80,
    "MOU": 170,
    "LID": 110,
    "EYL": 90,
    "EYR": 90,
    "EAL": 90,
    "EAR": 90,
}
pose_sleep = { # Each dictionary key = servo name, value = angle
#     "YAW": 88,
#     "RWH": 89,
    "ROL": 90,
    "PIT": 80,
    "MOU": 170,
    "LID": 30,
#     "EYL": 90,
#     "EYR": 90,
    "EAL": 150,
    "EAR": 30,
}

pose_base = { # Each dictionary key = servo name, value = angle
#     "YAW": 88,
#     "RWH": 89,
    "ROL": 90,
#     "PIT": 20,
#     "MOU": 170,
#     "LID": 130,
#     "EYL": 90,
#     "EYR": 90,
    "EAL": 130,
    "EAR": 70,
}

pose_speaking = { # Each dictionary key = servo name, value = angle
#     "YAW": 88,
#     "RWH": 89,
#     "ROL": 90,
#     "PIT": 20,
    "MOU": 10,
#     "LID": 130,
#     "EYL": 90,
#     "EYR": 90,
    "EAL": 130,
    "EAR": 70,
}

pose_stop_speaking = { # Each dictionary key = servo name, value = angle
#     "YAW": 88,
#     "RWH": 89,
#     "ROL": 90,
#     "PIT": 20,
    "MOU": 150,
#     "LID": 130,
#     "EYL": 90,
#     "EYR": 90,
    "EAL": 130,
    "EAR": 70,
}

pose_thinking_1 = { # Each dictionary key = servo name, value = angle
#     "YAW": 90,
#     "RWH": 90,
    "ROL": 130,
#     "PIT": 50,
    "MOU": 150,
#     "LID": 70,
#     "EYL": 90,
#     "EYR": 90,
    "EAL": 150,
    "EAR": 120,
}

pose_curious_2 = { # Each dictionary key = servo name, value = angle
#     "YAW": 90,
#     "RWH": 90,
    "ROL": 40,
#     "PIT": 10,
    "MOU": 160,
#     "LID": 130,
#     "EYL": 90,
#     "EYR": 90,
    "EAL": 60,
    "EAR": 60,
}

pose_map={
    "pose_calibrate": pose_calibrate,
    "pose_base": pose_base,
    "pose_thinking_1": pose_thinking_1,
    "pose_curious_2": pose_curious_2,
    "pose_sleep": pose_sleep,
    "pose_speaking": pose_speaking,
    "pose_stop_speaking": pose_stop_speaking,
}

#_________________#  #_________________#

#_________________# 眼球随机看（无 Grove 时的自动动画） #_________________#

def eye_random_look():
    """眼球随机左右看。Grove 有数据时跳过，由人脸追踪接管"""
    if grove_active:
        return
    global eye_last_change, eye_interval
    now = time.ticks_ms()
    if time.ticks_diff(now, eye_last_change) >= eye_interval:
        # 随机眼球目标角度：60°~120° 范围内（中位 90°，左右各偏 30°）
        target = 60 + urandom.getrandbits(6) % 61  # 60~120
        servos["EYL"].set_target(target)
        servos["EYR"].set_target(target)
        # 随机下次间隔 1.5~3.5 秒，看起来更自然
        eye_interval = 1500 + urandom.getrandbits(11) % 2001
        eye_last_change = now

def yaw_gentle_sway():
    """身体轻微摆动。Grove 有数据时跳过，由人脸追踪接管"""
    if grove_active:
        return
    global yaw_bool, yaw_last_toggle
    now = time.ticks_ms()
    if yaw_bool:
        servos["YAW"].set_target(80)   # 轻微偏左
    else:
        servos["YAW"].set_target(100)  # 轻微偏右
    if time.ticks_diff(now, yaw_last_toggle) >= 3000:  # 每 3 秒换方向
        yaw_bool = not yaw_bool
        yaw_last_toggle = now

#_________________# States (Mix of poses + animations) #_________________#



def state_startup():
    global new_state_flag
    if new_state_flag == True:
        apply_pose("pose_calibrate")
        new_state_flag = False
        
def state_calibrate():
    global new_state_flag
#     if new_state_flag == True:
    apply_pose("pose_calibrate")
    new_state_flag = False

def state_sleep():
    global new_state_flag
    if new_state_flag == True:
        apply_pose("pose_sleep")
        new_state_flag = False
        
def state_speaking():
    global new_state_flag
    if new_state_flag == True:
        apply_pose("pose_speaking")
        new_state_flag = False
    eye_random_look()
    yaw_gentle_sway()

def state_thinking():
    now=time.ticks_ms()
    global animation_bool_a, last_toggle_a, animation_bool_b, last_toggle_b, new_state_flag
    if new_state_flag == True:
        apply_pose("pose_base")
        new_state_flag = False
    if animation_bool_a == False:
        servos["PIT"].set_target(50)
    if animation_bool_a == True:
        servos["PIT"].set_target(120)
    if time.ticks_diff(now, last_toggle_a) >= 700:
        animation_bool_a = not animation_bool_a   # flip the boolean
        last_toggle_a = now
    eye_random_look()
        
        
def state_listen():
    global new_state_flag
    if new_state_flag == True:
        apply_pose("pose_curious_2")
        new_state_flag = False
    eye_random_look()
    yaw_gentle_sway()
        
def state_neutral():
    global new_state_flag
    if new_state_flag == True:
        apply_pose("pose_stop_speaking")
        new_state_flag = False
    eye_random_look()
        
def state_happy():
    now=time.ticks_ms()
    global animation_bool_a, last_toggle_a, animation_bool_b, last_toggle_b, new_state_flag
    if new_state_flag == True:
        apply_pose("pose_base")
        new_state_flag = False
    eye_random_look()
    if animation_bool_a == False:
        servos["ROL"].set_target(70)
    if animation_bool_a == True:
        servos["ROL"].set_target(110)
    if time.ticks_diff(now, last_toggle_a) >= 2000:
        animation_bool_a = not animation_bool_a   # flip the boolean
        last_toggle_a = now
        
    if animation_bool_b == False:
        servos["EAL"].set_target(100)
        servos["EAR"].set_target(60)
    if animation_bool_b == True:
        servos["EAL"].set_target(160)
        servos["EAR"].set_target(120)
    if time.ticks_diff(now, last_toggle_b) >= 1000:
        animation_bool_b = not animation_bool_b   # flip the boolean
        last_toggle_b = now
#         print("Toggled:", animation_bool_a)        

def state_limber_up():
    now=time.ticks_ms()
    global animation_bool_a, last_toggle_a, animation_bool_b, last_toggle_b, new_state_flag
    if new_state_flag == True:
        apply_pose("pose_base")
        new_state_flag = False
#     print("idle")
    if animation_bool_a == False:
#         servos["YAW"].set_target(50)
        servos["ROL"].set_target(70)
        servos["PIT"].set_target(70)
        servos["EAL"].set_target(100)
        servos["EYL"].set_target(60)
        servos["MOU"].set_target(20)
        servos["LID"].set_target(20)
        servos["EYR"].set_target(60)
        servos["EAL"].set_target(60)
        servos["EAR"].set_target(60)
    if animation_bool_a == True:
#         servos["YAW"].set_target(130)
        servos["ROL"].set_target(110)
        servos["PIT"].set_target(5)
        servos["EAL"].set_target(160)
        servos["EYL"].set_target(120)
        servos["MOU"].set_target(160)
        servos["LID"].set_target(160)
        servos["EYR"].set_target(120)
        servos["EAL"].set_target(120)
        servos["EAR"].set_target(120)
    if time.ticks_diff(now, last_toggle_a) >= 1000:
        animation_bool_a = not animation_bool_a   # flip the boolean
        last_toggle_a = now
        
# Note: system states:
#   - idle: sleeping, wakeword not initiated listening mode
#   - neutral: between finished speaking and begun listening. potentially use as marker for end of speech
#   - listening
#   - speaking
        
state_map={
    "neutral": state_neutral,
    "idle": state_sleep,
    "listening": state_listen,
    "state_startup": state_startup,
    "thinking": state_thinking,
    "speaking": state_speaking,
    "state_limber_up": state_limber_up,
    "happy": state_happy,
    "state_calibrate": state_calibrate
}

#_________________# 一次性动作系统 #_________________#

oneshot_active = False      # 是否正在执行一次性动作
oneshot_start_time = 0
oneshot_duration = 0        # 每段（半程）持续时间
oneshot_phase = 0           # 当前阶段：0,1,2,...,total_phases-1
oneshot_total_phases = 0    # 总阶段数（来回次数×2 + 1恢复阶段）
oneshot_pos_a = {}          # A 位置（摆动极限一侧）
oneshot_pos_b = {}          # B 位置（摆动极限另一侧）
saved_positions = {}        # 动作前各舵机位置快照

saved_speeds = {}           # 动作前各舵机的速度/加速度备份

def save_servo_positions():
    """保存当前所有舵机实际位置和运动参数"""
    global saved_positions, saved_speeds
    saved_positions = {}
    saved_speeds = {}
    for name, servo in servos.items():
        saved_positions[name] = servo.pos
        saved_speeds[name] = (servo.max_speed, servo.max_accel)

def restore_servo_positions():
    """恢复舵机到保存的位置和运动参数（只恢复被动作修改过的舵机）"""
    for name in oneshot_pos_a:
        if name in saved_positions and name in servos:
            servos[name].set_target(saved_positions[name])
        if name in saved_speeds and name in servos:
            servos[name].max_speed, servos[name].max_accel = saved_speeds[name]
            servos[name]._inv_2_accel = 1.0 / (2.0 * servos[name].max_accel)

def boost_servos(names):
    """临时提升指定舵机的速度和加速度，让动作又快又猛"""
    for name in names:
        if name in servos:
            servos[name].max_speed = 2000.0   # 原来 250，提升 8 倍
            servos[name].max_accel = 3000.0   # 原来 200，提升 15 倍
            servos[name]._inv_2_accel = 1.0 / (2.0 * 3000.0)

def start_oneshot(pos_a, pos_b, duration_ms=200, repeats=3):
    """
    启动多次来回摆动的一次性动作
    pos_a: {舵机名: 角度} — 摆动一侧极限
    pos_b: {舵机名: 角度} — 摆动另一侧极限（None 则用保存的原位）
    duration_ms: 每半程持续时间（越小越快）
    repeats: 来回摆动次数
    """
    global oneshot_active, oneshot_start_time, oneshot_duration
    global oneshot_phase, oneshot_total_phases
    global oneshot_pos_a, oneshot_pos_b

    if oneshot_active:
        return

    save_servo_positions()
    oneshot_pos_a = pos_a
    # pos_b 为 None 时用保存的原位作为 B 位置
    if pos_b is None:
        oneshot_pos_b = {}
        for name in pos_a:
            if name in saved_positions:
                oneshot_pos_b[name] = saved_positions[name]
    else:
        oneshot_pos_b = pos_b
    oneshot_duration = duration_ms
    oneshot_phase = 0
    oneshot_total_phases = repeats * 2 + 1
    oneshot_start_time = time.ticks_ms()
    oneshot_active = True

    # 临时提升涉及的舵机速度，让动作又快又猛
    boost_servos(pos_a.keys())
    if pos_b:
        boost_servos(pos_b.keys())

    # 第一阶段：去 A 位置
    for name, angle in pos_a.items():
        if name in servos:
            servos[name].set_target(angle)

def update_oneshot():
    """主循环中调用，管理多次来回摆动"""
    global oneshot_active, oneshot_phase, oneshot_start_time

    if not oneshot_active:
        return

    elapsed = time.ticks_diff(time.ticks_ms(), oneshot_start_time)
    if elapsed < oneshot_duration:
        return

    # 当前阶段完成，进入下一阶段
    oneshot_phase += 1
    oneshot_start_time = time.ticks_ms()

    if oneshot_phase >= oneshot_total_phases:
        # 所有阶段完成，恢复原位
        restore_servo_positions()
        oneshot_active = False
        return

    # 奇数阶段去 B，偶数阶段去 A
    side = "B" if (oneshot_phase % 2 == 1) else "A"
    targets = oneshot_pos_b if (oneshot_phase % 2 == 1) else oneshot_pos_a
    for name, angle in targets.items():
        if name in servos:
            servos[name].set_target(angle)

# --- 动作定义 ---
# duration_ms=每半程时间（越小越快），repeats=来回次数

def act_ear_left():
    # 左耳：全幅度快速来回 3 次（速度已 boost，250ms 足够跑满 90° 行程）
    start_oneshot({"EAL": 60}, {"EAL": 150}, 250, 3)

def act_ear_right():
    start_oneshot({"EAR": 120}, {"EAR": 30}, 250, 3)

def act_eye_left():
    start_oneshot({"EYL": 30, "EYR": 30}, {"EYL": 150, "EYR": 150}, 250, 3)

def act_eye_right():
    start_oneshot({"EYL": 150, "EYR": 150}, {"EYL": 30, "EYR": 30}, 250, 3)

def act_tilt_left():
    start_oneshot({"ROL": 30}, {"ROL": 120}, 300, 2)

def act_tilt_right():
    start_oneshot({"ROL": 120}, {"ROL": 30}, 300, 2)

def act_nod():
    # 大幅度快速点头 3 次
    start_oneshot({"PIT": 1}, {"PIT": 80}, 250, 3)

def act_mouth():
    start_oneshot({"MOU": 5}, {"MOU": 150}, 200, 3)

def act_blink():
    # 快速眨眼 3 次
    start_oneshot({"LID": 30}, {"LID": 110}, 150, 3)

def act_eyes_close():
    global oneshot_active
    if oneshot_active:
        return
    save_servo_positions()
    servos["LID"]._write_pwm(30)

def act_eyes_open():
    global oneshot_active
    if oneshot_active:
        return
    save_servo_positions()
    servos["LID"]._write_pwm(110)

def act_body_turn():
    start_oneshot({"YAW": 30}, {"YAW": 150}, 300, 2)

# --- 组合动作 ---

def act_ears_both():
    start_oneshot({"EAL": 60, "EAR": 120}, {"EAL": 150, "EAR": 30}, 250, 3)

def act_nod_blink():
    start_oneshot({"PIT": 1, "LID": 30}, {"PIT": 80, "LID": 110}, 250, 3)

def act_shake_head():
    start_oneshot({"ROL": 30}, {"ROL": 120}, 250, 3)

def act_look_around():
    start_oneshot({"EYL": 30, "EYR": 30, "YAW": 50}, {"EYL": 150, "EYR": 150, "YAW": 130}, 300, 2)

def act_cute():
    start_oneshot({"ROL": 50, "LID": 30, "EAL": 60}, {"ROL": 120, "LID": 110, "EAL": 150}, 250, 2)

def act_dance():
    start_oneshot({"YAW": 40, "ROL": 50, "EAL": 60, "EAR": 120}, {"YAW": 140, "ROL": 120, "EAL": 150, "EAR": 30}, 300, 3)

action_map = {
    # 单动作
    "act_ear_left":   act_ear_left,
    "act_ear_right":  act_ear_right,
    "act_eye_left":   act_eye_left,
    "act_eye_right":  act_eye_right,
    "act_tilt_left":  act_tilt_left,
    "act_tilt_right": act_tilt_right,
    "act_nod":        act_nod,
    "act_mouth":      act_mouth,
    "act_blink":      act_blink,
    "act_eyes_close": act_eyes_close,
    "act_eyes_open":  act_eyes_open,
    "act_body_turn":  act_body_turn,
    # 组合动作
    "act_ears_both":  act_ears_both,
    "act_nod_blink":  act_nod_blink,
    "act_shake_head": act_shake_head,
    "act_look_around": act_look_around,
    "act_cute":       act_cute,
    "act_dance":      act_dance,
}

#_________________#  #_________________#            





