use crate::base::body::{Body, Properties};
use crate::base::constants::{ASTRONAUMICAL_UNIT, STEPHAN_BOLTZMAN, TAU};
use crate::base::time::Time;
use crate::itertools::multizip;
use crate::toolbox::general::{Chrono, ExportFile};
use crate::toolbox::math::{newton_method, NewtonMethodArguments};
use crate::toolbox::matrix::pow;
use crate::{V3X, VX};
use indicatif::{ProgressBar, ProgressStyle};
use na::{DMatrix, Unit, U3};
use std::path::Path;

/// The world simulation
///
/// # Definition
///
/// This struct is the simulation manager. It takes the time manager, the targeted body of the
/// simulation and an optional secondary body for binary systems.
///
/// # Example
///
/// For single asteroid simulation:
///
/// ```ignore
/// use kalast::{Body, Time, World};
///
/// let time = Time::new(...); // you need to initialize Time
/// let mut body = Body::new(...); // you need to initialize Body
/// let mut world = World::new(time, body, None);
/// ```
pub struct World {
    /// Time manager.
    pub time: Time,
    /// Primary target of the simulation.
    pub body_target: Body,
    /// Secondary target of the simulation.
    pub body_secondary: Option<Body>,
}

/// # World simulation
impl World {
    /// World constructor from time and body.
    pub fn new(time: Time, body_target: Body, body_secondary: Option<Body>) -> Self {
        let mut world = Self {
            time,
            body_target,
            body_secondary,
        };
        world
            .body_target
            .properties
            .set_time_step(world.time.time_step());
        world
    }

    /// Start the simulation to compute surface temperatures.
    pub fn start(&mut self) {
        let mut ground_temperatures = DMatrix::<f64>::zeros(
            self.body_target.object.number_faces(),
            self.body_target.properties.ground_vector().nrows(),
        );
        let bar = ProgressBar::new(self.time.size() as u64);
        bar.set_style(
            ProgressStyle::default_bar()
                .template("{prefix} {bar:40.cyan/blue} [ELAPSED: {elapsed:>2}] [ETA: {eta:>2}]"),
        );
        bar.set_prefix(format!("KALAST>").as_str());
        // time loop
        let chrono = Chrono::tic();
        let previous_time = 0.;
        for time in self.time.data.iter() {
            let _time_step = time - previous_time;
            self.body_target.revolution_iteration();
            let solar_fluxes = self.body_target.solar_flux();
            // apply conduction to ground_temperatures
            self.temperatures_conduction(&mut ground_temperatures, &self.body_target.properties);
            // compute surface temperatures
            ground_temperatures.set_column(
                0,
                &newton_method(
                    self.body_target.temperatures().clone(),
                    newton_method_function,
                    newton_method_derivative,
                    CustomNewtonMethodArguments::new(
                        &ground_temperatures,
                        &solar_fluxes,
                        &self.body_target.properties,
                    ),
                ),
            );
            bar.inc(1);
        }
        self.body_target
            .temperatures()
            .copy_from(&ground_temperatures.column(0));
        bar.finish_and_clear();
        chrono.toc_str_print();
    }

    /// Write a minimal header and save surface temperatures
    pub fn save(&mut self, path: &'static Path) {
        let mut file = ExportFile::new(&path);
        file.add_to_header(format!("{} surface temperatures", self.body_target.name));
        file.add_to_header(format!(
            "Thermal intertia: {} J.K^{{-1}}.m^{{-2}}.s^{{-1/2}}",
            self.body_target.properties.thermal_inertia()
        ));
        file.add_to_header(format!(
            "Heliocentric distance: {} AU",
            self.body_target.distance_to_sun() / ASTRONAUMICAL_UNIT
        ));
        file.add_to_header(format!(
            "longitude [radian] | latitude [radian] | temperatures [K]",
        ));
        let mut data_to_write = V3X::<f64>::zeros(self.body_target.object.number_faces());
        data_to_write.set_row(0, &self.body_target.object.sphericals().row(0));
        data_to_write.set_row(1, &self.body_target.object.sphericals().row(1));
        data_to_write
            .row_mut(2)
            .tr_copy_from(&self.body_target.temperatures());
        for line in data_to_write.column_iter() {
            file.add_to_contents(format!("{:.3e} {:.3e} {:.3e}", line[0], line[1], line[2]));
        }
        file.write();
    }

    /// Apply 1D heat transfert conduction
    fn temperatures_conduction(
        &self,
        ground_temperatures: &mut DMatrix<f64>,
        properties: &Properties,
    ) {
        let ground_size = ground_temperatures.ncols();
        let ground_temperatures_old = ground_temperatures.clone();
        // define layers around
        let new_layer_iter = ground_temperatures
            .column_iter_mut()
            .enumerate()
            .filter(|&(i, _)| ![0, ground_size - 1].contains(&i))
            .map(|(_, e)| e);
        let layer_iter = ground_temperatures_old
            .column_iter()
            .enumerate()
            .filter(|&(i, _)| ![0, ground_size - 1].contains(&i))
            .map(|(_, e)| e);
        let layer_up_iter = ground_temperatures_old
            .column_iter()
            .enumerate()
            .filter(|&(i, _e)| i < ground_size - 2)
            .map(|(_, e)| e);
        let layer_down_iter = ground_temperatures_old.column_iter().skip(2);
        // inside ground conduction (without top layer)
        for (mut new_layer, layer, layer_up, layer_down) in
            multizip((new_layer_iter, layer_iter, layer_up_iter, layer_down_iter))
        {
            new_layer.copy_from(
                &(*properties.diffusivity_field() * (layer_up + layer_down)
                    + (1. - 2. * *properties.diffusivity_field()) * layer),
            );
        }
        // set bottom adiabatic
        ground_temperatures.set_column(
            ground_size - 1,
            &ground_temperatures_old.column(ground_size - 2),
        );
    }

    /// Compute the view factor describing the fraction of energy emitted from one facet toward
    /// another.
    ///
    /// ## Expression
    ///
    /// V_{ij}=\frac{a_i\cos\Theta_i\cos\Theta_j}{\pi r^2}
    pub fn view_factor(&self) -> DMatrix<f64> {
        let mut view_angles = DMatrix::<f64>::zeros(
            self.body_secondary.as_ref().unwrap().object.number_faces(),
            self.body_target.object.number_faces(),
        );
        let position_relative_bodies = self
            .body_target
            .position_relative_to_body(&self.body_secondary.as_ref().unwrap());
        for ((face_target, face_environment), view) in izip!(
            iproduct!(
                self.body_target.object.face_iter_all(),
                self.body_secondary.as_ref().unwrap().object.face_iter_all()
            ),
            view_angles.iter_mut()
        ) {
            let center_target = face_target.fixed_rows::<U3>(0);
            let normal_target = Unit::new_normalize(face_target.fixed_rows::<U3>(3).into_owned());
            let center_environment = face_environment.fixed_rows::<U3>(0);
            let normal_environment =
                Unit::new_normalize(face_environment.fixed_rows::<U3>(3).into_owned());
            let area_environment = face_target[9];
            let relative_position_faces =
                position_relative_bodies + center_environment - center_target;
            let direction_faces = Unit::new_normalize(relative_position_faces);
            let distance_faces = relative_position_faces.norm();
            let angle_target = normal_target.angle(&direction_faces);
            let angle_environment = normal_environment.angle(&-direction_faces);
            *view = 2. * area_environment * angle_target.cos() * angle_environment.cos()
                / (TAU * distance_faces.powi(2));
        }
        view_angles
    }
}

/// Struct that holds the parameters to be use for the Newton's method to compute the temperature
/// at the surface of the celestial body.
pub struct CustomNewtonMethodArguments<'a> {
    ground_temperatures: &'a DMatrix<f64>,
    solar_fluxes: &'a VX<f64>,
    properties: &'a Properties,
}

impl<'a> CustomNewtonMethodArguments<'a> {
    fn new(
        ground_temperatures: &'a DMatrix<f64>,
        solar_fluxes: &'a VX<f64>,
        properties: &'a Properties,
    ) -> Self {
        CustomNewtonMethodArguments {
            ground_temperatures,
            solar_fluxes,
            properties,
        }
    }
}

/// # Newton's method argument trait
///
/// Implement the trait [`NewtonMethodArguments`] to the struct that holds the paramaters in order to
/// be accepted by the Newton's method as a valid parameter.
impl<'a> NewtonMethodArguments for CustomNewtonMethodArguments<'a> {}

/// Implementation of the function of the Newton's method ton compute the temperatures at the
/// surface of the celestial body.
pub fn newton_method_function(
    value: &VX<f64>,
    newton_method_arguments: &CustomNewtonMethodArguments,
) -> VX<f64> {
    newton_method_arguments.solar_fluxes
        - newton_method_arguments.properties.emissivity() * STEPHAN_BOLTZMAN * pow(value, 4)
        + (-newton_method_arguments.ground_temperatures.column(2)
            + 4. * newton_method_arguments.ground_temperatures.column(1)
            - 3. * value)
            * newton_method_arguments.properties.conductivity()
            / (2. * newton_method_arguments.properties.ground_step())
}

/// Implementation of the derivative of the function of the Newton's method to compute the temperature
/// at the surface of the celestial body.
pub fn newton_method_derivative(
    value: &VX<f64>,
    newton_method_arguments: &CustomNewtonMethodArguments,
) -> VX<f64> {
    (-4. * newton_method_arguments.properties.emissivity() * STEPHAN_BOLTZMAN * pow(value, 3))
        .add_scalar(
            -3. * newton_method_arguments.properties.conductivity()
                / (2. * newton_method_arguments.properties.ground_step()),
        )
}
