//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      GUI/Model/Job/JobItem.cpp
//! @brief     Implements class JobItem.
//!
//! @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 "GUI/Model/Job/JobItem.h"
#include "Base/Axis/Frame.h"
#include "Base/Axis/Scale.h"
#include "Device/Data/Datafield.h"
#include "Device/Detector/IDetector.h"
#include "GUI/Model/Axis/AmplitudeAxisItem.h"
#include "GUI/Model/Axis/PointwiseAxisItem.h"
#include "GUI/Model/Beam/SourceItems.h"
#include "GUI/Model/Data/Data1DItem.h"
#include "GUI/Model/Data/Data2DItem.h"
#include "GUI/Model/Detector/DetectorItem.h"
#include "GUI/Model/File/DatafileItem.h"
#include "GUI/Model/Job/BatchInfo.h"
#include "GUI/Model/Job/JobStatus.h"
#include "GUI/Model/Job/JobWorker.h"
#include "GUI/Model/Job/ParameterTreeBuilder.h"
#include "GUI/Model/Mask/MaskUtil.h"
#include "GUI/Model/Mask/MasksSet.h"
#include "GUI/Model/Par/ParameterTreeItems.h"
#include "GUI/Model/Sample/SampleItem.h"
#include "GUI/Model/Sim/SimulationOptionsItem.h"
#include "GUI/Model/ToCore/SimulationToCore.h"
#include "GUI/Model/Tune/FitSuiteItem.h"
#include "GUI/Model/Util/Backup.h"
#include "GUI/Model/Util/Path.h"
#include "GUI/Model/Util/UtilXML.h"
#include "Sim/Simulation/ISimulation.h"

namespace {
namespace Tag {

const QString Activity("Activity");
const QString BatchInfo("BatchInfo");
const QString DatafileItem("DatafileItem");
const QString FitSuite("FitSuite");
const QString Instrument("Instrument");
const QString ParameterContainer("ParameterContainer");
const QString Sample("Sample");
const QString SimulatedData("SimulatedData");
const QString SimulationOptions("SimulationOptions");

} // namespace Tag

DataItem* newDataItem(size_t rank)
{
    if (rank == 1)
        return new Data1DItem;
    if (rank == 2)
        return new Data2DItem;
    ASSERT_NEVER;
}

} // namespace

JobItem::JobItem()
    : m_simulation_options_item(std::make_unique<SimulationOptionsItem>())
    , m_parameter_container(std::make_unique<ParameterContainerItem>())
    , m_sample_item(std::make_unique<SampleItem>())
    , m_batch_info(std::make_unique<BatchInfo>())
    , m_fit_suite_item(std::make_unique<FitSuiteItem>())
{
}

JobItem::JobItem(const SampleItem* sampleItem, const InstrumentItem* instrumentItem,
                 const DatafileItem* dfile_item, const SimulationOptionsItem* optionItem)
    : m_simulation_options_item(std::make_unique<SimulationOptionsItem>(*optionItem))
    , m_parameter_container(std::make_unique<ParameterContainerItem>())
    , m_sample_item(sampleItem->clone())
    , m_batch_info(std::make_unique<BatchInfo>())
    , m_fit_suite_item(std::make_unique<FitSuiteItem>())
{
    ASSERT(sampleItem);
    ASSERT(instrumentItem);

    m_instrument.setCertainItem(instrumentItem->clone());

    createParameterTree();
    parameterContainerItem()->addBackupValues("Initial state");

    createSimulatedDataItem();

    if (dfile_item) {
        copyDatafileItemIntoJob(dfile_item);
        adjustRealDataToJobInstrument();
        createDiffDataItem();
    }
}

JobItem::~JobItem() = default;

void JobItem::setFailed()
{
    if (DataItem* di = simulatedDataItem()) {
        if (Datafield* df = di->p_field())
            df->setAllTo(0.0);
        emit di->datafieldChanged();
    }
    batchInfo()->setStatus(JobStatus::Failed);
}

size_t JobItem::rank() const
{
    return instrumentItem()->detectorRank();
}

void JobItem::createParameterTree()
{
    ParameterTreeBuilder(this).build();
}

void JobItem::createSimulatedDataItem()
{
    ASSERT(!simulatedDataItem());
    m_simulated_data_item.reset(::newDataItem(rank()));
}

Data1DItem* JobItem::data1DItem()
{
    return dynamic_cast<Data1DItem*>(m_simulated_data_item.get());
}

Data2DItem* JobItem::data2DItem()
{
    return dynamic_cast<Data2DItem*>(m_simulated_data_item.get());
}

DataItem* JobItem::createDiffDataItem()
{
    ASSERT(!diffDataItem());
    m_diff_data_item.reset(::newDataItem(rank()));

    ASSERT(m_dfile_item);
    // use the same axes as simulated item (same as external data item if loading pre-22.0 project)
    const DataItem* source = m_simulated_data_item ? simulatedDataItem() : m_dfile_item->dataItem();
    GUI::Util::copyContents(source->axItemX(), m_diff_data_item->axItemX());
    if (rank() == 2)
        GUI::Util::copyContents(source->axItemY(), m_diff_data_item->axItemY());

    if (auto* spec_diff = dynamic_cast<Data1DItem*>(diffDataItem()))
        spec_diff->setDiffPlotStyle();
    else {
        auto* intensity_diff = dynamic_cast<Data2DItem*>(diffDataItem());
        ASSERT(intensity_diff);
        intensity_diff->setInterpolated(m_dfile_item->data2DItem()->isInterpolated());
    }
    return m_diff_data_item.get();
}

void JobItem::copyDatafileItemIntoJob(const DatafileItem* source)
{
    ASSERT(!dfileItem());
    ASSERT(source->rank() == rank());
    m_dfile_item.reset(source->clone());

    if (rank() == 1)
        m_dfile_item->data1DItem()->setRealPlotStyle();
    else {
        ASSERT(data2DItem());
        data2DItem()->setInterpolated(m_dfile_item->data2DItem()->isInterpolated());
    }
}

void JobItem::importMasksFromRealData()
{
    auto* iiI = dynamic_cast<Scatter2DInstrumentItem*>(instrumentItem());
    if (!iiI)
        return;

    // copy all masks from datafile to detector
    iiI->detectorItem()->setMasks(dfileItem()->data2DItem()->masks());

    // convert the copied masks to the coordinate system of the detector
    const Frame& orig_frame = dfileItem()->data2DItem()->c_field()->frame();
    const Frame det_frame = iiI->detectorItem()->createFrame();
    MaskUtil::convertMasks(iiI->detectorItem()->masks(), orig_frame, det_frame);
}

void JobItem::adjustRealDataToJobInstrument()
{
    if (const auto* spec_instr = dynamic_cast<const SpecularInstrumentItem*>(instrumentItem())) {
        const Datafield* old = m_dfile_item->data1DItem()->c_field();
        auto flatVector = old->flatVector();
        auto errorSigmas = old->errorSigmas();
        std::unique_ptr<Frame> frame = spec_instr->createFrame();

        // if needed, reverse values vector to match "reversed" frame argument axis
        BasicAxisItem* axis_item = spec_instr->scanItem()->currentAxisItem();
        if (auto pointwise_axis = dynamic_cast<const PointwiseAxisItem*>(axis_item)) {
            std::string unit_1 = pointwise_axis->axis()->unit();
            std::string unit_2 = frame->axis(0).unit();
            if ((unit_1 == "1/nm" && unit_2 == "nm") || (unit_1 == "nm" && unit_2 == "1/nm")) {
                std::reverse(flatVector.begin(), flatVector.end());
                std::reverse(errorSigmas.begin(), errorSigmas.end());
            }
        }
        m_dfile_item->data1DItem()->setDatafield({*frame, flatVector, errorSigmas});

    } else if (instrumentItem()->is<Scatter2DInstrumentItem>()) {
        importMasksFromRealData();
        applyMasksToRealDatafield(); // also crop Datafield to the region of interest
    }
}

void JobItem::applyMasksToRealDatafield()
{
    auto* iiI = dynamic_cast<Scatter2DInstrumentItem*>(instrumentItem());
    ASSERT(iiI);

    Data2DItem* data_item = dfileItem()->data2DItem();
    std::unique_ptr<Datafield> orig_field(data_item->c_field()->clone());
    std::unique_ptr<const IDetector> det = iiI->detectorItem()->createDetector();

    // (re)create zero-valued Datafield with the size of region of interest.
    data_item->setDatafield(det->clippedFrame());

    ASSERT(det->frame().rank() == 2);
    std::vector<size_t> ai = det->activeIndices();
    for (unsigned long i : ai)
        (*data_item->p_field())[i] = (*orig_field)[det->roiToFullIndex(i)];

    data_item->updateDataRange();
}

void JobItem::writeTo(QXmlStreamWriter* w) const
{
    XML::writeTaggedElement(w, Tag::SimulationOptions, *m_simulation_options_item);
    XML::writeTaggedElement(w, Tag::Instrument, m_instrument);
    XML::writeTaggedElement(w, Tag::Sample, *m_sample_item);
    XML::writeTaggedElement(w, Tag::ParameterContainer, *m_parameter_container);
    XML::writeTaggedElement(w, Tag::BatchInfo, *m_batch_info);
    XML::writeTaggedValue(w, Tag::Activity, m_activity);

    // simulated data
    if (m_simulated_data_item) {
        XML::writeTaggedElement(w, Tag::SimulatedData, *m_simulated_data_item);
    }

    // real item
    if (m_dfile_item) {
        XML::writeTaggedElement(w, Tag::DatafileItem, *m_dfile_item);
    }

    // fit suite
    if (m_fit_suite_item) {
        XML::writeTaggedElement(w, Tag::FitSuite, *m_fit_suite_item);
    }
}

void JobItem::readFrom(QXmlStreamReader* r)
{
    while (r->readNextStartElement()) {
        QString tag = r->name().toString();
        if (tag == Tag::SimulationOptions)
            XML::readTaggedElement(r, tag, *m_simulation_options_item);
        else if (tag == Tag::Instrument)
            XML::readTaggedElement(r, tag, m_instrument);
        else if (tag == Tag::ParameterContainer) {
            createParameterTree();
            XML::readTaggedElement(r, tag, *m_parameter_container);
        } else if (tag == Tag::Sample)
            XML::readTaggedElement(r, tag, *m_sample_item);
        else if (tag == Tag::Activity)
            m_activity = XML::readTaggedString(r, tag);
        else if (tag == Tag::BatchInfo)
            XML::readTaggedElement(r, tag, *m_batch_info);
        else if (tag == Tag::SimulatedData) {
            createSimulatedDataItem();
            XML::readTaggedElement(r, tag, *m_simulated_data_item);
        } else if (tag == Tag::DatafileItem) {
            m_dfile_item = std::make_unique<DatafileItem>();
            m_dfile_item->readFrom(r);
            createDiffDataItem()->alignXYranges(m_dfile_item->dataItem());
            XML::gotoEndElementOfTag(r, tag);
        } else if (tag == Tag::FitSuite)
            XML::readTaggedElement(r, tag, *m_fit_suite_item);
        else
            r->skipCurrentElement();
    }
}

void JobItem::saveDatafields(const QString& projectDir) const
{
    if (m_dfile_item)
        m_dfile_item->saveDatafield(projectDir);

    if (m_simulated_data_item)
        m_simulated_data_item->saveDatafield(projectDir);
}

void JobItem::loadDatafields(const QString& projectDir)
{
    if (m_dfile_item)
        m_dfile_item->loadDatafield(projectDir);

    if (m_simulated_data_item)
        m_simulated_data_item->loadDatafield(projectDir, rank());
}

//! Updates the name of file to store intensity data.
void JobItem::updateFileName()
{
    if (DataItem* item = simulatedDataItem())
        item->setFileName(GUI::Path::intensityDataFileName(batchInfo()->jobName(), "jobdata"));

    if (m_dfile_item)
        if (DataItem* item = m_dfile_item->dataItem())
            item->setFileName(GUI::Path::intensityDataFileName(batchInfo()->jobName(), "refdata"));
}

void JobItem::initWorker()
{
    ISimulation* simulation =
        GUI::ToCore::itemsToSimulation(sampleItem(), instrumentItem(), simulationOptionsItem())
            .release();
    m_worker = std::make_unique<JobWorker>(simulation);

    m_thread = std::make_unique<QThread>();
    m_worker->moveToThread(m_thread.get());

    connect(worker(), &JobWorker::started, this, &JobItem::onStartedJob);
    connect(worker(), &JobWorker::progressUpdate, this, &JobItem::onItemProgress);
    connect(worker(), &JobWorker::finished, this, &JobItem::onFinishedWork);

    connect(thread(), &QThread::started, worker(), &JobWorker::start);
    connect(thread(), &QThread::finished, this, &JobItem::onFinishedThread);
}

void JobItem::haltWorker()
{
    if (!m_worker)
        return;
    m_worker->terminate();
}

void JobItem::onStartedJob()
{
    batchInfo()->setProgress(0);
    batchInfo()->setStatus(JobStatus::Running);
    batchInfo()->setBeginTime(m_worker->simulationStart());
    batchInfo()->setEndTime(QDateTime());
}

void JobItem::onItemProgress()
{
    if (!worker())
        return;
    batchInfo()->setProgress(worker()->percentageDone());
    emit progressIncremented();
}

void JobItem::onFinishedWork()
{
    ASSERT(worker());
    batchInfo()->setEndTime(worker()->simulationEnd());

    // propagate status of runner
    if (isFailed(worker()->workerStatus()))
        batchInfo()->setComments(worker()->workerFailureMessage());
    else {
        batchInfo()->setComments("");
        ASSERT(worker()->workerResult());
        simulatedDataItem()->setDatafield(*worker()->workerResult());
        updateFileName();
    }
    batchInfo()->setStatus(worker()->workerStatus());

    // fix job progress (if job was successful, but due to wrong estimation, progress not 100%)
    if (isCompleted(batchInfo()->status()))
        batchInfo()->setProgress(100);

    // tell the thread to exit here (instead of connecting JobRunner::finished
    // to QThread::quit because of strange behaviour)
    ASSERT(thread());
    thread()->quit();

    emit jobFinished(this);

    worker()->disconnect();
    m_worker.release();
}

void JobItem::onFinishedThread()
{
    m_thread->deleteLater();
    m_thread.release();
}
