////////////////////////////////////////////////////////////////////////////////////////
//
//  Copyright 2023 OVITO GmbH, Germany
//
//  This file is part of OVITO (Open Visualization Tool).
//
//  OVITO is free software; you can redistribute it and/or modify it either under the
//  terms of the GNU General Public License version 3 as published by the Free Software
//  Foundation (the "GPL") or, at your option, under the terms of the MIT License.
//  If you do not alter this notice, a recipient may use your version of this
//  file under either the GPL or the MIT License.
//
//  You should have received a copy of the GPL along with this program in a
//  file LICENSE.GPL.txt.  You should have received a copy of the MIT License along
//  with this program in a file LICENSE.MIT.txt
//
//  This software is distributed on an "AS IS" basis, WITHOUT WARRANTY OF ANY KIND,
//  either express or implied. See the GPL or the MIT License for the specific language
//  governing rights and limitations.
//
////////////////////////////////////////////////////////////////////////////////////////

#include <ovito/core/Core.h>
#include <ovito/core/dataset/pipeline/PipelineCache.h>
#include <ovito/core/dataset/pipeline/PipelineNode.h>
#include <ovito/core/dataset/scene/Pipeline.h>
#include <ovito/core/dataset/DataSetContainer.h>
#include <ovito/core/dataset/data/TransformingDataVis.h>
#include <ovito/core/dataset/data/TransformedDataObject.h>
#include <ovito/core/app/Application.h>
#include <ovito/core/app/UserInterface.h>
#include <ovito/core/utilities/concurrent/Future.h>
#include <ovito/core/utilities/concurrent/TaskManager.h>

namespace Ovito {

/******************************************************************************
* Constructor.
******************************************************************************/
PipelineCache::PipelineCache(RefTarget* owner, bool includeVisElements) : _ownerObject(owner), _includeVisElements(includeVisElements)
{
}

/******************************************************************************
* Destructor.
******************************************************************************/
PipelineCache::~PipelineCache() // NOLINT
{
}

/******************************************************************************
* Starts a pipeline evaluation or returns a reference to an existing evaluation
* that is currently in progress.
******************************************************************************/
SharedFuture<PipelineFlowState> PipelineCache::evaluatePipeline(const PipelineEvaluationRequest& request)
{
    OVITO_ASSERT(!QCoreApplication::instance() || QThread::currentThread() == QCoreApplication::instance()->thread());
    OVITO_ASSERT(ExecutionContext::current().isValid());
    OVITO_ASSERT(Task::current());

    PipelineNode* pipelineNode = dynamic_object_cast<PipelineNode>(ownerObject());

    if(!isEnabled()) {
        if(pipelineNode)
            return pipelineNode->evaluateInternal(request);
        OVITO_ASSERT(false); // Cache may not be disabled for a whole pipeline.
    }

    if(_preparingEvaluation)
        return Future<PipelineFlowState>::createFailed(Pipeline::tr("A new pipeline evaluation is not permitted while another pipeline evaluation is already in progress. This error may be the result of an invalid user Python script invoking a function that is not permitted in this context."));

    // Update the times for which we should keep computed pipeline outputs.
    if(!_precomputeAllFrames)
        _requestedIntervals = request.cachingIntervals();
    else
        _requestedIntervals.add(TimeInterval::infinite());

    // Check if we can serve the request immediately by returning one of the cached states.
    for(const PipelineFlowState& state : _cachedStates) {
        if(state.stateValidity().contains(request.time())) {
            startFramePrecomputation(request);
            if(pipelineNode)
                return pipelineNode->postprocessCachedState(request, state);
            else
                return Future<PipelineFlowState>::createImmediateEmplace(state);
        }
    }

    // Check if there already is an evaluation in progress that is compatible with the new request.
    for(const EvaluationInProgress& evaluation : _evaluationsInProgress) {
        if(evaluation.validityInterval.contains(request.time())) {
            SharedFuture<PipelineFlowState> future = evaluation.future.lock();
            if(future.isValid() && !future.isCanceled()) {
                startFramePrecomputation(request);
                return future;
            }
        }
    }

    // To detect unexpected calls to invalidate() and reentrant function calls.
    _preparingEvaluation = true;
    try {
        SharedFuture<PipelineFlowState> future = evaluatePipelineImpl(request);

        // From now on, it is okay again to call invalidate().
        _preparingEvaluation = false;

        // Start the process of caching the pipeline results for all animation frames.
        startFramePrecomputation(request);

        return future;
    }
    catch(...) {
        _preparingEvaluation = false;
        throw;
    }
}

/******************************************************************************
* Starts a pipeline evaluation.
******************************************************************************/
SharedFuture<PipelineFlowState> PipelineCache::evaluatePipelineImpl(const PipelineEvaluationRequest& request)
{
    PipelineNode* pipelineNode = dynamic_object_cast<PipelineNode>(ownerObject());
    Pipeline* pipeline = !pipelineNode ? static_object_cast<Pipeline>(ownerObject()) : nullptr;
    OVITO_ASSERT(pipeline != nullptr || pipelineNode != nullptr);
    OVITO_ASSERT(pipeline != nullptr || _includeVisElements == false);

    SharedFuture<PipelineFlowState> future;
    TimeInterval preliminaryValidityInterval;

    if(!pipelineNode) {
        // Without a pipeline data source, the results will be an empty data collection.
        if(!pipeline->head())
            return Future<PipelineFlowState>::createImmediateEmplace(nullptr, PipelineStatus::Success);

        preliminaryValidityInterval = pipeline->head()->validityInterval(request);
        if(!_includeVisElements) {
            // When requesting the pipeline output without the effect of visualization elements,
            // delegate the evaluation to the head node of the pipeline.
            future = pipeline->head()->evaluate(request);
        }
        else {
            // When requesting the pipeline output with the effect of visualization elements,
            // delegate the evaluation to the pipeline's other cache.
            future = pipeline->evaluatePipeline(request);
        }
    }
    else {
        preliminaryValidityInterval = pipelineNode->validityInterval(request);
        try {
            future = pipelineNode->evaluateInternal(request);
        }
        catch(const Exception& ex) {
            if(request.throwOnError())
                throw;
            pipelineNode->setStatus(ex);
            future = Future<PipelineFlowState>::createImmediateEmplace(nullptr, pipelineNode->status());
        }
    }

    // Pre-register the evaluation operation.
    _evaluationsInProgress.push_front({ preliminaryValidityInterval });
    auto evaluation = _evaluationsInProgress.begin();
    OVITO_ASSERT(!evaluation->validityInterval.isEmpty());

    // When requesting the pipeline output with the effect of visualization elements,
    // let the visualization elements operate on the data collection.
    if(_includeVisElements) {
        future = future.then(*ownerObject(), [this, request, pipeline](const PipelineFlowState& state) {
            if(request.throwOnError() && state.status().type() == PipelineStatus::Error)
                throw Exception(state.status().text());
            // Give every visualization element the opportunity to apply an asynchronous data transformation.
            Future<PipelineFlowState> stateFuture;
            if(state) {
                for(const auto& dataObj : state.data()->objects()) {
                    for(DataVis* vis : dataObj->visElements()) {
                        // Let the PipelineSceneNode substitude the vis element with another one.
                        vis = pipeline->getReplacementVisElement(vis);
                        if(TransformingDataVis* transformingVis = dynamic_object_cast<TransformingDataVis>(vis)) {
                            if(transformingVis->isEnabled()) {
                                if(!stateFuture.isValid()) {
                                    stateFuture = transformingVis->transformData(request, dataObj, PipelineFlowState(state), _cachedTransformedDataObjects);
                                }
                                else {
                                    OORef<Pipeline> pipelineRef{pipeline}; // Used to keep the pipeline object alive.
                                    stateFuture = stateFuture.then(*transformingVis, [request, dataObj, transformingVis, this, pipeline = std::move(pipelineRef)](PipelineFlowState&& state) {
                                        if(request.throwOnError() && state.status().type() == PipelineStatus::Error)
                                            throw Exception(state.status().text());
                                        return transformingVis->transformData(request, dataObj, std::move(state), _cachedTransformedDataObjects);
                                    });
                                }
                            }
                        }
                    }
                }
            }
            if(!stateFuture.isValid()) {
                _cachedTransformedDataObjects.clear();
                stateFuture = Future<PipelineFlowState>::createImmediate(state);
            }
            else {
                // Cache the transformed data objects created by transforming visualization elements.
                stateFuture = stateFuture.then(*ownerObject(), [this, throwOnError = request.throwOnError()](PipelineFlowState&& state) {
                    if(throwOnError && state.status().type() == PipelineStatus::Error)
                        throw Exception(state.status().text());
                    cacheTransformedDataObjects(state);
                    return std::move(state);
                });
            }
            return stateFuture;
        });
    }

    // Store evaluation results in this cache.
    future = future.then(*ownerObject(), [this, pipeline, pipelineNode, evaluation, throwOnError = request.throwOnError()](PipelineFlowState state) {

        if(throwOnError && state.status().type() == PipelineStatus::Error)
            throw Exception(state.status().text());

        // Restrict the validity of the state.
        state.intersectStateValidity(evaluation->validityInterval);

        if(!state.stateValidity().isEmpty()) {
            // Let the cache decide whether the state should be stored or not.
            insertState(state);
            if(pipeline) {
                // Let the pipeline update its list of vis elements based on the new pipeline results.
                if(!_includeVisElements) {
                    // Only gather vis elements that are present in the pipeline at the animation time currently shown in the GUI.
                    std::optional<AnimationTime> time = currentAnimationTime();
                    if(time && state.stateValidity().contains(*time))
                        pipeline->updateVisElementList(state);
                }
            }
            else {
                // We also have a new preliminary state. Inform the upstream pipeline about it.
                if(pipelineNode->performPreliminaryUpdateAfterEvaluation()) {
                    std::optional<AnimationTime> time = currentAnimationTime();
                    if(time && state.stateValidity().contains(*time)) {
                        // Adopt the newly computed state as the current synchronous cache state.
                        _synchronousState = state;
                        _synchronousState.setStateValidity(TimeInterval::infinite());
                        pipelineNode->notifyDependents(ReferenceEvent::PreliminaryStateAvailable);
                    }
                }
            }
        }

        // Return state to caller.
        return std::move(state);
    });

    // Keep a weak reference to the future.
    evaluation->future = future;

    // Remove evaluation record from the list of ongoing evaluations once it is finished (successfully or not).
    future.finally(*ownerObject(), [this, evaluation](Task&) noexcept {
        cleanupEvaluation(evaluation);
    });

    OVITO_ASSERT(future.isValid());
    return future;
}

/******************************************************************************
* Removes an evaluation record from the list of evaluations currently in progress.
******************************************************************************/
void PipelineCache::cleanupEvaluation(std::forward_list<EvaluationInProgress>::iterator evaluation)
{
    OVITO_ASSERT(!_evaluationsInProgress.empty());
    OVITO_ASSERT(!QCoreApplication::instance() || QThread::currentThread() == QCoreApplication::instance()->thread());
    for(auto iter = _evaluationsInProgress.before_begin(), next = iter++; next != _evaluationsInProgress.end(); iter = next++) {
        if(next == evaluation) {
            _evaluationsInProgress.erase_after(iter);
            return;
        }
    }
    OVITO_ASSERT(false);
}

/******************************************************************************
* Inserts (or may reject) a pipeline state into the cache.
******************************************************************************/
void PipelineCache::insertState(const PipelineFlowState& state)
{
    OVITO_ASSERT(!QCoreApplication::instance() || QThread::currentThread() == QCoreApplication::instance()->thread());
    OVITO_ASSERT(!ownerObject()->isUndoRecording());
    OVITO_ASSERT(isEnabled());

    // Evict existing states from cache that do not overlap with the requested time intervals,
    // or which *do* overlap with the newly computed state and have now become outdated.
    _cachedStates.erase(std::remove_if(_cachedStates.begin(), _cachedStates.end(), [&](const PipelineFlowState& cachedState) {
        if(cachedState.stateValidity().overlap(state.stateValidity()))
            return true;
        return !std::any_of(_requestedIntervals.cbegin(), _requestedIntervals.cend(),
            std::bind(&TimeInterval::overlap, cachedState.stateValidity(), std::placeholders::_1));
    }), _cachedStates.end());

    // Decide whether to store the newly computed state in the cache or not.
    // To keep it, its validity interval must be overlapping with one of the requested time intervals.
    if(std::any_of(_requestedIntervals.cbegin(), _requestedIntervals.cend(),
            std::bind(&TimeInterval::overlap, state.stateValidity(), std::placeholders::_1))) {
        _cachedStates.push_back(state);
    }

    ownerObject()->notifyDependents(ReferenceEvent::PipelineCacheUpdated);
}

/******************************************************************************
* Performs a synchronous evaluation of the pipeline yielding a preliminary state.
******************************************************************************/
const PipelineFlowState& PipelineCache::evaluatePipelineSynchronous(const PipelineEvaluationRequest& request)
{
    OVITO_ASSERT(isEnabled());

    Pipeline* pipeline = static_object_cast<Pipeline>(ownerObject());

    // First, check if we can serve the request from the asynchronous evaluation cache.
    if(const PipelineFlowState& cachedState = getAt(request.time())) {
        if(cachedState.data() != _synchronousState.data()) {
            // Adopt the state from the asynchronous evaluation as new synchronous state
            // if it is valid at the current animation time being displayed in the GUI.
            std::optional<AnimationTime> time = currentAnimationTime();
            if(time && cachedState.stateValidity().contains(*time)) {
                _synchronousState = cachedState;
            }
        }
        return cachedState;
    }
    else {
        // Otherwise, try to serve the request from the synchronous evaluation cache.
        if(!_synchronousState.stateValidity().contains(request.time())) {

            // If no cached results are available, re-evaluate the pipeline.
            if(pipeline->head()) {
                // Adopt new state produced by the pipeline if it is not empty.
                // Otherwise stick with the old state from our own cache.
                UndoSuspender noUndo;
                if(PipelineFlowState newPreliminaryState = pipeline->head()->evaluateSynchronous(request)) {
                    _synchronousState = std::move(newPreliminaryState);

                    // Add the transformed data objects cached from the last pipeline evaluation.
                    if(_synchronousState) {
                        for(const auto& obj : _cachedTransformedDataObjects)
                            _synchronousState.addObject(obj);
                    }
                }
            }
            else {
                _synchronousState.reset();
            }

            // The preliminary state cache is time-independent.
            _synchronousState.setStateValidity(TimeInterval::infinite());
        }
        return _synchronousState;
    }
}

/******************************************************************************
* Performs a synchronous evaluation of a pipeline page yielding a preliminary state.
******************************************************************************/
PipelineFlowState PipelineCache::evaluatePipelineStageSynchronous(const PipelineEvaluationRequest& request)
{
    PipelineNode* pipelineNode = static_object_cast<PipelineNode>(ownerObject());

    if(!isEnabled()) {
        return pipelineNode->evaluateInternalSynchronous(request);
    }

    // First, check if we can serve the request from the asynchronous evaluation cache.
    if(const PipelineFlowState& cachedState = getAt(request.time())) {
        if(cachedState.data() != _synchronousState.data()) {
            // Adopt the state from the asynchronous evaluation as new synchronous state
            // if it is valid at the current animation time being displayed in the GUI.
            std::optional<AnimationTime> time = currentAnimationTime();
            if(time && cachedState.stateValidity().contains(*time)) {
                _synchronousState = cachedState;
            }
        }
        return cachedState;
    }
    else {
        // Otherwise, try to serve the request from the synchronous evaluation cache.
        if(!_synchronousState.stateValidity().contains(request.time())) {
            // If no cached results are available, re-evaluate the pipeline.
            // Adopt new state produced by the pipeline if it is not empty.
            // Otherwise stick with the old state from our own cache.
            UndoSuspender noUndo;
            if(PipelineFlowState newState = pipelineNode->evaluateInternalSynchronous(request)) {
                _synchronousState = std::move(newState);
            }

            // The preliminary state cache is time-independent.
            _synchronousState.setStateValidity(TimeInterval::infinite());
        }
    }
    return _synchronousState;
}

/******************************************************************************
* Marks the contents of the cache as outdated and throws away data that is no longer needed.
******************************************************************************/
void PipelineCache::invalidate(TimeInterval keepInterval, bool resetSynchronousCache)
{
    OVITO_ASSERT(!QCoreApplication::instance() || QThread::currentThread() == QCoreApplication::instance()->thread());
    if(_preparingEvaluation) {
        qWarning() << "Warning: Invalidating the pipeline cache while preparing the evaluation of the pipeline is not allowed. This error may be the result of an invalid user Python script invoking a function that is not permitted in this context.";
        return;
    }

    // Interrupt frame precomputation, which might be in progress.
    _precomputeFramesOperation.reset();
    _allFramesPrecomputed = false;

    // Reduce the validity of ongoing evaluations.
    for(EvaluationInProgress& evaluation : _evaluationsInProgress) {
        evaluation.validityInterval.intersect(keepInterval);
    }

    // Reduce the validity of the cached states. Throw away states that became completely invalid.
    for(PipelineFlowState& state : _cachedStates) {
        state.intersectStateValidity(keepInterval);
        if(state.stateValidity().isEmpty()) {
            state.reset();
        }
    }

    // Reduce the validity interval of the synchronous state cache.
    _synchronousState.intersectStateValidity(keepInterval);
    if(resetSynchronousCache && _synchronousState.stateValidity().isEmpty())
        _synchronousState.reset();

    if(resetSynchronousCache)
        _cachedTransformedDataObjects.clear();
}

/******************************************************************************
* Special method used by the FileSource class to replace the contents of the pipeline
* cache with a data collection modified by the user.
******************************************************************************/
void PipelineCache::overrideCache(const DataCollection* dataCollection, const TimeInterval& keepInterval)
{
    OVITO_ASSERT(dataCollection != nullptr);
    OVITO_ASSERT(!keepInterval.isEmpty());

    // Interrupt frame precomputation, which might be in progress.
    _precomputeFramesOperation.reset();
    _allFramesPrecomputed = false;

    // Reduce the validity of the cached states to the current animation time.
    // Throw away states that became completely invalid.
    // Replace the contents of the cache with the given data collection.
    for(PipelineFlowState& state : _cachedStates) {
        state.intersectStateValidity(keepInterval);
        if(state.stateValidity().isEmpty()) {
            state.reset();
        }
        else {
            state.setData(dataCollection);
        }
    }

    _synchronousState.setData(dataCollection);
}

/******************************************************************************
* Looks up the pipeline state for the given animation time.
******************************************************************************/
const PipelineFlowState& PipelineCache::getAt(AnimationTime time) const
{
    for(const PipelineFlowState& state : _cachedStates) {
        if(state.stateValidity().contains(time))
            return state;
    }
    static const PipelineFlowState emptyState;
    return emptyState;
}

/******************************************************************************
* Populates the internal cache with transformed data objects generated by
* transforming visual elements.
******************************************************************************/
void PipelineCache::cacheTransformedDataObjects(const PipelineFlowState& state)
{
    _cachedTransformedDataObjects.clear();
    if(state.data()) {
        for(const DataObject* o : state.data()->objects()) {
            if(const TransformedDataObject* transformedDataObject = dynamic_object_cast<TransformedDataObject>(o)) {
                _cachedTransformedDataObjects.push_back(transformedDataObject);
            }
        }
    }
}

/******************************************************************************
* Enables or disables the precomputation and caching of all frames of the animation.
******************************************************************************/
void PipelineCache::setPrecomputeAllFrames(bool enable)
{
    if(enable != _precomputeAllFrames) {
        _precomputeAllFrames = enable;
        if(!_precomputeAllFrames) {
            // Interrupt the precomputation process if it is currently in progress.
            _precomputeFramesOperation.reset();

            if(std::optional<AnimationTime> time = currentAnimationTime()) {
                // Throw away all precomputed data (except frame currently shown in the GUI) to reduce memory footprint.
                invalidate(TimeInterval(*time));
            }
            else {
                // Throw away all precomputed data to reduce memory footprint.
                invalidate();
            }
        }
    }
}

/******************************************************************************
* Starts the process of caching the pipeline results for all animation frames.
******************************************************************************/
void PipelineCache::startFramePrecomputation(const PipelineEvaluationRequest& request)
{
    OVITO_ASSERT(ExecutionContext::current().isValid());

    // Start the animation frame precomputation process if it has been activated.
    if(_precomputeAllFrames && !_precomputeFramesOperation.isValid() && !_allFramesPrecomputed) {
        // Create the async operation object that manages the frame precomputation.
        _precomputeFramesOperation = Promise<>::create<ProgressingTask>(true);

        // Show progress of the operation in the user interface by registering the asynchronous task.
        if(ExecutionContext::current().isValid())
            ExecutionContext::current().ui().taskManager().registerPromise(_precomputeFramesOperation);

        // Determine the number of frames that need to be precomputed.
        PipelineNode* pipelineNode = dynamic_object_cast<PipelineNode>(ownerObject());
        if(!pipelineNode)
            pipelineNode = static_object_cast<Pipeline>(ownerObject())->head();
        if(pipelineNode)
            _precomputeFramesOperation.setProgressMaximum(pipelineNode->numberOfSourceFrames());

        // Automatically reset the async operation object and the current frame precomputation when the
        // task gets canceled by the system.
        _precomputeFramesOperation.finally(*ownerObject(), [this](Task&) noexcept {
            _precomputeFrameFuture.reset();
            _precomputeFramesOperation.reset();
        });

        // Compute the first frame of the trajectory.
        precomputeNextAnimationFrame();
    }
}

/******************************************************************************
* Requests the next frame from the pipeline that needs to be precomputed.
******************************************************************************/
void PipelineCache::precomputeNextAnimationFrame()
{
    OVITO_ASSERT(_precomputeFramesOperation.isValid());
    OVITO_ASSERT(!_precomputeFramesOperation.isCanceled());

    // Determine the total number of animation frames.
    PipelineNode* pipelineNode = dynamic_object_cast<PipelineNode>(ownerObject());
    if(!pipelineNode)
        pipelineNode = static_object_cast<Pipeline>(ownerObject())->head();
    int numSourceFrames = pipelineNode ? pipelineNode->numberOfSourceFrames() : 0;

    // Determine what is the next animation frame that needs to be precomputed.
    int nextFrame = 0;
    AnimationTime nextFrameTime;
    while(nextFrame < numSourceFrames) {
        nextFrameTime = pipelineNode->sourceFrameToAnimationTime(nextFrame);
        const PipelineFlowState& state = getAt(nextFrameTime);
        if(!state) break;
        do {
            nextFrameTime = pipelineNode->sourceFrameToAnimationTime(++nextFrame);
        }
        while(state.stateValidity().contains(nextFrameTime) && nextFrame < numSourceFrames);
    }
    _precomputeFramesOperation.setProgressValue(nextFrame);
    _precomputeFramesOperation.setProgressText(Pipeline::tr("Caching trajectory (%1 frames remaining)").arg(numSourceFrames - nextFrame));
    if(nextFrame >= numSourceFrames) {
        // Precomputation of trajectory frames is complete.
        _precomputeFramesOperation.setFinished();
        OVITO_ASSERT(!_precomputeFrameFuture.isValid());
        _allFramesPrecomputed = true;
        return;
    }

    // Request the next frame from the input trajectory.
    _precomputeFrameFuture = evaluatePipeline(PipelineEvaluationRequest(nextFrameTime));

    // Wait until input frame is ready.
    _precomputeFrameFuture.finally(*ownerObject(), [this](Task& task) {
        try {
            // If the pipeline evaluation has been canceled for some reason, we interrupt the precomputation process.
            if(ownerObject()->isAboutToBeDeleted() || !_precomputeFramesOperation.isValid() || _precomputeFramesOperation.isFinished() || task.isCanceled()) {
                _precomputeFramesOperation.reset();
                OVITO_ASSERT(!_precomputeFrameFuture.isValid());
                return;
            }
            OVITO_ASSERT(_precomputeFrameFuture.isValid());

            // Schedule the pipeline evaluation at the next frame.
            precomputeNextAnimationFrame();
        }
        catch(const Exception&) {
            // In case of an error during pipeline evaluation or the unwrapping calculation,
            // abort the operation.
            _precomputeFramesOperation.setFinished();
        }
    });
}

/******************************************************************************
* Determines the current animation time shown in the GUI.
******************************************************************************/
std::optional<AnimationTime> PipelineCache::currentAnimationTime() const
{
    OVITO_ASSERT(ExecutionContext::current().isValid());

    if(AnimationSettings* anim = ExecutionContext::current().ui().datasetContainer().activeAnimationSettings())
        return anim->currentTime();
    return {};
}

}   // End of namespace
