#!/usr/bin/env python3
"""
Test Multi-Axis Homing with Nova API v2

This script demonstrates using the v2 API's ability to move all joints simultaneously
to achieve faster homing. Instead of moving one joint at a time, all joints move
toward their home position concurrently.
"""

import os
import sys
import time
import urllib.request
import json
from pathlib import Path

# Add parent directory to path
sys.path.insert(0, str(Path(__file__).parent.parent))

# Load .env.local
def load_env_file(filepath):
    """Simple .env file loader."""
    if not filepath.exists():
        return
    with open(filepath, 'r') as f:
        for line in f:
            line = line.strip()
            if line and not line.startswith('#') and '=' in line:
                key, value = line.split('=', 1)
                os.environ[key.strip()] = value.strip()

def get_robot_state():
    """Get current robot state from Nova API."""
    instance_url = os.getenv("NOVA_INSTANCE_URL")
    access_token = os.getenv("NOVA_ACCESS_TOKEN")
    cell_id = os.getenv("NOVA_CELL_ID", "cell")
    controller_id = os.getenv("NOVA_CONTROLLER_ID")
    motion_group_id = os.getenv("NOVA_MOTION_GROUP_ID")

    url = f"{instance_url}/api/v2/cells/{cell_id}/controllers/{controller_id}/motion-groups/{motion_group_id}/state"
    headers = {"Authorization": f"Bearer {access_token}"}
    req = urllib.request.Request(url, headers=headers)

    with urllib.request.urlopen(req, timeout=5) as response:
        data = json.loads(response.read())
        return data.get('joint_position', [])

def main():
    # Load .env.local
    env_path = Path(__file__).parent.parent / ".env.local"
    if not env_path.exists():
        print(f"Error: {env_path} not found")
        return 1

    print(f"Loading environment from {env_path}")
    load_env_file(env_path)

    # Import NovaJogger
    try:
        from robots.ur5.nova_jogger import NovaJogger
        print("✓ NovaJogger imported successfully")
    except ImportError as e:
        print(f"✗ Failed to import NovaJogger: {e}")
        return 1

    # Get config
    instance_url = os.getenv("NOVA_INSTANCE_URL")
    access_token = os.getenv("NOVA_ACCESS_TOKEN")
    cell_id = os.getenv("NOVA_CELL_ID", "cell")
    controller_id = os.getenv("NOVA_CONTROLLER_ID")
    motion_group_id = os.getenv("NOVA_MOTION_GROUP_ID")

    if not all([instance_url, access_token, controller_id, motion_group_id]):
        print("✗ Missing required environment variables")
        return 1

    # Define home position (typical UR5 home: all zeros)
    home_position = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]

    print("\n" + "=" * 70)
    print("Multi-Axis Homing Test with Nova API v2")
    print("=" * 70)
    print(f"Instance: {instance_url}")
    print(f"Controller: {controller_id}")
    print(f"Motion Group: {motion_group_id}")
    print(f"Home Position: {home_position}")
    print()

    # Get initial position
    print("[1] Getting initial robot position...")
    try:
        initial_joints = get_robot_state()
        print(f"Initial position: {[f'{j:.3f}' for j in initial_joints]}")
    except Exception as e:
        print(f"✗ Failed to get robot state: {e}")
        return 1

    # Calculate deltas and velocities
    deltas = [home_position[i] - initial_joints[i] for i in range(6)]
    print(f"Deltas from home: {[f'{d:.3f}' for d in deltas]}")

    # Calculate proportional velocities (all joints reach home at approximately same time)
    max_delta = max(abs(d) for d in deltas)
    if max_delta < 0.001:
        print("\n✓ Already at home position!")
        return 0

    # Use proportional control: faster for joints further from home
    base_velocity = 0.1  # rad/s
    velocities = [base_velocity * (d / max_delta) if abs(d) > 0.001 else 0.0 for d in deltas]

    print(f"Calculated velocities: {[f'{v:.3f}' for v in velocities]} rad/s")
    print(f"Estimated time to home: {max_delta / base_velocity:.1f} seconds")
    print()

    # Create jogger instance
    jogger = NovaJogger(
        instance_url=instance_url,
        access_token=access_token,
        cell_id=cell_id,
        controller_id=controller_id,
        motion_group_id=motion_group_id
    )

    try:
        # Connect
        print("[2] Connecting to Nova API v2...")
        if not jogger.connect():
            print("✗ Failed to connect")
            return 1
        print("✓ Connected\n")

        # Start multi-axis homing
        print("[3] Starting multi-axis homing (all joints moving simultaneously)...")
        if not jogger.start_multi_joint_jog(velocities):
            print("✗ Failed to start multi-joint jog")
            return 1
        print("✓ Homing in progress")

        # Monitor progress
        print("\nProgress:")
        start_time = time.time()
        last_print = 0

        while True:
            time.sleep(0.2)
            elapsed = time.time() - start_time

            # Print status every 0.5 seconds
            if elapsed - last_print >= 0.5:
                try:
                    current_joints = get_robot_state()
                    current_deltas = [abs(home_position[i] - current_joints[i]) for i in range(6)]
                    max_current_delta = max(current_deltas)

                    print(f"  t={elapsed:.1f}s | Max delta: {max_current_delta:.4f} rad | " +
                          f"Joints: {[f'{j:.3f}' for j in current_joints]}")

                    # Check if close enough to home
                    if max_current_delta < 0.01:  # Within 0.01 rad of home
                        print(f"\n✓ Reached home position in {elapsed:.1f} seconds!")
                        break

                    last_print = elapsed

                except Exception as e:
                    print(f"  Warning: Failed to get state: {e}")

            # Safety timeout
            if elapsed > 30:
                print("\n! Timeout after 30 seconds")
                break

        # Stop
        print("\n[4] Stopping...")
        if not jogger.stop():
            print("✗ Failed to stop")
            return 1
        print("✓ Stopped")

        # Get final position
        time.sleep(0.5)
        final_joints = get_robot_state()
        final_deltas = [abs(home_position[i] - final_joints[i]) for i in range(6)]

        print("\n" + "=" * 70)
        print("HOMING RESULTS")
        print("=" * 70)
        print(f"Initial position: {[f'{j:.3f}' for j in initial_joints]}")
        print(f"Final position:   {[f'{j:.3f}' for j in final_joints]}")
        print(f"Home position:    {[f'{j:.3f}' for j in home_position]}")
        print(f"Final deltas:     {[f'{d:.4f}' for d in final_deltas]}")
        print(f"Max delta:        {max(final_deltas):.4f} rad")
        print()

        if max(final_deltas) < 0.05:
            print("✓ HOMING SUCCESSFUL - Robot is at home position")
        else:
            print("⚠ HOMING INCOMPLETE - Robot not fully at home")

        # Disconnect
        print("\n[5] Disconnecting...")
        jogger.disconnect()
        print("✓ Disconnected")

        print("\n" + "=" * 70)
        print("✓ MULTI-AXIS HOMING TEST COMPLETE")
        print("=" * 70)
        print("\nKey advantages of v2 API multi-axis homing:")
        print("- All joints move simultaneously (parallel motion)")
        print("- Faster homing compared to sequential joint movement")
        print("- Smoother motion with proportional velocity control")
        print("- More efficient use of robot capabilities")

        return 0

    except Exception as e:
        print(f"\n✗ Test failed: {type(e).__name__}: {e}")
        import traceback
        traceback.print_exc()
        try:
            jogger.disconnect()
        except:
            pass
        return 1


if __name__ == "__main__":
    sys.exit(main())