use crate::base::body::Properties;
use crate::base::constants::{ASTRONAUMICAL_UNIT, STEPHAN_BOLTZMAN, TAU};
use crate::base::environment_system::{BinaryEnvironmentSystem, EnvironmentSystem};
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::{matrix_flat_to_vector, pow};
use crate::{V3, V3X, VX};
use console::{style, Emoji};
use indicatif::{ProgressBar, ProgressStyle};
use na::DMatrix;
use std::path::Path;

/// Emoji Indicatif step [1/3]
static MILKY_WAY: Emoji<'_, '_> = Emoji("🌌  ", "");
/// Emoji Indicatif step [2/3]
static PLANET: Emoji<'_, '_> = Emoji("🪐  ", "");
/// Emoji Indicatif step [3/3]
static SAVE: Emoji<'_, '_> = Emoji("💾  ", "");

/// The world simulation.
///
/// # Definition
///
/// This struct is the simulation manager. It takes the time manager and the environment system of
/// the simulation.
///
/// # Example
///
/// ```ignore
/// use kalast::{EnvironmentSystem, Time, World};
///
/// // You need to initialize Time
/// let time = Time::new(...);
///
/// // You need to initialize the environment system.
/// let system = EnvironmentSystem::new(...);
///
/// // Create the world.
/// let mut world = World::new(time, system);
/// ```
pub struct World<S> {
    /// Time manager.
    pub time: Time,
    /// Environment system, can be: [`SingleBody`][crate::base::environment_system::SingleBody],
    /// [`BinarySystem`][crate::base::environment_system::BinarySystem].
    pub environment_system: S,
    /// Position of the Sun.
    _position_sun: V3<f64>,
    /// Whether the sun should be activated or not.
    _no_sun: bool,
}

/// # World simulation
impl<S: EnvironmentSystem> World<S> {
    /// World constructor from time and body.
    pub fn new(time: Time, environment_system: S) -> Self {
        // Indicatif
        println!(
            "{} {}Initialize environment system {}...",
            style("[1/3]").bold().dim(),
            MILKY_WAY,
            style(environment_system.name()).bold(),
        );
        // Initiate the world.
        let mut world = Self {
            time,
            environment_system,
            _position_sun: V3::zeros(),
            _no_sun: false,
        };
        // Complete the initialization of the world, like giving the time step of the simulation to
        // the targeted body for it to compute advance ground properties.
        if !world._no_sun {
            world
                .environment_system
                .body_mut()
                .compute_global_equilibrium_temperature(&world._position_sun);
        }
        world
            .environment_system
            .body_mut()
            .share_time_step(world.time.time_step());
        world.environment_system.recompute();
        world.time.recompute_to_match_last_revolution(
            world
                .environment_system
                .body()
                .properties
                .revolution_period(),
        );
        world
    }

    /// Start the simulation to compute surface temperatures.
    pub fn start(&mut self) {
        // Indicatif.
        println!(
            "{} {}Compute surface flux of {}...",
            style("[2/3]").bold().dim(),
            PLANET,
            style(&self.environment_system.body().name).bold(),
        );
        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}]"),
        );
        // let properties = self.environment_system.body().properties.clone();
        // crate::debug!("rot angle: {}", properties.rotation_angle());
        // Time loop.
        let chrono = Chrono::tic();
        for (index, time) in self.time.data().clone().iter().enumerate().skip(1) {
            let time_step = self.time.time_step_backward(index);
            self.environment_system
                .body_mut()
                .properties
                .set_time_step(time_step);
            crate::debug!(
                "index: {:3}/{:3}, time: {:5}/{:5}, time step: {:3}, rotang: {}",
                index,
                self.time.size() - 1,
                time,
                self.time.last(),
                time_step,
                self.environment_system.body().properties.rotation_angle()
            );
            self.environment_system.orbit_and_revolution_iteration();
            self.compute_and_apply_heat_flux_at_boundaries_conditions();
            self.ground_conduction();
            self.save_daily_if_last_revolution(index, *time);
            // Prepare the next iteration.
            bar.inc(1);
        }
        bar.finish_and_clear();
        println!("{:10}Elapsed time: {}", "", chrono.toc_str());
    }

    /// Compute surface and ground limit heat fluxes and set conditions at boundaries.
    fn compute_and_apply_heat_flux_at_boundaries_conditions(&mut self) {
        let number_faces = self.environment_system.body().object.number_faces();
        let ground_size = self
            .environment_system
            .body()
            .properties
            .ground_vector()
            .len();
        // Proceed some necessary recomputations, e.g. view factors.
        self.environment_system.recompute();

        // Surface temperatures.
        let surface_temperatures;
        if self.environment_system.body().fixed_surface_temperatures() {
            surface_temperatures = self.environment_system.body().surface_temperatures_as_vec()
        } else {
            // Compute heat flux at the surface.
            self.environment_system
                .compute_surface_heat_flux(self._no_sun, &self._position_sun);
            // Apply surface heat flux.
            surface_temperatures = newton_method(
                self.environment_system.body().surface_temperatures_as_vec(),
                newton_method_function,
                newton_method_derivative,
                CustomNewtonMethodArguments::new(
                    self.environment_system.body().ground_temperatures(),
                    self.environment_system.body().surface_heat_flux(),
                    &self.environment_system.body().properties,
                ),
            );
        }
        // Apply.
        self.environment_system
            .body_mut()
            .set_surface_temperatures(surface_temperatures);

        // Ground limit temperatures.
        if ground_size > 2 {
            let ground_limit_temperatures;
            if self
                .environment_system
                .body()
                .fixed_ground_limit_temperatures()
            {
                ground_limit_temperatures = matrix_flat_to_vector(
                    &self
                        .environment_system
                        .body()
                        .ground_temperatures()
                        .column(ground_size - 1),
                );
            } else {
                // Compute heat flux at the ground limit.
                let ground_limit_heat_flux = VX::<f64>::zeros(number_faces);
                // Apply ground limit heat flux.
                let body = self.environment_system.body();
                let properties = &body.properties;
                ground_limit_temperatures = -properties.ground_step() / properties.conductivity()
                    * ground_limit_heat_flux
                    + body.ground_temperatures().column(ground_size - 2);
            }
            // Apply.
            self.environment_system
                .body_mut()
                .ground_temperatures_mut()
                .set_column(ground_size - 1, &ground_limit_temperatures);
        }
    }

    /// Apply 1D heat conduction
    ///
    /// # Expression
    ///
    /// u\left(x,t+\Delta t\right)=u\left(x,t\right)+\frac{\alpha\Delta t}{\Delta x^2}\left(u\left(x-\Delta x,t\right)-2u\left(x,t\right)+u\left(x+\Delta x,t\right)\right)
    fn ground_conduction(&mut self) {
        let ground_size = self
            .environment_system
            .body()
            .properties
            .ground_vector()
            .len();
        if ground_size == 1 {
            return;
        }

        // Define 1D neighbouring layers:
        //     + layer:       temperatures at ([x_1, x_{n-1}], t_{i  })
        //     + layer_up:    temperatures at ([x_0, x_{n-2}], t_{i  })
        //     + layer_down:  temperatures at ([x_2, x_{n  }], t_{i  })
        //     + new_layer:   temperatures at ([x_1, x_{n-1}], t_{i+1})
        let diffusivity_field = *self
            .environment_system
            .body()
            .properties
            .diffusivity_field();
        let ground_temperatures = self.environment_system.body().ground_temperatures().clone();
        let layer_iter = ground_temperatures
            .column_iter()
            .enumerate()
            .filter(|&(i, _)| ![0, ground_size - 1].contains(&i))
            .map(|(_, e)| e);
        let layer_up_iter = ground_temperatures
            .column_iter()
            .enumerate()
            .filter(|&(i, _e)| i < ground_size - 2)
            .map(|(_, e)| e);
        let layer_down_iter = ground_temperatures.column_iter().skip(2);
        let new_layer_iter = self
            .environment_system
            .body_mut()
            .ground_temperatures_mut()
            .column_iter_mut()
            .enumerate()
            .filter(|&(i, _)| ![0, ground_size - 1].contains(&i))
            .map(|(_, e)| e);

        // Apply 1D ground conduction.
        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(&(layer + diffusivity_field * (layer_up - 2. * layer + layer_down)));
        }
    }

    /// Compute the global temperature at equilibrium of the targeted body.
    pub fn compute_global_equilibrium_temperature(&self) -> f64 {
        self.environment_system
            .compute_global_equilibrium_temperature(&self._position_sun)
    }

    /// Save the temperatures at the surface if it is the last revolution.
    fn save_daily_if_last_revolution(&mut self, index: usize, time: f64) {
        if time
            >= self.time.duration()
                - self
                    .environment_system
                    .body()
                    .properties
                    .revolution_period()
        {
            if self.environment_system.body().has_last_revolution_started() {
                self.environment_system
                    .body_mut()
                    .set_daily_surface_temperatures();
            } else {
                self.environment_system
                    .body_mut()
                    .initiate_daily_record(self.time.size() - index);
            }
        }
    }

    /// Get the position of the Sun.
    pub fn position_sun(&self) -> V3<f64> {
        self._position_sun
    }

    /// Set the position of the Sun.
    pub fn set_position_sun(&mut self, position_sun: V3<f64>) {
        self._position_sun = position_sun
    }

    /// Set whether the Sun should be activated.
    pub fn no_sun(&mut self, _no_sun: bool) {
        self._no_sun = _no_sun;
    }

    /// Write a minimal header and save surface temperatures.
    pub fn save<P: AsRef<Path>>(&mut self, path: P) {
        // Indicatif
        println!(
            "{} {}Saving surface temperatures to {}...",
            style("[3/3]").bold().dim(),
            SAVE,
            style(path.as_ref().to_str().unwrap()).bold(),
        );
        let mut file = ExportFile::new(path);
        file.add_to_header(format!(
            "{} surface temperatures",
            self.environment_system.body().name
        ));
        file.add_to_header(format!(
            "Thermal intertia: {:0} J.K^{{-1}}.m^{{-2}}.s^{{-1/2}}",
            self.environment_system.body().properties.thermal_inertia()
        ));
        file.add_to_header(format!(
            "Heliocentric distance: {:3} AU",
            self.environment_system
                .body()
                .distance_to(&self._position_sun)
                / ASTRONAUMICAL_UNIT
        ));
        file.add_to_header("longitude [radian] | latitude [radian] | temperatures [K]".to_string());
        let mut data_to_write =
            V3X::<f64>::zeros(self.environment_system.body().object.number_faces());
        data_to_write.set_row(
            0,
            &self.environment_system.body().object.sphericals().row(0),
        );
        data_to_write.set_row(
            1,
            &self.environment_system.body().object.sphericals().row(1),
        );
        data_to_write
            .row_mut(2)
            .tr_copy_from(&self.environment_system.body().surface_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();
    }

    /// Write a minimal header and save self view factor relative to a specific face.
    pub fn save_self_view_factor(&mut self, path: &'static Path, face_reference_index: usize) {
        // Indicatif
        println!(
            "{} {}Saving self view factor to {}...",
            style("[3/3]").bold().dim(),
            SAVE,
            style(path.to_str().unwrap()).bold(),
        );
        let mut file = ExportFile::new(&path);
        file.add_to_header(format!(
            "{} self view factor relative to face n°{}",
            self.environment_system.body().name,
            face_reference_index
        ));
        file.add_to_header("longitude [radian] | latitude [radian] | view factor []".to_string());
        let mut data_to_write =
            V3X::<f64>::zeros(self.environment_system.body().object.number_faces());
        data_to_write.set_row(
            0,
            &self.environment_system.body().object.sphericals().row(0),
        );
        data_to_write.set_row(
            1,
            &self.environment_system.body().object.sphericals().row(1),
        );
        data_to_write.row_mut(2).tr_copy_from(
            &self
                .environment_system
                .view_factor_self()
                .column(face_reference_index),
        );
        for line in data_to_write.column_iter() {
            file.add_to_contents(format!("{:.3e} {:.3e} {:.3e}", line[0], line[1], line[2]));
        }
        file.write();
    }

    /// Write a minimal header and save view factor relative to a specific face.
    pub fn save_view_factor(&mut self, path: &'static Path, face_reference_index: usize)
    where
        S: BinaryEnvironmentSystem,
    {
        // Indicatif
        println!(
            "{} {}Saving self view factor to {}...",
            style("[3/3]").bold().dim(),
            SAVE,
            style(path.to_str().unwrap()).bold(),
        );
        let mut file = ExportFile::new(&path);
        file.add_to_header(format!(
            "{} view factor relative to face n°{} of {}",
            self.environment_system.primary().name,
            face_reference_index,
            self.environment_system.secondary().name,
        ));
        file.add_to_header("longitude [radian] | latitude [radian] | view factor []".to_string());
        let mut data_to_write =
            V3X::<f64>::zeros(self.environment_system.primary().object.number_faces());
        data_to_write.set_row(
            0,
            &self.environment_system.primary().object.sphericals().row(0),
        );
        data_to_write.set_row(
            1,
            &self.environment_system.primary().object.sphericals().row(1),
        );
        data_to_write.row_mut(2).tr_copy_from(
            &self
                .environment_system
                .view_factor()
                .column(face_reference_index),
        );
        for line in data_to_write.column_iter() {
            file.add_to_contents(format!("{:.3e} {:.3e} {:.3e}", line[0], line[1], line[2]));
        }
        file.write();
    }

    /// Save last revolution daily temperature of a single facet, i.e. the hour angle.
    pub fn save_daily<P: AsRef<Path>>(&mut self, path: P) {
        // Indicatif
        println!(
            "{} {}Saving daily surface temperatures to {}...",
            style("[3/3]").bold().dim(),
            SAVE,
            style(path.as_ref().to_str().unwrap()).bold(),
        );
        let mut file = ExportFile::new(path);
        file.add_to_header(format!(
            "{} daily surface temperatures",
            self.environment_system.body().name
        ));
        file.add_to_header(format!(
            "Thermal intertia: {:0} J.K^{{-1}}.m^{{-2}}.s^{{-1/2}}",
            self.environment_system.body().properties.thermal_inertia()
        ));
        file.add_to_header(format!(
            "Heliocentric distance: {:3} AU",
            self.environment_system
                .body()
                .distance_to(&self._position_sun)
                / ASTRONAUMICAL_UNIT
        ));
        file.add_to_header("hour angle [rad] | temperature [K]".to_string());
        let size = self
            .environment_system
            .body()
            .daily_surface_temperatures()
            .nrows();
        let mut hour_angles = VX::from_fn(size, |index, _| index as f64 / size as f64 * TAU);
        for (hour_angle, temperature) in izip!(
            hour_angles.column_mut(0).iter_mut(),
            self.environment_system
                .body_mut()
                .daily_surface_temperatures_mut()
                .iter_mut()
        ) {
            file.add_to_contents(format!("{:.3e} {:.3e}", hour_angle, temperature));
        }
        file.write();
    }
}

/// 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> {
    /// Constructor.
    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 to compute the temperatures at the
/// surface of the celestial body.
pub fn newton_method_function(
    value: &VX<f64>,
    newton_method_arguments: &CustomNewtonMethodArguments,
) -> VX<f64> {
    let mut result = newton_method_arguments.solar_fluxes
        - newton_method_arguments.properties.emissivity() * STEPHAN_BOLTZMAN * pow(value, 4);
    // If the body ground vector size is only the surface, consider the 2nd and 3rd layer the same
    // as the surface
    if newton_method_arguments.properties.ground_vector().len() > 2 {
        result += (-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());
    }
    result
}

/// 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> {
    let mut result =
        -4. * newton_method_arguments.properties.emissivity() * STEPHAN_BOLTZMAN * pow(value, 3);
    if newton_method_arguments.properties.ground_vector().len() > 2 {
        result = result.add_scalar(
            -3. * newton_method_arguments.properties.conductivity()
                / (2. * newton_method_arguments.properties.ground_step()),
        );
    }
    result
}
