#include "JacobiSmoother.H"
#include "JacobiSmootherF.H"
#include "lduMatrixSolutionCache.H"

namespace Foam
{
    defineTypeNameAndDebug(JacobiSmoother, 0);

    gpulduMatrix::smoother::addsymMatrixConstructorToTable<JacobiSmoother>
        addJacobiSmootherSymMatrixConstructorToTable_;

    gpulduMatrix::smoother::addasymMatrixConstructorToTable<JacobiSmoother>
        addJacobiSmootherAsymMatrixConstructorToTable_;
}

Foam::JacobiSmoother::JacobiSmoother
(
    const word& fieldName,
    const gpulduMatrix& matrix,
    const FieldField<gpuField, scalar>& interfaceBouCoeffs,
    const FieldField<gpuField, scalar>& interfaceIntCoeffs,
    const lduInterfacegpuFieldPtrsList& interfaces
)
:
    gpulduMatrix::smoother
    (
        fieldName,
        matrix,
        interfaceBouCoeffs,
        interfaceIntCoeffs,
        interfaces
    ),
    omega_(0.9)
{
    //solverControls.readIfPresent("omega", omega_);
}


void Foam::JacobiSmoother::smooth
(
    scalargpuField& psi,
    const scalargpuField& source,
    const direction cmpt,
    const label nSweeps
) const
{
    scalargpuField Apsi(lduMatrixSolutionCache::first(psi.size()),psi.size());
    scalargpuField sourceTmp(lduMatrixSolutionCache::second(source.size()),source.size());

    bool fastPath = lduMatrixSolutionCache::favourSpeed >= 2;//||
                    //(lduMatrixSolutionCache::favourSpeed && ( matrix_.coarsestLevel() || ! matrix_.level()));

    const labelgpuList& l = fastPath?
                            matrix_.lduAddr().gpuOwnerSortAddr():
                            matrix_.lduAddr().lowerAddr();
    const labelgpuList& u = matrix_.lduAddr().upperAddr();

    const labelgpuList& ownStart = matrix_.lduAddr().ownerStartAddr();
    const labelgpuList& losortStart = matrix_.lduAddr().losortStartAddr();
    const labelgpuList& losort = matrix_.lduAddr().losortAddr();

    const scalargpuField& Lower = fastPath?
                                  matrix_.gpuLowerSort():
                                  matrix_.gpuLower();

    const scalargpuField& Upper = matrix_.gpuUpper();
    const scalargpuField& Diag = matrix_.gpuDiag();

    const label startRequest = Pstream::nRequests();

    for (label sweep=0; sweep<nSweeps; sweep++)
    {
        sourceTmp = source;

        matrix_.initMatrixInterfaces
        (
            false,
            interfaceBouCoeffs_,
            interfaces_,
            psi,
            sourceTmp,
            cmpt
        );

        matrix_.updateMatrixInterfaces
        (
            false,
            interfaceBouCoeffs_,
            interfaces_,
            psi,
            sourceTmp,
            cmpt,
            startRequest
        );

        if(fastPath)
        {

            thrust::transform
            (
                thrust::make_counting_iterator(0),
                thrust::make_counting_iterator(0)+psi.size(),
                Apsi.begin(),
                JacobiSmootherFunctor<true,3>
                (
                    omega_,
                    psi.data(),
                    Diag.data(),
                    sourceTmp.data(),
                    Lower.data(),
                    Upper.data(),
                    l.data(),
                    u.data(),
                    ownStart.data(),
                    losortStart.data(),
                    losort.data()
                )
            );

        }
        else
        {

            thrust::transform
            (
                thrust::make_counting_iterator(0),
                thrust::make_counting_iterator(0)+psi.size(),
                Apsi.begin(),
                JacobiSmootherFunctor<false,3>
                (
                    omega_,
                    psi.data(),
                    Diag.data(),
                    sourceTmp.data(),
                    Lower.data(),
                    Upper.data(),
                    l.data(),
                    u.data(),
                    ownStart.data(),
                    losortStart.data(),
                    losort.data()
                )
            );
        }

        psi = Apsi;
    }
}

void Foam::JacobiSmoother::scalarSmooth
(
    scalargpuField& psi,
    const scalargpuField& Source,
    const direction cmpt,
    const label nSweeps
) const
{} 
