/////////////////////////////////////////////////////////////////////////////////
//
//  Copyright (C) 2020-     Equinor ASA
//
//  ResInsight is free software: you can redistribute it and/or modify
//  it under the terms of the GNU General Public License as published by
//  the Free Software Foundation, either version 3 of the License, or
//  (at your option) any later version.
//
//  ResInsight is distributed in the hope that it will be useful, but WITHOUT ANY
//  WARRANTY; without even the implied warranty of MERCHANTABILITY or
//  FITNESS FOR A PARTICULAR PURPOSE.
//
//  See the GNU General Public License at <http://www.gnu.org/licenses/gpl.html>
//  for more details.
//
/////////////////////////////////////////////////////////////////////////////////

#include "RigFemPartResultCalculatorPoreCompressibility.h"

#include "RiaEclipseUnitTools.h"

#include "RigFemPart.h"
#include "RigFemPartCollection.h"
#include "RigFemPartResultsCollection.h"
#include "RigFemResultAddress.h"
#include "RigFemScalarResultFrames.h"

#include "Riu3DMainWindowTools.h"

#include "cafProgressInfo.h"

#include <QString>

//--------------------------------------------------------------------------------------------------
///
//--------------------------------------------------------------------------------------------------
RigFemPartResultCalculatorPoreCompressibility::RigFemPartResultCalculatorPoreCompressibility( RigFemPartResultsCollection& collection )
    : RigFemPartResultCalculator( collection )
{
}

//--------------------------------------------------------------------------------------------------
///
//--------------------------------------------------------------------------------------------------
RigFemPartResultCalculatorPoreCompressibility::~RigFemPartResultCalculatorPoreCompressibility()
{
}

//--------------------------------------------------------------------------------------------------
///
//--------------------------------------------------------------------------------------------------
bool RigFemPartResultCalculatorPoreCompressibility::isMatching( const RigFemResultAddress& resVarAddr ) const
{
    return ( resVarAddr.fieldName == "COMPRESSIBILITY" && ( resVarAddr.componentName == "PORE" || resVarAddr.componentName == "VERTICAL" ||
                                                            resVarAddr.componentName == "VERTICAL-RATIO" ) );
}

//--------------------------------------------------------------------------------------------------
///
//--------------------------------------------------------------------------------------------------
RigFemScalarResultFrames* RigFemPartResultCalculatorPoreCompressibility::calculate( int partIndex, const RigFemResultAddress& resAddr )
{
    caf::ProgressInfo stepCountProgress( static_cast<size_t>( m_resultCollection->timeStepCount() ) * 7, "Calculating Pore Compressibility" );

    auto loadFrameLambda = [&]( RigFemResultAddress addr, const QString& errMsg = "" ) -> RigFemScalarResultFrames*
    {
        auto task   = stepCountProgress.task( QString( "Loading %1: %2" )
                                                .arg( QString::fromStdString( addr.fieldName ) )
                                                .arg( QString::fromStdString( addr.componentName ) ),
                                            m_resultCollection->timeStepCount() );
        auto result = m_resultCollection->findOrLoadScalarResult( partIndex, addr );
        if ( result->frameData( 0, 0 ).empty() )
        {
            if ( !errMsg.isEmpty() ) Riu3DMainWindowTools::reportAndShowWarning( "Required data missing", errMsg );
            return nullptr;
        }
        return result;
    };

    RigFemScalarResultFrames* srcPORDataFrames = loadFrameLambda( RigFemResultAddress( RIG_NODAL, "POR-Bar", "" ) );

    // Volumetric Strain
    RigFemScalarResultFrames* srcEVDataFrames = loadFrameLambda( RigFemResultAddress( resAddr.resultPosType, "NE", "EV" ) );

    // Vertical Strain
    RigFemScalarResultFrames* verticalStrainDataFrames = loadFrameLambda( RigFemResultAddress( resAddr.resultPosType, "NE", "E33" ) );

    // Biot porelastic coefficient (alpha)
    RigFemScalarResultFrames* biotCoefficient = nullptr;
    if ( !m_resultCollection->biotResultAddress().isEmpty() )
    {
        biotCoefficient = loadFrameLambda( RigFemResultAddress( RIG_ELEMENT, m_resultCollection->biotResultAddress().toStdString(), "" ) );
    }

    QString youngsErrMsg = QString( "Failed to compute %1\n" ).arg( QString::fromStdString( resAddr.componentName ) );
    youngsErrMsg += "Missing Young's Modulus element data (MODULUS)";

    RigFemScalarResultFrames* youngsModuliFrames = loadFrameLambda( RigFemResultAddress( RIG_ELEMENT, "MODULUS", "" ), youngsErrMsg );
    if ( !youngsModuliFrames ) return nullptr;

    QString poissonError = QString( "Failed to compute %1\n" ).arg( QString::fromStdString( resAddr.componentName ) );
    poissonError += "Missing Poisson Ratio element data (RATIO)";

    RigFemScalarResultFrames* poissonRatioFrames = loadFrameLambda( RigFemResultAddress( RIG_ELEMENT, "RATIO", "" ), poissonError );
    if ( !poissonRatioFrames ) return nullptr;

    RigFemScalarResultFrames* voidRatioFrames = loadFrameLambda( RigFemResultAddress( resAddr.resultPosType, "VOIDR", "" ) );

    RigFemScalarResultFrames* poreCompressibilityFrames =
        m_resultCollection->createScalarResult( partIndex, RigFemResultAddress( resAddr.resultPosType, resAddr.fieldName, "PORE" ) );

    RigFemScalarResultFrames* verticalCompressibilityFrames =
        m_resultCollection->createScalarResult( partIndex, RigFemResultAddress( resAddr.resultPosType, resAddr.fieldName, "VERTICAL" ) );
    RigFemScalarResultFrames* verticalCompressibilityRatioFrames =
        m_resultCollection->createScalarResult( partIndex, RigFemResultAddress( resAddr.resultPosType, resAddr.fieldName, "VERTICAL-RATIO" ) );

    const RigFemPart* femPart = m_resultCollection->parts()->part( partIndex );
    float             inf     = std::numeric_limits<float>::infinity();

    int refStepIdx, refFrameIdx;
    std::tie( refStepIdx, refFrameIdx ) = m_resultCollection->referenceStepAndFrameIndex();

    const int timeSteps = srcEVDataFrames->timeStepCount();
    for ( int stepIdx = 0; stepIdx < timeSteps; stepIdx++ )
    {
        auto task = stepCountProgress.task( QString( "Step %1" ).arg( stepIdx ) );

        const int frameCount = srcEVDataFrames->frameCount( stepIdx );
        for ( int fIdx = 0; fIdx < frameCount; fIdx++ )
        {
            const std::vector<float>& evData                      = srcEVDataFrames->frameData( stepIdx, fIdx );
            const std::vector<float>& referenceEvData             = srcEVDataFrames->frameData( refStepIdx, refFrameIdx );
            const std::vector<float>& verticalStrainData          = verticalStrainDataFrames->frameData( stepIdx, fIdx );
            const std::vector<float>& referenceVerticalStrainData = verticalStrainDataFrames->frameData( refStepIdx, refFrameIdx );

            const std::vector<float>& youngsModuliData      = youngsModuliFrames->frameData( stepIdx, fIdx );
            const std::vector<float>& poissonRatioData      = poissonRatioFrames->frameData( stepIdx, fIdx );
            const std::vector<float>& voidRatioData         = voidRatioFrames->frameData( 0, 0 );
            const std::vector<float>& referencePorFrameData = srcPORDataFrames->frameData( refStepIdx, refFrameIdx );
            const std::vector<float>& porFrameData          = srcPORDataFrames->frameData( stepIdx, fIdx );

            std::vector<float>& poreCompressibilityFrameData          = poreCompressibilityFrames->frameData( stepIdx, fIdx );
            std::vector<float>& verticalCompressibilityFrameData      = verticalCompressibilityFrames->frameData( stepIdx, fIdx );
            std::vector<float>& verticalCompressibilityRatioFrameData = verticalCompressibilityRatioFrames->frameData( stepIdx, fIdx );

            size_t valCount = evData.size();
            poreCompressibilityFrameData.resize( valCount );
            verticalCompressibilityFrameData.resize( valCount );
            verticalCompressibilityRatioFrameData.resize( valCount );

            int elementCount = femPart->elementCount();

            std::vector<float> biotData;
            if ( biotCoefficient )
            {
                biotData = biotCoefficient->frameData( stepIdx, fIdx );
                if ( !m_resultCollection->isValidBiotData( biotData, elementCount ) )
                {
                    m_resultCollection->deleteResult( resAddr );
                    return nullptr;
                }
            }

#pragma omp parallel for
            for ( int elmIdx = 0; elmIdx < elementCount; ++elmIdx )
            {
                RigElementType elmType = femPart->elementType( elmIdx );

                int elmNodeCount = RigFemTypes::elementNodeCount( femPart->elementType( elmIdx ) );

                if ( elmType == HEX8P )
                {
                    for ( int elmNodIdx = 0; elmNodIdx < elmNodeCount; ++elmNodIdx )
                    {
                        size_t elmNodResIdx = femPart->elementNodeResultIdx( elmIdx, elmNodIdx );
                        if ( elmNodResIdx < evData.size() )
                        {
                            if ( ( fIdx == refFrameIdx ) && ( stepIdx == refStepIdx ) )
                            {
                                // The time step and the reference time step are the same: results undefined
                                poreCompressibilityFrameData[elmNodResIdx]          = inf;
                                verticalCompressibilityFrameData[elmNodResIdx]      = inf;
                                verticalCompressibilityRatioFrameData[elmNodResIdx] = inf;
                            }
                            else
                            {
                                // Use biot coefficient for all timesteps
                                double biotCoefficient = 1.0;
                                if ( biotData.empty() )
                                {
                                    biotCoefficient = m_resultCollection->biotFixedFactor();
                                }
                                else
                                {
                                    // Use coefficient from element property table
                                    biotCoefficient = biotData[elmIdx];
                                }

                                int nodeIdx = femPart->nodeIdxFromElementNodeResultIdx( elmNodResIdx );

                                // Calculate bulk modulus for solids (grains).
                                // Incoming unit for Young's Modulus is GPa: convert to Pa.
                                double poissonRatio     = poissonRatioData[elmIdx];
                                double youngsModuli     = RiaEclipseUnitTools::gigaPascalToPascal( youngsModuliData[elmIdx] );
                                double bulkModulusFrame = youngsModuli / ( 3.0 * ( 1.0 - 2.0 * poissonRatio ) );
                                double bulkModulus      = bulkModulusFrame / ( 1.0 - biotCoefficient );

                                // Calculate initial porosity (always from geostatic timestep)
                                double voidr    = voidRatioData[elmNodResIdx];
                                double porosity = voidr / ( 1.0 + voidr );

                                // Calculate difference in pore pressure between reference state and this state,
                                // and convert unit from Bar to Pascal.
                                double referencePorePressure = referencePorFrameData[nodeIdx];
                                double framePorePressure     = porFrameData[nodeIdx];
                                double deltaPorePressure = RiaEclipseUnitTools::barToPascal( framePorePressure - referencePorePressure );

                                // Calculate pore compressibility
                                double poreCompressibility = inf;
                                if ( deltaPorePressure != 0.0 && porosity != 0.0 )
                                {
                                    double deltaEv      = evData[elmNodResIdx] - referenceEvData[elmNodResIdx];
                                    poreCompressibility = -( biotCoefficient * deltaEv ) / ( deltaPorePressure * porosity );
                                    // Guard against divide by zero: second term can be ignored when bulk modulus is
                                    // zero, which can happens when biot coefficient is 1.0
                                    if ( biotCoefficient != 1.0 && porosity != 1.0 )
                                    {
                                        poreCompressibility += ( 1.0 / bulkModulus ) * ( biotCoefficient / porosity - 1.0 );
                                    }
                                }
                                // Convert from 1/Pa to 1/GPa
                                poreCompressibilityFrameData[elmNodResIdx] = poreCompressibility * 1.0e9;

                                double verticalCompressibility      = inf;
                                double verticalCompressibilityRatio = inf;

                                if ( biotCoefficient != 0.0 && deltaPorePressure != 0.0 )
                                {
                                    double deltaStrain = verticalStrainData[elmNodResIdx] - referenceVerticalStrainData[elmNodResIdx];

                                    // Calculate vertical compressibility (unit: 1/Pa)
                                    verticalCompressibility = -deltaStrain / ( biotCoefficient * deltaPorePressure );

                                    // Calculate vertical compressibility ratio
                                    verticalCompressibilityRatio = ( verticalCompressibility * youngsModuli * ( 1.0 - poissonRatio ) ) /
                                                                   ( ( 1.0 + poissonRatio ) * ( 1.0 - 2.0 * poissonRatio ) );
                                }

                                // Convert from 1/Pa to 1/GPa
                                verticalCompressibilityFrameData[elmNodResIdx]      = verticalCompressibility * 1.0e9;
                                verticalCompressibilityRatioFrameData[elmNodResIdx] = verticalCompressibilityRatio;
                            }
                        }
                    }
                }
                else
                {
                    for ( int elmNodIdx = 0; elmNodIdx < elmNodeCount; ++elmNodIdx )
                    {
                        size_t elmNodResIdx = femPart->elementNodeResultIdx( elmIdx, elmNodIdx );
                        if ( elmNodResIdx < poreCompressibilityFrameData.size() )
                        {
                            poreCompressibilityFrameData[elmNodResIdx] = inf;
                        }
                    }
                }
            }
        }
    }
    RigFemScalarResultFrames* requestedResultFrames = m_resultCollection->findOrLoadScalarResult( partIndex, resAddr );
    return requestedResultFrames;
}
