#!/usr/bin/env pybricks-micropython

from math import cos, pi

from pybricks.hubs import EV3Brick
from pybricks.ev3devices import Motor
from pybricks.parameters import Port, Stop
from pybricks.tools import wait, StopWatch
from pybricks.messaging import BluetoothMailboxClient, TextMailbox


# ============================================================
# MASTER BRICK PROGRAM
# ============================================================
#
# Run slave_main.py FIRST.
# Then run this master program.
#
# Master:
#   - connects to slave,
#   - sends full slave-side movement plan,
#   - waits for READY,
#   - sends GO,
#   - executes master-side local joints,
#   - waits for slave DONE.
#
# This avoids sending per-joint commands during the movement itself.


# ============================================================
# BLUETOOTH SETTINGS
# ============================================================

# Change this to the Bluetooth name of your slave EV3.
# By default many EV3 bricks are named 'ev3dev'.
SLAVE_BRICK_NAME = 'ev3armslave'

# Optional delay after sending GO before master starts its own local movement.
#
# Bluetooth mailboxes are not hard real-time. If your measurements show that
# master starts earlier than slave, increase this value slightly.
#
# Start with 0. Later calibrate using camera/audio/LED timestamps.
MASTER_START_DELAY_AFTER_GO_MS = 0


# ============================================================
# MASTER JOINT PORTS
# ============================================================

BASE_PORT = Port.A
SHOULDER_LEFT_PORT = Port.B
SHOULDER_RIGHT_PORT = Port.C
ELBOW_PORT = Port.D


# ============================================================
# MOTION SETTINGS
# ============================================================

DEFAULT_SPEED = 200
PROFILE_UPDATE_MS = 20
POLL_MS = 10

DEFAULT_EASING = 'linear'
CLAMP_EASING_OUTPUT = True

DEADBAND_MOTOR_DEG = 1.5
TRACK_TARGET_MIN_UPDATE_MOTOR_DEG = 1.0

SOUND_ENABLED = True


# ============================================================
# MASTER JOINT CONFIGURATION
# ============================================================

GEAR_RATIOS = {
    'base': 7.5,
    'shoulder': 7.5,
    'elbow': 5.0,
}

JOINT_LIMITS = {
    'base': None,
    'shoulder': None,
    'elbow': None,
}

SHOULDER_SIGNS = [1, 1]

STOP_MODE_BY_JOINT = {
    'base': Stop.BRAKE,
    'shoulder': Stop.BRAKE,
    'elbow': Stop.BRAKE,
}


# ============================================================
# HARDWARE / BLUETOOTH
# ============================================================

ev3 = EV3Brick()

base_motor = Motor(BASE_PORT)
shoulder_left_motor = Motor(SHOULDER_LEFT_PORT)
shoulder_right_motor = Motor(SHOULDER_RIGHT_PORT)
elbow_motor = Motor(ELBOW_PORT)

client = BluetoothMailboxClient()
cmd_box = TextMailbox('command', client)
status_box = TextMailbox('status', client)


# ============================================================
# UTILITIES
# ============================================================

def clamp(value, min_value, max_value):
    if value < min_value:
        return min_value
    if value > max_value:
        return max_value
    return value


def abs_value(value):
    if value < 0:
        return -value
    return value


def beep_ok():
    if SOUND_ENABLED:
        ev3.speaker.beep(frequency=1000, duration=80)


def beep_error():
    if SOUND_ENABLED:
        ev3.speaker.beep(frequency=220, duration=200)
        wait(80)
        ev3.speaker.beep(frequency=220, duration=200)


def beep_start_marker():
    if SOUND_ENABLED:
        ev3.speaker.beep(frequency=1400, duration=60)


# ============================================================
# EASING FUNCTIONS
# ============================================================

def ease_linear(x):
    return x


def ease_in_out_sine(x):
    return -(cos(pi * x) - 1) / 2


def ease_in_out_quad(x):
    if x < 0.5:
        return 2 * x * x
    return 1 - ((-2 * x + 2) ** 2) / 2


def ease_in_out_cubic(x):
    if x < 0.5:
        return 4 * x * x * x
    return 1 - ((-2 * x + 2) ** 3) / 2


def ease_out_quad(x):
    return 1 - (1 - x) * (1 - x)


EASINGS = {
    'linear': ease_linear,
    'easeInOutSine': ease_in_out_sine,
    'easeInOutQuad': ease_in_out_quad,
    'easeInOutCubic': ease_in_out_cubic,
    'easeOutQuad': ease_out_quad,
}


def get_easing(name):
    if name is None or name == '':
        name = DEFAULT_EASING

    if name in EASINGS:
        return EASINGS[name]

    return EASINGS[DEFAULT_EASING]


# ============================================================
# JOINT CLASS
# ============================================================

class Joint:
    """
    Logical master-side joint.

    Supports one or more physical motors.

    Coordinate convention:
        motor_angle = joint_angle * gear_ratio
    """

    def __init__(self, name, motors, signs, gear_ratio, joint_limits, stop_mode):
        self.name = name
        self.motors = motors
        self.signs = signs
        self.gear_ratio = gear_ratio
        self.joint_limits = joint_limits
        self.stop_mode_value = stop_mode

    def motor_angle(self):
        total = 0
        count = len(self.motors)

        for i in range(count):
            total += self.motors[i].angle() * self.signs[i]

        return total / count

    def joint_angle(self):
        return self.motor_angle() / self.gear_ratio

    def joint_to_motor(self, joint_angle):
        return joint_angle * self.gear_ratio

    def motor_to_joint(self, motor_angle):
        return motor_angle / self.gear_ratio

    def reset_joint_angle(self, joint_angle):
        motor_angle = self.joint_to_motor(joint_angle)

        for i in range(len(self.motors)):
            self.motors[i].reset_angle(motor_angle * self.signs[i])

    def check_joint_limit(self, joint_target):
        if self.joint_limits is None:
            return

        min_angle = self.joint_limits[0]
        max_angle = self.joint_limits[1]

        if joint_target < min_angle or joint_target > max_angle:
            raise ValueError(
                '{} target {} outside limits {}..{}'.format(
                    self.name, joint_target, min_angle, max_angle
                )
            )

    def physical_target(self, logical_motor_target, motor_index):
        return int(logical_motor_target * self.signs[motor_index])

    def track_motor_target(self, logical_motor_target):
        joint_target = self.motor_to_joint(logical_motor_target)
        self.check_joint_limit(joint_target)

        for i in range(len(self.motors)):
            self.motors[i].track_target(
                self.physical_target(logical_motor_target, i)
            )

    def apply_stop(self):
        for motor in self.motors:
            if self.stop_mode_value == Stop.HOLD:
                motor.hold()
            elif self.stop_mode_value == Stop.COAST:
                motor.stop()
            else:
                motor.brake()


# ============================================================
# MASTER JOINTS
# ============================================================

base = Joint(
    'base',
    [base_motor],
    [1],
    GEAR_RATIOS['base'],
    JOINT_LIMITS['base'],
    STOP_MODE_BY_JOINT['base']
)

shoulder = Joint(
    'shoulder',
    [shoulder_left_motor, shoulder_right_motor],
    SHOULDER_SIGNS,
    GEAR_RATIOS['shoulder'],
    JOINT_LIMITS['shoulder'],
    STOP_MODE_BY_JOINT['shoulder']
)

elbow = Joint(
    'elbow',
    [elbow_motor],
    [1],
    GEAR_RATIOS['elbow'],
    JOINT_LIMITS['elbow'],
    STOP_MODE_BY_JOINT['elbow']
)

JOINTS = {
    'base': base,
    'shoulder': shoulder,
    'elbow': elbow,
}

JOINT_ORDER = ['base', 'shoulder', 'elbow']


# ============================================================
# COMMAND HELPERS
# ============================================================

def ABS_J(joint, angle):
    return {
        'joint': joint,
        'unit': 'joint',
        'mode': 'abs',
        'value': angle
    }


def REL_J(joint, delta):
    return {
        'joint': joint,
        'unit': 'joint',
        'mode': 'rel',
        'value': delta
    }


def ABS_M(joint, angle):
    return {
        'joint': joint,
        'unit': 'motor',
        'mode': 'abs',
        'value': angle
    }


def REL_M(joint, delta):
    return {
        'joint': joint,
        'unit': 'motor',
        'mode': 'rel',
        'value': delta
    }


def STEP(label, master_cmds, slave_cmds, duration_ms, easing='linear', pause_ms=1000):
    return {
        'label': label,
        'master_cmds': master_cmds,
        'slave_cmds': slave_cmds,
        'duration_ms': duration_ms,
        'easing': easing,
        'pause_ms': pause_ms
    }


# ============================================================
# LOCAL MASTER PLANNING
# ============================================================

def validate_parallel(commands):
    seen = []

    for cmd in commands:
        name = cmd['joint']

        if name in seen:
            raise ValueError('Joint repeated in one parallel group: {}'.format(name))

        seen.append(name)


def plan_master_commands(commands):
    validate_parallel(commands)

    plans = []

    for cmd in commands:
        joint = JOINTS[cmd['joint']]
        unit = cmd['unit']
        mode = cmd['mode']
        value = cmd['value']

        start_motor = joint.motor_angle()
        start_joint = joint.joint_angle()

        if mode == 'abs':
            if unit == 'motor':
                target_motor = value
                target_joint = joint.motor_to_joint(target_motor)
            elif unit == 'joint':
                target_joint = value
                target_motor = joint.joint_to_motor(target_joint)
            else:
                raise ValueError('Unknown unit: {}'.format(unit))

        elif mode == 'rel':
            if unit == 'motor':
                target_motor = start_motor + value
                target_joint = joint.motor_to_joint(target_motor)
            elif unit == 'joint':
                target_joint = start_joint + value
                target_motor = joint.joint_to_motor(target_joint)
            else:
                raise ValueError('Unknown unit: {}'.format(unit))

        else:
            raise ValueError('Unknown mode: {}'.format(mode))

        joint.check_joint_limit(target_joint)

        delta_motor = target_motor - start_motor
        skip = abs_value(delta_motor) <= DEADBAND_MOTOR_DEG

        plans.append({
            'joint': joint,
            'start_motor': start_motor,
            'target_motor': target_motor,
            'delta_motor': delta_motor,
            'last_sent': None,
            'skip': skip
        })

    return plans


def maybe_track(plan, target, force):
    joint = plan['joint']
    last = plan['last_sent']

    if force or last is None:
        joint.track_motor_target(target)
        plan['last_sent'] = target
        return

    if abs_value(target - last) >= TRACK_TARGET_MIN_UPDATE_MOTOR_DEG:
        joint.track_motor_target(target)
        plan['last_sent'] = target


def execute_master_plan_timed(plans, duration_ms, easing_name):
    if duration_ms <= 0:
        return True

    easing = get_easing(easing_name)
    active = []

    for plan in plans:
        if not plan['skip']:
            active.append(plan)

    if len(active) == 0:
        return True

    timer = StopWatch()

    while timer.time() < duration_ms:
        ratio = timer.time() / duration_ms
        ratio = clamp(ratio, 0, 1)

        eased = easing(ratio)

        if CLAMP_EASING_OUTPUT:
            eased = clamp(eased, 0, 1)

        for plan in active:
            target = plan['start_motor'] + plan['delta_motor'] * eased
            maybe_track(plan, target, False)

        wait(PROFILE_UPDATE_MS)

    for plan in active:
        maybe_track(plan, plan['target_motor'], True)

    for plan in active:
        plan['joint'].apply_stop()

    return True


def reset_master_as_zero():
    for name in JOINT_ORDER:
        JOINTS[name].reset_joint_angle(0)


def print_master_states(prefix):
    print(prefix)

    for name in JOINT_ORDER:
        joint = JOINTS[name]
        print(
            '{}: motor={} joint={}'.format(
                name, joint.motor_angle(), joint.joint_angle()
            )
        )


# ============================================================
# SLAVE MESSAGE ENCODING
# ============================================================

def encode_slave_commands(commands):
    if len(commands) == 0:
        return 'NONE'

    parts = []

    for cmd in commands:
        parts.append(
            '{},{},{},{}'.format(
                cmd['joint'],
                cmd['unit'],
                cmd['mode'],
                cmd['value']
            )
        )

    return ';'.join(parts)


def send_slave_prepare(step_id, duration_ms, easing, slave_cmds):
    payload = encode_slave_commands(slave_cmds)

    msg = 'PREP|{}|{}|{}|{}'.format(
        step_id,
        duration_ms,
        easing,
        payload
    )

    print('MASTER TX: {}'.format(msg))
    cmd_box.send(msg)

    status_box.wait()
    reply = status_box.read()

    print('MASTER RX: {}'.format(reply))

    expected = 'READY|{}'.format(step_id)

    if reply != expected:
        ev3.screen.clear()
        ev3.screen.print('SLAVE PREP ERR')
        ev3.screen.print(reply[:16])
        ev3.screen.print(reply[16:32])
        raise RuntimeError('Unexpected slave reply: {}'.format(reply))


def send_slave_go(step_id):
    msg = 'GO|{}'.format(step_id)
    print('MASTER TX: {}'.format(msg))
    cmd_box.send(msg)


def wait_slave_done(step_id):
    status_box.wait()
    reply = status_box.read()

    print('MASTER RX: {}'.format(reply))

    expected = 'DONE|{}'.format(step_id)

    if reply != expected:
        ev3.screen.clear()
        ev3.screen.print('SLAVE RUN ERR')
        ev3.screen.print(reply[:16])
        ev3.screen.print(reply[16:32])
        raise RuntimeError('Unexpected slave completion reply: {}'.format(reply))


# ============================================================
# COORDINATED STEP EXECUTION
# ============================================================

def run_coordinated_step(step_id, step):
    label = step['label']
    master_cmds = step['master_cmds']
    slave_cmds = step['slave_cmds']
    duration_ms = step['duration_ms']
    easing = step['easing']
    pause_ms = step['pause_ms']

    ev3.screen.clear()
    ev3.screen.print('STEP {}'.format(step_id))
    ev3.screen.print(label)

    print('')
    print('=' * 70)
    print('COORDINATED STEP {}: {}'.format(step_id, label))
    print('=' * 70)

    master_plans = plan_master_commands(master_cmds)

    # Prepare slave before movement.
    if len(slave_cmds) > 0:
        send_slave_prepare(step_id, duration_ms, easing, slave_cmds)

    # Start marker.
    beep_start_marker()

    # Send GO to slave, then start local motion.
    if len(slave_cmds) > 0:
        send_slave_go(step_id)

    if MASTER_START_DELAY_AFTER_GO_MS > 0:
        wait(MASTER_START_DELAY_AFTER_GO_MS)

    execute_master_plan_timed(master_plans, duration_ms, easing)

    # Wait for slave completion.
    if len(slave_cmds) > 0:
        wait_slave_done(step_id)

    print_master_states('MASTER after step:')

    beep_ok()

    if pause_ms > 0:
        wait(pause_ms)


# ============================================================
# TEST SCENARIO INVOLVING BOTH BRICKS
# ============================================================
#
# Master joints:
#   base, shoulder, elbow
#
# Slave joints:
#   elbow_rotate, wrist_flexion_extension, wrist_pronation_supination, gripper
#
# Rename slave joints after you define their real mechanical meaning.

SCENARIO = [

    STEP(
        'Test one slave joint only',
        master_cmds=[
            ABS_J('base', 0),
        ],
        slave_cmds=[
            ABS_J('elbow_rotate', 0),
        ],
        duration_ms=1000,
        easing='linear',
        pause_ms=1000
    ),


    STEP(
        'All joints zero',
        master_cmds=[
            ABS_J('base', 0),
            ABS_J('shoulder', 0),
            ABS_J('elbow', 0),
        ],
        slave_cmds=[
            ABS_J('elbow_rotate', 0),
            ABS_J('wrist_flexion_extension', 0),
            ABS_J('wrist_pronation_supination', 0),
            ABS_J('gripper', 0),
        ],
        duration_ms=3000,
        easing='linear',
        pause_ms=1000
    ),

    STEP(
        'Base + elbow_rotate synchronized',
        master_cmds=[
            ABS_J('base', 15),
        ],
        slave_cmds=[
            ABS_J('elbow_rotate', 15),
        ],
        duration_ms=3000,
        easing='easeInOutSine',
        pause_ms=1000
    ),

    STEP(
        'Shoulder + elbow + two slave joints',
        master_cmds=[
            ABS_J('shoulder', 8),
            ABS_J('elbow', 12),
        ],
        slave_cmds=[
            ABS_J('wrist_flexion_extension', 10),
            ABS_J('wrist_pronation_supination', -10),
        ],
        duration_ms=4000,
        easing='easeInOutQuad',
        pause_ms=1000
    ),

    STEP(
        'All master and slave joints small negative pose',
        master_cmds=[
            ABS_J('base', -12),
            ABS_J('shoulder', -8),
            ABS_J('elbow', -12),
        ],
        slave_cmds=[
            ABS_J('elbow_rotate', -10),
            ABS_J('wrist_flexion_extension', -10),
            ABS_J('wrist_pronation_supination', 10),
            ABS_J('gripper', 5),
        ],
        duration_ms=4000,
        easing='easeInOutCubic',
        pause_ms=1000
    ),

    STEP(
        'Relative coordinated movement',
        master_cmds=[
            REL_J('base', 10),
            REL_J('elbow', 8),
        ],
        slave_cmds=[
            REL_J('elbow_rotate', 8),
            REL_J('gripper', -5),
        ],
        duration_ms=3000,
        easing='easeOutQuad',
        pause_ms=1000
    ),

    STEP(
        'Return all joints to zero',
        master_cmds=[
            ABS_J('base', 0),
            ABS_J('shoulder', 0),
            ABS_J('elbow', 0),
        ],
        slave_cmds=[
            ABS_J('elbow_rotate', 0),
            ABS_J('wrist_flexion_extension', 0),
            ABS_J('wrist_pronation_supination', 0),
            ABS_J('gripper', 0),
        ],
        duration_ms=4000,
        easing='easeInOutSine',
        pause_ms=1000
    ),
]


# ============================================================
# MAIN PROGRAM
# ============================================================

ev3.screen.clear()
ev3.screen.print('MASTER')
ev3.screen.print('Connecting')

print('MASTER: connecting to slave {}'.format(SLAVE_BRICK_NAME))
client.connect(SLAVE_BRICK_NAME)
print('MASTER: connected')

ev3.screen.clear()
ev3.screen.print('MASTER')
ev3.screen.print('Connected')

reset_master_as_zero()
print_master_states('MASTER initial zero:')

try:
    for i in range(len(SCENARIO)):
        run_coordinated_step(i, SCENARIO[i])

except Exception as e:
    print('MASTER ERROR: {}'.format(e))
    beep_error()
    cmd_box.send('STOP')

ev3.screen.clear()
ev3.screen.print('MASTER')
ev3.screen.print('DONE')

print('MASTER: scenario complete')
beep_ok()