#!/usr/bin/env python3
"""Position Estimator - Fuses sensors for local position estimate."""

import rclpy
from rclpy.node import Node
from geometry_msgs.msg import PoseStamped, TwistStamped, PoseWithCovarianceStamped
from sensor_msgs.msg import Imu
from nav_msgs.msg import Odometry
import numpy as np
from scipy.spatial.transform import Rotation


class PositionEstimator(Node):
    
    def __init__(self):
        super().__init__('position_estimator')
        
        self.declare_parameter('fusion_method', 'weighted_average')
        self.declare_parameter('vo_weight', 0.6)
        self.declare_parameter('optical_flow_weight', 0.3)
        self.declare_parameter('imu_weight', 0.1)
        self.declare_parameter('update_rate', 50.0)
        
        self.fusion_method = self.get_parameter('fusion_method').value
        self.vo_weight = self.get_parameter('vo_weight').value
        self.of_weight = self.get_parameter('optical_flow_weight').value
        self.imu_weight = self.get_parameter('imu_weight').value
        
        self.position = np.zeros(3)
        self.velocity = np.zeros(3)
        self.orientation = np.array([0.0, 0.0, 0.0, 1.0])
        
        self.vo_pose = None
        self.of_velocity = None
        self.imu_data = None
        
        self.prev_time = None
        
        self.vo_sub = self.create_subscription(
            PoseStamped, '/uav/visual_odometry/pose', self.vo_callback, 10)
        
        self.of_sub = self.create_subscription(
            TwistStamped, '/uav/optical_flow/velocity', self.of_callback, 10)
        
        self.imu_sub = self.create_subscription(
            Imu, '/uav/imu/data', self.imu_callback, 10)
        
        self.pose_pub = self.create_publisher(
            PoseWithCovarianceStamped, '/uav/local_position/pose', 10)
        
        self.odom_pub = self.create_publisher(
            Odometry, '/uav/local_position/odom', 10)
        
        update_rate = self.get_parameter('update_rate').value
        self.timer = self.create_timer(1.0 / update_rate, self.update_estimate)
        
        self.get_logger().info(f'Position Estimator Started - {self.fusion_method}')
    
    def vo_callback(self, msg):
        self.vo_pose = msg
    
    def of_callback(self, msg):
        self.of_velocity = msg
    
    def imu_callback(self, msg):
        self.imu_data = msg
    
    def update_estimate(self):
        current_time = self.get_clock().now()
        
        if self.prev_time is not None:
            dt = (current_time.nanoseconds - self.prev_time.nanoseconds) / 1e9
            if dt <= 0:
                dt = 0.02
        else:
            dt = 0.02
        
        self.prev_time = current_time
        
        if self.fusion_method == 'weighted_average':
            self.weighted_average_fusion(dt)
        else:
            self.simple_fusion(dt)
        
        self.publish_estimate()
    
    def weighted_average_fusion(self, dt):
        total_weight = 0.0
        weighted_pos = np.zeros(3)
        weighted_vel = np.zeros(3)
        
        if self.vo_pose is not None:
            vo_pos = np.array([
                self.vo_pose.pose.position.x,
                self.vo_pose.pose.position.y,
                self.vo_pose.pose.position.z
            ])
            weighted_pos += self.vo_weight * vo_pos
            total_weight += self.vo_weight
            
            self.orientation = np.array([
                self.vo_pose.pose.orientation.x,
                self.vo_pose.pose.orientation.y,
                self.vo_pose.pose.orientation.z,
                self.vo_pose.pose.orientation.w
            ])
        
        if self.of_velocity is not None:
            of_vel = np.array([
                self.of_velocity.twist.linear.x,
                self.of_velocity.twist.linear.y,
                0.0
            ])
            weighted_vel += self.of_weight * of_vel
        
        if self.imu_data is not None:
            imu_accel = np.array([
                self.imu_data.linear_acceleration.x,
                self.imu_data.linear_acceleration.y,
                self.imu_data.linear_acceleration.z - 9.81
            ])
            self.velocity += imu_accel * dt * self.imu_weight
        
        if total_weight > 0:
            self.position = weighted_pos / total_weight
        
        if self.of_velocity is not None:
            self.velocity = weighted_vel
    
    def simple_fusion(self, dt):
        if self.vo_pose is not None:
            self.position = np.array([
                self.vo_pose.pose.position.x,
                self.vo_pose.pose.position.y,
                self.vo_pose.pose.position.z
            ])
    
    def publish_estimate(self):
        pose_msg = PoseWithCovarianceStamped()
        pose_msg.header.stamp = self.get_clock().now().to_msg()
        pose_msg.header.frame_id = 'odom'
        
        pose_msg.pose.pose.position.x = float(self.position[0])
        pose_msg.pose.pose.position.y = float(self.position[1])
        pose_msg.pose.pose.position.z = float(self.position[2])
        
        pose_msg.pose.pose.orientation.x = float(self.orientation[0])
        pose_msg.pose.pose.orientation.y = float(self.orientation[1])
        pose_msg.pose.pose.orientation.z = float(self.orientation[2])
        pose_msg.pose.pose.orientation.w = float(self.orientation[3])
        
        pose_msg.pose.covariance = [
            0.1, 0, 0, 0, 0, 0,
            0, 0.1, 0, 0, 0, 0,
            0, 0, 0.1, 0, 0, 0,
            0, 0, 0, 0.05, 0, 0,
            0, 0, 0, 0, 0.05, 0,
            0, 0, 0, 0, 0, 0.05
        ]
        
        self.pose_pub.publish(pose_msg)
        
        odom_msg = Odometry()
        odom_msg.header = pose_msg.header
        odom_msg.child_frame_id = 'base_link'
        odom_msg.pose = pose_msg.pose
        
        odom_msg.twist.twist.linear.x = float(self.velocity[0])
        odom_msg.twist.twist.linear.y = float(self.velocity[1])
        odom_msg.twist.twist.linear.z = float(self.velocity[2])
        
        self.odom_pub.publish(odom_msg)


def main(args=None):
    rclpy.init(args=args)
    node = PositionEstimator()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.shutdown()


if __name__ == '__main__':
    main()