/*---------------------------------------------------------------------------*\
  =========                 |
  \\      /  F ield         | OpenFOAM: The Open Source CFD Toolbox
   \\    /   O peration     |
    \\  /    A nd           | www.openfoam.com
     \\/     M anipulation  |
-------------------------------------------------------------------------------
    Copyright (C) 2011-2017 OpenFOAM Foundation
    Copyright (C) 2018-2021 OpenCFD Ltd.
-------------------------------------------------------------------------------
License
    This file is part of OpenFOAM.

    OpenFOAM 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.

    OpenFOAM 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
    for more details.

    You should have received a copy of the GNU General Public License
    along with OpenFOAM.  If not, see <http://www.gnu.org/licenses/>.

\*---------------------------------------------------------------------------*/

#include "gpuMRFZone.H"
#include "gpufvMesh.H"
#include "volgpuFields.H"
#include "surfacegpuFields.H"
#include "gpufvMatrices.H"
#include "faceSet.H"
#include "geometricOneField.H"
#include "syncTools.H"

// * * * * * * * * * * * * * * Static Data Members * * * * * * * * * * * * * //

namespace Foam
{
    defineTypeNameAndDebug(gpuMRFZone, 0);

    struct MRFZoneCorrectVelocityFunctor : public std::unary_function<vector, vector>
    {
        const vector omega;
        const vector origin;

        MRFZoneCorrectVelocityFunctor(vector _omega, vector _origin) : omega(_omega), origin(_origin) {}
        __host__ __device__
            vector
            operator()(const vector &x)
        {
            return omega ^ (x - origin);
        }
    };

    template <bool rhs>
    struct MRFZoneAddCoriollisFunctor : public std::binary_function<vector, thrust::tuple<scalar, vector>, vector>
    {
        const vector omega;

        MRFZoneAddCoriollisFunctor(vector _omega) : omega(_omega) {}

        __host__ __device__
            vector
            operator()(const vector &Usource, const thrust::tuple<scalar, vector> &t)
        {
            vector delta = thrust::get<0>(t) * (omega ^ thrust::get<1>(t));
            if (rhs)
                return Usource + delta;
            else
                return Usource - delta;
        }
    };

    template <bool rhs>
    struct MRFZoneAddRhoCoriollisFunctor : public std::binary_function<vector, thrust::tuple<scalar, scalar, vector>, vector>
    {
        const vector omega;

        MRFZoneAddRhoCoriollisFunctor(vector _omega) : omega(_omega) {}

        __host__ __device__
            vector
            operator()(const vector &Usource, const thrust::tuple<scalar, scalar, vector> &t)
        {
            vector delta = thrust::get<0>(t) * thrust::get<1>(t) * (omega ^ thrust::get<2>(t));
            if (rhs)
                return Usource + delta;
            else
                return Usource - delta;
        }
    };
}

// * * * * * * * * * * * * * Private Member Functions  * * * * * * * * * * * //

void Foam::gpuMRFZone::setMRFFaces()
{
    fvMesh &mesh2 = mesh_.hostmesh();

    const polyBoundaryMesh &patches = mesh2.boundaryMesh();

    // Type per face:
    //  0:not in zone
    //  1:moving with frame
    //  2:other
    labelList faceType(mesh2.nFaces(), Zero);

    // Determine faces in cell zone
    // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    // (without constructing cells)

    const labelList &own = mesh2.faceOwner();
    const labelList &nei = mesh2.faceNeighbour();

    // Cells in zone
    boolList zoneCell(mesh2.nCells(), false);

    if (cellZoneID_ != -1)
    {
        const labelList &cellLabels = mesh2.cellZones()[cellZoneID_];
        forAll(cellLabels, i)
        {
            zoneCell[cellLabels[i]] = true;
        }
    }

    label nZoneFaces = 0;

    for (label facei = 0; facei < mesh2.nInternalFaces(); facei++)
    {
        if (zoneCell[own[facei]] || zoneCell[nei[facei]])
        {
            faceType[facei] = 1;
            nZoneFaces++;
        }
    }

    labelHashSet excludedPatches(excludedPatchLabels_);

    forAll(patches, patchi)
    {
        const polyPatch &pp = patches[patchi];

        if (pp.coupled() || excludedPatches.found(patchi))
        {
            forAll(pp, i)
            {
                label facei = pp.start() + i;

                if (zoneCell[own[facei]])
                {
                    faceType[facei] = 2;
                    nZoneFaces++;
                }
            }
        }
        else if (!isA<emptyPolyPatch>(pp))
        {
            forAll(pp, i)
            {
                label facei = pp.start() + i;

                if (zoneCell[own[facei]])
                {
                    faceType[facei] = 1;
                    nZoneFaces++;
                }
            }
        }
    }

    // Synchronize the faceType across processor patches
    syncTools::syncFaceList(mesh2, faceType, maxEqOp<label>());

    // Now we have for faceType:
    //  0   : face not in cellZone
    //  1   : internal face or normal patch face
    //  2   : coupled patch face or excluded patch face

    // Sort into lists per patch.
    labelList internalFaces2;
    internalFaces2.setSize(mesh2.nFaces());
    
    label nInternal = 0;

    for (label facei = 0; facei < mesh2.nInternalFaces(); facei++)
    {
        if (faceType[facei] == 1)
        {
            internalFaces2[nInternal++] = facei;
        }
    }
    internalFaces2.setSize(nInternal);

    internalFaces_ = internalFaces2;

    labelList nIncludedFaces(patches.size(), Zero);
    labelList nExcludedFaces(patches.size(), Zero);

    forAll(patches, patchi)
    {
        const polyPatch &pp = patches[patchi];

        forAll(pp, patchFacei)
        {
            label facei = pp.start() + patchFacei;

            if (faceType[facei] == 1)
            {
                nIncludedFaces[patchi]++;
            }
            else if (faceType[facei] == 2)
            {
                nExcludedFaces[patchi]++;
            }
        }
    }

    includedFaces_.setSize(patches.size());
    excludedFaces_.setSize(patches.size());
    forAll(nIncludedFaces, patchi)
    {
        includedFaces_[patchi].setSize(nIncludedFaces[patchi]);
        excludedFaces_[patchi].setSize(nExcludedFaces[patchi]);
    }
    nIncludedFaces = 0;
    nExcludedFaces = 0;

    forAll(patches, patchi)
    {
        const polyPatch &pp = patches[patchi];

        labelList includedFaces2(includedFaces_[patchi].size());
        labelList excludedFaces2(excludedFaces_[patchi].size());

        forAll(pp, patchFacei)
        {
            label facei = pp.start() + patchFacei;

            if (faceType[facei] == 1)
            {
                includedFaces2[nIncludedFaces[patchi]++] = patchFacei;
            }
            else if (faceType[facei] == 2)
            {
                excludedFaces2[nExcludedFaces[patchi]++] = patchFacei;
            }

            includedFaces_[patchi] = includedFaces2;
            excludedFaces_[patchi] = excludedFaces2;
        }
    }

    //if (debug)
    //{
    //    faceSet internalFaces(mesh_, "internalFaces", internalFaces_);
    //    Pout << "Writing " << internalFaces.size()
    //         << " internal faces in MRF zone to faceSet "
    //         << internalFaces.name() << endl;
    //    internalFaces.write();
    //
    //    faceSet MRFFaces(mesh_, "includedFaces", 100);
    //    forAll(includedFaces_, patchi)
    //    {
    //        forAll(includedFaces_[patchi], i)
    //        {
    //            label patchFacei = includedFaces_[patchi][i];
    //            MRFFaces.insert(patches[patchi].start() + patchFacei);
    //        }
    //    }
    //    Pout << "Writing " << MRFFaces.size()
    //         << " patch faces in MRF zone to faceSet "
    //         << MRFFaces.name() << endl;
    //    MRFFaces.write();
    //
    //    faceSet excludedFaces(mesh_, "excludedFaces", 100);
    //    forAll(excludedFaces_, patchi)
    //    {
    //        forAll(excludedFaces_[patchi], i)
    //        {
    //            label patchFacei = excludedFaces_[patchi][i];
    //            excludedFaces.insert(patches[patchi].start() + patchFacei);
    //        }
    //    }
    //    Pout << "Writing " << excludedFaces.size()
    //         << " faces in MRF zone with special handling to faceSet "
    //         << excludedFaces.name() << endl;
    //    excludedFaces.write();
    //}
}

// * * * * * * * * * * * * * * * * Constructors  * * * * * * * * * * * * * * //

Foam::gpuMRFZone::gpuMRFZone(
    const word &name,
    const gpufvMesh &mesh,
    const dictionary &dict,
    const word &cellZoneName)
    : mesh_(mesh),
      name_(name),
      coeffs_(dict),
      active_(true),
      cellZoneName_(cellZoneName),
      cellZoneID_(-1),
      excludedPatchNames_(wordRes()),
      origin_(Zero),
      axis_(Zero),
      omega_(nullptr)
{
    read(dict);
}

// * * * * * * * * * * * * * * * Member Functions  * * * * * * * * * * * * * //

Foam::vector Foam::gpuMRFZone::Omega() const
{
    return omega_->value(mesh_.time().timeOutputValue()) * axis_;
}

void Foam::gpuMRFZone::addCoriolis(
    const volVectorgpuField &U,
    volVectorgpuField &ddtU) const
{
    if (cellZoneID_ == -1)
    {
        return;
    }

    const labelgpuList &cells = mesh_.hostmesh().cellZones()[cellZoneID_].getList();
    vectorgpuField &ddtUc = ddtU.primitiveFieldRef();
    const vectorgpuField &Uc = U;

    const vector Omega = this->Omega();

    thrust::transform(
        thrust::make_permutation_iterator(
            ddtUc.begin(),
            cells.begin()),
        thrust::make_permutation_iterator(
            ddtUc.begin(),
            cells.end()),
        thrust::make_transform_iterator(
            thrust::make_permutation_iterator(
                Uc.begin(),
                cells.end()),
            crossOperatorSFFunctor<vector,vector,vector>(Omega)),
        thrust::make_permutation_iterator(
            ddtUc.begin(),
            cells.begin()),
        addOperatorFunctor<vector,vector,vector>());
}

void Foam::gpuMRFZone::addCoriolis(gpufvVectorMatrix &UEqn, const bool rhs) const
{
    if (cellZoneID_ == -1)
    {
        return;
    }

    const labelgpuList &cells = mesh_.hostmesh().cellZones()[cellZoneID_].getList();
    const scalargpuField &V = mesh_.V().field();
    vectorgpuField &Usource = UEqn.source();
    
    const vectorgpuField &U = UEqn.psi().field();

    const vector Omega = this->Omega();

    if (rhs)
    {
        thrust::transform(
            thrust::make_permutation_iterator(
                Usource.begin(),
                cells.begin()),
            thrust::make_permutation_iterator(
                Usource.begin(),
                cells.end()),
            thrust::make_zip_iterator(thrust::make_tuple(
                thrust::make_permutation_iterator(
                    V.begin(),
                    cells.begin()),
                thrust::make_permutation_iterator(
                    U.begin(),
                    cells.begin()))),
            thrust::make_permutation_iterator(
                Usource.begin(),
                cells.begin()),
            MRFZoneAddCoriollisFunctor<true>(Omega));
    }
    else
    {
        thrust::transform(
            thrust::make_permutation_iterator(
                Usource.begin(),
                cells.begin()),
            thrust::make_permutation_iterator(
                Usource.begin(),
                cells.end()),
            thrust::make_zip_iterator(thrust::make_tuple(
                thrust::make_permutation_iterator(
                    V.begin(),
                    cells.begin()),
                thrust::make_permutation_iterator(
                    U.begin(),
                    cells.begin()))),
            thrust::make_permutation_iterator(
                Usource.begin(),
                cells.begin()),
            MRFZoneAddCoriollisFunctor<false>(Omega));
    }
}

void Foam::gpuMRFZone::addCoriolis(
    const volScalargpuField &rho,
    gpufvVectorMatrix &UEqn,
    const bool rhs) const
{
    if (cellZoneID_ == -1)
    {
        return;
    }

    const labelgpuList &cells = mesh_.hostmesh().cellZones()[cellZoneID_].getList();
    const scalargpuField &V = mesh_.V().field();
    vectorgpuField &Usource = UEqn.source();
    const vectorgpuField &U = UEqn.psi().field();

    const vector Omega = this->Omega();

    if (rhs)
    {
        thrust::transform(
            thrust::make_permutation_iterator(
                Usource.begin(),
                cells.begin()),
            thrust::make_permutation_iterator(
                Usource.begin(),
                cells.end()),
            thrust::make_zip_iterator(thrust::make_tuple(
                thrust::make_permutation_iterator(
                    V.begin(),
                    cells.begin()),
                thrust::make_permutation_iterator(
                    rho.field().begin(),
                    cells.begin()),
                thrust::make_permutation_iterator(
                    U.begin(),
                    cells.begin()))),
            thrust::make_permutation_iterator(
                Usource.begin(),
                cells.begin()),
            MRFZoneAddRhoCoriollisFunctor<true>(Omega));
    }
    else
    {
        thrust::transform(
            thrust::make_permutation_iterator(
                Usource.begin(),
                cells.begin()),
            thrust::make_permutation_iterator(
                Usource.begin(),
                cells.end()),
            thrust::make_zip_iterator(thrust::make_tuple(
                thrust::make_permutation_iterator(
                    V.begin(),
                    cells.begin()),
                thrust::make_permutation_iterator(
                    rho.field().begin(),
                    cells.begin()),
                thrust::make_permutation_iterator(
                    U.begin(),
                    cells.begin()))),
            thrust::make_permutation_iterator(
                Usource.begin(),
                cells.begin()),
            MRFZoneAddRhoCoriollisFunctor<false>(Omega));
    }
}

void Foam::gpuMRFZone::makeRelative(volVectorgpuField &U) const
{
    if (cellZoneID_ == -1)
    {
        return;
    }

    const volVectorgpuField &C = mesh_.C();

    const vector Omega = this->Omega();

    const labelgpuList &cells = mesh_.hostmesh().cellZones()[cellZoneID_].getList();

    thrust::transform(
        thrust::make_permutation_iterator(
            U.field().begin(),
            cells.begin()),
        thrust::make_permutation_iterator(
            U.field().begin(),
            cells.end()),
        thrust::make_transform_iterator(
            thrust::make_permutation_iterator(
                C.field().begin(),
                cells.begin()),
            MRFZoneCorrectVelocityFunctor(Omega, origin_)),
        thrust::make_permutation_iterator(
            U.field().begin(),
            cells.begin()),
        subtractOperatorFunctor<vector,vector,vector>());

    // Included patches

    //volVectorField::Boundary& Ubf = U.boundaryFieldRef();

    forAll(includedFaces_, patchi)
    {
        const labelgpuList &faces = includedFaces_[patchi];
        thrust::fill(
            thrust::make_permutation_iterator(
                U.boundaryFieldRef()[patchi].begin(),
                faces.begin()),
            thrust::make_permutation_iterator(
                U.boundaryFieldRef()[patchi].begin(),
                faces.end()),
            vector(Zero)); // change 1
    }

    // Excluded patches
    forAll(excludedFaces_, patchi)
    {
        const vectorgpuField &patchC = mesh_.C().boundaryField()[patchi];
        const vectorgpuField &pfld = U.boundaryField()[patchi];
        const labelgpuList &faces = excludedFaces_[patchi];


        thrust::transform(
            thrust::make_permutation_iterator(
                pfld.begin(),
                faces.begin()),
            thrust::make_permutation_iterator(
                pfld.begin(),
                faces.end()),
            thrust::make_transform_iterator(
                thrust::make_permutation_iterator(
                    patchC.begin(),
                    faces.begin()),
                MRFZoneCorrectVelocityFunctor(Omega, origin_)),
            thrust::make_permutation_iterator(
                U.boundaryFieldRef()[patchi].begin(),
                faces.begin()),
            subtractOperatorFunctor<vector,vector,vector>());
    }
}

void Foam::gpuMRFZone::makeRelative(surfaceScalargpuField &phi) const
{
    makeRelativeRhoFlux(geometricOneField(), phi);
}

void Foam::gpuMRFZone::makeRelative(FieldField<fvsPatchgpuField, scalar> &phi) const
{
    makeRelativeRhoFlux(oneFieldField(), phi);
}

void Foam::gpuMRFZone::makeRelative(gpuField<scalar> &phi, const label patchi) const
{
    makeRelativeRhoFlux(oneField(), phi, patchi);
}

void Foam::gpuMRFZone::makeRelative(
    const surfaceScalargpuField &rho,
    surfaceScalargpuField &phi) const
{
    makeRelativeRhoFlux(rho, phi); 
}

void Foam::gpuMRFZone::makeAbsolute(volVectorgpuField &U) const
{
    if (cellZoneID_ == -1)
    {
        return;
    }

    const volVectorgpuField &C = mesh_.C();

    const vector Omega = this->Omega();

    const labelgpuList &cells = mesh_.hostmesh().cellZones()[cellZoneID_].getList();

    thrust::transform(
        thrust::make_permutation_iterator(
            U.field().begin(),
            cells.begin()),
        thrust::make_permutation_iterator(
            U.field().begin(),
            cells.end()),
        thrust::make_transform_iterator(
            thrust::make_permutation_iterator(
                C.field().begin(),
                cells.begin()),
            MRFZoneCorrectVelocityFunctor(Omega, origin_)),
        thrust::make_permutation_iterator(
            U.field().begin(),
            cells.begin()),
        addOperatorFunctor<vector,vector,vector>());

    // Included patches
    // volVectorField::Boundary &Ubf = U.boundaryFieldRef();

    forAll(includedFaces_, patchi)
    {
        const vectorgpuField &patchC = C.boundaryField()[patchi];

        vectorgpuField &pfld = U.boundaryFieldRef()[patchi];
        const labelgpuList &faces = includedFaces_[patchi];
        
        thrust::transform
        (
            thrust::make_permutation_iterator
            (
                patchC.begin(),
                faces.begin()
            ),
            thrust::make_permutation_iterator
            (
                patchC.begin(),
                faces.end()
            ),
            thrust::make_permutation_iterator
            (
                pfld.begin(),
                faces.begin()
            ),
            MRFZoneCorrectVelocityFunctor(Omega,origin_)
        );
    }

    // Excluded patches
    forAll(excludedFaces_, patchi)
    {
        const vectorgpuField &patchC = mesh_.C().boundaryField()[patchi];

        vectorgpuField &pfld = U.boundaryFieldRef()[patchi];
        const labelgpuList &faces = excludedFaces_[patchi];

        thrust::transform
        (
            thrust::make_permutation_iterator
            (
                pfld.begin(),
                faces.begin()
            ),
            thrust::make_permutation_iterator
            (
                pfld.begin(),
                faces.end()
            ),
            thrust::make_transform_iterator
            (
                thrust::make_permutation_iterator
                (
                    patchC.begin(),
                    faces.begin()
                ),
                MRFZoneCorrectVelocityFunctor(Omega,origin_)
            ),
            thrust::make_permutation_iterator
            (
                U.boundaryFieldRef()[patchi].begin(),
                faces.begin()
            ),
            addOperatorFunctor<vector,vector,vector>()
        );
    }
}

void Foam::gpuMRFZone::makeAbsolute(surfaceScalargpuField &phi) const
{
    makeAbsoluteRhoFlux(geometricOneField(), phi);
}

void Foam::gpuMRFZone::makeAbsolute(
    const surfaceScalargpuField &rho,
    surfaceScalargpuField &phi) const
{
    makeAbsoluteRhoFlux(rho, phi);
}

void Foam::gpuMRFZone::correctBoundaryVelocity(volVectorgpuField &U) const
{
    if (!active_)
    {
        return;
    }

    const vector Omega = this->Omega();

    // Included patches
    volVectorgpuField::Boundary &Ubf = U.boundaryFieldRef();

    forAll(includedFaces_, patchi)
    {
        const vectorgpuField &patchC = mesh_.Cf().boundaryField()[patchi];

        vectorgpuField pfld(U.boundaryFieldRef()[patchi]);
        const labelgpuList &faces = includedFaces_[patchi];

        thrust::transform(
            thrust::make_permutation_iterator(
                patchC.begin(),
                faces.begin()),
            thrust::make_permutation_iterator(
                patchC.begin(),
                faces.end()),
            thrust::make_permutation_iterator(
                pfld.begin(),
                faces.begin()),
            MRFZoneCorrectVelocityFunctor(Omega, origin_));

        Ubf[patchi] == pfld;
    }
}

void Foam::gpuMRFZone::writeData(Ostream &os) const
{
    os << nl;
    os.beginBlock(name_);

    os.writeEntry("active", active_);
    os.writeEntry("cellZone", cellZoneName_);
    os.writeEntry("origin", origin_);
    os.writeEntry("axis", axis_);
    omega_->writeData(os);

    if (excludedPatchNames_.size())
    {
        os.writeEntry("nonRotatingPatches", excludedPatchNames_);
    }

    os.endBlock();
}

bool Foam::gpuMRFZone::read(const dictionary &dict)
{
    coeffs_ = dict;

    coeffs_.readIfPresent("active", active_);

    if (!active_)
    {
        cellZoneID_ = -1;
        return true;
    }

    coeffs_.readIfPresent("nonRotatingPatches", excludedPatchNames_);

    origin_ = coeffs_.get<vector>("origin");
    axis_ = coeffs_.get<vector>("axis").normalise();
    omega_.reset(Function1<scalar>::New("omega", coeffs_, &mesh_.hostmesh()));

    const word oldCellZoneName = cellZoneName_;
    if (cellZoneName_ == word::null)
    {
        coeffs_.readEntry("cellZone", cellZoneName_);
    }
    else
    {
        coeffs_.readIfPresent("cellZone", cellZoneName_);
    }

    if (cellZoneID_ == -1 || oldCellZoneName != cellZoneName_)
    {
        cellZoneID_ = mesh_.hostmesh().cellZones().findZoneID(cellZoneName_);

        const labelHashSet excludedPatchSet(
            mesh_.hostmesh().boundaryMesh().patchSet(excludedPatchNames_));

        excludedPatchLabels_.setSize(excludedPatchSet.size());

        label i = 0;
        for (const label patchi : excludedPatchSet)
        {
            excludedPatchLabels_[i++] = patchi;
        }

        bool cellZoneFound = (cellZoneID_ != -1);

        reduce(cellZoneFound, orOp<bool>());

        if (!cellZoneFound)
        {
            FatalErrorInFunction
                << "cannot find MRF cellZone " << cellZoneName_
                << exit(FatalError);
        }

        setMRFFaces();
    }

    return true;
}


void Foam::gpuMRFZone::update()
{
    if (mesh_.hostmesh().topoChanging())
    {
        setMRFFaces();
    }
}

// ************************************************************************* //
