use std::iter::zip;

use argmin::{
    core::{CostFunction, Error, Executor, Gradient, State, TerminationReason},
    solver::{linesearch::MoreThuenteLineSearch, quasinewton::LBFGS},
};
use finitediff::vec;
use geographiclib_rs::{DirectGeodesic, Geodesic, InverseGeodesic};
use log::debug;
use serde::Serialize;
use thiserror::Error;
#[cfg(feature = "bindings")]
use ts_rs::TS;

use crate::geo::{Degrees, FloatMeters};

#[derive(Error, Debug)]
#[cfg_attr(feature = "bindings", derive(TS))]
#[cfg_attr(feature = "bindings", ts(export))]
pub enum OptimizationError {
    #[error("Could not create solver: {0}")]
    SolverCreationError(String),

    #[error("Task optimization did not converge")]
    NoConvergence,

    #[error("Invalid optimization argument")]
    InvalidArgument,

    #[error("Invalid task: {0}")]
    InvalidTask(String),
}

#[derive(Serialize)]
#[cfg_attr(feature = "bindings", derive(TS))]
#[cfg_attr(feature = "bindings", ts(export))]
pub struct Point {
    pub lat: Degrees,
    pub lon: Degrees,
}

#[derive(Serialize)]
#[cfg_attr(feature = "bindings", derive(TS))]
#[cfg_attr(feature = "bindings", ts(export))]
pub struct Route {
    pub distance: FloatMeters,
    pub points: Vec<Point>,
}

// Task optimization is computed piecewise.
// The pre speed-section is optional because the takeoff is not used for navigation and can be omitted from the task.
// However, the goal is mandatory so the post speed-section is not optional.
#[derive(Serialize)]
#[cfg_attr(feature = "bindings", derive(TS))]
#[cfg_attr(feature = "bindings", ts(export))]
pub struct TaskOptimizationResult {
    pub pre_speed_section: Option<Route>,
    pub speed_section: Route,
    pub post_speed_section: Option<Route>,
}

impl TaskOptimizationResult {
    pub fn total_distance(&self) -> FloatMeters {
        self.pre_speed_section.as_ref().map_or(0., |r| r.distance)
            + self.speed_section.distance
            + self.post_speed_section.as_ref().map_or(0., |r| r.distance)
    }
}

#[derive(Clone, Copy)]
pub struct CompactTurnpoint {
    pub lat: Degrees,
    pub lon: Degrees,
    pub radius: FloatMeters,
}

#[derive(Clone)]
pub struct CompactTask {
    pub turnpoints: Vec<CompactTurnpoint>,
    pub geodesic: Geodesic,
}

impl CompactTask {
    fn get_route_points_from_angles(&self, angles: &[Degrees]) -> Vec<Point> {
        zip(&self.turnpoints, angles)
            .map(|(tp, theta)| {
                let (lat, lon, _) = self.geodesic.direct(tp.lat, tp.lon, *theta, tp.radius);
                Point { lat, lon }
            })
            .collect()
    }

    fn get_distance_from_points(&self, route_points: &[Point]) -> FloatMeters {
        route_points
            .windows(2)
            .map(|window| {
                let p1 = &window[0];
                let p2 = &window[1];
                let distance: FloatMeters = self.geodesic.inverse(p1.lat, p1.lon, p2.lat, p2.lon);
                distance
            })
            .sum()
    }

    pub fn route(&self, angles: &[Degrees]) -> Route {
        let points = self.get_route_points_from_angles(angles);
        let distance = self.get_distance_from_points(&points);

        Route { distance, points }
    }

    pub fn distance(&self, angles: &[Degrees]) -> FloatMeters {
        let points = self.get_route_points_from_angles(angles);
        self.get_distance_from_points(&points)
    }
}

impl CostFunction for CompactTask {
    type Param = Vec<Degrees>;
    type Output = FloatMeters;

    fn cost(&self, p: &Self::Param) -> Result<Self::Output, Error> {
        Ok(self.distance(p))
    }
}

impl Gradient for CompactTask {
    type Param = Vec<Degrees>;
    type Gradient = Vec<Degrees>;

    fn gradient(&self, p: &Self::Param) -> Result<Self::Gradient, Error> {
        vec::central_diff(&|x: &Vec<Degrees>| Ok(self.distance(x)))(p)
    }
}

pub fn find_shortest_route(task: CompactTask, n_tries: u32) -> Result<Route, OptimizationError> {
    if n_tries == 0 {
        return Err(OptimizationError::InvalidArgument);
    }

    let mut best_route_so_far: Option<Route> = None;

    for i in 0..n_tries {
        let p_init = task
            .turnpoints
            .iter()
            .map(|_| rand::random_range(-180.0..180.0))
            .collect();

        let linesearch = MoreThuenteLineSearch::new();
        let solver = LBFGS::new(linesearch, 10)
            .with_tolerance_cost(0.001)
            .map_err(|e| OptimizationError::SolverCreationError(e.to_string()))?;

        debug!(
            "Starting L-BFGS run {}/{} with p_init: {:?}",
            i + 1,
            n_tries,
            p_init
        );

        match Executor::new(task.clone(), solver)
            .configure(|state| state.param(p_init).max_iters(1000).target_cost(0.0))
            .run()
        {
            Ok(run_result) => {
                let state = run_result.state();
                debug!(
                    "Run {}/{}: Terminated with {:?}. Reason: {:?}. Iterations: {}. Cost: {:.4e}. Time: {:?}",
                    i + 1,
                    n_tries,
                    state.get_termination_status(),
                    state.get_termination_reason(),
                    state.get_iter(),
                    state.get_best_cost(),
                    state.get_time()
                );

                if state.get_termination_reason() == Some(&TerminationReason::SolverConverged) {
                    if let Some(best_param_for_run) = state.get_best_param() {
                        let current_route = task.route(best_param_for_run);
                        if best_route_so_far
                            .as_ref()
                            .is_none_or(|best| current_route.distance < best.distance)
                        {
                            debug!(
                                "Run {}/{}: New best route found with distance {:.2}m",
                                i + 1,
                                n_tries,
                                current_route.distance
                            );
                            best_route_so_far = Some(current_route);
                        }
                    }
                }
            }
            Err(e) => {
                debug!("Run {}/{}: Optimization failed: {}", i + 1, n_tries, e);
            }
        }
    }

    best_route_so_far.ok_or(OptimizationError::NoConvergence)
}

#[cfg(test)]
mod tests {
    use approx::assert_abs_diff_eq;

    use crate::task::Task;

    #[test]
    fn test_find_shortest_route() {
        let contents = include_str!("../data/2025-pwca-algodonales-t2.xctsk");
        let task = Task::from_xctsk_contents(contents).unwrap();
        let opti = task.find_shortest_route().unwrap();
        assert_eq!(opti.speed_section.distance.round(), 68479.0);
        assert_abs_diff_eq!(
            opti.pre_speed_section.as_ref().unwrap().distance.round(),
            190.0,
            epsilon = 10.0
        );
        assert_eq!(
            opti.post_speed_section.as_ref().unwrap().distance.round(),
            2000.0
        );
        assert_abs_diff_eq!(opti.total_distance().round(), 70670.0, epsilon = 10.0);
    }
}
