nmWTAI-Platform/Src/nmNum/nmSubWxs/nmPebiResultAnimationExporter.cpp

#include "nmPebiResultAnimationExporter.h"

#include "nmFFmpegVideoWriter.h"
#include "nmPebiResultDisplayStateBuilder.h"
#include "nmPebiResultPropertyCatalog.h"
#include "nmPebiResultSnapshot.h"
#include "nmPebiResultSnapshotRenderAdapter.h"
#include "nmPebiResultWellAnnotationLayer.h"
#include "nmVTKMeshLayer.h"
#include "nmVTKPlanarScaleGridLayer.h"
#include "nmVTKScalarBarLayer.h"
#include "nmVTKScene.h"

#include <QCoreApplication>
#include <QDir>
#include <QEventLoop>
#include <QFile>
#include <QFileInfo>
#include <QUuid>

#include <vtkActor.h>
#include <vtkAlgorithm.h>
#include <vtkAlgorithmOutput.h>
#include <vtkCamera.h>
#include <vtkCellData.h>
#include <vtkDataArray.h>
#include <vtkDataSet.h>
#include <vtkImageData.h>
#include <vtkLookupTable.h>
#include <vtkMath.h>
#include <vtkProperty.h>
#include <vtkRenderer.h>
#include <vtkRenderWindow.h>
#include <vtkType.h>
#include <vtkWindowToImageFilter.h>

#include <limits.h>
#include <string>

#if defined(Q_OS_WIN)
#include <windows.h>
#endif

namespace {

const char* const EXPORT_MESH_LAYER_ID = "pebi.export.mesh";
const char* const EXPORT_SCALAR_BAR_LAYER_ID = "pebi.export.scalar_bar";
const char* const EXPORT_SCALE_GRID_LAYER_ID = "pebi.export.scale_grid";
const char* const EXPORT_WELLS_LAYER_ID = "pebi.export.wells";

QString exporterText(const char* pText)
{
    return QCoreApplication::translate(
                "nmPebiResultAnimationExporter", pText);
}

bool getAdapterBounds(nmPebiResultSnapshotRenderAdapter& oAdapter,
                      vtkIdType nExpectedCellCount,
                      double dBounds[6])
{
    vtkAlgorithmOutput* pOutput = oAdapter.getOutputPort();
    if(pOutput == nullptr || pOutput->GetProducer() == nullptr ||
       oAdapter.getCellCount() != nExpectedCellCount ||
       nExpectedCellCount <= 0) {
        return false;
    }
    pOutput->GetProducer()->Update(pOutput->GetIndex());
    vtkDataSet* pDataSet = vtkDataSet::SafeDownCast(
                pOutput->GetProducer()->GetOutputDataObject(
                    pOutput->GetIndex()));
    if(pDataSet == nullptr ||
       pDataSet->GetNumberOfCells() != nExpectedCellCount) {
        return false;
    }
    pDataSet->GetBounds(dBounds);
    if(!vtkMath::AreBoundsInitialized(dBounds)) {
        return false;
    }
    for(int nIndex = 0; nIndex < 6; ++nIndex) {
        if(!vtkMath::IsFinite(dBounds[nIndex])) {
            return false;
        }
    }
    return dBounds[0] < dBounds[1] &&
            dBounds[2] < dBounds[3] &&
            dBounds[4] <= dBounds[5];
}

QString createTemporaryPath(const QString& sTargetPath)
{
    // 临时文件保留目标扩展名，Muxer才能按最终容器格式生成完整文件。
    const QFileInfo oTargetInfo(sTargetPath);
    const QString sToken = QUuid::createUuid().toString()
            .remove('{').remove('}');
    return QDir(oTargetInfo.absolutePath()).filePath(
                QString::fromLatin1(".%1.nm-export-%2.%3")
                .arg(oTargetInfo.completeBaseName())
                .arg(sToken)
                .arg(oTargetInfo.suffix()));
}

bool isNonEmptyFile(const QString& sFilePath)
{
    const QFileInfo oFileInfo(sFilePath);
    return oFileInfo.isFile() && oFileInfo.size() > 0;
}

bool publishTemporaryFile(const QString& sTemporaryPath,
                          const QString& sTargetPath,
                          QString* pError)
{
    if(!isNonEmptyFile(sTemporaryPath)) {
        if(pError != nullptr) {
            *pError = exporterText(
                        "Unable to publish the completed video file.");
        }
        return false;
    }

    // 编码与发布是两个事务：只有非空临时视频完成后才允许触碰原目标。
    const QFileInfo oTargetInfo(sTargetPath);
    const QString sBackupPath = QDir(oTargetInfo.absolutePath()).filePath(
                QString::fromLatin1(".%1.nm-backup-%2.%3")
                .arg(oTargetInfo.completeBaseName())
                .arg(QUuid::createUuid().toString().remove('{').remove('}'))
                .arg(oTargetInfo.suffix()));
#if defined(Q_OS_WIN)
    const std::wstring wsTarget = QDir::toNativeSeparators(
                QFileInfo(sTargetPath).absoluteFilePath()).toStdWString();
    const std::wstring wsTemporary = QDir::toNativeSeparators(
                QFileInfo(sTemporaryPath).absoluteFilePath()).toStdWString();
    const std::wstring wsBackup = QDir::toNativeSeparators(
                QFileInfo(sBackupPath).absoluteFilePath()).toStdWString();
    if(QFileInfo(sTargetPath).exists()) {
        QFile::remove(sBackupPath);
        // ReplaceFileW原子替换目标，并让系统把旧文件保存到唯一备份路径。
        if(ReplaceFileW(
                    wsTarget.c_str(), wsTemporary.c_str(),
                    wsBackup.c_str(), REPLACEFILE_WRITE_THROUGH,
                    nullptr, nullptr) != FALSE) {
            QFile::remove(sBackupPath);
            return true;
        }

        // ReplaceFileW失败时Windows仍可能已经生成备份，必须先恢复原文件。
        if(QFileInfo(sBackupPath).exists()) {
            if(MoveFileExW(
                   wsBackup.c_str(), wsTarget.c_str(),
                   MOVEFILE_REPLACE_EXISTING |
                   MOVEFILE_WRITE_THROUGH) == FALSE &&
               QFileInfo(sBackupPath).exists()) {
                if(pError != nullptr) {
                    *pError = exporterText(
                                "Unable to publish the completed video file.") +
                            QString::fromLatin1(" ") +
                            exporterText(
                                "The original file backup was retained at \"%1\".")
                            .arg(QDir::toNativeSeparators(sBackupPath));
                }
                return false;
            }
        }
    } else {
        if(MoveFileExW(
                    wsTemporary.c_str(), wsTarget.c_str(),
                    MOVEFILE_WRITE_THROUGH) != FALSE) {
            return true;
        }
    }
    if(pError != nullptr) {
        *pError = exporterText(
                    "Unable to publish the completed video file.");
    }
    return false;
#else
    if(QFileInfo(sTargetPath).exists()) {
        QFile::remove(sBackupPath);
        if(!QFile::rename(sTargetPath, sBackupPath)) {
            if(pError != nullptr) {
                *pError = exporterText(
                            "Unable to publish the completed video file.");
            }
            return false;
        }
        if(QFile::rename(sTemporaryPath, sTargetPath)) {
            QFile::remove(sBackupPath);
            return true;
        }
        if(!QFile::rename(sBackupPath, sTargetPath) &&
           QFileInfo(sBackupPath).exists() && pError != nullptr) {
            *pError = exporterText(
                        "Unable to publish the completed video file.") +
                    QString::fromLatin1(" ") +
                    exporterText(
                        "The original file backup was retained at \"%1\".")
                    .arg(QDir::toNativeSeparators(sBackupPath));
            return false;
        }
        if(pError != nullptr) {
            *pError = exporterText(
                        "Unable to publish the completed video file.");
        }
        return false;
    }
    if(QFile::rename(sTemporaryPath, sTargetPath)) {
        return true;
    }
    if(pError != nullptr) {
        *pError = exporterText(
                    "Unable to publish the completed video file.");
    }
    return false;
#endif
}

bool addLayerOrDelete(nmVTKScene& oScene, nmVTKRenderLayer* pLayer)
{
    // addLayer成功后Scene取得所有权；失败时仍由调用侧负责销毁。
    if(pLayer != nullptr && oScene.addLayer(pLayer)) {
        return true;
    }
    delete pLayer;
    return false;
}

class OffscreenRenderWindowGuard
{
public:
    OffscreenRenderWindowGuard(vtkRenderWindow* pRenderWindow,
                               vtkRenderer* pRenderer)
        : m_pRenderWindow(pRenderWindow),
          m_pRenderer(pRenderer)
    {
    }

    ~OffscreenRenderWindowGuard()
    {
        if(m_pRenderWindow != nullptr) {
            if(m_pRenderer != nullptr) {
                m_pRenderWindow->RemoveRenderer(m_pRenderer);
            }
            m_pRenderWindow->Finalize();
        }
    }

private:
    vtkRenderWindow* m_pRenderWindow;
    vtkRenderer* m_pRenderer;
};

}

nmPebiResultAnimationExporter::Request::Request()
    : m_bScalarBarVisible(true),
      m_bScaleGridVisible(false),
      m_bWellNamesVisible(true),
      m_bGradientBackground(false),
      m_pCamera(nullptr)
{
    m_dBackground[0] = m_dBackground[1] = m_dBackground[2] = 0.0;
    m_dBackground2[0] = m_dBackground2[1] = m_dBackground2[2] = 0.0;
}

nmPebiResultAnimationExporter::nmPebiResultAnimationExporter(
        QObject* pParent)
    : QObject(pParent),
      m_bCancelRequested(false)
{
}

nmPebiResultAnimationExporter::~nmPebiResultAnimationExporter()
{
}

nmPebiResultAnimationExporter::ExportResult
nmPebiResultAnimationExporter::exportAnimation(
        const Request& oRequest,
        QString* pError)
{
    if(pError != nullptr) {
        pError->clear();
    }
    m_bCancelRequested = false;
    const int nAvailableFrameCount = oRequest.m_pSnapshot.isNull()
            ? 0 : oRequest.m_pSnapshot->getPropertyFrameCount(
                oRequest.m_sPropertyId);
    nmPebiResultPropertyDescriptor oDescriptor;
    if(oRequest.m_pSnapshot.isNull() ||
       !oRequest.m_pSnapshot->isComplete() ||
       !nmPebiResultPropertyCatalog::findDescriptor(
           oRequest.m_sPropertyId, oDescriptor) ||
       oDescriptor.m_eTimeType != NM_PebiResultProperty_Dynamic ||
       oDescriptor.m_eAssociation !=
           NM_PebiResultPropertyAssociation_CellData ||
       !oRequest.m_oExportSettings.isValid(nAvailableFrameCount) ||
       oRequest.m_sTargetFilePath.isEmpty() ||
       oRequest.m_pCamera == nullptr) {
        if(pError != nullptr) {
            *pError = exporterText("The animation export request is invalid.");
        }
        return Export_Failed;
    }

    const QString sExtension = QFileInfo(
                oRequest.m_sTargetFilePath).suffix().toLower();
    if(!nmFFmpegVideoWriter::isSupportedExtension(sExtension)) {
        if(pError != nullptr) {
            *pError = exporterText("The selected video format is unsupported.");
        }
        return Export_Failed;
    }

    // 屏幕场景与离屏场景共用同一构造器，避免颜色规则出现两套实现。
    nmPebiResultDisplayState oDisplayState;
    if(!nmPebiResultDisplayStateBuilder::build(
           oRequest.m_pSnapshot, oRequest.m_sPropertyId,
           oRequest.m_oDisplaySettings, oDisplayState, pError)) {
        return Export_Failed;
    }

    const QVector<int> vecFrameIndices =
            oRequest.m_oExportSettings.buildFrameIndices();
    // 复制属性数组后逐帧更新，导出过程不会切换屏幕Scene中的当前帧。
    nmPebiResultSnapshotRenderAdapter oAdapter;
    if(vecFrameIndices.isEmpty() ||
       !oAdapter.initialize(
           oRequest.m_pSnapshot,
           nmPebiResultSnapshotRenderAdapter::CopiedPropertyArray) ||
       !oAdapter.setPropertyFrame(
           oRequest.m_sPropertyId, vecFrameIndices.first())) {
        if(pError != nullptr) {
            *pError = exporterText("Unable to prepare the animation result data.");
        }
        return Export_Failed;
    }

    double dBounds[6] = { 0.0, 0.0, 0.0, 0.0, 0.0, 0.0 };
    if(!getAdapterBounds(
           oAdapter, oRequest.m_pSnapshot->getResultCellCount(), dBounds)) {
        if(pError != nullptr) {
            *pError = exporterText("The animation render data failed validation.");
        }
        return Export_Failed;
    }

    vtkSmartPointer<vtkRenderer> pRenderer =
            vtkSmartPointer<vtkRenderer>::New();
    vtkSmartPointer<vtkRenderWindow> pRenderWindow =
            vtkSmartPointer<vtkRenderWindow>::New();
    pRenderWindow->OffScreenRenderingOn();
    pRenderWindow->SetSize(
                oRequest.m_oExportSettings.m_nOutputWidth,
                oRequest.m_oExportSettings.m_nOutputHeight);
    pRenderWindow->AddRenderer(pRenderer);
    OffscreenRenderWindowGuard oRenderWindowGuard(
                pRenderWindow, pRenderer);
    pRenderer->SetBackground(
                oRequest.m_dBackground[0],
                oRequest.m_dBackground[1],
                oRequest.m_dBackground[2]);
    pRenderer->SetBackground2(
                oRequest.m_dBackground2[0],
                oRequest.m_dBackground2[1],
                oRequest.m_dBackground2[2]);
    pRenderer->SetGradientBackground(
                oRequest.m_bGradientBackground ? 1 : 0);
    // 保留用户视角；后续只更新ClippingRange，不执行ResetCamera适应窗口。
    pRenderer->GetActiveCamera()->DeepCopy(oRequest.m_pCamera);

    nmVTKScene oScene(pRenderer);
    nmVTKMeshLayer* pMeshLayer = new nmVTKMeshLayer(
                QString::fromLatin1(EXPORT_MESH_LAYER_ID));
    pMeshLayer->setDisplayMode(
                nmVTKMeshLayer::DisplayMode_SurfaceWithEdges);
    pMeshLayer->getActor()->GetProperty()->SetEdgeColor(0.0, 0.0, 0.0);
    pMeshLayer->getActor()->GetProperty()->LightingOff();
    if(!pMeshLayer->setInputConnection(
           oAdapter.getOutputPort(),
           oAdapter.getCurrentScalarArrayName(),
           nmVTKMeshLayer::ScalarAssociation_CellData,
           oDisplayState.m_pLookupTable)) {
        delete pMeshLayer;
        if(pError != nullptr) {
            *pError = exporterText("Unable to create the animation mesh layer.");
        }
        return Export_Failed;
    }
    pMeshLayer->setScalarVisibility(true);
    pMeshLayer->setVisible(true);
    if(!addLayerOrDelete(oScene, pMeshLayer)) {
        if(pError != nullptr) {
            *pError = exporterText("Unable to create the animation mesh layer.");
        }
        return Export_Failed;
    }

    nmVTKScalarBarLayer* pScalarBarLayer = new nmVTKScalarBarLayer(
                QString::fromLatin1(EXPORT_SCALAR_BAR_LAYER_ID));
    pScalarBarLayer->setLookupTable(oDisplayState.m_pLookupTable);
    pScalarBarLayer->setTitle(oDisplayState.m_sScalarBarTitle);
    pScalarBarLayer->setLabelFormat(QString::fromLatin1("%.4g"));
    pScalarBarLayer->setViewportSize(
                oRequest.m_oExportSettings.m_nOutputWidth,
                oRequest.m_oExportSettings.m_nOutputHeight);
    pScalarBarLayer->setNumberOfLabels(
                oDisplayState.m_bSingleValueMode
                ? 0 : pScalarBarLayer->getResponsiveLabelCount());
    pScalarBarLayer->setConstantValueLabel(
                oDisplayState.m_bSingleValueMode,
                oDisplayState.m_bSingleValueMode
                ? QString::number(
                    oDisplayState.m_dBusinessRange[0], 'g', 6)
                : QString());
    pScalarBarLayer->setVisible(oRequest.m_bScalarBarVisible);
    if(!addLayerOrDelete(oScene, pScalarBarLayer)) {
        if(pError != nullptr) {
            *pError = exporterText("Unable to create the animation scalar bar.");
        }
        return Export_Failed;
    }

    nmVTKPlanarScaleGridLayer* pScaleGridLayer =
            new nmVTKPlanarScaleGridLayer(
                QString::fromLatin1(EXPORT_SCALE_GRID_LAYER_ID));
    if(!pScaleGridLayer->setBounds(
           dBounds[0], dBounds[1], dBounds[2], dBounds[3], dBounds[4])) {
        delete pScaleGridLayer;
        if(pError != nullptr) {
            *pError = exporterText("Unable to create the animation scale grid.");
        }
        return Export_Failed;
    }
    pScaleGridLayer->setVisible(oRequest.m_bScaleGridVisible);
    if(!addLayerOrDelete(oScene, pScaleGridLayer)) {
        if(pError != nullptr) {
            *pError = exporterText("Unable to create the animation scale grid.");
        }
        return Export_Failed;
    }

    nmPebiResultWellAnnotationLayer* pWellLayer =
            new nmPebiResultWellAnnotationLayer(
                QString::fromLatin1(EXPORT_WELLS_LAYER_ID));
    if(!pWellLayer->prepareCandidate(oRequest.m_pSnapshot)) {
        delete pWellLayer;
        if(pError != nullptr) {
            *pError = exporterText("Unable to create the animation well annotations.");
        }
        return Export_Failed;
    }
    pWellLayer->commitCandidate();
    pWellLayer->setVisible(oRequest.m_bWellNamesVisible);
    if(!addLayerOrDelete(oScene, pWellLayer)) {
        if(pError != nullptr) {
            *pError = exporterText("Unable to create the animation well annotations.");
        }
        return Export_Failed;
    }

    pRenderer->ResetCameraClippingRange(dBounds);

    const QString sTemporaryPath = createTemporaryPath(
                oRequest.m_sTargetFilePath);
    nmFFmpegVideoWriter::VideoSettings oVideoSettings;
    oVideoSettings.m_sTemporaryFilePath = sTemporaryPath;
    oVideoSettings.m_sFileExtension = sExtension;
    oVideoSettings.m_nWidth =
            oRequest.m_oExportSettings.m_nOutputWidth;
    oVideoSettings.m_nHeight =
            oRequest.m_oExportSettings.m_nOutputHeight;
    oVideoSettings.m_nFramesPerSecond =
            oRequest.m_oExportSettings.m_nFramesPerSecond;
    nmFFmpegVideoWriter oWriter;
    if(!oWriter.open(oVideoSettings, pError)) {
        return Export_Failed;
    }

    vtkSmartPointer<vtkWindowToImageFilter> pCapture =
            vtkSmartPointer<vtkWindowToImageFilter>::New();
    pCapture->SetInput(pRenderWindow);
    pCapture->SetInputBufferTypeToRGB();
    pCapture->ReadFrontBufferOff();
    pCapture->ShouldRerenderOff();
    emit progressChanged(0, vecFrameIndices.size());
    QCoreApplication::processEvents(QEventLoop::AllEvents);
    if(m_bCancelRequested) {
        oWriter.cancel();
        return Export_Cancelled;
    }

    for(int nOutputIndex = 0;
        nOutputIndex < vecFrameIndices.size(); ++nOutputIndex) {
        if(!oAdapter.setPropertyFrame(
               oRequest.m_sPropertyId, vecFrameIndices[nOutputIndex])) {
            oWriter.cancel();
            if(pError != nullptr) {
                *pError = exporterText("Unable to load an animation frame.");
            }
            return Export_Failed;
        }

        pRenderWindow->Render();
        pCapture->Modified();
        pCapture->Update();
        vtkImageData* pImage = pCapture->GetOutput();
        int pDimensions[3] = { 0, 0, 0 };
        vtkIdType pIncrements[3] = { 0, 0, 0 };
        if(pImage != nullptr) {
            pImage->GetDimensions(pDimensions);
            pImage->GetIncrements(pIncrements);
        }
        unsigned char* pPixels = pImage == nullptr ? nullptr :
                static_cast<unsigned char*>(pImage->GetScalarPointer());
        if(pImage == nullptr || pPixels == nullptr ||
           pImage->GetScalarType() != VTK_UNSIGNED_CHAR ||
           pImage->GetNumberOfScalarComponents() != 3 ||
           pDimensions[0] != oVideoSettings.m_nWidth ||
           pDimensions[1] != oVideoSettings.m_nHeight ||
           pIncrements[1] < oVideoSettings.m_nWidth * 3 ||
           pIncrements[1] > INT_MAX ||
           !oWriter.writeRgb24Frame(
               pPixels, static_cast<int>(pIncrements[1]), true, pError)) {
            oWriter.cancel();
            if(pError != nullptr && pError->isEmpty()) {
                *pError = exporterText("Unable to capture an animation frame.");
            }
            return Export_Failed;
        }

        // 单帧完成编码后再处理事件，避免取消或其他事件造成半帧状态。
        emit progressChanged(nOutputIndex + 1, vecFrameIndices.size());
        QCoreApplication::processEvents(QEventLoop::AllEvents);
        if(m_bCancelRequested) {
            oWriter.cancel();
            return Export_Cancelled;
        }
    }

    // Writer只负责完成临时视频；成功后再单独发布到用户目标路径。
    if(!oWriter.finish(pError)) {
        return Export_Failed;
    }
    if(!publishTemporaryFile(
           sTemporaryPath, oRequest.m_sTargetFilePath, pError)) {
        QFile::remove(sTemporaryPath);
        if(pError != nullptr && pError->isEmpty()) {
            *pError = exporterText("Unable to publish the completed video file.");
        }
        return Export_Failed;
    }
    return Export_Succeeded;
}

void nmPebiResultAnimationExporter::cancel()
{
    m_bCancelRequested = true;
}