use std::collections::BTreeSet;
use std::time::Instant;
use bevy::app::App;
use bevy::color::palettes::basic::{RED, WHITE};
use bevy::color::palettes::css::LIMEGREEN;
use bevy::math::Vec3Swizzles;
use good_lp::{default_solver, variable, Expression, ProblemVariables, Solution, SolutionStatus, SolverModel, Variable};
use leafwing_input_manager::prelude::ActionState;
use crate::attachment::Parts;
use crate::client::input::ClientAction;
use crate::ecs::thruster::{PartThrusters, Thruster, ThrusterOfPart};
use crate::prelude::*;
use crate::client::input::util::ActionStateExt;
use crate::ecs::Me;
use crate::thrust::ThrustSolution;

pub fn client_thrusters_plugin(app: &mut App) {
    app
        .insert_resource(ThrusterDebugRes(false))
        .insert_resource(ThrustSolution {
            thrusters_on: BTreeSet::default(),
            converged: true,
        })
        .add_systems(Update, draw_thruster_debug)
        .add_systems(Update, solve_thrust);
}

#[derive(Resource, Deref)]
pub struct ThrusterDebugRes(pub bool);

fn draw_thruster_debug(
    thruster_debug_res: Res<ThrusterDebugRes>,
    thrusters: Query<(&Thruster, Entity, &GlobalTransform)>,
    thrust_solution: Res<ThrustSolution>,
    mut gizmos: Gizmos,
) {
    if !thruster_debug_res.0 { return };
    for thruster in thrusters {
        // Draw white if it's just a thruster, bright green if it's in the current thrust solution
        let mut color = if thrust_solution.thrusters_on.contains(&thruster.1) {
            LIMEGREEN
        } else {
            WHITE
        };
        // Exception: if the thrust solution failed to converge, RED
        if !thrust_solution.converged {
            color = RED;
        }
        let rescaled_thrust_vector = thruster.0.thrust_vector / 200.0;
        gizmos.arrow_2d(
            thruster.2.translation().xy(),
            thruster.2.translation().xy() + thruster.2.rotation().mul_vec3(rescaled_thrust_vector.extend(0.0)).xy(),
            color
        );
    }
}

// TODO: split this into two passes
/// The thrust solver!
/// This is an annoyingly complicated function...
fn solve_thrust(
    me: Query<(Option<&Parts>, &GlobalTransform, Entity), With<Me>>,
    parts: Query<&PartThrusters>,
    thrusters: Query<(&Thruster, &GlobalTransform)>,
    input: Res<ActionState<ClientAction>>,
    mut solution: ResMut<ThrustSolution>,
    mut events: MessageWriter<ThrustSolution>,
) {
    if !(
        input.button_changed(&ClientAction::ThrustForward)
            || input.button_changed(&ClientAction::ThrustBackward)
            || input.button_changed(&ClientAction::TorqueCw)
            || input.button_changed(&ClientAction::TorqueCcw)
            || input.button_changed(&ClientAction::ThrustRight)
            || input.button_changed(&ClientAction::ThrustLeft)
    ) { return; /* no changes, existing thrust solution is valid */ }

    trace!("input changed, recalculating thrust solution");
    let start = Instant::now();
    solution.thrusters_on.clear();
    solution.converged = false;

    // we need to find our entire ship
    let Ok((our_parts, hearty_transform, hearty)) = me.single() else {
        error!("could not solve for thrust: hearty does not exist?");
        error!("failed to solve for thrust after {}ms", start.elapsed().as_millis());
        return;
    };

    // determine our target vector:
    // unit vector in the intended direction of movement

    // Z-axis torque: this cursed thing is apparently standard
    // +Z == counterclockwise/ccw
    // -Z == clockwise/cw


    /*

    Background info:

    The thrust solver operates in two passes.

    Pass 1: confusingly called the "Thrust" pass, is responsible for cartesian thrust vectoring
    Pass 2: appropriately named the "Torque" pass, is soley responsible for torque

    In the end, the results of both passes are added together and set as the thrust solution.

    This is why this whole function operates in duplicate.

     */

    let mut target_unit_vector = Vec3::ZERO; // target vector for thrust pass
    let mut target_torque_vector = Vec3::ZERO; // target vector for torque pass

    let mut anything_pressed = false; // are we going to do anything?

    // +y
    if input.pressed(&ClientAction::ThrustForward) {
        anything_pressed = true;
        target_unit_vector += hearty_transform.rotation() * Vec3::new(0.0, 1.0, 0.0);
    }
    // -y
    if input.pressed(&ClientAction::ThrustBackward) {
        anything_pressed = true;
        target_unit_vector += hearty_transform.rotation() * Vec3::new(0.0, -1.0, 0.0);
    }
    // +x
    if input.pressed(&ClientAction::ThrustRight) {
        anything_pressed = true;
        target_unit_vector += hearty_transform.rotation() * Vec3::new(1.0, 0.0, 0.0);
    }
    // -x
    if input.pressed(&ClientAction::ThrustLeft) {
        anything_pressed = true;
        target_unit_vector += hearty_transform.rotation() * Vec3::new(-1.0, 0.0, 0.0);
    }
    // cw => -z
    if input.pressed(&ClientAction::TorqueCw) {
        anything_pressed = true;
        target_torque_vector += Vec3::new(0.0, 0.0, -1.0);
    }
    // ccw => +z
    if input.pressed(&ClientAction::TorqueCcw) {
        anything_pressed = true;
        target_torque_vector += Vec3::new(0.0, 0.0, 1.0);
    }

    if !anything_pressed {
        trace!("no buttons are pressed; zeroing thrust solution");
        trace!("solved thrust in {}ms", start.elapsed().as_millis());
        solution.converged = true;
        events.write(solution.clone()); // send our solution to the server, to be applied
        return;
    }

    // Normalize the target vectors.
    // The thrust solver operates purely based on direction;
    // it does not care about the strength of the thrusters.
    // Thus, we normalize everything;
    // including the target unit vectors, which are rotated with respect to Hearty
    if target_unit_vector != Vec3::ZERO {
        target_unit_vector = target_unit_vector.normalize();
    }
    // TODO(core): the torque vector should already be normalized, but the solver breaks without this.
    // TODO(core): Investigate
    if target_torque_vector != Vec3::ZERO {
        target_torque_vector = target_torque_vector.normalize();
    }

    // Determine all parts on the ship. It contains at least Hearty...
    let mut all_parts = vec![hearty];
    if let Some(parts) = our_parts {
        // and if we have an &Parts, all the attached parts too
        all_parts.extend(parts.iter());
    }

    // collect all thrusters on our ship, and figure out their thrust vectors
    let mut all_thrusters = vec![];

    for part in &all_parts {
        let Ok(part_thrusters) = parts.get(*part) else {
            continue; // This part has no thrusters
        };
        for thruster_id in &**part_thrusters {
            let Ok((thruster, thruster_transform)) = thrusters.get(*thruster_id) else {
                warn!("issue while solving for thrust: thruster {:?} of part {:?} does not exist? skipping...", *thruster_id, *part);
                continue;
            };

            // determine the thruster force in world space
            let thruster_vector = thruster_transform.rotation().mul_vec3(thruster.thrust_vector.extend(0.0)).xy();

            // determine our xy offset from hearty
            let relative_translation = thruster_transform.translation().xy() - hearty_transform.translation().xy();

            // Magic torque equation: I stole this from avian's code
            // The only difference is that like everything else, it's all normalized
            // I haven't the faintest idea what this actually does
            // No touchy
            let thruster_torque = relative_translation.normalize().extend(0.0).cross(thruster_vector.normalize().extend(0.0)).z;

            // Although all the numbers going in were normalized, the torque output is in different
            // units and is wacky. Re-normalize it to a set of expected values, since this is all
            // direction based anyway.
            let renormalized_thruster_torque = if thruster_torque.abs() < 0.1 {
                0.0 // This thruster's effect is small enough to be ignored
            } else if thruster_torque < 0.0 {
                -1.0 // if it's negative, force to -1
            } else {
                1.0 // if it's positive, force to +1
            };

            // TODO(core): remove overly verbose debug logging
            trace!("thruster: {:?} {}({})", thruster_vector, thruster_torque, renormalized_thruster_torque);

            // Then, push all this data for the next section to deal with.
            all_thrusters.push((thruster_id,
                                thruster_vector.extend(0.0),
                                Vec3::new(0.0, 0.0, renormalized_thruster_torque)
            ));
        }
    }

    /*
    Why are we normalizing everything, you may ask?

    A: The thrust solver concerns itself only with direction. It is intended to be a more dynamic
    alternative to the standard "assume the ship is the structure it should be and guess the thrust
    offsets from that" approach, mostly because I didn't want to implement that, and this seemed
    more fun.

    Also, it makes the solver converge faster, because reasons.
     */

    // calculate thrust ~~and torque~~ values

    /*
    Consult the paper for more information.

    Recall that we're optimizing an equation of form i_0 * x_0 + i_1 * x_1 + i_2 * x_2 ... i_n * x_n
    "Coefficients" are i_0 ... i_n, and can be precomputed, and x_0 ... x_n is the "decision variables"
     */

    // TODO(core): Remove overly verbose debug logging
    trace!("found {} thrusters, computing coefficients", all_thrusters.len());

    if all_thrusters.len() == 0 {
        trace!("there are no thrusters; zeroing thrust solution");
        trace!("solved thrust in {}ms", start.elapsed().as_millis());
        solution.converged = true;
        events.write(solution.clone()); // send our solution to the server, to be applied
        return;
    }

    // TODO(core): Remove overly verbose debug logging
    for thruster in &all_thrusters {
        trace!("thruster on ship: {:?}", thruster);
    }

    let coefficients = all_thrusters.iter()
        .map(|u| {
            // TODO(core): Remove overly verbose debug logging
            trace!("{} dot {}, {} dot {}", target_unit_vector, u.1.normalize(), target_torque_vector, u.2.normalize());
            // Computes both system coefficients, for simplicity
            (
                target_unit_vector.dot(u.1.normalize()), // Thrust coefficient
                target_torque_vector.dot(u.2.normalize()) // Torque coefficient
            )
        })
        .map(|u| {
            // TODO(core): Remove overly verbose debug logging
            trace!("=> {}, {}", u.0, u.1);

            // improve reliability:
            // if thrust coefficient is <0.1, zap it entirely (this thruster is not helping)
            // This is done elsewhere for torque, so pass it (u.1) through unchanged
            // TODO(core): figure out how to make this adjustable
            if u.0.abs() < 0.1 {
                (0.0, u.1)
            } else {
                (u.0, u.1)
            }
        })
        .map(|u| {
            // Sometimes NaN shows up. This just means zero. I hate math.
            (
                if u.0.is_nan() { 0.0 } else { u.0 },
                if u.1.is_nan() { 0.0 } else { u.1 },
            )
        })
        .collect::<Vec<_>>();

    trace!("preparing models");
    /* The Model is the actual solver. Currently using clarabel, but this could change. */
    let mut thrust_variables = ProblemVariables::new();
    let mut torque_variables = ProblemVariables::new();

    // add variables to problem
    // Iterate through each of our variables (thrusters) and add them to the model.
    // This will be used next to create the actual problem.
    let variables = coefficients.iter()
        .map(|u| {
            // We need to return these handles later to get the values back
            (
                (
                    u.0 as f64,
                    thrust_variables.add(variable().min(0.0).max(1.0).initial(u.0))
                ),
                (
                    u.1 as f64,
                    torque_variables.add(variable().min(0.0).max(1.0).initial(u.1))
                ),
            )
        })
        .collect::<Vec<_>>();

    // Calculate the actual problem; this is a bounded sum
    let thrust_problem: Expression = variables.iter().map(|u| u.0.0 * u.0.1).sum();
    let torque_problem: Expression = variables.iter().map(|u| u.1.0 * u.1.1).sum();

    trace!("prepared {} variables; solving", variables.len());

    // now, we run the actual solver!

    trace!("starting thrust solve @ {:?}", start.elapsed());

    let thrust_solution = match thrust_variables.maximise(thrust_problem).using(default_solver).solve() {
        Ok(soln) => soln,
        Err(e) => {
            error!("failed to solve for thrust: {}", e.to_string());
            error!("failed to solve for thrust after {}ms", start.elapsed().as_millis());
            return;
        }
    };

    // did the solution converge?
    match thrust_solution.status() {
        SolutionStatus::Optimal => {}, // yay!
        SolutionStatus::TimeLimit => {
            warn!("thrust solver failed to converge, hit time limit")
        }
        SolutionStatus::GapLimit => {
            warn!("thrust solver failed to converge, hit gap limit")
        }
    }

    trace!("finished thrust solve @ {:?}", start.elapsed());
    trace!("starting torque solve @ {:?}", start.elapsed());

    let torque_solution = match torque_variables.maximise(torque_problem).using(default_solver).solve() {
        Ok(soln) => soln,
        Err(e) => {
            error!("failed to solve for torque: {}", e.to_string());
            error!("failed to solve for torque after {}ms", start.elapsed().as_millis());
            return;
        }
    };

    // did the solution converge?
    match torque_solution.status() {
        SolutionStatus::Optimal => {}, // yay!
        SolutionStatus::TimeLimit => {
            warn!("torque solver failed to converge, hit time limit")
        }
        SolutionStatus::GapLimit => {
            warn!("torque solver failed to converge, hit gap limit")
        }
    }

    trace!("finished torque solve @ {:?}ms", start.elapsed());

    trace!("found thrust+torque solution!");

    // Finally, extract the info out of the models and compile it into a cohesive ThrustSolution.

    let mut new_soln = ThrustSolution {
        thrusters_on: BTreeSet::default(),
        converged: true
    };

    for thruster in all_thrusters.iter().enumerate() {
        // TODO(core): Remove overly verbose debug logging
        trace!("thrust solution: thruster #{} ({:?}): {} @ coeff {}", thruster.0, thruster.1.0, thrust_solution.value(variables[thruster.0].0.1), coefficients[thruster.0].0);
        trace!("torque solution: thruster #{} ({:?}): {} @ coeff {}", thruster.0, thruster.1.0, torque_solution.value(variables[thruster.0].1.1), coefficients[thruster.0].1);
        // TODO(core): make this more easily adjustable

        // Currently, we only turn on a thruster if it's variable value (think weight in a weighted sum)
        // is above 80%.
        // The solver seems to be picking 0.0 or 1.0 in all circumstances anyway, but just in case.

        if thrust_solution.value(variables[thruster.0].0.1) > 0.8 || torque_solution.value(variables[thruster.0].1.1) > 0.8 {
            new_soln.thrusters_on.insert(*thruster.1.0);
        }
    }

    let elapsed = start.elapsed();
    debug!(?elapsed, ?target_unit_vector, ?target_torque_vector, "solved for thrust and torque");
    *solution = new_soln; // save it to the Resource for use on the client...
    events.write(solution.clone()); // ...then send it to the server!
    return;
}