use crate::base::body::{Body, Properties};
use crate::base::constants::{STEPHAN_BOLTZMAN, TAU};
use crate::{V3, VX};
use na::{DMatrix, Unit, U3};
use std::path::Path;

/// Generic trait for environment system, can be: [`SingleBody`], [`BinarySystem`].
pub trait EnvironmentSystem {
    /// Targeted body name.
    fn name(&self) -> &String;

    /// Targeted body getter.
    fn body(&self) -> &Body;

    /// Targeted body mutable getter.
    fn body_mut(&mut self) -> &mut Body;

    /// Get the number of bodies of the environment.
    fn number_bodies(&self) -> usize;

    /// Recompute memoized fields.
    fn recompute(&mut self);

    /// Proceed to one iteration in the orbit and in revolution of the bodies in the system.
    fn orbit_and_revolution_iteration(&mut self);

    /// Compute the global temperature at equilibrium of the targeted body.
    fn compute_global_equilibrium_temperature(&self, position_sun: &V3<f64>) -> f64 {
        self.body().global_equilibrium_temperature(position_sun)
    }

    /// Compute the surface heat flux of the targeted body.
    fn compute_surface_heat_flux(&mut self, no_sun: bool, position_sun: &V3<f64>);

    /// The view factor of the target body with itself, describing the fraction of energy emitted
    /// from one facet toward another.
    ///
    /// ## Definition
    ///
    /// The viewing factor is set to zero if the line of sight is intercepted by topography (TODO)
    /// or if any of the angles between the facet outward and the line between facet centers is
    /// superior to 90 degrees.
    ///
    /// Note that when the distance between the facets is lower than the square root of the area
    /// of the facet of the other body (this is the case for neighboring facets), the view factor
    /// risk to become unphysically large. Therefore, when the view factor is greater than 0.1, we
    /// can subdivide the facet into multiple sub-facets and recompute the view factor.
    ///
    /// ## Expression
    ///
    /// The view factor for two well-separated facets i (facets of the target) and j (facets of the
    /// other body) is given by
    ///
    /// V_{ij}=\frac{a_j\cos\Theta_i\cos\Theta_j}{\pi r^2}
    ///
    /// where a_j is the surface area of facet j, \Theta_i and \Thera_j are the angles between the
    /// facets outward normals and the line between their centers, and r is the distance between
    /// the facets.
    fn view_factor_self(&self) -> &DMatrix<f64>;

    /// Set the view factor of the target body with itself.
    fn set_view_factor_self(&mut self, view_factor: DMatrix<f64>);

    /// Compute the view factor of the target body with itself.
    fn compute_view_factor_self(&mut self) {
        let mut view_factor = DMatrix::<f64>::zeros(
            self.body().object.number_faces(),
            self.body().object.number_faces(),
        );
        // Cartesian product of secondary's and primary's faces.
        for ((face_target, face_other), view) in izip!(
            iproduct!(
                self.body().object.face_iter_all(),
                self.body().object.face_iter_all()
            ),
            view_factor.iter_mut()
        ) {
            let normal_target = face_target.fixed_rows::<U3>(3);
            let normal_other = face_other.fixed_rows::<U3>(3);
            let center_target = face_target.fixed_rows::<U3>(0);
            let center_other = face_other.fixed_rows::<U3>(0);
            if center_target == center_other {
                continue;
            }
            let area_other = face_other[9];
            let relative_position_faces = center_other - center_target;
            let direction_faces = relative_position_faces.normalize();
            let distance_faces = relative_position_faces.norm();
            if distance_faces < area_other.sqrt() {
                continue;
            }
            let angle_target = normal_target.angle(&direction_faces);
            let angle_other = normal_other.angle(&-direction_faces);
            if angle_target > TAU / 4.0 || angle_other > TAU / 4.0 {
                continue;
            }
            *view = 2. * area_other * angle_target.cos() * angle_other.cos()
                / (TAU * distance_faces.powi(2));
        }
        self.set_view_factor_self(view_factor);
    }

    /// Change whether the thermophysical model includes mutual and self.
    fn complete_model(&mut self, complete: bool);

    /// Change whether the ground temperatures are initialized at the global equilibrium
    /// temperatures.
    fn initialize_at_global_equilibrium(&mut self, init_at_global: bool) {
        self.body_mut()
            .set_initialize_at_global_equilibrium(init_at_global);
    }
}

/// Environment system of a single celestial body.
///
/// # Example
///
/// ```
/// use kalast::{SingleBody, Properties, V3, ASTRONAUMICAL_UNIT, HOUR};
///
/// let mut system = SingleBody::new(
///     "Dimorphos",
///     "rsc/obj/dimorphos.obj",
///     V3::new(0.0, 1.0, 0.0) * ASTRONAUMICAL_UNIT * 1.664,
///     Properties::new(
///         11.92 * HOUR, // rotation period
///         11.92 * HOUR, // revolution period
///         0.0,          // obliquity 162.0 * TAU / 360.
///         500.0,        // thermal inertia
///         2146.0,       // density
///         600.0,        // heat capacity
///         0.07,         // albedo
///         0.9,          // emissivity
///     ),
/// );
/// ```
pub struct SingleBody {
    /// The body.
    _body: Body,
    /// The view factor of the target body with itself, describing the fraction of energy emitted
    /// from one facet toward another.
    _view_factor_self: Option<DMatrix<f64>>,
    /// Whether the mutual and self heating are turned on or off.
    _complete_model: bool,
}

impl SingleBody {
    /// Create a body.
    pub fn new<P: AsRef<Path>>(
        name: &'static str,
        path: P,
        position: V3<f64>,
        properties: Properties,
    ) -> Self {
        Self {
            _body: Body::new(name, position, path, properties),
            _view_factor_self: None,
            _complete_model: true,
        }
    }

    /// Create a body from an empty mesh.
    pub fn empty_object(name: &'static str, position: V3<f64>, properties: Properties) -> Self {
        Self {
            _body: Body::empty(name, position, properties),
            _view_factor_self: None,
            _complete_model: true,
        }
    }

    /// See [`self view factor getter`][EnvironmentSystem::view_factor_self]
    fn view_factor_self(&self) -> &DMatrix<f64> {
        self._view_factor_self.as_ref().unwrap()
    }

    /// See [`self view factor setter`][EnvironmentSystem::set_view_factor_self]
    fn set_view_factor_self(&mut self, view_factor: DMatrix<f64>) {
        self._view_factor_self = Some(view_factor);
    }

    /// Compute the diffuse thermal self heating.
    ///
    /// ## Expression
    ///
    /// W_{s,j}=\sum_{i\ne j}^{N}V_{ij}\frac{F_s A\cos\varsigma_j\left(t\right)}{r_H^2\left(t\right)}
    fn diffuse_solar_self_heating(&self, position_sun: &V3<f64>) -> VX<f64> {
        let mut diffuse_solar_radiation = VX::zeros(self.body().object.number_faces());
        let direct_solar_flux_from_other = self.body().direct_solar_flux(position_sun)
            * self.body().properties.albedo()
            / (1. - self.body().properties.albedo());
        for (flux, view_factor) in izip!(
            diffuse_solar_radiation.iter_mut(),
            self.view_factor_self().column_iter()
        ) {
            *flux = (view_factor.component_mul(&direct_solar_flux_from_other)).sum();
        }
        diffuse_solar_radiation
    }

    /// Compute the direct thermal self heating.
    ///
    /// ## Expression
    ///
    /// u_{s,j}=\sum_{i\ne j}^{N}V_{ij}\epsilon\sigma T_{s,i}^4
    fn direct_thermal_self_heating(&self) -> VX<f64> {
        let mut direct_thermal_heating = VX::zeros(self.body().object.number_faces());
        let max_temperature = self.body().surface_temperatures().max();
        for (flux, view_factor) in izip!(
            direct_thermal_heating.iter_mut(),
            self.view_factor_self().column_iter()
        ) {
            *flux = (view_factor
                * self.body().properties.emissivity()
                * STEPHAN_BOLTZMAN
                * max_temperature.powi(4))
            .sum();
        }
        direct_thermal_heating
    }
}

impl EnvironmentSystem for SingleBody {
    /// Targeted body name.
    fn name(&self) -> &String {
        &self._body.name
    }

    /// Targeted body getter.
    fn body(&self) -> &Body {
        &self._body
    }

    /// Targeted body mutable getter.
    fn body_mut(&mut self) -> &mut Body {
        &mut self._body
    }

    /// Get the number of bodies of the environment.
    fn number_bodies(&self) -> usize {
        1
    }

    /// Recompute memoized fields.
    fn recompute(&mut self) {
        if self._view_factor_self.as_ref().is_none() && self._complete_model {
            self.compute_view_factor_self();
        }
    }

    /// Proceed to one iteration in the orbit and in revolution of the bodies in the system.
    fn orbit_and_revolution_iteration(&mut self) {
        self.body_mut().revolution_iteration();
    }

    /// Compute the surface heat flux of the targeted body.
    ///
    /// ## Expression
    ///
    /// U=\frac{F_s\left(1-A\right)\cos\varsigma\left(t\right)}{r_H^2\left(t\right)}
    fn compute_surface_heat_flux(&mut self, no_sun: bool, position_sun: &V3<f64>) {
        let mut flux;
        if no_sun {
            flux = VX::zeros(self.body().object.number_faces());
        } else {
            flux = self.body().direct_solar_flux(position_sun);
            if self._complete_model {
                flux += self.diffuse_solar_self_heating(position_sun)
                    + self.direct_thermal_self_heating()
            }
        }
        self.body_mut().set_surface_heat_flux(flux);
    }

    /// See [`self view factor getter`][EnvironmentSystem::view_factor_self]
    fn view_factor_self(&self) -> &DMatrix<f64> {
        self.view_factor_self()
    }

    /// See [`self view factor setter`][EnvironmentSystem::view_factor_self]
    fn set_view_factor_self(&mut self, view_factor: DMatrix<f64>) {
        self.set_view_factor_self(view_factor)
    }

    /// Change whether the thermophysical model includes mutual and self.
    fn complete_model(&mut self, complete: bool) {
        self._complete_model = complete;
    }
}

/// Environment system of two celestial bodies.
///
/// # Definition
///
/// By default, the primary body is the target of the computation. You can change the target to be
/// the secondary with [`set_target_secondary`][BinarySystem::set_target_secondary].
///
/// # Example
///
/// ```
/// use kalast::{BinarySystem, Properties, V3, ASTRONAUMICAL_UNIT, YEAR, HOUR};
///
/// let mut binary_system = BinarySystem::new(
///     "Didymos system",
///     V3::new(0.0, 1.0, 0.0) * ASTRONAUMICAL_UNIT * 1.664,
/// )
/// .with_primary(
///     "rsc/obj/didymos.obj",
///     Properties::new(
///         1.0 * YEAR, // rotation period
///         2.7 * HOUR, // revolution period
///         0.0,        // obliquity
///         100.0,      // thermal inertia
///         2146.0,     // density
///         600.0,      // heat capacity
///         0.1,        // albedo
///         0.9,        // emissivity
///     ),
/// )
/// .with_secondary(
///     "rsc/obj/dimorphos_light.obj",
///     Properties::new(
///         11.92 * HOUR, // rotation period
///         11.92 * HOUR, // revolution period
///         0.0,          // obliquity 162.0 * TAU / 360.
///         500.0,        // thermal inertia
///         2146.0,       // density
///         600.0,        // heat capacity
///         0.07,         // albedo
///         0.9,          // emissivity
///     ),
///     -V3::new(1.0, 0.0, 0.0) * 1e3, // relative position
/// );
///
/// binary_system.set_target_secondary(true);
/// ```
pub struct BinarySystem {
    /// The name of the binary system of celestial bodies.
    pub name: String,
    /// The cartesian position of the center of mass of the binary system.
    pub position: V3<f64>,
    /// The primary body of the system.
    _primary: Option<Body>,
    /// The secondary body of the system.
    _secondary: Option<Body>,
    /// Wheter the primary body is the target of the computation or not.
    _is_primary_target: bool,
    /// The view factor describing the fraction of energy emitted from one facet of the secondary
    /// toward another of the primary.
    _view_factor: Option<DMatrix<f64>>,
    /// The view factor describing the fraction of energy emitted from one facet of the secondary
    /// toward another of the primary.
    _view_factor_self: Option<DMatrix<f64>>,
    /// Whether the mutual and self heating are turned on or off.
    _complete_model: bool,
}

impl BinarySystem {
    /// Initiate the binary system as build design.
    pub fn new(name: &str, position: V3<f64>) -> Self {
        Self {
            name: String::from(name),
            position,
            _primary: None,
            _secondary: None,
            _is_primary_target: true,
            _view_factor: None,
            _view_factor_self: None,
            _complete_model: true,
        }
    }

    /// Set the primary body of the system in the build design.
    pub fn with_primary<P: AsRef<Path>>(mut self, path: P, properties: Properties) -> Self {
        self._primary = Some(Body::new(
            format!("{} primary", self.name).as_str(),
            self.position,
            path,
            properties,
        ));
        self
    }

    /// Set the primary body of the system in the build design.
    pub fn with_secondary<P: AsRef<Path>>(
        mut self,
        path: P,
        properties: Properties,
        relative_position: V3<f64>,
    ) -> Self {
        self._secondary = Some(Body::new(
            format!("{} secondary", self.name).as_str(),
            self.position + relative_position,
            path,
            properties,
        ));
        self
    }

    /// Primary getter.
    pub fn primary(&self) -> &Body {
        self._primary.as_ref().unwrap()
    }

    /// Primary mutable getter.
    pub fn primary_mut(&mut self) -> &mut Body {
        self._primary.as_mut().unwrap()
    }

    /// Secondary getter.
    pub fn secondary(&self) -> &Body {
        self._secondary.as_ref().unwrap()
    }

    /// Secondary mutable getter.
    pub fn secondary_mut(&mut self) -> &mut Body {
        self._secondary.as_mut().unwrap()
    }

    /// Change whether the primary is the target of the computation or not: `true` to the set the
    /// secondary as the target. In case you want to switch back to the primary as the target, use
    /// `false`.
    pub fn set_target_secondary(&mut self, is_target_secondary: bool) {
        self._is_primary_target = !is_target_secondary;
    }

    /// Relative position between bodies.
    pub fn relative_position(&self) -> V3<f64> {
        self.primary().position - self.secondary().position
    }

    /// Distance between bodies.
    pub fn distance_between(&self) -> f64 {
        (self.primary().position - self.secondary().position).norm()
    }

    /// Direction from secondary to primary.
    pub fn direction_from_secondary_to_primary(&self) -> Unit<V3<f64>> {
        Unit::new_normalize(self.primary().position - self.secondary().position)
    }

    /// The view factor describing the fraction of energy emitted from one facet of the secondary
    /// toward another of the primary.
    fn view_factor(&self) -> &DMatrix<f64> {
        self._view_factor.as_ref().unwrap()
    }

    /// Compute the view factor describing the fraction of energy emitted from one facet of the
    /// secondary toward another of the primary.
    ///
    /// ## Definition
    ///
    /// The viewing factor is null if the line of sight is intercepted by topography (TODO) or if
    /// any of the angles between the facet outward and the line between facet centers is superior
    /// to 90 degrees.
    /// Rows are the faces of the primary and columns are the faces of the secondary.
    /// Meaning, if you read this matrix column by column, you get the view of the face n°(indice
    /// of the column) of the target body with regard to all the faces of the primary body.
    ///
    /// ## Expression
    ///
    /// V_{ij}=\frac{a_i\cos\Theta_i\cos\Theta_j}{\pi r^2}
    fn compute_view_factor(&mut self) {
        let mut view_factor = DMatrix::<f64>::zeros(
            self.primary().object.number_faces(),
            self.secondary().object.number_faces(),
        );
        let position_relative = self.relative_position();
        // Cartesian product of secondary's and primary's faces.
        for ((face_target, face_other), view) in izip!(
            iproduct!(
                self.secondary().object.face_iter_all(),
                self.primary().object.face_iter_all()
            ),
            view_factor.iter_mut()
        ) {
            let normal_target = face_target.fixed_rows::<U3>(3);
            let normal_other = face_other.fixed_rows::<U3>(3);
            let center_target = face_target.fixed_rows::<U3>(0);
            let center_other = face_other.fixed_rows::<U3>(0);
            let area_other = face_other[9];
            let relative_position_faces = position_relative + center_other - center_target;
            let direction_faces = relative_position_faces.normalize();
            let distance_faces = relative_position_faces.norm();
            let angle_target = normal_target.angle(&direction_faces);
            let angle_other = normal_other.angle(&-direction_faces);
            if angle_target > TAU / 4.0 || angle_other > TAU / 4.0 {
                continue;
            }
            *view = 2. * area_other * angle_target.cos() * angle_other.cos()
                / (TAU * distance_faces.powi(2));
        }
        self._view_factor = Some(view_factor);
    }

    /// See [`self view factor getter`][EnvironmentSystem::view_factor_self]
    fn view_factor_self(&self) -> &DMatrix<f64> {
        self._view_factor_self.as_ref().unwrap()
    }

    /// See [`self view factor setter`][EnvironmentSystem::set_view_factor_self]
    fn set_view_factor_self(&mut self, view_factor: DMatrix<f64>) {
        self._view_factor_self = Some(view_factor);
    }

    /// Compute the diffuse solar radial from primary.
    ///
    /// ## Expression
    ///
    /// W_{p,j}=\sum_{i\ne j}^{N}V_{ij}\frac{F_s A\cos\varsigma_j\left(t\right)}{r_H^2\left(t\right)}
    fn diffuse_solar_radiation_from_primary(&self, position_sun: &V3<f64>) -> VX<f64> {
        let mut diffuse_solar_radiation = VX::zeros(self.secondary().object.number_faces());
        let direct_solar_flux_from_other = self.primary().direct_solar_flux(position_sun)
            * self.primary().properties.albedo()
            / (1. - self.primary().properties.albedo());
        for (flux, view_factor) in izip!(
            diffuse_solar_radiation.iter_mut(),
            self.view_factor().column_iter()
        ) {
            *flux = (view_factor.component_mul(&direct_solar_flux_from_other)).sum();
        }
        diffuse_solar_radiation
    }

    /// Compute the direct thermal heating from the primary.
    ///
    /// ## Expression
    ///
    /// u_{p,j}=\sum_{i\ne j}^{N}V_{ij}\epsilon\sigma T_{s,i}^4
    fn direct_thermal_heating_from_primary(&self) -> VX<f64> {
        let mut direct_thermal_heating = VX::zeros(self.secondary().object.number_faces());
        let primary_temperature: f64 = self.secondary().surface_temperatures().max();
        for (flux, view_factor) in izip!(
            direct_thermal_heating.iter_mut(),
            self.view_factor().column_iter()
        ) {
            *flux = (view_factor
                * self.secondary().properties.emissivity()
                * STEPHAN_BOLTZMAN
                * primary_temperature.powi(4))
            .sum();
        }
        direct_thermal_heating
    }

    /// Compute the diffuse thermal self heating.
    ///
    /// ## Expression
    ///
    /// W_{s,j}=\sum_{i\ne j}^{N}V_{ij}\frac{F_s A\cos\varsigma_j\left(t\right)}{r_H^2\left(t\right)}
    fn diffuse_solar_self_heating(&self, position_sun: &V3<f64>) -> VX<f64> {
        let mut diffuse_solar_radiation = VX::zeros(self.secondary().object.number_faces());
        let direct_solar_flux_from_other = self.secondary().direct_solar_flux(position_sun)
            * self.secondary().properties.albedo()
            / (1. - self.secondary().properties.albedo());
        for (flux, view_factor) in izip!(
            diffuse_solar_radiation.iter_mut(),
            self.view_factor_self().column_iter()
        ) {
            *flux = (view_factor.component_mul(&direct_solar_flux_from_other)).sum();
        }
        diffuse_solar_radiation
    }

    /// Compute the direct thermal self heating.
    ///
    /// ## Expression
    ///
    /// u_{s,j}=\sum_{i\ne j}^{N}V_{ij}\epsilon\sigma T_{s,i}^4
    fn direct_thermal_self_heating(&self) -> VX<f64> {
        let mut direct_thermal_heating = VX::zeros(self.secondary().object.number_faces());
        let secondary_temperature: f64 = self.secondary().surface_temperatures().max();
        for (flux, view_factor) in izip!(
            direct_thermal_heating.iter_mut(),
            self.view_factor_self().column_iter()
        ) {
            *flux = (view_factor
                * self.secondary().properties.emissivity()
                * STEPHAN_BOLTZMAN
                * secondary_temperature.powi(4))
            .sum();
        }
        direct_thermal_heating
    }
}

impl EnvironmentSystem for BinarySystem {
    /// Targeted body name.
    fn name(&self) -> &String {
        &self.name
    }
    /// Targeted body getter.
    fn body(&self) -> &Body {
        self.secondary()
    }

    /// Targeted body mutable getter.
    fn body_mut(&mut self) -> &mut Body {
        self.secondary_mut()
    }

    /// Get the number of bodies of the environment.
    fn number_bodies(&self) -> usize {
        2
    }

    /// Recompute memoized fields.
    fn recompute(&mut self) {
        if self._view_factor.as_ref().is_none() && self._complete_model {
            self.compute_view_factor();
        }
        if self._view_factor_self.as_ref().is_none() && self._complete_model {
            self.compute_view_factor_self();
        }
    }

    /// Proceed to one iteration in the orbit and in revolution of the bodies in the system.
    fn orbit_and_revolution_iteration(&mut self) {
        self.primary_mut().revolution_iteration();
        self.secondary_mut().revolution_iteration();
    }

    /// Compute the surface heat flux of the targeted body.
    fn compute_surface_heat_flux(&mut self, no_sun: bool, position_sun: &V3<f64>) {
        let mut flux;
        if no_sun {
            flux = VX::zeros(self.body().object.number_faces());
        } else {
            flux = self.body().direct_solar_flux(position_sun);
            if self._complete_model {
                flux += self.diffuse_solar_radiation_from_primary(position_sun)
                    + self.direct_thermal_heating_from_primary()
                    + self.diffuse_solar_self_heating(position_sun)
                    + self.direct_thermal_self_heating()
            }
        }
        self.body_mut().set_surface_heat_flux(flux);
    }

    /// See [`self view factor getter`][EnvironmentSystem::view_factor_self]
    fn view_factor_self(&self) -> &DMatrix<f64> {
        self.view_factor_self()
    }

    /// See [`self view factor setter`][EnvironmentSystem::view_factor_self]
    fn set_view_factor_self(&mut self, view_factor: DMatrix<f64>) {
        self.set_view_factor_self(view_factor)
    }

    /// Change whether the thermophysical model includes mutual and self.
    fn complete_model(&mut self, complete: bool) {
        self._complete_model = complete;
    }
}

/// Trait for binary environment system.
pub trait BinaryEnvironmentSystem {
    /// Primary getter.
    fn primary(&self) -> &Body;

    /// Secondary getter.
    fn secondary(&self) -> &Body;

    /// The view factor of the secondary body with the primary describing the fraction of energy
    /// emitted from one facet toward another.
    fn view_factor(&self) -> &DMatrix<f64>;
}

impl BinaryEnvironmentSystem for BinarySystem {
    /// Primary getter.
    fn primary(&self) -> &Body {
        self.primary()
    }

    /// Secondary getter.
    fn secondary(&self) -> &Body {
        self.secondary()
    }

    /// The view factor of the secondary body with the primary describing the fraction of energy
    /// emitted from one facet toward another.
    fn view_factor(&self) -> &DMatrix<f64> {
        self.view_factor()
    }
}
