# RobotState/walking.py
import time
import numpy as np
import config as cfg
import kinematics as kin
import DataTypes as dt
from RobotState.RobotState import RobotState

class WalkingState(RobotState):

    def on_enter(self, ctx):
        self.t = 0.0
        self.tick_duration = 1 / cfg.standard_tickpersec
        self.ticks = cfg.standard_duration * cfg.standard_tickpersec

        self.current_pos_copy = ctx.current_pos.copy()
        self.dirmov = [0.0, 0.0]

    def update(self, ctx, intent, dt):
        # 🛑 Transition check
        if not intent.walk:
            return "idle"

        self.dirmov = [intent.move_x, intent.move_y]

        t = self.t / self.ticks
        tick_pos_temp = []

        for leg_id in range(6):
            if ctx.leg_state[leg_id] == "drag":
                tick_pos_temp.append(
                    self.current_pos_copy[leg_id]
                    + (ctx.center_points[leg_id] - self.current_pos_copy[leg_id]) * t
                )

            elif ctx.leg_state[leg_id] == "step":
                mid = [
                    (self.current_pos_copy[leg_id][0] + ctx.center_points[leg_id][0]) / 2,
                    (self.current_pos_copy[leg_id][1] + ctx.center_points[leg_id][1]) / 2,
                    ctx.center_points[leg_id][2] + cfg.step_height,
                ]

                def bez(a, b, c):
                    return (1 - t)**2 * a + 2*(1 - t)*t*b + t*t*c

                x = bez(self.current_pos_copy[leg_id][0], mid[0], ctx.center_points[leg_id][0])
                y = bez(self.current_pos_copy[leg_id][1], mid[1], ctx.center_points[leg_id][1])
                z = bez(self.current_pos_copy[leg_id][2], mid[2], ctx.center_points[leg_id][2])

                tick_pos_temp.append([x, y, z])

        tick_pos = dt.PosArray(tick_pos_temp)
        target_rad = kin.ikpyInverse(tick_pos)

        ctx.current_pos = tick_pos
        ctx.current_rad = target_rad

        if ctx.robotCommunication:
            ctx.robotCommunication.send_command(target_rad)

        self.t += 1
        if self.t >= self.ticks:
            self.t = 0
            ctx.leg_state = (
                np.array(["drag","step","drag","step","drag","step"])
                if ctx.leg_state[0] == "step"
                else np.array(["step","drag","step","drag","step","drag"])
            )

        return None