#!/usr/bin/env python
# -*- coding: utf-8 -*-

import rospy
from clover import srv
from std_srvs.srv import Trigger
import math
import csv

rospy.init_node('flight')

# Define service proxies
get_telemetry = rospy.ServiceProxy('get_telemetry', srv.GetTelemetry)
navigate = rospy.ServiceProxy('navigate', srv.Navigate)
land = rospy.ServiceProxy('land', Trigger)

def wait_arrival(tolerance=0.2, debug=False, log_writer=None):
    """
    Waits until the drone reaches the target position within a certain tolerance.
    Logs telemetry data to the provided CSV writer during the wait.
    """
    rate = rospy.Rate(5)  # 5 Hz
    while not rospy.is_shutdown():
        telem = get_telemetry(frame_id='navigate_target')
        distance = math.sqrt(telem.x ** 2 + telem.y ** 2 + telem.z ** 2)

        if debug:
            rospy.loginfo(f"Current distance to target: {distance}")

        if log_writer:
            global_telem = get_telemetry(frame_id='map')
            log_writer.writerow([rospy.get_time(), global_telem.x, global_telem.y, global_telem.z, global_telem.yaw])

        if distance < tolerance:
            break
        rate.sleep()
    rospy.sleep(2)

def navigate_and_wait(x, y, z, frame_id='body', auto_arm=False, debug=False, log_writer=None):
    """
    Sends the drone to the specified coordinates and waits until arrival.
    """
    navigate(x=x, y=y, z=z, frame_id=frame_id, auto_arm=auto_arm)
    if debug:
        rospy.loginfo(f"Navigating to: x={x}, y={y}, z={z}, frame_id={frame_id}")
    wait_arrival(debug=debug, log_writer=log_writer)

# Main mission loop
for flight_num in range(1, 4):
    rospy.loginfo(f"Starting flight test {flight_num}")

    filename = f'flight_log_{flight_num}.csv'
    with open(filename, 'w', newline='') as csvfile:
        log_writer = csv.writer(csvfile)
        log_writer.writerow(['time', 'x', 'y', 'z', 'yaw'])  # CSV header

        # Flight mission
        # 1. Take off to 1 meter
        navigate_and_wait(x=0, y=0, z=1, frame_id='body', auto_arm=True, debug=True, log_writer=log_writer)
        rospy.loginfo("Reached altitude of 1 meter")

        # 2. Fly to (1, 0, 1)
        navigate_and_wait(x=1, y=0, z=1, frame_id='map', debug=True, log_writer=log_writer)
        rospy.loginfo("Reached point (1, 0, 1)")

        # 3. Fly to (1, 1, 1)
        navigate_and_wait(x=1, y=1, z=1, frame_id='map', debug=True, log_writer=log_writer)
        rospy.loginfo("Reached point (1, 1, 1)")

        # 4. Fly to (0, 1, 1)
        navigate_and_wait(x=0, y=1, z=1, frame_id='map', debug=True, log_writer=log_writer)
        rospy.loginfo("Reached point (0, 1, 1)")

        # 5. Land at (0, 0, 0)
        navigate_and_wait(x=0, y=0, z=0.5, frame_id='map', debug=True, log_writer=log_writer)
        land()
        rospy.loginfo("Landed at (0, 0, 0)")

    rospy.loginfo(f"Completed flight test {flight_num}")
    rospy.sleep(5)

# Analyze flight data
for flight_num in range(1, 4):
    filename = f'flight_log_{flight_num}.csv'
    rospy.loginfo(f"Analyzing flight data from {filename}")

    with open(filename, 'r') as csvfile:
        log_reader = csv.reader(csvfile)
        next(log_reader)  # Skip header

        times = []
        xs = []
        ys = []
        zs = []
        yaws = []

        for row in log_reader:
            times.append(float(row[0]))
            xs.append(float(row[1]))
            ys.append(float(row[2]))
            zs.append(float(row[3]))
            yaws.append(float(row[4]))

        if times:
            total_time = times[-1] - times[0]
            rospy.loginfo(f"Flight {flight_num} total time: {total_time:.2f} seconds")

            total_distance = 0
            for i in range(1, len(xs)):
                dx = xs[i] - xs[i-1]
                dy = ys[i] - ys[i-1]
                dz = zs[i] - zs[i-1]
                total_distance += math.sqrt(dx**2 + dy**2 + dz**2)
            rospy.loginfo(f"Flight {flight_num} total distance: {total_distance:.2f} meters")

            average_speed = total_distance / total_time if total_time > 0 else 0
            rospy.loginfo(f"Flight {flight_num} average speed: {average_speed:.2f} m/s")
        else:
            rospy.loginfo(f"No data recorded for flight {flight_num}")

rospy.loginfo("Mission completed")