use crate::base::constants::{
    ASTRONAUMICAL_UNIT, EQUATOR_THRESHOLD, NUMBER_THERMAL_SKIN_DEPTH, NUMERICAL_STABILITY,
    SOLAR_FLUX, TAU,
};
use crate::base::object_3d::Object3D;
use crate::toolbox::matrix::{clip, dot_vectors, linspace};
use crate::{V3, V3X, VX};
use itertools::zip;
use na::{Rotation3, Unit};
use std::path::Path;

/// Message for expected error when calling a field of Properties that has not been computed.
const EXPECT_TIME_STEP_NOT_SET: &str = "The time step of the simulation has not been shared yet with the properties of the celestial body. Once the World is initialized, the time step is shared and you can call this variable.";

/// A celestial body representation.
///
/// # Definition
///
/// [`Body`] is defined by its name, its position with respect to the Sun, the shape model of its object file, the physical properties of its ground layers and its surface temperatures
///
/// # Example
///
/// ```no_run
/// use kalast::{Properties, Body, ASTRONAUMICAL_UNIT, HOUR, V3};
/// use std::path::Path;
///
/// let mut body = Body::new(
///     "Dimorphos", // name
///     V3::new(0., 1., 0.) * ASTRONAUMICAL_UNIT, // distance to Sun
///     Path::new("path/to/dimorphos.obj"), // 3D object path
///     Properties::new(
///         11.92 * HOUR, // rotation period
///         11.92 * HOUR, // revolution period
///         0.,           // obliquity
///         500.,         // thermal inertia
///         2146.,        // density
///         600.,         // heat capacity
///         0.07,         // albedo
///         0.9,          // emissivity
///     ),
/// );
/// ```
pub struct Body {
    /// Name of the celestial body.
    pub name: &'static str,
    /// Position of the celestial body with respect to the Sun (in meters).
    pub position: V3<f64>,
    /// Path to the 3D object file for the shape model of the celestial body.
    pub object: Object3D,
    /// Ground properties of the celestial body (completed with time).
    pub properties: Properties,
    /// Temperatures at the surface of the celestial body (witout the mask of the faces of the
    /// 3D object applied).
    _temperatures: VX<f64>,
    /// Temperatures at the surface of the celestial body (with the mask of the faces of the
    /// 3D object applied).
    _masked_temperatures: Option<VX<f64>>,
}

impl Body {
    /// Celestial body constructor.
    pub fn new(
        name: &'static str,
        position: V3<f64>,
        path: &'static Path,
        properties: Properties,
    ) -> Self {
        let object = Object3D::new(path);
        let number_faces_raw = object.number_faces_raw();
        let mut body = Self {
            name,
            position,
            object,
            properties,
            _temperatures: VX::zeros(number_faces_raw),
            _masked_temperatures: None,
        };
        body.initiate_obliquity();
        body
    }

    /// Temperatures getter that applies 3D object faces' mask.
    ///
    /// + return masked temperatures if available
    /// + fill masked temperatures with unmasked temperatures if available and masked temperatures
    /// unavailable
    pub fn temperatures(&mut self) -> &mut VX<f64> {
        if self._masked_temperatures.is_none() {
            let mut temperatures = VX::<f64>::zeros(self.object.number_faces());
            for (temperature, (_, temperature_ref)) in zip(
                temperatures.iter_mut(),
                self._temperatures
                    .iter()
                    .enumerate()
                    .filter(|&(i, _)| self.object.faces_mask()[i]),
            ) {
                *temperature = temperature_ref.clone();
            }
            self._masked_temperatures = Some(temperatures);
        }
        self._masked_temperatures.as_mut().unwrap()
    }

    /// Temperatures setter that shall match 3D object faces' mask.
    ///
    /// Also update unmasked temperatures at masked indices, for the computed temperatures
    /// with mask will not be lost in case the memoized masked temperatures is reset.
    pub fn set_temperatures(&mut self, temperatures: VX<f64>) {
        // Update masked temperatures
        self._masked_temperatures = Some(temperatures.clone());
        // Update unmasked temperatures at masked indices
        let mut unmasked_temperatures = self._temperatures.clone();
        for (temperature, (_, temperature_ref)) in zip(
            unmasked_temperatures.iter_mut(),
            temperatures
                .iter()
                .enumerate()
                .filter(|&(i, _)| self.object.faces_mask()[i]),
        ) {
            *temperature = temperature_ref.clone();
        }
        self._temperatures = unmasked_temperatures;
    }

    /// Iterate in the revolution of the celestial body around its spin axis over a time step.
    pub fn revolution_iteration(&mut self) {
        let old_normals = self.object.normals().into_owned();
        self.object
            .normals_mut()
            .copy_from(&(self.properties.rotation_matrix() * old_normals));
    }

    /// Compute solar flux at the surface of the celestial body. This fixes the surface flux
    /// conditions.
    pub fn solar_flux(&mut self) -> VX<f64> {
        clip(
            &(SOLAR_FLUX * (1. - self.properties.albedo()) * self.illumination_cosine()
                / (self.distance_to_sun() / ASTRONAUMICAL_UNIT).powi(2)),
            Some(0.),
            None,
        )
    }

    /// Compute the illumation angles for each face at the surface of the celestial body (angle
    /// between the normal of each face and the direction of the Sun).
    pub fn illumination_cosine(&mut self) -> VX<f64> {
        let facets_directions_to_sun = self.facets_directions_to_sun();
        dot_vectors(
            &-&facets_directions_to_sun,
            &self.object.normals().into_owned(),
        )
    }

    /// Compute the direction of the Sun for each face at the surface of the celestial body.
    pub fn facets_directions_to_sun(&mut self) -> V3X<f64> {
        let position_sun = V3::<f64>::new(0., 0., 0.);
        let mut directions = self.object.centers().into_owned();
        for mut face in directions.column_iter_mut() {
            face.copy_from(&(&self.position + &face - &position_sun).normalize());
        }
        directions
    }

    /// Compute the distance of the celestial body to the Sun.
    pub fn distance_to_sun(&self) -> f64 {
        self.position.norm()
    }

    /// Compute the distance to another body.
    pub fn distance_to_body(&self, other: &Self) -> f64 {
        (self.position_relative_to_body(other)).norm()
    }

    /// Compute the distance to another body.
    pub fn position_relative_to_body(&self, other: &Self) -> V3<f64> {
        other.position - self.position
    }

    /// Compute the direction to another body.
    pub fn direction_to_body(&self, other: &Self) -> Unit<V3<f64>> {
        Unit::new_normalize(self.position_relative_to_body(other))
    }

    /// Update the view with the indices of the faces that are located at the equator.
    pub fn set_faces_mask_equator(&mut self) {
        let mask = VX::from_iterator(
            self.object.number_faces_raw(),
            self.object
                .sphericals()
                .row(1)
                .iter()
                .map(|latitude| latitude.abs() < EQUATOR_THRESHOLD),
        );
        self.object.set_faces_mask(mask);
        self.reset_memoized_fields();
    }

    /// Reset all memoized fields.
    fn reset_memoized_fields(&mut self) {
        self._masked_temperatures = None;
        self.object.recompute_faces();
    }

    /// Initiate the obliquity by tilting the revolution axis
    fn initiate_obliquity(&mut self) {
        let old_vertices = self.object.vertices().clone();
        self.object
            .set_vertices(self.properties.obliquity_initiation_rotation_matrix() * old_vertices);
    }
}

/// Ground properties of celestial bodies
///
/// # Definition
///
/// [`Properties`] are initialized
///
/// # Example
///
/// ```
/// use kalast::{Properties, HOUR};
///
/// let properties = Properties::new(
///     11.92 * HOUR, // rotation period
///     11.92 * HOUR, // revolution period
///     0.,           // obliquity
///     500.,         // thermal inertia
///     2146.,        // density
///     600.,         // heat capacity
///     0.07,         // albedo
///     0.9,          // emissivity
/// );
/// ```
pub struct Properties {
    /// Time to complete a rotation of the celestial body around its primary body (corresponds to a body year), given in seconds.
    _rotation_period: f64,
    /// Time to complete a rotation of the celestial around its spin axis (corresponds to a day), given in seconds.
    _revolution_period: f64,
    /// The obliquity defines the inclination of the spin axis, given in radians.
    _obliquity: f64,
    /// The thermal inertia characterize the sensitivity to temperature changes, given in J.K^{-1}.m^{-2}.s^-{-1/2}.
    _thermal_inertia: f64,
    /// Material density of the surface, given in kg.m^{-3}.
    _density: f64,
    /// Heat capacity of the surface, given in J.K^{-1}.
    _heat_capacity: f64,
    /// The surface albedo defines the capacity to reflect the light.
    _albedo: f64,
    /// Surface emissivity.
    _emissivity: f64,
    /// The number of thermal skin depth to consider before adiabatic condition in the ground.
    _number_thermal_skin_depth: usize,
    /// Time step of the simulation.
    _time_step: Option<f64>,
    /// The surface conductivity, given in J.kg.m^{-1}.s^{-1}.K^{-2}.
    _conductivity: Option<f64>,
    /// The surface diffusivity, given in m^{2}.K^{-1}.s^{-1}.
    _diffusivity: Option<f64>,
    /// The thermal skin depth, given in m.K^{-1/2}.
    _thermal_skin_depth: Option<f64>,
    /// The ground depth is the number of thermal skin depth time the thermal skin depth.
    _ground_depth: Option<f64>,
    /// The ground step, given in meters.
    _ground_step: Option<f64>, // unit: m
    /// The vector of 1D grounds layers defined by the ground steps.
    _ground_vector: Option<VX<f64>>,
    /// The surface diffusivity times the space and time steps, given in K^{-1}.
    _diffusivity_field: Option<f64>,
    /// The axis of the revolution rotation.
    _rotation_axis: Option<Unit<V3<f64>>>,
    /// The angle for a revolution of one time step, given in radians.
    _rotation_angle: Option<f64>,
    /// The rotation matrix for one time step built from the axis and angle of rotation.
    _rotation_matrix: Option<Rotation3<f64>>,
    /// The axis for initial obliquity incination.
    _obliquity_initiation_axis: Option<Unit<V3<f64>>>,
    /// The rotation matrix for initial obliquity incination.
    _obliquity_initiation_rotation_matrix: Option<Rotation3<f64>>,
}

impl Properties {
    /// Basic properties constructor.
    pub fn new(
        rotation_period: f64,
        revolution_period: f64,
        obliquity: f64,
        thermal_inertia: f64,
        density: f64,
        heat_capacity: f64,
        albedo: f64,
        emissivity: f64,
    ) -> Self {
        let mut prop = Self {
            _rotation_period: rotation_period,
            _revolution_period: revolution_period,
            _obliquity: obliquity,
            _thermal_inertia: thermal_inertia,
            _density: density,
            _heat_capacity: heat_capacity,
            _albedo: albedo,
            _emissivity: emissivity,
            _number_thermal_skin_depth: NUMBER_THERMAL_SKIN_DEPTH,
            _time_step: None,
            _conductivity: None,
            _diffusivity: None,
            _thermal_skin_depth: None,
            _ground_depth: None,
            _ground_step: None,
            _ground_vector: None,
            _diffusivity_field: None,
            _rotation_axis: None,
            _rotation_angle: None,
            _rotation_matrix: None,
            _obliquity_initiation_axis: None,
            _obliquity_initiation_rotation_matrix: None,
        };
        prop.compute_obliquity_initiation_axis();
        prop.compute_obliquity_initiation_rotation_matrix();
        prop
    }

    /// Rotation period getter.
    pub fn rotation_period(&self) -> f64 {
        self._rotation_period
    }

    /// Revolution period getter.
    pub fn revolution_period(&self) -> f64 {
        self._revolution_period
    }

    /// Obliquity getter.
    pub fn obliquity(&self) -> f64 {
        self._obliquity
    }

    /// Thermal inertia getter.
    pub fn thermal_inertia(&self) -> f64 {
        self._thermal_inertia
    }

    /// Density getter.
    pub fn density(&self) -> f64 {
        self._density
    }

    /// Heat capacity getter.
    pub fn heat_capacity(&self) -> f64 {
        self._heat_capacity
    }

    /// Albedo getter.
    pub fn albedo(&self) -> f64 {
        self._albedo
    }

    /// Emissivity getter.
    pub fn emissivity(&self) -> f64 {
        self._emissivity
    }

    /// Number of thermal skin depth getter.
    pub fn number_thermal_skin_depth(&self) -> usize {
        self._number_thermal_skin_depth
    }

    /// Number of thermal skin depth setter.
    pub fn set_number_thermal_skin_depth(&mut self, number_thermal_skin_depth: usize) {
        self._number_thermal_skin_depth = number_thermal_skin_depth;
        self._ground_depth = None;
        self.compute_ground_depth();
    }

    /// Time step getter. Fails if the time step of the simulation has not been shared with this
    /// struct.
    pub fn time_step(&self) -> f64 {
        self._time_step.expect(EXPECT_TIME_STEP_NOT_SET)
    }

    /// Time step setter.
    pub fn set_time_step(&mut self, time_step: f64) {
        self._time_step = Some(time_step);
        self.compute_properties();
    }

    /// Compute properties not initialized.
    pub fn compute_properties(&mut self) {
        self.compute_conductivity();
        self.compute_diffusivity();
        self.compute_thermal_skin_depth();
        self.compute_ground_depth();
        self.compute_ground_step();
        self.compute_ground_vector();
        self.compute_diffusivity_field();
        self.compute_rotation_axis();
        self.compute_rotation_angle();
        self.compute_rotation_matrix();
    }

    /// Conductivity getter.
    pub fn conductivity(&self) -> f64 {
        self._conductivity.unwrap()
    }

    /// Compute conductivity.
    pub fn compute_conductivity(&mut self) {
        if self._conductivity.is_none() {
            self._conductivity =
                Some(self.thermal_inertia().powi(2) / (self.density() * self.heat_capacity()));
        }
    }

    /// Diffusivity getter.
    pub fn diffusivity(&self) -> f64 {
        self._diffusivity.expect(EXPECT_TIME_STEP_NOT_SET)
    }

    /// Compute diffusivity.
    pub fn compute_diffusivity(&mut self) {
        if self._diffusivity.is_none() {
            self._diffusivity = Some(self.conductivity() / (self.density() * self.heat_capacity()));
        }
    }

    /// Thermal skin depth getter.
    pub fn thermal_skin_depth(&self) -> f64 {
        self._thermal_skin_depth.expect(EXPECT_TIME_STEP_NOT_SET)
    }

    /// Compute thermal skin depth.
    pub fn compute_thermal_skin_depth(&mut self) {
        if self._thermal_skin_depth.is_none() {
            self._thermal_skin_depth =
                Some((self.diffusivity() * TAU / 2. * self.revolution_period()).sqrt());
        }
    }

    /// Ground depth getter.
    pub fn ground_depth(&self) -> f64 {
        self._ground_depth.expect(EXPECT_TIME_STEP_NOT_SET)
    }

    /// Compute depth memoized.
    pub fn compute_ground_depth(&mut self) {
        if self._ground_depth.is_none() {
            self._ground_depth =
                Some(self.thermal_skin_depth() * self.number_thermal_skin_depth() as f64);
        }
    }

    /// Ground step getter.
    pub fn ground_step(&self) -> f64 {
        self._ground_step.expect(EXPECT_TIME_STEP_NOT_SET)
    }

    /// Compute ground step.
    pub fn compute_ground_step(&mut self) {
        if self._ground_step.is_none() {
            self._ground_step =
                Some((self.time_step() * self.diffusivity() / NUMERICAL_STABILITY).sqrt());
        }
    }

    /// Ground vector getter as reference.
    pub fn ground_vector(&self) -> &VX<f64> {
        self._ground_vector
            .as_ref()
            .expect(EXPECT_TIME_STEP_NOT_SET)
    }

    /// Compute ground vector.
    pub fn compute_ground_vector(&mut self) {
        if self._ground_vector.is_none() {
            self._ground_vector = Some(linspace(0., self.ground_depth(), self.ground_step()));
        }
    }

    /// Diffusivity field getter.
    pub fn diffusivity_field(&self) -> &f64 {
        self._diffusivity_field
            .as_ref()
            .expect(EXPECT_TIME_STEP_NOT_SET)
    }

    /// Compute diffusivity field.
    pub fn compute_diffusivity_field(&mut self) {
        if self._diffusivity_field.is_none() {
            self._diffusivity_field =
                Some(self.diffusivity() * self.time_step() / self.ground_step().powi(2));
        }
    }

    /// Rotation axis getter as reference.
    pub fn rotation_axis(&self) -> &Unit<V3<f64>> {
        self._rotation_axis
            .as_ref()
            .expect(EXPECT_TIME_STEP_NOT_SET)
    }

    /// Compute rotation axis.
    pub fn compute_rotation_axis(&mut self) {
        if self._rotation_axis.is_none() {
            self._rotation_axis = Some(Unit::new_normalize(V3::new(
                0.,
                -self.obliquity().sin(),
                self.obliquity().cos(),
            )));
        }
    }

    /// Rotation angle getter.
    pub fn rotation_angle(&self) -> f64 {
        self._rotation_angle.expect(EXPECT_TIME_STEP_NOT_SET)
    }

    /// Compute rotation angle.
    pub fn compute_rotation_angle(&mut self) {
        if self._rotation_angle.is_none() {
            self._rotation_angle = Some(TAU * self.time_step() / self.revolution_period());
        }
    }

    /// Rotation matrix getter as reference.
    pub fn rotation_matrix(&self) -> &Rotation3<f64> {
        self._rotation_matrix
            .as_ref()
            .expect(EXPECT_TIME_STEP_NOT_SET)
    }

    /// Compute rotation matrix.
    pub fn compute_rotation_matrix(&mut self) {
        if self._rotation_matrix.is_none() {
            self._rotation_matrix = Some(Rotation3::from_axis_angle(
                self.rotation_axis(),
                self.rotation_angle(),
            ));
        }
    }

    /// Obliquity initiation axis getter as reference.
    pub fn obliquity_initiation_axis(&self) -> &Unit<V3<f64>> {
        self._obliquity_initiation_axis
            .as_ref()
            .expect("Obliquity initiation axis not set.")
    }

    /// Compute obliquity initiation axis.
    pub fn compute_obliquity_initiation_axis(&mut self) {
        if self._obliquity_initiation_axis.is_none() {
            self._obliquity_initiation_axis = Some(Unit::new_normalize(V3::new(1., 0., 0.)));
        }
    }

    /// Obliquity initiation rotation matrix getter as reference.
    pub fn obliquity_initiation_rotation_matrix(&self) -> &Rotation3<f64> {
        self._obliquity_initiation_rotation_matrix
            .as_ref()
            .expect("Obliquity initiation rotation matrix no set.")
    }

    /// Compute obliquity initiation rotation matrix.
    pub fn compute_obliquity_initiation_rotation_matrix(&mut self) {
        if self._obliquity_initiation_rotation_matrix.is_none() {
            self._obliquity_initiation_rotation_matrix = Some(Rotation3::from_axis_angle(
                self.obliquity_initiation_axis(),
                self.obliquity(),
            ));
        }
    }
}
