/*---------------------------------------------------------------------------*\
  =========                 |
  \\      /  F ield         | OpenFOAM: The Open Source CFD Toolbox
   \\    /   O peration     | faSavageHutterFOAM
    \\  /    A nd           | Copyright (C) 2017 Matthias Rauter
     \\/     M anipulation  |
-------------------------------------------------------------------------------
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/>.

Application
    faSavageHutterFoam

Description
    A depth-integrated solver for shallow granular flows.
    The solver is based on the Finite Area Method.
    Model derivation and description:
    Rauter and Tukovic (submitted to Computer & Fluids):
    A finite area scheme for shallow granular flows on three-dimensional
    surfaces

Author
    Matthias Rauter matthias.rauter@uibk.ac.at

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

#include "fvCFD.H"
#include "faCFD.H"
#include "constraints.H"
#include "frictionModel.H"
#include "entrainmentModel.H"
#include "depositionModel.H"

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

int main(int argc, char *argv[])
{
#   include "setRootCase.H"
#   include "createTime.H"
#   include "createMesh.H"
#   include "createFaMesh.H"
#   include "readGravitationalAcceleration.H"
#   include "createFaFields.H"
#   include "readTransportProperties.H"
#   include "createFvFields.H"
#   include "createTimeControls.H"

    Info << "\nStarting time loop\n" << endl;

#       include "readSolutionControls.H"

        Info << endl
             << "Numerical settings" << endl
             << "    max number of iterations " << nCorr << endl
             << "    min number of iterations " << minCorr << endl
             << "    TOL h " << hResidualMax << endl
             << "    TOL Us " << UsResidualMax << endl << endl;

    if (initDeltaT)
    {
        Info << "Initializing Delta T" << endl;
#       include "readTimeControls.H"
#       include "surfaceCourantNo.H"
        runTime.setDeltaT(
            min(maxCo/(CoNum + SMALL)*runTime.deltaT().value(), maxDeltaT)
        );
    }

    bool final = false;

    while (runTime.run())
    {
#       include "readSolutionControls.H"
#       include "readTimeControls.H"
#       include "surfaceCourantNo.H"
#       include "setDeltaT.H"

        runTime++;

        Info << "Time = " << runTime.timeName() << nl << endl;

        for (int iCorr = 0; iCorr < nCorr; iCorr++)
        {

//            phi2s = (fac::interpolate(h)*fac::interpolate(Us) & aMesh.Le());

#           include "calcBasalstress.H"

            const areaVectorField & tauSc = friction->tauSc();
            const areaScalarField & tauSp = friction->tauSp();

            faVectorMatrix UsEqn
            (
                fam::ddt(h, Us)
              + xi*fam::div(phi2s, Us)
              + tauSc
              + fam::Sp(tauSp, Us)
             ==
                gs*h
              - fac::grad(pb*h/(2.*friction->rho()))
              //- gn*h*entrainment->gradhentrain()
            );

            if(explicitDryAreas)
            {
                scalarField &diag = UsEqn.diag();
                vectorField &source = UsEqn.source();
                scalarField &upper = UsEqn.upper();
                scalarField &lower = UsEqn.lower();

                vector forcedValue(0,0,0);
                label dryCellCount = 0;

                forAll(diag, i)
                {

                    const label startFaceOwn = aMesh.lduAddr().ownerStartAddr()[i];
                    const label endFaceOwn = aMesh.lduAddr().ownerStartAddr()[i + 1];

                    const label startFaceNbr = aMesh.lduAddr().losortStartAddr()[i];
                    const label endFaceNbr = aMesh.lduAddr().losortStartAddr()[i + 1];

                    const unallocLabelList& owner = aMesh.lduAddr().lowerAddr();
                    const unallocLabelList& neighbour = aMesh.lduAddr().upperAddr();
                    const unallocLabelList& losort = aMesh.lduAddr().losortAddr();


                    if (diag[i] < 1e-8)
                    {
                        UsEqn.eliminatedEqns().insert(i);

                        diag[i] = 1;
                        source[i] = forcedValue;

                        label faceIndex = startFaceOwn;
                        while (faceIndex < endFaceOwn)
                        {
                            source[neighbour[faceIndex]] -= lower[faceIndex]*forcedValue;
                            lower[faceIndex] = 0.0;
                            faceIndex++;
                        }

                        faceIndex = startFaceNbr;
                        while (faceIndex < endFaceNbr)
                        {
                            source[owner[losort[faceIndex]]] -= upper[losort[faceIndex]]*forcedValue;
                            upper[losort[faceIndex]] = 0.0;
                            faceIndex++;
                        }

                        dryCellCount++;
                    }
                }
                if (dryCellCount  > 0)
                {
                    Info << "Number of dry cells = " << dryCellCount << endl;
                }
            }

            if (!final)
                UsEqn.relax();

            lduSolverPerformance UsResidual = solve(UsEqn);

            tau = (tauSc + tauSp*Us)*friction->rho();
            phis = (fac::interpolate(Us) & aMesh.Le());

            const areaScalarField & Sm = entrainment->Sm();
            const areaScalarField & Sd = deposition->Sd();


            h.storePrevIter();

            faScalarMatrix hEqn
            (
                fam::ddt(h)
              + fam::div(phis, h)
             ==
                Sm
              - fam::Sp
                (
                    Sd/(h + dimensionedScalar("small", dimLength, SMALL)),
                    h
                )
            );

            //hEqn.relax();
            lduSolverPerformance hResidual = hEqn.solve();


            phi2s = hEqn.flux();

            if (!final)
                h.relax();


            // Bind h
            if (bindHeight)
            {
                h.internalField() = max
                (
                    h.internalField(),
                    hmin.value()
                );
            }

            faScalarMatrix hentrainEqn
            (
                fam::ddt(hentrain)
              ==
                Sd
                - fam::Sp
                (
                    Sm/(hentrain + dimensionedScalar("small", dimLength, SMALL)),
                    hentrain
                )
            );

            hentrainEqn.solve();


            //h.correctBoundaryConditions();
            //Us.correctBoundaryConditions();

            if (final)
            {
                Info << "reached residual in h = "
                     << hResidual.initialResidual()
                     << " < " << hResidualMax
                     << " and in Us = "
                     << UsResidual.initialResidual()
                     << " < " << UsResidualMax
                     << ", stopping loop!" << endl;
                final = false;
                break;
            }
            if (hResidual.initialResidual() < hResidualMax &&  UsResidual.initialResidual() < UsResidualMax &&  iCorr >= minCorr-1)
            {
                final = true;
            }



        }

        if (runTime.outputTime())
        {
            vsm.mapToVolume(Us, U.boundaryField());
            vsm.mapToVolume(h, H.boundaryField());
            vsm.mapToVolume(pb, Pb.boundaryField());
            vsm.mapToVolume(tau, Tau.boundaryField());
            vsm.mapToVolume(hentrain, Hentrain.boundaryField());

            runTime.write();
        }


        Info << "ExecutionTime = "
            << scalar(runTime.elapsedCpuTime())
            << " s\n" << endl << endl;
    }

    return(0);
}

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