#include "RigGeoMechBoreHoleStressCalculator.h"

//==================================================================================================
/// Internal root finding class to find a Well Pressure that gives:
/// a) a zero SigmaT for estimating the fracture gradient.
/// b) a solution to the Stassi-d'Alia failure criterion for estimating the shear failure gradient.
//==================================================================================================

//--------------------------------------------------------------------------------------------------
///
//--------------------------------------------------------------------------------------------------
RigGeoMechBoreHoleStressCalculator::RigGeoMechBoreHoleStressCalculator( const caf::Ten3d& tensor,
                                                                        double            porePressure,
                                                                        double            poissonRatio,
                                                                        double            uniaxialCompressiveStrength,
                                                                        int               nThetaSubSamples )
    : m_tensor( tensor )
    , m_porePressure( porePressure )
    , m_poissonRatio( poissonRatio )
    , m_uniaxialCompressiveStrength( uniaxialCompressiveStrength )
    , m_nThetaSubSamples( nThetaSubSamples )
{
    calculateStressComponents();
}

//--------------------------------------------------------------------------------------------------
/// Simple bisection method for now
//--------------------------------------------------------------------------------------------------
double RigGeoMechBoreHoleStressCalculator::solveFractureGradient( double* thetaOut )
{
    MemberFunc fn = &RigGeoMechBoreHoleStressCalculator::sigmaTMinOfMin;
    return solveSecant( fn, thetaOut );
}

//--------------------------------------------------------------------------------------------------
///
//--------------------------------------------------------------------------------------------------
double RigGeoMechBoreHoleStressCalculator::solveStassiDalia( double* thetaOut )
{
    MemberFunc fn = &RigGeoMechBoreHoleStressCalculator::stassiDalia;
    return solveSecant( fn, thetaOut );
}

//--------------------------------------------------------------------------------------------------
/// Bi-section root finding method: https://en.wikipedia.org/wiki/Bisection_method
/// Used as fall-back in case the secant method doesn't converge.
//--------------------------------------------------------------------------------------------------
double RigGeoMechBoreHoleStressCalculator::solveBisection( double minPw, double maxPw, MemberFunc fn, double* thetaOut )
{
    const int    N       = 50;
    const double epsilon = 1.0e-10;

    double theta = 0.0;

    std::pair<double, double> largestNegativeValue( 0.0, -std::numeric_limits<double>::infinity() );
    std::pair<double, double> smallestPositiveValue( 0.0, std::numeric_limits<double>::infinity() );

    for ( int i = 0; i <= N; ++i )
    {
        double pw   = minPw + ( maxPw - minPw ) * i / static_cast<double>( N );
        double f_pw = ( this->*fn )( pw, &theta );
        if ( f_pw >= 0.0 && f_pw < smallestPositiveValue.second )
        {
            smallestPositiveValue = std::make_pair( pw, f_pw );
        }
        if ( f_pw < 0.0 && f_pw > largestNegativeValue.second )
        {
            largestNegativeValue = std::make_pair( pw, f_pw );
        }
    }

    // TODO: Provide a warning if there was no solution to the equation
    if ( largestNegativeValue.second == -std::numeric_limits<double>::infinity() )
    {
        // No solution. Function is always positive. Pick smallest value.
        return smallestPositiveValue.first;
    }
    if ( smallestPositiveValue.second == std::numeric_limits<double>::infinity() )
    {
        // No solution. Function is always negative. Pick largest value.
        return largestNegativeValue.first;
    }
    minPw               = largestNegativeValue.first;
    double minPwFuncVal = largestNegativeValue.second;
    maxPw               = smallestPositiveValue.first;

    double range = std::abs( maxPw - minPw );

    int i = 0;
    for ( ; i <= N && range > m_porePressure * epsilon; ++i )
    {
        double midPw        = ( minPw + maxPw ) * 0.5;
        double midPwFuncVal = ( this->*fn )( midPw, &theta );
        if ( midPwFuncVal * minPwFuncVal < 0.0 )
        {
            maxPw = midPw;
        }
        else
        {
            minPw        = midPw;
            minPwFuncVal = midPwFuncVal;
        }
        range = std::abs( maxPw - minPw );
    }
    CVF_ASSERT( i < N ); // Otherwise it hasn't converged

    if ( thetaOut )
    {
        *thetaOut = theta;
    }

    // Return average of minPw and maxPw.
    return 0.5 * ( maxPw + minPw );
}

//--------------------------------------------------------------------------------------------------
/// Secant root finding method: https://en.wikipedia.org/wiki/Secant_method
/// Basically a Newton's method using finite differences for the derivative.
//--------------------------------------------------------------------------------------------------
double RigGeoMechBoreHoleStressCalculator::solveSecant( MemberFunc fn, double* thetaOut )
{
    const double epsilon = 1.0e-10;
    const int    N       = 50;
    double       theta   = 0.0;

    double x_0  = 0.0;
    double f_x0 = ( this->*fn )( x_0, &theta );
    double x_1  = m_porePressure;
    double f_x1 = ( this->*fn )( x_1, &theta );
    double x    = m_porePressure;
    double f_x  = 0.0;
    int    i    = 0;
    for ( ; i <= N && std::abs( f_x1 - f_x0 ) > epsilon; ++i )
    {
        x   = x_1 - f_x1 * ( x_1 - x_0 ) / ( f_x1 - f_x0 );
        f_x = ( this->*fn )( x, &theta );
        if ( std::abs( f_x ) < epsilon * m_porePressure ) break;

        // Update iteration variables
        x_0  = x_1;
        f_x0 = f_x1;
        x_1  = x;
        f_x1 = f_x;
    }

    if ( i == 0 || i == N || std::abs( f_x ) > epsilon * m_porePressure )
    {
        // Fallback to bisection if secant doesn't converge or converged to a wrong solution.
        return solveBisection( 0.0, m_porePressure * 2.0, fn, thetaOut );
    }

    if ( thetaOut )
    {
        *thetaOut = theta;
    }
    return x;
}

//--------------------------------------------------------------------------------------------------
///
//--------------------------------------------------------------------------------------------------
double RigGeoMechBoreHoleStressCalculator::sigmaTMinOfMin( double wellPressure, double* thetaAtMin ) const
{
    CVF_ASSERT( thetaAtMin );
    double sigma_t_min_min = std::numeric_limits<double>::max();
    for ( const cvf::Vec4d& stressComponentsForAngle : m_stressComponents )
    {
        // Perform all these internal calculations in double to reduce significance errors
        double        sigma_theta = stressComponentsForAngle[1] - wellPressure;
        const double& sigma_z     = stressComponentsForAngle[2];
        double        tauSqrx4    = std::pow( stressComponentsForAngle[3], 2 ) * 4.0;
        double sigma_t_min = 0.5 * ( ( sigma_z + sigma_theta ) - std::sqrt( std::pow( sigma_z - sigma_theta, 2 ) + tauSqrx4 ) ) - m_porePressure;
        if ( sigma_t_min < sigma_t_min_min )
        {
            sigma_t_min_min = sigma_t_min;
            *thetaAtMin     = stressComponentsForAngle[0];
        }
    }
    return sigma_t_min_min;
}

//--------------------------------------------------------------------------------------------------
///
//--------------------------------------------------------------------------------------------------
double RigGeoMechBoreHoleStressCalculator::stassiDalia( double wellPressure, double* thetaAtMin ) const
{
    CVF_ASSERT( thetaAtMin );
    double minStassiDalia = std::numeric_limits<double>::max();
    for ( const cvf::Vec4d& stressComponentsForAngle : m_stressComponents )
    {
        double        sigma_theta = stressComponentsForAngle[1] - wellPressure;
        const double& sigma_z     = stressComponentsForAngle[2];
        double        tauSqrx4    = std::pow( stressComponentsForAngle[3], 2 ) * 4.0;

        double sigma_1 = wellPressure - m_porePressure;
        double sigma_2 = 0.5 * ( ( sigma_z + sigma_theta ) + std::sqrt( std::pow( sigma_z - sigma_theta, 2 ) + tauSqrx4 ) ) - m_porePressure;
        double sigma_3 = 0.5 * ( ( sigma_z + sigma_theta ) - std::sqrt( std::pow( sigma_z - sigma_theta, 2 ) + tauSqrx4 ) ) - m_porePressure;

        double stassiDalia = std::pow( sigma_1 - sigma_2, 2 ) + std::pow( sigma_2 - sigma_3, 2 ) + std::pow( sigma_1 - sigma_3, 2 ) -
                             2 * m_uniaxialCompressiveStrength * ( sigma_1 + sigma_2 + sigma_3 );

        if ( stassiDalia < minStassiDalia )
        {
            minStassiDalia = stassiDalia;
            *thetaAtMin    = stressComponentsForAngle[0];
        }
    }
    return minStassiDalia;
}

//--------------------------------------------------------------------------------------------------
///
//--------------------------------------------------------------------------------------------------
void RigGeoMechBoreHoleStressCalculator::calculateStressComponents()
{
    m_stressComponents.reserve( m_nThetaSubSamples );

    for ( int i = 0; i < m_nThetaSubSamples; ++i )
    {
        double     theta                    = ( i * cvf::PI_F ) / ( m_nThetaSubSamples - 1.0 );
        cvf::Vec4d stressComponentsForAngle = calculateStressComponentsForSegmentAngle( theta );
        m_stressComponents.push_back( stressComponentsForAngle );
    }
}

//--------------------------------------------------------------------------------------------------
///
//--------------------------------------------------------------------------------------------------
cvf::Vec4d RigGeoMechBoreHoleStressCalculator::calculateStressComponentsForSegmentAngle( double theta ) const
{
    cvf::Vec4d stressComponents;

    const double& sx  = m_tensor[caf::Ten3d::SXX];
    const double& sy  = m_tensor[caf::Ten3d::SYY];
    const double& sz  = m_tensor[caf::Ten3d::SZZ];
    const double& txy = m_tensor[caf::Ten3d::SXY];
    const double& txz = m_tensor[caf::Ten3d::SZX];
    const double& tyz = m_tensor[caf::Ten3d::SYZ];

    stressComponents[0] = theta;
    stressComponents[1] = sx + sy - 2 * ( sx - sy ) * cos( 2 * theta ) - 4 * txy * sin( 2 * theta );
    stressComponents[2] = sz - m_poissonRatio * ( 2 * ( sx - sy ) * cos( 2 * theta ) + 4 * txy * sin( 2 * theta ) );
    stressComponents[3] = 2 * ( tyz * cos( theta ) - txz * sin( theta ) );

    return stressComponents;
}