// Copyright (c) 2020-2022  David Sorokin <david.sorokin@gmail.com>, based in Yoshkar-Ola, Russia
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at https://mozilla.org/MPL/2.0/.

use std::sync::Arc;
use std::default::Default;
use std::result;

#[cfg(feature="branch_mode")]
use std::io;

#[cfg(feature="branch_mode")]
use std::io::*;

use dvcompute_dist::simulation;
use dvcompute_dist::simulation::*;
use dvcompute_dist::simulation::generator::*;
use dvcompute_dist::simulation::composite::*;
use dvcompute_dist::simulation::simulation::*;
use dvcompute_dist::simulation::event::*;
use dvcompute_dist::simulation::observable::disposable::*;
use dvcompute_results_dist::simulation::results;
use dvcompute_results_dist::simulation::results::*;
use dvcompute_results_dist::simulation::results::locale::*;

#[cfg(feature="branch_mode")]
use rayon::prelude::*;

#[cfg(feature="dist_mode")]
use dvcompute_dist::simulation::comm::context::*;

#[cfg(feature="dist_mode")]
use dvcompute_dist::simulation::comm::pid::*;

#[cfg(feature="dist_mode")]
use dvcompute_dist::simulation::comm::lp::*;

#[cfg(feature="dist_mode")]
use dvcompute_dist::simulation::comm::time::*;

#[cfg(feature="dist_mode")]
use dvcompute_network::network::*;

/// Default renderers.
pub mod rendering;

/// Default view instances.
pub mod view;

/// It defines the simulation experiment.
#[derive(Clone)]
pub struct Experiment {

    /// The simulation specs for the experiment.
    pub specs: Specs,

    /// How the results must be transformed before rendering.
    pub transform: Arc<dyn Fn() -> ResultTransform + Sync + Send>,

    /// Specifies a locale.
    pub locale: ResultLocale,

    /// How many simulation runs should be launched.
    pub run_count: usize,

    /// The experiment title.
    pub title: String,

    /// The experiment description.
    pub description: Option<String>,

    /// Whether the process of generating the results is verbose.
    pub verbose: bool,

    /// The number of threads used when running the simulation experiment (can be ignored).
    pub num_threads: Option<usize>
}

impl Default for Experiment {

    fn default() -> Self {
        Self {
            specs: Specs {
                start_time: 0.0,
                stop_time: 10.0,
                dt: 0.01,
                generator_type: GeneratorType::Simple
            },
            transform: Arc::new(|| { ResultTransform::new(|x| { results::Result::Ok(x.clone()) }) }),
            locale: ResultLocale::En,
            run_count: 1,
            title: String::from("Simulation Experiment"),
            description: None,
            verbose: true,
            num_threads: None
        }
    }
}

impl Experiment {

    /// Run the simulation experiment with the specified executor.
    #[cfg(any(feature="branch_mode", feature="branch_wasm_mode"))]
    pub fn run<I, R, F, E, M>(&self,
        generators: I,
        rendering: R,
        simulation: F,
        executor: E) -> simulation::Result<()>
            where
                I: IntoIterator<Item = Box<dyn ExperimentGenerator<R>>>,
                R: ExperimentRendering + Send + 'static,
                F: FnOnce() -> M + Sync + Send + Clone + 'static,
                M: Simulation<Item = ResultSet> + Clone + 'static,
                E: ExperimentExecutor
    {
        let specs = self.specs.clone();
        let launcher = move |comp: SimulationBox<()>, run_index, run_count| {
            comp.run_by_index(specs, run_index, run_count)
        };

        self.run_with_executor_and_launcher(generators, rendering, simulation, executor, launcher)
    }

    /// Run the simulation experiment.
    #[cfg(feature="dist_mode")]
    pub fn run<'a, I, R, F, E>(&self,
        generators: I,
        rendering: R,
        simulation: F,
        executor: E) -> simulation::Result<()>
            where
                I: IntoIterator<Item = Box<dyn ExperimentGenerator<R>>>,
                R: ExperimentRendering + Send + 'static,
                F: FnOnce(&LogicalProcessContext, ExperimentCont) -> simulation::Result<()> + Sync + Send + Clone + 'static,
                E: ExperimentExecutor
    {
        let specs = self.specs.clone();
        let launcher = move |ctx: &LogicalProcessContext, comp: SimulationBox<()>, run_index, run_count| {
            comp.run_by_index(specs, ctx, run_index, run_count)
        };

        self.run_with_executor_and_launcher(generators, rendering, simulation, executor, launcher)
    }

    /// Run the simulation experiment with the specified executor and launcher.
    #[cfg(any(feature="branch_mode", feature="branch_wasm_mode"))]
    fn run_with_executor_and_launcher<I, R, F, E, L, M>(&self,
        generators: I,
        rendering: R,
        simulation: F,
        executor: E,
        launcher: L) -> simulation::Result<()>
            where
                I: IntoIterator<Item = Box<dyn ExperimentGenerator<R>>>,
                R: ExperimentRendering + Send + 'static,
                F: FnOnce() -> M + Sync + Send + Clone + 'static,
                M: Simulation<Item = ResultSet> + Clone + 'static,
                E: ExperimentExecutor,
                L: FnOnce(SimulationBox<()>, usize, usize) -> simulation::Result<()> + Sync + Send + Clone + 'static
    {
        let env = rendering.prepare(self)?;
        let reporters: Vec<_> = generators.into_iter()
            .map(|x| { x.report(&rendering, self, &env) })
            .collect();
        let reporters = Arc::new(reporters);
        for reporter in reporters.iter() {
            reporter.initialise()?;
        }
        let simulate: Vec<_> = (0 .. self.run_count)
            .map(|run_index| {
                let run_count  = self.run_count;
                let simulation = simulation.clone();
                let launcher   = launcher.clone();
                let reporters  = reporters.clone();
                let f: Box<dyn Fn() -> simulation::Result<()> + Sync + Send> = {
                    Box::new(move || {
                        let simulation = simulation.clone();
                        let launcher   = launcher.clone();
                        let reporters  = reporters.clone();
                        let comp = {
                            ResultPredefinedObservableSet::new()
                                .flat_map(move |predefined_observables| {
                                    simulation().flat_map(move |results| {
                                        let d = ExperimentData { results, predefined_observables };
                                        let comps: Vec<_> = reporters.iter()
                                            .map(|x| { x.simulate(&d) })
                                            .collect();
                                        composite_sequence_(comps)
                                            .run(empty_disposable())
                                            .run_in_start_time_by(false)
                                            .flat_map(move |((), fs)| {
                                                return_event(())
                                                    .run_in_stop_time()
                                                    .finally({
                                                        fs.into_event()
                                                            .run_in_stop_time()
                                                    })
                                            })
                                    })
                                })
                                .into_boxed()
                        };

                        launcher(comp, run_index, run_count)
                    })
                };

                f
            })
            .collect();
        let x = executor.execute(simulate);
        for reporter in reporters.iter() {
            reporter.finalise()?;
        }
        match x {
            result::Result::Ok(a) => {
                rendering.render(&self, &reporters, &env)?;
                rendering.on_completed(&self, &env)?;
                result::Result::Ok(a)
            },
            result::Result::Err(e) => {
                rendering.on_failed(&self, &env, &e)?;
                result::Result::Err(e)
            }
        }
    }

    /// Run the simulation experiment with the specified executor and launcher.
    #[cfg(feature="dist_mode")]
    fn run_with_executor_and_launcher<I, R, F, E, L>(&self,
        generators: I,
        rendering: R,
        simulation: F,
        executor: E,
        launcher: L) -> simulation::Result<()>
            where
                I: IntoIterator<Item = Box<dyn ExperimentGenerator<R>>>,
                R: ExperimentRendering + Send + 'static,
                F: FnOnce(&LogicalProcessContext, ExperimentCont) -> simulation::Result<()> + Sync + Send + Clone + 'static,
                E: ExperimentExecutor,
                L: FnOnce(&LogicalProcessContext, SimulationBox<()>, usize, usize) -> simulation::Result<()> + Sync + Send + Clone + 'static
    {
        let env = rendering.prepare(self)?;
        let reporters: Vec<_> = generators.into_iter()
            .map(|x| { x.report(&rendering, self, &env) })
            .collect();
        let reporters = Arc::new(reporters);
        for reporter in reporters.iter() {
            reporter.initialise()?;
        }
        let simulate: Vec<_> = (0 .. self.run_count)
            .map(|run_index| {
                let run_count  = self.run_count;
                let simulation = simulation.clone();
                let launcher   = launcher.clone();
                let reporters  = reporters.clone();
                let f: Box<dyn Fn(&LogicalProcessContext) -> simulation::Result<()> + Sync + Send> = {
                    Box::new(move |ctx| {
                        let simulation = simulation.clone();
                        let launcher   = launcher.clone();
                        let reporters  = reporters.clone();
                        simulation(ctx, Box::new(move |ctx, results| {
                            let comp = {
                                results
                                    .flat_map(move |results| {
                                        ResultPredefinedObservableSet::new()
                                            .flat_map(move |predefined_observables| {
                                                let d = ExperimentData { results, predefined_observables };
                                                let comps: Vec<_> = reporters.iter()
                                                    .map(|x| { x.simulate(&d) })
                                                    .collect();
                                                composite_sequence_(comps)
                                                    .run(empty_disposable())
                                                    .run_in_start_time_by(false)
                                                    .flat_map(move |((), fs)| {
                                                        return_event(())
                                                            .run_in_stop_time()
                                                            .finally({
                                                                fs.into_event()
                                                                    .run_in_stop_time()
                                                            })
                                                    })
                                            })
                                            .into_boxed()
                                    })
                                    .into_boxed()
                            };

                            launcher(ctx, comp, run_index, run_count)
                        }))
                    })
                };

                f
            })
            .collect();
        let x = executor.execute(simulate);
        for reporter in reporters.iter() {
            reporter.finalise()?;
        }
        match x {
            result::Result::Ok(a) => {
                rendering.render(&self, &reporters, &env)?;
                rendering.on_completed(&self, &env)?;
                result::Result::Ok(a)
            },
            result::Result::Err(e) => {
                rendering.on_failed(&self, &env, &e)?;
                result::Result::Err(e)
            }
        }
    }

    /// Run the time server for the specified number of run counts.
    #[cfg(feature="dist_mode")]
    pub fn run_time_server<F>(network: F, run_count: usize, ps: TimeServerParameters)
        where F: Fn(usize) -> NetworkSupport
    {
        for run_index in 0 .. run_count {
            let mut network = network(run_index);
            let init_quorum = (network.size() - 1) as usize;
            let ps = ps.clone();

            network.barrier();
            TimeServer::run(network, init_quorum, ps);
        }
    }
}

/// The continuation of simulation experiment.
#[cfg(feature="dist_mode")]
pub type ExperimentCont = Box<dyn FnOnce(&LogicalProcessContext, SimulationBox<ResultSet>) -> simulation::Result<()>>;

/// It allows rendering the simulation results in an arbitrary way.
pub trait ExperimentRendering {

    /// Defines a context used when rendering the experiment.
    type ExperimentContext: Send;

    /// Defines the experiment environment.
    type ExperimentEnvironment: Send;

    /// Prepare before rendering the experiment.
    fn prepare(&self, experiment: &Experiment) -> simulation::Result<Self::ExperimentEnvironment>;

    /// Render the experiment after the simulation is finished, example,
    /// create the `index.html` file in the specified directory.
    fn render(&self, experiment: &Experiment, reporters: &Vec<Box<dyn ExperimentReporter<Self> + Sync + Send>>, env: &Self::ExperimentEnvironment) -> simulation::Result<()>
        where Self: Sized;

    /// It is called when the experiment has been completed.
    fn on_completed(&self, experiment: &Experiment, env: &Self::ExperimentEnvironment) -> simulation::Result<()>;

    /// It is called when the experiment has been failed.
    fn on_failed(&self, experiment: &Experiment, env: &Self::ExperimentEnvironment, err: &simulation::error::Error) -> simulation::Result<()>;
}

/// Defines a view in which the simulation results should be saved.
/// You should extend this type class to define your own views such
/// as the PDF document.
pub trait ExperimentView<R: ExperimentRendering> {

    /// Get the view of the corresponding results.
    fn view(self) -> Box<dyn ExperimentGenerator<R>>;
}

/// This is a generator of the reporter with the specified rendering backend.
pub trait ExperimentGenerator<R: ExperimentRendering> {

    /// Create the result reporter.
    fn report(&self, rendering: &R, experiment: &Experiment, env: &R::ExperimentEnvironment) -> Box<dyn ExperimentReporter<R> + Sync + Send>;
}

/// It describes the source simulation data used in the experiment.
#[derive(Clone)]
pub struct ExperimentData {

    /// The simulation results used in the experiment.
    pub results: ResultSet,

    /// The predefined signals provided by every model.
    pub predefined_observables: ResultPredefinedObservableSet
}

/// Defines what creates the simulation reports by the specified renderer.
pub trait ExperimentReporter<R: ExperimentRendering> {

    /// Initialise the reporting before the simulation runs are started.
    fn initialise(&self) -> simulation::Result<()>;

    /// Finalise the reporting after all simulation runs are finished.
    fn finalise(&self) -> simulation::Result<()>;

    /// Start the simulation run in the start time.
    fn simulate(&self, xs: &ExperimentData) -> CompositeBox<()>;

    /// The context used by the renderer.
    fn context(&self) -> R::ExperimentContext;
}

/// It executes the simulation experiment.
pub trait ExperimentExecutor {

    /// Execute the simulation models.
    #[cfg(any(feature="branch_mode", feature="branch_wasm_mode"))]
    fn execute(self, models: Vec<Box<dyn Fn() -> simulation::Result<()> + Sync + Send>>) -> simulation::Result<()>;

    /// Execute the simulation models.
    #[cfg(feature="dist_mode")]
    fn execute(self, models: Vec<Box<dyn Fn(&LogicalProcessContext) -> simulation::Result<()> + Sync + Send>>) -> simulation::Result<()>;
}

/// The prebuit simulation experiment executors.
#[cfg(feature="branch_mode")]
pub enum BasicExperimentExecutor {

    /// Sequentially launch each experiment after experiment.
    Seq,

    /// Lauch experiments in parallel.
    Par
}

/// The prebuit simulation experiment executors.
#[cfg(feature="branch_wasm_mode")]
pub enum BasicExperimentExecutor {

    /// Sequentially launch each experiment after experiment.
    Seq
}

/// The prebuit simulation experiment executors.
#[cfg(feature="branch_wasm_mode")]
pub enum BasicExperimentExecutor {

    /// Sequentially launch each experiment after experiment.
    Seq
}

#[cfg(feature="branch_mode")]
impl Default for BasicExperimentExecutor {

    fn default() -> Self {
        BasicExperimentExecutor::Par
    }
}

#[cfg(feature="branch_wasm_mode")]
impl Default for BasicExperimentExecutor {

    fn default() -> Self {
        BasicExperimentExecutor::Seq
    }
}

#[cfg(feature="branch_wasm_mode")]
impl Default for BasicExperimentExecutor {

    fn default() -> Self {
        BasicExperimentExecutor::Seq
    }
}

#[cfg(feature="branch_mode")]
impl ExperimentExecutor for BasicExperimentExecutor {

    fn execute(self, models: Vec<Box<dyn Fn() -> simulation::Result<()> + Sync + Send>>) -> simulation::Result<()> {
        match self {
            BasicExperimentExecutor::Seq => {
                for x in models {
                    match x() {
                        result::Result::Ok(()) => continue,
                        result::Result::Err(e) => return result::Result::Err(e)
                    }
                }

                result::Result::Ok(())
            },
            BasicExperimentExecutor::Par => {
                models.par_iter()
                    .map(|x| {
                        match x() {
                            result::Result::Ok(()) => {},
                            result::Result::Err(e) => {
                                let _ = writeln!(io::stderr(), "Error: {:?}", e);
                            }
                        }
                        0
                    })
                    .sum::<usize>();

                result::Result::Ok(())
            }
        }
    }
}

#[cfg(feature="branch_wasm_mode")]
impl ExperimentExecutor for BasicExperimentExecutor {

    fn execute(self, models: Vec<Box<dyn Fn() -> simulation::Result<()> + Sync + Send>>) -> simulation::Result<()> {
        match self {
            BasicExperimentExecutor::Seq => {
                for x in models {
                    match x() {
                        result::Result::Ok(()) => continue,
                        result::Result::Err(e) => return result::Result::Err(e)
                    }
                }

                result::Result::Ok(())
            }
        }
    }
}

/// The logical process executor.
#[cfg(feature="dist_mode")]
pub struct LogicalProcessExecutor {

    /// Get the network support by the specified run index.
    pub network: Box<dyn Fn(usize) -> NetworkSupport>,

    /// The time server logical process identifier.
    pub time_server_id: LogicalProcessId,

    /// The logical process parameters.
    pub ps: LogicalProcessParameters
}

#[cfg(feature="dist_mode")]
impl ExperimentExecutor for LogicalProcessExecutor {

    fn execute(self, models: Vec<Box<dyn Fn(&LogicalProcessContext) -> simulation::Result<()> + Sync + Send>>) -> simulation::Result<()> {
        for (x, run_index) in models.into_iter().zip(0..) {
            let mut network = (self.network)(run_index);
            let x = move |ctx: &LogicalProcessContext| {
                match x(ctx) {
                    result::Result::Ok(()) => {},
                    result::Result::Err(e) => {
                        log::error!("{:?}", e);
                    }
                }
            };
            let time_server_id = self.time_server_id.clone();
            let ps = self.ps.clone();

            network.barrier();
            LogicalProcess::run(network, time_server_id, ps, x);
        }

        result::Result::Ok(())
    }
}