#!/usr/bin/env python3
"""
LeRobot AI Controller
Imitation Learning + Computer Vision
Author: Global Byte Shop Thailand
"""

import cv2
import numpy as np
import torch
import torch.nn as nn
import rospy
from sensor_msgs.msg import Image, Imu
from std_msgs.msg import Float32MultiArray
from cv_bridge import CvBridge
import serial
import time
import json
from collections import deque
import threading

class LeRobotController:
    def __init__(self):
        # Initialize ROS node
        rospy.init_node('lerobot_controller', anonymous=True)
        
        # Hardware connections
        self.arduino = serial.Serial('/dev/ttyACM0', 115200, timeout=1)
        self.bridge = CvBridge()
        
        # AI Models
        self.vision_model = self.load_vision_model()
        self.policy_model = self.load_policy_model()
        
        # Robot state
        self.current_pose = np.zeros(6)  # 6 servo positions
        self.target_pose = np.zeros(6)
        self.sensor_data = {}
        
        # Data buffers for learning
        self.observation_buffer = deque(maxlen=1000)
        self.action_buffer = deque(maxlen=1000)
        
        # ROS subscribers
        self.image_sub = rospy.Subscriber('/camera/image_raw', Image, self.image_callback)
        self.imu_sub = rospy.Subscriber('/imu/data', Imu, self.imu_callback)
        
        # Control parameters
        self.control_frequency = 50  # Hz
        self.vision_frequency = 10   # Hz
        
        print("🤖 LeRobot Controller initialized successfully!")
        
    def load_vision_model(self):
        """Load YOLO v8 model for object detection"""
        try:
            from ultralytics import YOLO
            model = YOLO('yolov8n.pt')  # Nano version for speed
            print("✅ Vision model loaded")
            return model
        except Exception as e:
            print(f"❌ Failed to load vision model: {e}")
            return None
    
    def load_policy_model(self):
        """Load imitation learning policy network"""
        class PolicyNetwork(nn.Module):
            def __init__(self, input_size=512, hidden_size=256, output_size=6):
                super(PolicyNetwork, self).__init__()
                self.fc1 = nn.Linear(input_size, hidden_size)
                self.fc2 = nn.Linear(hidden_size, hidden_size)
                self.fc3 = nn.Linear(hidden_size, output_size)
                self.relu = nn.ReLU()
                self.tanh = nn.Tanh()
                
            def forward(self, x):
                x = self.relu(self.fc1(x))
                x = self.relu(self.fc2(x))
                x = self.tanh(self.fc3(x))  # Output range [-1, 1]
                return x
        
        model = PolicyNetwork()
        try:
            model.load_state_dict(torch.load('lerobot_policy.pth', map_location='cpu'))
            model.eval()
            print("✅ Policy model loaded")
        except:
            print("⚠️ Policy model not found, using random initialization")
        
        return model
    
    def image_callback(self, msg):
        """Process camera images for computer vision"""
        try:
            # Convert ROS image to OpenCV
            cv_image = self.bridge.imgmsg_to_cv2(msg, "bgr8")
            
            # Object detection
            if self.vision_model:
                results = self.vision_model(cv_image, verbose=False)
                
                # Extract detected objects
                objects = []
                for result in results:
                    boxes = result.boxes
                    if boxes is not None:
                        for box in boxes:
                            x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
                            conf = box.conf[0].cpu().numpy()
                            cls = int(box.cls[0].cpu().numpy())
                            
                            if conf > 0.5:  # Confidence threshold
                                objects.append({
                                    'class': cls,
                                    'confidence': conf,
                                    'bbox': [x1, y1, x2, y2],
                                    'center': [(x1+x2)/2, (y1+y2)/2]
                                })
                
                self.sensor_data['objects'] = objects
                
                # Visual feedback (optional)
                self.draw_detections(cv_image, objects)
                
        except Exception as e:
            print(f"❌ Image processing error: {e}")
    
    def imu_callback(self, msg):
        """Process IMU data for balance and orientation"""
        self.sensor_data['imu'] = {
            'orientation': [
                msg.orientation.x,
                msg.orientation.y,
                msg.orientation.z,
                msg.orientation.w
            ],
            'angular_velocity': [
                msg.angular_velocity.x,
                msg.angular_velocity.y,
                msg.angular_velocity.z
            ],
            'linear_acceleration': [
                msg.linear_acceleration.x,
                msg.linear_acceleration.y,
                msg.linear_acceleration.z
            ]
        }
    
    def get_ultrasonic_distance(self):
        """Read ultrasonic sensor data from Arduino"""
        try:
            self.arduino.write(b'GET_DISTANCE\n')
            response = self.arduino.readline().decode().strip()
            distance = float(response)
            return distance
        except:
            return 999.0  # Max distance if error
    
    def create_observation_vector(self):
        """Create observation vector for AI model"""
        obs = np.zeros(512)  # Fixed size observation
        
        # Current servo positions (6 values)
        obs[0:6] = self.current_pose
        
        # IMU data (10 values)
        if 'imu' in self.sensor_data:
            imu = self.sensor_data['imu']
            obs[6:10] = imu['orientation']
            obs[10:13] = imu['angular_velocity']
            obs[13:16] = imu['linear_acceleration']
        
        # Object detection features (up to 10 objects, 5 features each)
        if 'objects' in self.sensor_data:
            objects = self.sensor_data['objects'][:10]  # Limit to 10 objects
            for i, obj in enumerate(objects):
                start_idx = 16 + i * 5
                obs[start_idx] = obj['class'] / 80.0  # Normalize class ID
                obs[start_idx + 1] = obj['confidence']
                obs[start_idx + 2:start_idx + 4] = np.array(obj['center']) / 640.0  # Normalize coordinates
                obs[start_idx + 4] = (obj['bbox'][2] - obj['bbox'][0]) / 640.0  # Normalize width
        
        # Distance sensor
        distance = self.get_ultrasonic_distance()
        obs[66] = min(distance / 200.0, 1.0)  # Normalize to [0, 1]
        
        # Previous actions (6 values)
        obs[67:73] = self.target_pose
        
        return obs
    
    def predict_action(self, observation):
        """Use AI model to predict next action"""
        if self.policy_model is None:
            # Fallback: simple reactive behavior
            return self.reactive_behavior(observation)
        
        try:
            obs_tensor = torch.FloatTensor(observation).unsqueeze(0)
            with torch.no_grad():
                action = self.policy_model(obs_tensor)
            return action.squeeze().numpy()
        except Exception as e:
            print(f"❌ Policy prediction error: {e}")
            return self.reactive_behavior(observation)
    
    def reactive_behavior(self, observation):
        """Simple reactive behavior as fallback"""
        action = np.zeros(6)
        
        # Simple head tracking
        if 'objects' in self.sensor_data and len(self.sensor_data['objects']) > 0:
            obj = self.sensor_data['objects'][0]  # Track first object
            center_x = obj['center'][0]
            
            # Pan head to track object
            if center_x < 280:  # Left side
                action[0] = 0.3  # Turn head left
            elif center_x > 360:  # Right side
                action[0] = -0.3  # Turn head right
        
        # Obstacle avoidance
        distance = observation[66] * 200.0  # Denormalize
        if distance < 30:  # Too close
            action[1] = -0.5  # Move arm back
            action[2] = -0.5
        
        return action
    
    def send_servo_commands(self, actions):
        """Send servo commands to Arduino"""
        try:
            # Convert actions [-1, 1] to servo angles [0, 180]
            servo_angles = ((actions + 1) * 90).astype(int)
            servo_angles = np.clip(servo_angles, 0, 180)
            
            # Create command string
            cmd = f"SERVO:{','.join(map(str, servo_angles))}\n"
            self.arduino.write(cmd.encode())
            
            # Update current pose
            self.current_pose = actions
            
        except Exception as e:
            print(f"❌ Servo command error: {e}")
    
    def collect_demonstration_data(self, duration=60):
        """Collect demonstration data for imitation learning"""
        print(f"🎓 Collecting demonstration data for {duration} seconds...")
        print("Use keyboard to control the robot:")
        print("WASD: Head movement, IJKL: Left arm, Arrow keys: Right arm")
        
        import keyboard
        start_time = time.time()
        
        while time.time() - start_time < duration:
            # Get current observation
            obs = self.create_observation_vector()
            
            # Get human input
            action = np.zeros(6)
            
            # Head control (servo 0)
            if keyboard.is_pressed('a'):
                action[0] = -0.5
            elif keyboard.is_pressed('d'):
                action[0] = 0.5
            
            # Left arm control (servos 1, 2)
            if keyboard.is_pressed('i'):
                action[1] = 0.5
            elif keyboard.is_pressed('k'):
                action[1] = -0.5
            if keyboard.is_pressed('j'):
                action[2] = -0.5
            elif keyboard.is_pressed('l'):
                action[2] = 0.5
            
            # Right arm control (servos 3, 4)
            if keyboard.is_pressed('up'):
                action[3] = 0.5
            elif keyboard.is_pressed('down'):
                action[3] = -0.5
            if keyboard.is_pressed('left'):
                action[4] = -0.5
            elif keyboard.is_pressed('right'):
                action[4] = 0.5
            
            # Store demonstration data
            if np.any(action != 0):  # Only store when there's action
                self.observation_buffer.append(obs.copy())
                self.action_buffer.append(action.copy())
            
            # Execute action
            self.send_servo_commands(action)
            time.sleep(1.0 / self.control_frequency)
        
        print(f"✅ Collected {len(self.observation_buffer)} demonstration samples")
        self.save_demonstration_data()
    
    def save_demonstration_data(self):
        """Save collected demonstration data"""
        data = {
            'observations': [obs.tolist() for obs in self.observation_buffer],
            'actions': [act.tolist() for act in self.action_buffer]
        }
        
        with open('demonstration_data.json', 'w') as f:
            json.dump(data, f)
        
        print("💾 Demonstration data saved to demonstration_data.json")
    
    def train_policy_model(self, epochs=100):
        """Train the policy model using collected data"""
        if len(self.observation_buffer) < 100:
            print("❌ Not enough demonstration data for training")
            return
        
        print(f"🎓 Training policy model with {len(self.observation_buffer)} samples...")
        
        # Prepare training data
        X = torch.FloatTensor([obs for obs in self.observation_buffer])
        y = torch.FloatTensor([act for act in self.action_buffer])
        
        # Training setup
        optimizer = torch.optim.Adam(self.policy_model.parameters(), lr=0.001)
        criterion = nn.MSELoss()
        
        # Training loop
        for epoch in range(epochs):
            optimizer.zero_grad()
            predictions = self.policy_model(X)
            loss = criterion(predictions, y)
            loss.backward()
            optimizer.step()
            
            if epoch % 10 == 0:
                print(f"Epoch {epoch}, Loss: {loss.item():.4f}")
        
        # Save trained model
        torch.save(self.policy_model.state_dict(), 'lerobot_policy.pth')
        print("✅ Policy model training completed and saved!")
    
    def draw_detections(self, image, objects):
        """Draw detection results on image"""
        for obj in objects:
            x1, y1, x2, y2 = map(int, obj['bbox'])
            conf = obj['confidence']
            cls = obj['class']
            
            # Draw bounding box
            cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
            
            # Draw label
            label = f"Class {cls}: {conf:.2f}"
            cv2.putText(image, label, (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
        
        # Show image (optional, for debugging)
        cv2.imshow('LeRobot Vision', image)
        cv2.waitKey(1)
    
    def run_autonomous_mode(self):
        """Run robot in autonomous mode"""
        print("🤖 Starting autonomous mode...")
        rate = rospy.Rate(self.control_frequency)
        
        while not rospy.is_shutdown():
            try:
                # Get current observation
                observation = self.create_observation_vector()
                
                # Predict action using AI
                action = self.predict_action(observation)
                
                # Execute action
                self.send_servo_commands(action)
                
                # Store data for continuous learning
                if len(self.observation_buffer) > 0:  # Only if we have reference data
                    self.observation_buffer.append(observation)
                    self.action_buffer.append(action)
                
                rate.sleep()
                
            except KeyboardInterrupt:
                print("🛑 Stopping autonomous mode...")
                break
            except Exception as e:
                print(f"❌ Autonomous mode error: {e}")
                time.sleep(0.1)
    
    def run_interactive_mode(self):
        """Run robot in interactive demonstration mode"""
        print("🎮 Interactive mode - Use keyboard to control robot")
        print("Commands:")
        print("  'c' - Collect demonstration data")
        print("  't' - Train policy model")
        print("  'a' - Run autonomous mode")
        print("  'q' - Quit")
        
        while True:
            cmd = input("Enter command: ").lower()
            
            if cmd == 'c':
                duration = int(input("Duration (seconds): "))
                self.collect_demonstration_data(duration)
            elif cmd == 't':
                epochs = int(input("Training epochs (default 100): ") or "100")
                self.train_policy_model(epochs)
            elif cmd == 'a':
                self.run_autonomous_mode()
            elif cmd == 'q':
                break
            else:
                print("Unknown command")
        
        print("👋 Goodbye!")

def main():
    try:
        robot = LeRobotController()
        robot.run_interactive_mode()
    except rospy.ROSInterruptException:
        pass
    except KeyboardInterrupt:
        print("\n🛑 Program interrupted by user")
    except Exception as e:
        print(f"❌ Fatal error: {e}")

if __name__ == '__main__':
    main()