// Copyright (C) 2024, NTNU
// Author: Jarle Vinje Kramer <jarlekramer@gmail.com; jarle.a.kramer@ntnu.no>
// License: GPL v3.0 (see separate file LICENSE or https://www.gnu.org/licenses/gpl-3.0.html)


use super::*;

use crate::line_force_model::LineForceModel;

/// This code block contains the logic to update the wake structure
impl DynamicWake {
    /// Function that updates the wake data before the solver executes a new time step. The main job
    /// is to ensure that geometry of the wake and the strength behind the first panels is correct.
    ///
    /// # Steps that are performed in this function
    /// Since the last time step, the line force model might have moved, and the wake points have
    /// streamed downstream. That means that the following must happen:
    /// - Synchronize the geometry for the first points in the wake to the geometry of the line
    /// force model.
    /// - Stream the old wake points downstream.
    /// - Shift the strength of the panels downstream.
    ///
    /// # Arguments
    /// * `time_step` - The current value of the time step
    /// * `line_force_model` - The line force model that the wake is based on
    /// * `wake_points_freestream` - The freestream velocity at the points in the wake for the
    /// current time step
    pub fn update_before_solving(
        &mut self,
        time_step: Float,
        line_force_model: &LineForceModel,
        felt_span_points_freestream: &[SpatialVector]
    ) {
        self.update_wake_points_before_solving(
            time_step,
            line_force_model,
            felt_span_points_freestream
        );

        self.update_panel_data();
        self.stream_strength_values_downstream();
    }

    /// Update the wake geometry and strength based on the final solution at a time step.
    ///
    /// This will:
    /// 1) stream the wake points downstream
    /// 2) stream the strength downstream
    pub fn update_after_solving(
        &mut self,
        new_circulation_strength: &[Float],
        wake_points_freestream: &[SpatialVector],
    ) {
        self.update_wing_strength(new_circulation_strength);

        self.update_velocity_at_points(wake_points_freestream);

        self.number_of_time_steps_completed += 1;
    }

    /// Update the wake points by streaming them downstream.
    ///
    /// The first and second "rows" - meaning the wing geometries and the first row of wake points -
    /// are treaded as special cases. The rest are moved based on the euler method
    pub fn update_wake_points_before_solving(
        &mut self,
        time_step: Float,
        line_force_model: &LineForceModel,
        felt_span_points_freestream: &[SpatialVector]
    ) {
        self.synchronize_first_points_to_wing_geometry(line_force_model);

        self.stream_free_wake_points_based_on_stored_velocity(time_step);

        if self.settings.shape_damping_factor > 0.0 {
            self.move_first_free_wake_points_with_damping(line_force_model, felt_span_points_freestream);
        } else {
            self.move_first_free_wake_points_no_damping(line_force_model, felt_span_points_freestream);
        }

        self.move_last_wake_points(felt_span_points_freestream);
    }

    /// Takes a line force vector as input, that might have a different position and orientation
    /// than the previous time step, and updates the first points in the wake to match the new
    /// geometry.
    ///
    /// # Argument
    /// * `line_force_model` - The line force model that the wake is based on
    pub fn synchronize_first_points_to_wing_geometry(
        &mut self,
        line_force_model: &LineForceModel
    ) {
        let nr_span_points = line_force_model.span_points_global.len();

        for i in 0..nr_span_points {
            self.points[i] = line_force_model.span_points_global[i];
        }
    }

    /// Recalculates the panel data based on the current geometry of the wake.
    pub fn update_panel_data(&mut self) {
        for i in 0..self.indices.nr_panels() {
            let (stream_index, span_index) = self.indices.reverse_panel_index(i);

            let panel_points = self.panel_points(stream_index, span_index);

            self.panels[i] = Panel::new(
                panel_points,
                self.potential_theory_settings.far_field_ratio,
                self.panels_viscous_core_length[i]
            );

        }
    }

    /// Computes the velocity at all the wake points.
    ///
    /// The velocity is calculated as the sum of the freestream velocity and the induced velocity.
    /// However, if the settings contains and end-index for the induced velocities, the induced
    /// velocities can be neglected for the last panels. This is useful for speeding up simulations.
    ///
    /// # Argument
    /// * `wake_points_freestream` - A vector containing the freestream velocity at the wake points
    pub fn update_velocity_at_points(&mut self, wake_points_freestream: &[SpatialVector]) {
        for i in 0..self.points.len() {
            self.velocity_at_points[i] = wake_points_freestream[i];
        }

        let end_index = self.settings.end_index_induced_velocities_on_wake.min(
            self.points.len()
        );

        if end_index > 0 && self.number_of_time_steps_completed > 2 {
            let u_i_calc: Vec<SpatialVector> = self.induced_velocities(&self.points[0..end_index]);

            for i in 0..end_index {
                self.velocity_at_points[i] += u_i_calc[i];
            }
        }
    }

    /// Update the strength of the wake panels closest to the wing geometry.
    ///
    /// This is the same as updating the circulation strength on the first panels in the wake.
    pub fn update_wing_strength(&mut self, new_circulation_strength: &[Float]) {
        for i in 0..new_circulation_strength.len() {
            self.strengths[i] = new_circulation_strength[i];
        }
    }

    /// Moves the first wake points after the wing geometry itself.
    ///
    /// In general, the principle is that the first free wake points are moved from the wing
    /// geometry and then *either* in the direction of the chord vector or the velocity vector.
    fn move_first_free_wake_points_no_damping(
        &mut self,
        line_force_model: &LineForceModel,
        felt_span_points_freestream: &[SpatialVector]
    ) {
        let nr_first_wake_points = self.indices.nr_points_along_span;

        let mut direction_vectors: Vec<SpatialVector> = Vec::with_capacity(nr_first_wake_points);

        if self.settings.use_chord_direction {
            for i in 0..nr_first_wake_points{
                direction_vectors.push(
                    line_force_model.chord_vectors_global_at_span_points[i].normalize()
                )
            }
        } else {
            for i in 0..nr_first_wake_points {
                direction_vectors.push(
                    felt_span_points_freestream[i].normalize()
                )
            }
        }

        let first_panel_length = self.representative_chord_length * self.settings.first_panel_relative_length;

        for i in 0..self.indices.nr_points_along_span {
            let estimated_new_wake_point = self.points[i] + direction_vectors[i] * first_panel_length;

            self.points[i + self.indices.nr_points_along_span] = estimated_new_wake_point;
        }
    }

    /// Moves the first wake points after the wing geometry itself.
    ///
    /// In general, the principle is that the first free wake points are moved from the wing
    /// geometry and then *either* in the direction of the chord vector or the velocity vector.
    fn move_first_free_wake_points_with_damping(
        &mut self,
        line_force_model: &LineForceModel,
        felt_span_points_freestream: &[SpatialVector]
    ) {
        let old_start_index = self.indices.nr_points_along_span;
        let old_end_index   = 2 * self.indices.nr_points_along_span;

        let nr_first_wake_points = self.indices.nr_points_along_span;

        let old_wake_points = self.points[old_start_index..old_end_index].to_vec();

        let mut direction_vectors: Vec<SpatialVector> = Vec::with_capacity(nr_first_wake_points);

        if self.settings.use_chord_direction {
            for i in 0..nr_first_wake_points{
                direction_vectors.push(
                    line_force_model.chord_vectors_global_at_span_points[i].normalize()
                )
            }
        } else {
            for i in 0..nr_first_wake_points {
                direction_vectors.push(
                    felt_span_points_freestream[i].normalize()
                )
            }
        }

        let first_panel_length = self.representative_chord_length * self.settings.first_panel_relative_length;

        for i in 0..self.indices.nr_points_along_span {
            let estimated_new_wake_point = self.points[i] + direction_vectors[i] * first_panel_length;

            self.points[i + self.indices.nr_points_along_span] =
                old_wake_points[i] * self.settings.shape_damping_factor +
                estimated_new_wake_point * (1.0 - self.settings.shape_damping_factor);
        }
    }

    /// Moves the last points in the wake.
    ///
    /// The length of the last panel is determined based on the `last_panel_relative_length`
    /// parameter in the settings, plus the chord length. The direction of the last panel is taken
    /// to be the same as the direction between the previous two points in the wake.
    ///
    /// # Arguments
    /// * `line_force_model_data` - The line force model data that the should use for the update
    pub fn move_last_wake_points(
        &mut self,
        felt_span_points_freestream: &[SpatialVector],
    ) {
        let nr_points = self.points.len();
        let nr_points_along_span = self.indices.nr_points_along_span;

        let multi_panel_wake = self.indices.nr_panels_per_line_element > 1;

        let start_index_last = nr_points - nr_points_along_span;
        let start_index_previous = start_index_last - nr_points_along_span;

        let change_directions: Vec<SpatialVector> = if multi_panel_wake {
            (0..self.indices.nr_points_along_span).map(|i| {
                let previous_point        = self.points[start_index_previous + i];
                let second_previous_point = self.points[start_index_previous - nr_points_along_span + i];

                (previous_point - second_previous_point).normalize()
            }).collect()
        } else {
            felt_span_points_freestream.iter()
                .map(|v| v.normalize())
                .collect()
        };

        for i in 0..self.indices.nr_points_along_span {
            let change_vector = self.settings.last_panel_relative_length * self.representative_chord_length * change_directions[i];

            self.points[start_index_last + i] = self.points[start_index_previous + i] + change_vector;
        }
    }

    /// Stream all free wake points based on the Euler method.
    fn stream_free_wake_points_based_on_stored_velocity(&mut self, time_step: Float) {
        for i_stream in (2..self.indices.nr_points_per_line_element).rev() {
            for i_span in 0..self.indices.nr_points_along_span {
                let previous_flat_index = self.indices.point_index(i_stream - 1, i_span);
                let current_flat_index  = self.indices.point_index(i_stream, i_span);

                let previous_wake_point = self.points[previous_flat_index];
                let previous_velocity = self.velocity_at_points[previous_flat_index];

                let integrated_point = previous_wake_point + time_step * previous_velocity;

                if self.settings.shape_damping_factor > 0.0 {
                    let current_wake_point = self.points[current_flat_index];

                    self.points[current_flat_index] = current_wake_point * self.settings.shape_damping_factor +
                        integrated_point * (1.0 - self.settings.shape_damping_factor);
                } else {
                    self.points[current_flat_index] = integrated_point;
                }
            }
        }
    }

    /// Shifts the strength of the panels downstream.
    pub fn stream_strength_values_downstream(&mut self) {
        for i_stream in (1..self.indices.nr_panels_per_line_element).rev() {
            for i_span in 0..self.indices.nr_panels_along_span {
                let current_index  = self.indices.panel_index(i_stream, i_span);
                let previous_index = self.indices.panel_index(i_stream - 1, i_span);

                self.strengths[current_index] = self.strengths[previous_index];
            }
        }
    }
}
