/****************************************************************************
 *
 *   Copyright (c) 2024 PX4 Development Team. All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in
 *    the documentation and/or other materials provided with the
 *    distribution.
 * 3. Neither the name PX4 nor the names of its contributors may be
 *    used to endorse or promote products derived from this software
 *    without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
 * FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
 * COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
 * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS
 * OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED
 * AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
 * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
 * ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 *
 ****************************************************************************/

/**
 * @file offline_task_planner.cpp
 *
 * Task planner implementation using nearest neighbor algorithm
 */

#include "offline_task_planner.hpp"
#include <px4_platform_common/log.h>
#include <math.h>

namespace offline_task_manager
{

/* Earth radius in meters */
static constexpr double EARTH_RADIUS_M = 6371000.0;

/* Degrees to radians */
static double deg2rad(double deg)
{
	return deg * M_PI / 180.0;
}

float haversine_distance(double lat1, double lon1, double lat2, double lon2)
{
	double d_lat = deg2rad(lat2 - lat1);
	double d_lon = deg2rad(lon2 - lon1);

	double a = sin(d_lat / 2.0) * sin(d_lat / 2.0) +
		   cos(deg2rad(lat1)) * cos(deg2rad(lat2)) *
		   sin(d_lon / 2.0) * sin(d_lon / 2.0);

	double c = 2.0 * atan2(sqrt(a), sqrt(1.0 - a));

	return (float)(EARTH_RADIUS_M * c);
}

/* Find nearest unselected task using nearest neighbor algorithm */
static int find_nearest_task(
	const OfflineTask *tasks,
	int task_count,
	double current_lat,
	double current_lon,
	bool *selected_mask)
{
	float min_distance = -1.0f;
	int nearest_idx = -1;

	for (int i = 0; i < task_count; i++) {
		if (selected_mask[i]) {
			continue;
		}

		float dist = haversine_distance(
			     current_lat, current_lon,
			     tasks[i].lat, tasks[i].lon);

		if (nearest_idx < 0 || dist < min_distance) {
			min_distance = dist;
			nearest_idx = i;
		}
	}

	return nearest_idx;
}

/* Nearest neighbor planning for a subset of tasks */
static int nearest_neighbor_plan(
	const OfflineTask *source_tasks,
	int source_count,
	double start_lat,
	double start_lon,
	OfflineTask *planned_tasks,
	int max_planned,
	float &total_distance)
{
	if (source_count <= 0 || max_planned <= 0) {
		return 0;
	}

	/* Track which tasks are selected */
	bool *selected = new bool[source_count];
	if (!selected) {
		return -1;
	}

	for (int i = 0; i < source_count; i++) {
		selected[i] = false;
	}

	double current_lat = start_lat;
	double current_lon = start_lon;
	total_distance = 0.0f;
	int planned_count = 0;

	/* Greedy nearest neighbor selection */
	while (planned_count < source_count && planned_count < max_planned) {
		int nearest_idx = find_nearest_task(
				  source_tasks, source_count,
				  current_lat, current_lon,
				  selected);

		if (nearest_idx < 0) {
			break;
		}

		/* Calculate distance to this task */
		float dist = haversine_distance(
			     current_lat, current_lon,
			     source_tasks[nearest_idx].lat,
			     source_tasks[nearest_idx].lon);

		/* Add to plan */
		planned_tasks[planned_count] = source_tasks[nearest_idx];
		total_distance += dist;

		PX4_DEBUG("planned task %d: id=%d dist=%.1fm (lat=%.6f lon=%.6f)",
			  planned_count,
			  source_tasks[nearest_idx].task_id,
			  dist,
			  source_tasks[nearest_idx].lat,
			  source_tasks[nearest_idx].lon);

		/* Update current position to this task */
		current_lat = source_tasks[nearest_idx].lat;
		current_lon = source_tasks[nearest_idx].lon;
		selected[nearest_idx] = true;
		planned_count++;
	}

	delete[] selected;

	PX4_INFO("nearest neighbor planned %d tasks, total distance: %.1fm",
		 planned_count, (double)total_distance);

	return planned_count;
}

int plan_offline_tasks(
	const OfflineTask *input_tasks,
	int input_count,
	double start_lat,
	double start_lon,
	float battery_percent,
	float low_batt_threshold,
	OfflineTask *planned_tasks,
	int max_planned_tasks,
	float &total_distance_m)
{
	if (!input_tasks || !planned_tasks || input_count <= 0 || max_planned_tasks <= 0) {
		PX4_ERR("invalid planner parameters");
		return -1;
	}

	PX4_INFO("planning %d tasks from (%.6f, %.6f), battery: %.1f%%",
		 input_count, start_lat, start_lon, (double)battery_percent);

	/* Determine which tasks to include based on battery */
	bool low_battery = (battery_percent < low_batt_threshold);

	if (low_battery) {
		PX4_WARN("low battery mode: only urgent tasks will be planned");
	}

	/* Separate urgent and normal tasks */
	OfflineTask *urgent_tasks = new OfflineTask[input_count];
	OfflineTask *normal_tasks = new OfflineTask[input_count];

	if (!urgent_tasks || !normal_tasks) {
		PX4_ERR("failed to allocate task buffers");
		delete[] urgent_tasks;
		delete[] normal_tasks;
		return -1;
	}

	int urgent_count = 0;
	int normal_count = 0;

	for (int i = 0; i < input_count; i++) {
		bool is_urgent = (input_tasks[i].must_first || input_tasks[i].priority == 5);

		if (is_urgent) {
			urgent_tasks[urgent_count++] = input_tasks[i];
			PX4_DEBUG("task %d is urgent (must_first=%s priority=%d)",
				  input_tasks[i].task_id,
				  input_tasks[i].must_first ? "true" : "false",
				  input_tasks[i].priority);
		} else if (!low_battery) {
			/* Only include normal tasks if not in low battery mode */
			normal_tasks[normal_count++] = input_tasks[i];
		}
	}

	PX4_INFO("urgent tasks: %d, normal tasks: %d", urgent_count, normal_count);

	total_distance_m = 0.0f;
	int total_planned = 0;
	double current_lat = start_lat;
	double current_lon = start_lon;

	/* Plan urgent tasks first using nearest neighbor */
	if (urgent_count > 0) {
		PX4_INFO("planning urgent tasks...");

		OfflineTask *planned_urgent = new OfflineTask[urgent_count];
		if (!planned_urgent) {
			PX4_ERR("failed to allocate urgent plan buffer");
			delete[] urgent_tasks;
			delete[] normal_tasks;
			return -1;
		}

		float urgent_distance = 0.0f;
		int planned_urgent_count = nearest_neighbor_plan(
			   urgent_tasks, urgent_count,
			   current_lat, current_lon,
			   planned_urgent, urgent_count,
			   urgent_distance);

		/* Copy urgent tasks to output */
		for (int i = 0; i < planned_urgent_count && total_planned < max_planned_tasks; i++) {
			planned_tasks[total_planned++] = planned_urgent[i];
		}

		total_distance_m += urgent_distance;

		/* Update current position to last urgent task */
		if (planned_urgent_count > 0) {
			current_lat = planned_urgent[planned_urgent_count - 1].lat;
			current_lon = planned_urgent[planned_urgent_count - 1].lon;
		}

		delete[] planned_urgent;
	}

	/* Plan normal tasks if not in low battery mode */
	if (normal_count > 0 && !low_battery && total_planned < max_planned_tasks) {
		PX4_INFO("planning normal tasks...");

		int remaining_slots = max_planned_tasks - total_planned;
		OfflineTask *planned_normal = new OfflineTask[normal_count];
		if (!planned_normal) {
			PX4_ERR("failed to allocate normal plan buffer");
			delete[] urgent_tasks;
			delete[] normal_tasks;
			return total_planned;
		}

		float normal_distance = 0.0f;
		int planned_normal_count = nearest_neighbor_plan(
			   normal_tasks, normal_count,
			   current_lat, current_lon,
			   planned_normal, remaining_slots,
			   normal_distance);

		/* Copy normal tasks to output */
		for (int i = 0; i < planned_normal_count && total_planned < max_planned_tasks; i++) {
			planned_tasks[total_planned++] = planned_normal[i];
		}

		total_distance_m += normal_distance;

		delete[] planned_normal;
	}

	/* Cleanup */
	delete[] urgent_tasks;
	delete[] normal_tasks;

	PX4_INFO("planning complete: %d tasks, total distance: %.1fm",
		 total_planned, (double)total_distance_m);

	/* Print planned sequence */
	for (int i = 0; i < total_planned; i++) {
		PX4_INFO("  [%d] task_id=%d (%.6f, %.6f) alt=%.1f priority=%d",
			 i,
			 planned_tasks[i].task_id,
			 planned_tasks[i].lat,
			 planned_tasks[i].lon,
			 (double)planned_tasks[i].alt,
			 planned_tasks[i].priority);
	}

	return total_planned;
}

} // namespace offline_task_manager
