// 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)

//! Functionality for representing wings as "line objects", without any assumptions about how
//! lift-induced velocities are estimated. In other words, this part is common for all methods
//! available in the library, and therefore the foundation of all simulations.

use std::ops::Range;

use stormath::{
    type_aliases::Float,
    consts::PI,
    spatial_vector::SpatialVector,
    spatial_vector::transformations::RotationType,
    statistics::mean, 
    interpolation::linear_interpolation,
    rigid_body_motion::RigidBodyMotion
};

pub mod builder;
pub mod data_access;
pub mod data_update;
pub mod force_calculations;
pub mod value_mapping;
pub mod span_line;
pub mod input_power;

pub mod corrections;
pub mod prelude;

#[cfg(test)]
mod tests;

use crate::common_utils::prelude::*;
use crate::section_models::SectionModel;
use crate::controller::output::ControllerOutput;

use corrections::{
    circulation::CirculationCorrection,
    angle_of_attack::AngleOfAttackCorrection,
};
use builder::single_wing::SingleWing;
use span_line::*;

use input_power::InputPowerModel;

#[derive(Clone, Debug)]
/// The struct holds variables for a model that calculate the forces on wings, under the assumption
/// that they can be represented as a set of line elements. The intended use is in lifting line and
/// actuator line simulations.
pub struct LineForceModel {
    /// Vector of line segments that defines the span geometry of the wings. Each have its own start
    /// and end point, to allow for uncoupled analysis
    pub span_lines_local: Vec<SpanLine>,
    /// Vectors representing both the chord length and the direction of the chord for each span line
    pub chord_vectors_local_not_rotated: Vec<SpatialVector>,
    /// The length of the chord vectors, stored as it is needed for several calculations
    pub chord_lengths: Vec<Float>,
    /// Two dimensional models for lift and drag coefficients for each wing in the model
    pub section_models: Vec<SectionModel>,
    /// Indices used to sort different wings from each other.
    pub wing_indices: Vec<Range<usize>>,
    /// A vector that contains booleans that indicate whether the circulation should be zero at the
    /// ends or not. The variables are mainly used when using corrections for the circulation 
    /// distribution. The vector is structured as follows:
    /// - The first index is the wing index
    /// - The second index is the end index, where 0 means that start of the wing and 1 means the
    /// end
    /// - When the boolean is true, it is a warning to any correction method that that it is not 
    /// correct to assume zero circulation at that end.
    pub non_zero_circulation_at_ends: Vec<[bool; 2]>,
    /// Density used in force calculations
    pub density: Float,
    /// Optional correction that can be applied to the estimated circulation strength.
    pub circulation_correction: CirculationCorrection,
    /// Optional correction for the angle of attack
    pub angle_of_attack_correction: AngleOfAttackCorrection,
    /// The coordinate system to generate the output in. Variants consists of Global and Body.
    pub output_coordinate_system: CoordinateSystem,
    /// Rigid body motion of the line force model
    pub rigid_body_motion: RigidBodyMotion,
    /// Vector used to store local angles for each wing. This can be used to rotate the wing along
    /// the span axis during a dynamic simulation. The typical example is changing the angle of
    /// attack on a wing sail due to changing apparent wind conditions.
    pub local_wing_angles: Vec<Float>,
    /// The local chord vector after applying the local wing angles, but before applying the rigid
    /// body rotation.
    pub chord_vectors_local: Vec<SpatialVector>,
    /// The global chord vectors, i.e., after applying the rigid body rotation and local wing angles
    pub chord_vectors_global: Vec<SpatialVector>,
    /// The chord vectors at the span points, interpolated and extrapolated from the global chord vectors
    pub chord_vectors_global_at_span_points: Vec<SpatialVector>,
    /// The global span lines, i.e., after applying the rigid body transformations
    pub span_lines_global: Vec<SpanLine>,
    /// The global span points, i.e., after applying the rigid body transformations
    pub span_points_global: Vec<SpatialVector>,
    /// The global ctrl points, i.e., after applying the rigid body transformations
    pub ctrl_points_global: Vec<SpatialVector>,
    /// The spanwise distance from the center of each wing for each control point
    pub ctrl_point_spanwise_distance: Vec<Float>,
    /// The non-dimensional spanwise distance from the center of each wing for each control point
    pub ctrl_point_spanwise_distance_non_dimensional: Vec<Float>,
    /// Effective values for the ctrl point spanwise distance when applying prescribed circulation
    pub ctrl_point_spanwise_distance_circulation_model: Vec<Float>,
    /// Models for estimating the input energy for the sails
    pub input_power_models: Vec<InputPowerModel>,
}

impl Default for LineForceModel {
    fn default() -> Self {
        Self::new(Self::default_density())
    }
}

impl LineForceModel {
    /// Default density for air at sea level in kg/m^3
    pub fn default_density() -> Float {
        1.225
    }

    /// Creates a new empty line force model. Wings can be added using the 
    /// [LineForceModel::add_wing] function.
    pub fn new(density: Float) -> LineForceModel {
        Self {
            span_lines_local: Vec::new(),
            chord_vectors_local_not_rotated: Vec::new(),
            chord_lengths: Vec::new(),
            section_models: Vec::new(),
            wing_indices: Vec::new(),
            rigid_body_motion: RigidBodyMotion::default(),
            local_wing_angles: Vec::new(),
            non_zero_circulation_at_ends: Vec::new(),
            density,
            circulation_correction: Default::default(),
            angle_of_attack_correction: Default::default(),
            output_coordinate_system: CoordinateSystem::Global,
            chord_vectors_local: Vec::new(),
            chord_vectors_global: Vec::new(),
            chord_vectors_global_at_span_points: Vec::new(),
            span_lines_global: Vec::new(),
            span_points_global: Vec::new(),
            ctrl_points_global: Vec::new(),
            ctrl_point_spanwise_distance: Vec::new(),
            ctrl_point_spanwise_distance_non_dimensional: Vec::new(),
            ctrl_point_spanwise_distance_circulation_model: Vec::new(),
            input_power_models: Vec::new(),
        }
    }

    /// Adds a new wing to the complete model. This involves appending the span lines, chord vectors,
    /// and section models to the existing vectors, and adding the indices of the new wing to the
    /// wing indices vector.
    pub fn add_wing(&mut self, wing: &SingleWing) {
        let start_index = if self.span_lines_local.is_empty() {
            0
        } else {
            self.wing_indices.last().unwrap().end
        };

        let end_index = start_index + wing.span_lines_local.len();

        self.wing_indices.push(Range {
            start: start_index,
            end: end_index,
        });

        for line in &wing.span_lines_local {
            self.span_lines_local.push(line.clone());
            self.span_lines_global.push(line.clone());
        }

        for chord_vector in &wing.chord_vectors_local {
            self.chord_vectors_local_not_rotated.push(*chord_vector);
            self.chord_vectors_local.push(*chord_vector);
            self.chord_vectors_global.push(*chord_vector);
        }

        for chord_length in &wing.chord_lengths {
            self.chord_lengths.push(*chord_length);
        }

        self.section_models.push(wing.section_model.clone());

        self.local_wing_angles.push(0.0);
        self.non_zero_circulation_at_ends.push(wing.non_zero_circulation_at_ends);

        self.input_power_models.push(wing.input_power_model.clone());

        self.update_global_data_representations();
    }

    

    /// Finds the wing index from the global index of a line element.
    pub fn wing_index_from_global(&self, global_index: usize) -> usize {
        for (wing_index, wing_indices) in self.wing_indices.iter().enumerate() {
            if wing_indices.contains(&global_index) {
                return wing_index;
            }
        }

        panic!("Wing index not found. The global index is not part of any wing")
    }

    /// Returns the local index of a single wing. For instance, the first line element of each wing
    /// will have a local index of 0, etc. This function is primarily used to map a global index to
    /// a local index used as input to sectional model where the properties vary for each line
    /// element.
    pub fn local_index_from_global(&self, global_index: usize) -> usize {
        for wing_indices in &self.wing_indices {
            if wing_indices.contains(&global_index) {
                return global_index - wing_indices.start;
            }
        }

        panic!("Local index not found. The global index is not part of any wing")
    }

    /// Removes the velocity in the span direction from the input velocity vector.
    pub fn remove_span_velocity(
        &self,
        velocity: &[SpatialVector],
        input_coordinate_system: CoordinateSystem,
    ) -> Vec<SpatialVector> {
        let span_lines = match input_coordinate_system {
            CoordinateSystem::Global => &self.span_lines_global,
            CoordinateSystem::Body => &self.span_lines_local,
        };

        velocity
            .iter()
            .zip(span_lines.iter())
            .map(|(vel, line)| {
                let span_velocity = vel.project(line.relative_vector());

                *vel - span_velocity
            })
            .collect()
    }

    /// Calculates the wake angle behind each line element.
    pub fn wake_angles(&self, velocity: &[SpatialVector]) -> Vec<Float> {
        (0..self.nr_span_lines())
            .map(|index| {
                let wing_index = self.wing_index_from_global(index);

                match &self.section_models[wing_index] {
                    SectionModel::Foil(_) => 0.0,
                    SectionModel::VaryingFoil(_) => 0.0,
                    SectionModel::RotatingCylinder(cylinder) => cylinder.wake_angle(
                        self.chord_lengths[index],
                        velocity[index].length(),
                    ),
                    SectionModel::EffectiveWindSensor => 0.0,
                }
            })
            .collect()
    }    
    
    /// Shorthand for quickly calculating the typical force factor used when presenting
    /// non-dimensional forces from a simulation (i.e., lift and drag coefficients)
    pub fn total_force_factor(&self, freestream_velocity: Float) -> Float {
        0.5 * self.density * freestream_velocity.powi(2) * self.total_projected_area()
    }
}
