import time
import sys
import struct
from Logger import Logs

from unitree_sdk2py.core.channel import ChannelSubscriber, ChannelFactoryInitialize
from unitree_sdk2py.idl.default import unitree_hg_msg_dds__LowState_
from unitree_sdk2py.idl.unitree_hg.msg.dds_ import LowState_

# Initialize Logger
controller_logs = Logs()
controller_logs.LogEngine("G1_Logs", "Wireless_Controller_keys.log")



"""_summary_

1. Posture Switch (Modifier: Hold L2)

These commands change the fundamental physical state or stance of the robot.

    L2 + R2: Debug Mode

    L2 + Y: Zero Torque / Zero Moment Mode

    L2 + B: ① Damping Mode

    L2 + UP: ② Lock Stand / Locked Standing

    L2 + LEFT: ④ Seated Mode

    L2 + X: ⑤ Lying and Standing (Lie -> Stand)

    L2 + A: ⑥ Squat Switch (Squat <-> Stand)

2. Interactional Functions (Modifier: Hold SELECT or Double-Click)

These are social or "demo" movements.

    SELECT + A: Handshake

    SELECT + Y: Wave Hand

    SELECT + X: Reject / Turn Around and Waves Hand

    SELECT + B: Face Wave

    Double-Click Y: Right Kiss

    Double-Click X: Left Kiss

    Double-Click A: Clap

    Double-Click B: Right Raise

    Double-Click UP: Flex Left

    Double-Click DOWN: Flex Right

    Double-Click LEFT: Flex Both

3. Motion & Run Control (Modifier: Hold R2 or START)

Commands related to locomotion and body leaning.

    R2 + DOWN: Slow Running

    R2 + UP: Fast Running

    R2 + A: ③ Running Mode

    R2 + B: ⑧ Climb Mode

    START + UP: Forward Lean

    START + DOWN: Backward Lean

4. Main Operation & Special Modes (Modifier: Hold R1)

Controls for complex body structures and specialized routines.

    R1 + X: ⑦ Main Operation Control (1-DOF Waist)

    R1 + Y: ⑧ Main Operation Control (3-DOF Waist)

    R1 + SELECT: Kung Fu

    R1 + A: Dance 1

    R1 + B: Dance 2

    R1 + X (alt): Dance 3

    R1 + Y (alt): Roll Up

5. System & Offset Settings

Technical adjustments for calibration and UI feedback.

    Click START: Stand / Standing

    Double-Click START: Keep Stepping (Not Recommended)

    Double-Click L1: High Speed Mode

    Double-Click L2: Low Speed Mode

    F1 (Pressed 3 times): Sound/Vibration Switch

    F3 (Pressed 3 times): Sound/Vibration Toggle

    Hold R1 + Click →: Left Offset

    Hold R1 + Click ←: Right Offset

    Hold R1 + Click ↓: Forward Offset

    Hold R1 + Click ↑: Backward Offset

6. Joysticks (Rockers)

    Left Rocker (Lx, Ly): Translational movement (Forward/Backward/Sideways).

    Right Rocker (Rx, Ry): Rotational movement (Yaw) or limb control in Operation Mode.
    
    
    
"""



class unitreeRemoteController:
    def __init__(self):
        # Joysticks
        self.Lx = 0; self.Rx = 0; self.Ry = 0; self.Ly = 0
        # Buttons
        self.L1 = 0; self.L2 = 0; self.R1 = 0; self.R2 = 0
        self.A = 0; self.B = 0; self.X = 0; self.Y = 0
        self.Up = 0; self.Down = 0; self.Left = 0; self.Right = 0
        self.Select = 0; self.F1 = 0; self.F3 = 0; self.Start = 0
       
    def parse_botton(self, data1, data2):
        self.R1 = (data1 >> 0) & 1; self.L1 = (data1 >> 1) & 1
        self.Start = (data1 >> 2) & 1; self.Select = (data1 >> 3) & 1
        self.R2 = (data1 >> 4) & 1; self.L2 = (data1 >> 5) & 1
        self.F1 = (data1 >> 6) & 1; self.F3 = (data1 >> 7) & 1
        self.A = (data2 >> 0) & 1; self.B = (data2 >> 1) & 1
        self.X = (data2 >> 2) & 1; self.Y = (data2 >> 3) & 1
        self.Up = (data2 >> 4) & 1; self.Right = (data2 >> 5) & 1
        self.Down = (data2 >> 6) & 1; self.Left = (data2 >> 7) & 1

    def parse_key(self, data):
        offsets = [4, 8, 12, 20] # Lx, Rx, Ry, Ly
        self.Lx, self.Rx, self.Ry, self.Ly = [struct.unpack('<f', data[o:o+4])[0] for o in offsets]

    def parse(self, remoteData):
        self.parse_key(remoteData)
        self.parse_botton(remoteData[2], remoteData[3])

    def get_state(self):
        return self.__dict__.copy()

class Custom:
    def __init__(self):
        self.remoteController = unitreeRemoteController()
        self.last_state = None
        self.click_times = {} 
        self.CLICK_THRESHOLD = 0.35 

    def Init(self):
        self.print_usage_keys()
        self.lowstate_subscriber = ChannelSubscriber("rt/lf/lowstate", LowState_)
        self.lowstate_subscriber.Init(self.LowStateMessageHandler, 10)

    def print_usage_keys(self):
        usage = """
------------------------------------------------------------------
G1 REMOTE USAGE KEYS:
[L2]      : Modifier for Posture (Squat, Sit, Stand, Damping)
[R1]      : Modifier for Operation (Waist, Dance, Offset Comp)
[R2]      : Modifier for Run Control (Fast, Slow, Climb, Running)
[SELECT]  : Modifier for Interaction (Wave, Handshake, Reject)
[START]   : Standing / Lean Control (Hold for Lean)
------------------------------------------------------------------
"""
        controller_logs.print_and_log(usage, "info")

    def LowStateMessageHandler(self, msg: LowState_):
        self.remoteController.parse(msg.wireless_remote)
        curr = self.remoteController.get_state()

        # Detect Rising Edge (Initial Button Press)
        if self.last_state:
            for key, val in curr.items():
                if key in ['Lx', 'Ly', 'Rx', 'Ry']: continue # Ignore joysticks for binary press log
                if val == 1 and self.last_state[key] == 0:
                    controller_logs.print_and_log(f"[PRESS] {key} button", "info")
                    self.handle_multi_click(key, curr)

        # Detect State Changes for Combinations/Joysticks
        if curr != self.last_state:
            self.process_actions(curr)
            self.last_state = curr

    def handle_multi_click(self, btn, state):
        now = time.time()
        prev_clicks, last_time = self.click_times.get(btn, (0, 0))
        count = prev_clicks + 1 if now - last_time < self.CLICK_THRESHOLD else 1
        self.click_times[btn] = (count, now)

        # Double Click Actions
        if count == 2:
            if btn == 'L1': controller_logs.print_and_log("[ACTION] High Speed Mode (Double-Click L1)", "info")
            if btn == 'L2': controller_logs.print_and_log("[ACTION] Low Speed Mode (Double-Click L2)", "info")
            if btn == 'Start': controller_logs.print_and_log("[ACTION] Keep Stepping (Double-Click Start - Not Recommended)", "warning")
            # Interactional Multi-clicks
            if not any([state['Select'], state['L2'], state['R1']]):
                mapping = {'Y':'Right Kiss', 'X':'Left Kiss', 'A':'Clap', 'B':'Right Raise', 
                           'Up':'Flex Left', 'Down':'Flex Right', 'Left':'Flex Both'}
                if btn in mapping: controller_logs.print_and_log(f"[ACTION] {mapping[btn]} (Double-Click {btn})", "info")
        
        # Triple Click System Settings
        if count == 3:
            if btn == 'F1': controller_logs.print_and_log("[ACTION] Sound/Vibration Switch (3x F1)", "info")
            if btn == 'F3': controller_logs.print_and_log("[ACTION] Sound/Vibration Toggle (3x F3)", "info")

    def process_actions(self, s):
        # 1. Posture Switches (L2 Modifier)
        if s['L2']:
            mapping = {'R2':'Debug Mode', 'Y':'Zero Torque', 'B':'Damping Mode', 
                       'Up':'Lock Stand', 'Left':'Seated Mode', 'X':'Lying/Standing', 'A':'Squat Switch'}
            for btn, act in mapping.items():
                if s[btn]: controller_logs.print_and_log(f"[ACTION] {act} (L2 + {btn})", "info"); return

        # 2. Interactive (Select Modifier)
        if s['Select']:
            mapping = {'A':'Handshake', 'Y':'Wave Hand', 'X':'Reject / Turn Around', 'B':'Face Wave'}
            for btn, act in mapping.items():
                if s[btn]: controller_logs.print_and_log(f"[ACTION] {act} (Select + {btn})", "info"); return

        # 3. Run Control (R2 Modifier)
        if s['R2']:
            mapping = {'Down':'Slow Running', 'Up':'Fast Running', 'A':'Running Mode', 'B':'Climb Mode'}
            for btn, act in mapping.items():
                if s[btn]: controller_logs.print_and_log(f"[ACTION] {act} (R2 + {btn})", "info"); return

        # 4. Operation / Dance / Offset (R1 Modifier)
        if s['R1']:
            if s['Select']: controller_logs.print_and_log("[ACTION] Kung Fu (R1 + Select)", "info"); return
            mapping = {'X':'Main Op (1-DOF) / Dance 3', 'Y':'Main Op (3-DOF) / Roll Up', 'A':'Dance 1', 'B':'Dance 2',
                       'Right':'Left Offset', 'Left':'Right Offset', 'Down':'Forward Offset', 'Up':'Backward Offset'}
            for btn, act in mapping.items():
                if s[btn]: controller_logs.print_and_log(f"[ACTION] {act} (R1 + {btn})", "info"); return

        # 5. Lean Control (Start Modifier)
        if s['Start']:
            if s['Up']: controller_logs.print_and_log("[ACTION] Forward Lean (Start + Up)", "info"); return
            if s['Down']: controller_logs.print_and_log("[ACTION] Backward Lean (Start + Down)", "info"); return
            # Default single click stand
            if not self.last_state['Start']:
                controller_logs.print_and_log("[ACTION] Standing (Click Start)", "info")

        # 6. Joystick check (Deadzone 0.1)
        for stick in ['Lx', 'Ly', 'Rx', 'Ry']:
            if abs(s[stick]) > 0.1:
                controller_logs.print_and_log(f"[STICK] {stick}: {s[stick]:.2f}", "debug")
                break

if __name__ == '__main__':
    if len(sys.argv) > 1: ChannelFactoryInitialize(0, sys.argv[1])
    else: ChannelFactoryInitialize(0)

    custom = Custom()
    custom.Init()
    controller_logs.print_and_log("G1 Logger Ready. Monitoring inputs...", "info")
    while True: time.sleep(0.01)