//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Sim/Simulation/ISimulation.cpp
//! @brief     Implements interface ISimulation.
//!
//! @homepage  http://www.bornagainproject.org
//! @license   GNU General Public License v3 or higher (see COPYING)
//! @copyright Forschungszentrum Jülich GmbH 2018
//! @authors   Scientific Computing Group at MLZ (see CITATION, AUTHORS)
//
//  ************************************************************************************************

#include "Sim/Simulation/ISimulation.h"
#include "Base/Progress/ProgressHandler.h"
#include "Base/Util/Assert.h"
#include "Base/Util/StringUtil.h"
#include "Device/Data/Datafield.h"
#include "Param/Distrib/DistributionHandler.h"
#include "Resample/Option/SimulationOptions.h"
#include "Resample/Processed/ReSample.h"
#include "Sample/Multilayer/Sample.h"
#include "Sim/Background/IBackground.h"
#include <algorithm>
#include <gsl/gsl_errno.h>
#include <iomanip>
#include <iostream>
#include <mutex>
#include <thread>

namespace {

size_t indexStep(size_t total_size, size_t n_handlers)
{
    ASSERT(total_size > 0);
    ASSERT(n_handlers > 0);
    size_t result = total_size / n_handlers;
    return total_size % n_handlers ? ++result : result;
}

size_t startIndex(size_t n_handlers, size_t current_handler, size_t n_elements)
{
    const size_t handler_size = indexStep(n_elements, n_handlers);
    const size_t start_index = current_handler * handler_size;
    if (start_index >= n_elements)
        return n_elements;
    return start_index;
}

size_t batchSize(size_t n_handlers, size_t current_handler, size_t n_elements)
{
    const size_t handler_size = indexStep(n_elements, n_handlers);
    const size_t start_index = current_handler * handler_size;
    if (start_index >= n_elements)
        return 0;
    return std::min(handler_size, n_elements - start_index);
}

} // namespace


//  ************************************************************************************************
//  class implementation
//  ************************************************************************************************

ISimulation::ISimulation(const Sample& sample)
    : m_sample(sample.clone())
    , m_options(std::make_unique<SimulationOptions>())
    , m_distribution_handler(std::make_unique<DistributionHandler>())
    , m_progress(std::make_unique<ProgressHandler>())
{
    ASSERT(m_sample);
}

ISimulation::~ISimulation() = default;

//... Setters:

void ISimulation::setBackground(const IBackground& bg)
{
    m_background.reset(bg.clone());
}

void ISimulation::subscribe(const std::function<bool(size_t)>& inform)
{
    ASSERT(m_progress);
    m_progress->subscribe(inform);
}

//! Initializes a progress monitor that prints to stdout.
void ISimulation::setTerminalProgressMonitor()
{
#ifndef SILENT_PROGRESS
    subscribe([](size_t percentage_done) -> bool {
        if (percentage_done < 100)
            std::cout << std::setprecision(2) << "\r... " << percentage_done << "%" << std::flush;
        else // wipe out
            std::cout << "\r... 100%\n";
        return true;
    });
#endif
}

const SimulationOptions& ISimulation::options() const
{
    ASSERT(m_options);
    return *m_options;
}

//... Executor:

//! Runs simulation with possible averaging over parameter distributions; returns result.
Datafield ISimulation::simulate()
{
    ASSERT(m_sample);
    const std::string errs = m_sample->validateAmbientSubstrate();
    if (!errs.empty())
        throw std::runtime_error("Invalid sample model: " + errs + ".");

    gsl_set_error_handler_off();

    prepareSimulation();
    m_cache = std::vector<double>(nOutChannels(), 0.);

    const ReSample re_sample = ReSample::make(*m_sample, options(), force_polarized());

    const size_t total_size = nElements();
    if (total_size == 0)
        throw std::runtime_error("No output pixels. All masked? Invalid axis?");
    size_t n_combinations = distributionHandler().nParamSamples();

    m_progress->reset();
    m_progress->setExpectedNTicks(n_combinations * total_size);

    // restrict calculation to current batch // TODO: clarify
    const size_t n_batches = m_options->getNumberOfBatches();
    const size_t current_batch = m_options->getCurrentBatch();
    const size_t batch_start = startIndex(n_batches, current_batch, total_size);
    const size_t batch_size = batchSize(n_batches, current_batch, total_size);
    ASSERT(batch_size);

    if (n_combinations == 1) {
        runSingleSimulation(re_sample, batch_start, batch_size, 1.);
    } else {
        // only GISAS
        initDistributionHandler();
        for (size_t index = 0; index < n_combinations; ++index) {
            double weight = distributionHandler().setParameterValues(index);
            runSingleSimulation(re_sample, batch_start, batch_size, weight);
        }
    }

    return packResult();
}

//... Getters:

std::vector<const INode*> ISimulation::nodeChildren() const
{
    std::vector<const INode*> result;
    if (m_sample)
        result << m_sample.get();
    return result;
}

const Sample* ISimulation::sample() const
{
    return m_sample.get();
}

const IBackground* ISimulation::background() const
{
    return m_background.get();
}

const std::vector<ParameterDistribution>& ISimulation::paramDistributions() const
{
    return m_distribution_handler->paramDistributions();
}

SimulationOptions& ISimulation::options()
{
    ASSERT(m_options);
    return *m_options;
}

//... Protected accessors:

ProgressHandler& ISimulation::progress()
{
    ASSERT(m_progress);
    return *m_progress;
}

DistributionHandler& ISimulation::distributionHandler()
{
    ASSERT(m_distribution_handler);
    return *m_distribution_handler;
}

//... Private executor:

//! Runs a single simulation with fixed parameter values.
//! If desired, the simulation is run in several threads.
void ISimulation::runSingleSimulation(const ReSample& re_sample, size_t batch_start,
                                      size_t batch_size, double weight)
{
    initScanElementVector();

    const size_t n_threads = m_options->getNumberOfThreads();
    ASSERT(n_threads > 0);

    if (n_threads == 1) {
        // Run computation in current thread.
        try {
            for (size_t i = 0; i < batch_size; ++i) {
                if (!m_progress->alive())
                    break;
                runComputation(re_sample, batch_start + i, weight);
            }
        } catch (const std::exception& ex) {
            throw std::runtime_error(std::string("Unexpected error in simulation:\n") + ex.what());
        }

    } else {
        // Launch computation threads.
        std::vector<std::unique_ptr<std::thread>> threads;
        std::vector<std::string> failure_messages;
        std::mutex mutex;
        for (size_t i_thread = 0; i_thread < n_threads; ++i_thread) {
            const size_t thread_start = batch_start + startIndex(n_threads, i_thread, batch_size);
            const size_t thread_size = batchSize(n_threads, i_thread, batch_size);
            if (thread_size == 0)
                break;
            threads.emplace_back(new std::thread(
                [this, &re_sample, &weight, &failure_messages, &mutex, thread_start, thread_size] {
                    try {
                        for (size_t i = 0; i < thread_size; ++i) {
                            if (!m_progress->alive())
                                break;
                            runComputation(re_sample, thread_start + i, weight);
                        }
                    } catch (const std::exception& ex) {
                        mutex.lock();
                        if (std::find(failure_messages.begin(), failure_messages.end(), ex.what())
                            == failure_messages.end())
                            failure_messages.emplace_back(ex.what());
                        mutex.unlock();
                    }
                }));
        }

        // Wait for threads to complete.
        for (auto& thread : threads)
            thread->join();

        // Check successful completion.
        if (!failure_messages.empty())
            throw std::runtime_error("Unexpected error(s) in simulation thread(s):\n"
                                     + Base::String::join(failure_messages, "\n"));
    }
}