// -*- mode: C++; tab-width: 4; indent-tabs-mode: nil; c-basic-offset: 4 -*-
// vi: set et ts=4 sw=4 sts=4:
/*
  This file is part of the Open Porous Media project (OPM).

  OPM 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 2 of the License, or
  (at your option) any later version.

  OPM 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 OPM.  If not, see <http://www.gnu.org/licenses/>.

  Consult the COPYING file in the top-level source directory of this
  module for the precise wording of the license and the list of
  copyright holders.
*/
/**
 * \file
 *
 * \copydoc Opm::TemperatureModel
 */
#ifndef OPM_TEMPERATURE_MODEL_HPP
#define OPM_TEMPERATURE_MODEL_HPP

#include <opm/common/OpmLog/OpmLog.hpp>

#include <opm/models/parallel/gridcommhandles.hh>
#include <opm/models/utils/propertysystem.hh>

#include <opm/simulators/flow/BlackoilModelParameters.hpp>
#include <opm/simulators/flow/GenericTemperatureModel.hpp>
#include <opm/simulators/linalg/findOverlapRowsAndColumns.hpp>
#include <opm/simulators/aquifers/AquiferGridUtils.hpp>
#include <opm/models/discretization/common/tpfalinearizer.hh>
#include <opm/simulators/linalg/istlsparsematrixadapter.hh>

#include <algorithm>
#include <cassert>
#include <cmath>
#include <cstddef>
#include <string>
#include <vector>

namespace Opm::Properties {

template<class TypeTag, class MyTypeTag>
struct EnableTemperatureModel {
    using type = UndefinedProperty;
};

} // namespace Opm::Properties

namespace Opm {

template<typename Scalar, typename IndexTraits> class WellState;


template <class TypeTag>
class BlackOilEnergyIntensiveQuantitiesTemp
{
    using FluidSystem = GetPropType<TypeTag, Properties::FluidSystem>;
    using Scalar = GetPropType<TypeTag, Properties::Scalar>;
    using SolidEnergyLaw = GetPropType<TypeTag, Properties::SolidEnergyLaw>;
    using ThermalConductionLaw = GetPropType<TypeTag, Properties::ThermalConductionLaw>;
    using Indices = GetPropType<TypeTag, Properties::Indices>;
    using Problem = GetPropType<TypeTag, Properties::Problem>;



    enum { enableBrine = getPropValue<TypeTag, Properties::EnableBrine>() };
    enum { enableVapwat = getPropValue<TypeTag, Properties::EnableVapwat>() };
    enum { enableDisgasInWater = getPropValue<TypeTag, Properties::EnableDisgasInWater>() };
    enum { enableSaltPrecipitation = getPropValue<TypeTag, Properties::EnableSaltPrecipitation>() };
    enum { enableSolvent = getPropValue<TypeTag, Properties::EnableSolvent>() };
    enum { numPhases = getPropValue<TypeTag, Properties::NumPhases>() };
    static constexpr bool compositionSwitchEnabled = Indices::compositionSwitchIdx >= 0;

    public:
    using EvaluationTemp = DenseAd::Evaluation<Scalar, 1>;
    using FluidStateTemp = BlackOilFluidState<EvaluationTemp,
                                              FluidSystem,
                                              true, //store temperature
                                              true, // store enthalpy
                                              compositionSwitchEnabled,
                                              enableVapwat,
                                              enableBrine,
                                              enableSaltPrecipitation,
                                              enableDisgasInWater,
                                              enableSolvent,
                                              Indices::numPhases>;



    void updateTemperature_(const Problem& problem, unsigned globalDofIdx, unsigned timeIdx)
    {
        const EvaluationTemp T = EvaluationTemp::createVariable(problem.temperature(globalDofIdx, timeIdx), 0);
        fluidState_.setTemperature(T);
    }

    void updateEnergyQuantities_(const Problem& problem,
                                 const unsigned globalSpaceIdx,
                                 const unsigned timeIdx)
    {
        // compute the specific enthalpy of the fluids, the specific enthalpy of the rock
        // and the thermal conductivity coefficients
        for (int phaseIdx = 0; phaseIdx < numPhases; ++ phaseIdx) {
            if (!FluidSystem::phaseIsActive(phaseIdx)) {
                continue;
            }

            const auto& h = FluidSystem::enthalpy(fluidState_, phaseIdx, problem.pvtRegionIndex(globalSpaceIdx));
            fluidState_.setEnthalpy(phaseIdx, h);
        }

        const auto& solidEnergyLawParams = problem.solidEnergyLawParams(globalSpaceIdx, timeIdx);
        rockInternalEnergy_ = SolidEnergyLaw::solidInternalEnergy(solidEnergyLawParams, fluidState_);

        const auto& thermalConductionLawParams = problem.thermalConductionLawParams(globalSpaceIdx, timeIdx);
        totalThermalConductivity_ = ThermalConductionLaw::thermalConductivity(thermalConductionLawParams, fluidState_);

        // Retrieve the rock fraction from the problem
        // Usually 1 - porosity, but if pvmult is used to modify porosity
        // we will apply the same multiplier to the rock fraction
        // i.e. pvmult*(1 - porosity) and thus interpret multpv as a volume
        // multiplier. This is to avoid negative rock volume for pvmult*porosity > 1
        rockFraction_ = problem.rockFraction(globalSpaceIdx, timeIdx);
    }

    const EvaluationTemp& rockInternalEnergy() const
    { return rockInternalEnergy_; }

    const EvaluationTemp& totalThermalConductivity() const
    { return totalThermalConductivity_; }

    const Scalar& rockFraction() const
    { return rockFraction_; }

    const FluidStateTemp& fluidStateTemp() const
    { return fluidState_; }

    template <class FluidState>
    void setFluidState(const FluidState& fs)
    {
        // copy the needed part of the fluid state
        fluidState_.setPvtRegionIndex(fs.pvtRegionIndex());
        fluidState_.setRs(getValue(fs.Rs()));
        fluidState_.setRv(getValue(fs.Rv()));
        for (int phaseIdx = 0; phaseIdx < numPhases; ++ phaseIdx) {
            fluidState_.setPressure(phaseIdx, getValue(fs.pressure(phaseIdx)));
            fluidState_.setDensity(phaseIdx, getValue(fs.density(phaseIdx)));
            fluidState_.setSaturation(phaseIdx, getValue(fs.saturation(phaseIdx)));
            fluidState_.setInvB(phaseIdx, getValue(fs.invB(phaseIdx)));
        }
    }

protected:
    EvaluationTemp rockInternalEnergy_;
    EvaluationTemp totalThermalConductivity_;
    FluidStateTemp fluidState_;
    Scalar rockFraction_;
};

/*!
 * \ingroup BlackOilSimulator
 *
 * \brief A class which handles sequential implicit solution of the energy equation as specified in by TEMP
 */
template <class TypeTag, bool enableTempV = getPropValue<TypeTag, Properties::EnergyModuleType>() == EnergyModules::SequentialImplicitThermal >
class TemperatureModel : public GenericTemperatureModel<GetPropType<TypeTag, Properties::Grid>,
                                              GetPropType<TypeTag, Properties::GridView>,
                                              GetPropType<TypeTag, Properties::DofMapper>,
                                              GetPropType<TypeTag, Properties::Stencil>,
                                              GetPropType<TypeTag, Properties::FluidSystem>,
                                              GetPropType<TypeTag, Properties::Scalar>>
{
    using BaseType = GenericTemperatureModel<GetPropType<TypeTag, Properties::Grid>,
                                        GetPropType<TypeTag, Properties::GridView>,
                                        GetPropType<TypeTag, Properties::DofMapper>,
                                        GetPropType<TypeTag, Properties::Stencil>,
                                        GetPropType<TypeTag, Properties::FluidSystem>,
                                        GetPropType<TypeTag, Properties::Scalar>>;
    using Simulator = GetPropType<TypeTag, Properties::Simulator>;
    using GridView = GetPropType<TypeTag, Properties::GridView>;
    using Grid = GetPropType<TypeTag, Properties::Grid>;
    using Scalar = GetPropType<TypeTag, Properties::Scalar>;
    using Stencil = GetPropType<TypeTag, Properties::Stencil>;
    using FluidSystem = GetPropType<TypeTag, Properties::FluidSystem>;
    using ElementContext = GetPropType<TypeTag, Properties::ElementContext>;
    using Indices = GetPropType<TypeTag, Properties::Indices>;
    using LocalResidual = GetPropType<TypeTag, Properties::LocalResidual>;
    using IntensiveQuantities = GetPropType<TypeTag, Properties::IntensiveQuantities>;
    using EnergyModule = BlackOilEnergyModule<TypeTag>;
    using IndexTraits = typename FluidSystem::IndexTraitsType;
    using WellStateType = WellState<Scalar, IndexTraits>;

    using IntensiveQuantitiesTemp = BlackOilEnergyIntensiveQuantitiesTemp<TypeTag>;
    using FluidStateTemp = typename IntensiveQuantitiesTemp::FluidStateTemp;
    using Evaluation = typename IntensiveQuantitiesTemp::EvaluationTemp;
    using EnergyMatrix = typename BaseType::EnergyMatrix;
    using EnergyVector = typename BaseType::EnergyVector;
    using MatrixBlockTemp = typename BaseType::MatrixBlockTemp;
    using SpareMatrixEnergyAdapter = typename Linear::IstlSparseMatrixAdapter<MatrixBlockTemp>;


    //using ResidualNBInfo = typename LocalResidual::ResidualNBInfo;
    struct ResidualNBInfo
    {
        double faceArea;
        double inAlpha;
        double outAlpha;
    };

    using NeighborInfoCPU = NeighborInfoStruct<ResidualNBInfo, MatrixBlockTemp>;

    enum { numEq = getPropValue<TypeTag, Properties::NumEq>() };
    enum { numPhases = FluidSystem::numPhases };
    enum { waterPhaseIdx = FluidSystem::waterPhaseIdx };
    enum { oilPhaseIdx = FluidSystem::oilPhaseIdx };
    enum { gasPhaseIdx = FluidSystem::gasPhaseIdx };
    static constexpr int temperatureIdx = 0;

public:
    explicit TemperatureModel(Simulator& simulator)
        : BaseType(simulator.vanguard().gridView(),
                   simulator.vanguard().eclState(),
                   simulator.vanguard().cartesianIndexMapper(),
                   simulator.model().dofMapper())
        , simulator_(simulator)
    {}

    void init()
    {
        const unsigned int numCells = simulator_.model().numTotalDof();
        this->doInit(numCells);

        if (!this->doTemp())
            return;

        // we need the storage term at start of the iteration (timeIdx = 1)
        storage1_.resize(numCells);

        // set the initial temperature
        for (unsigned globI = 0; globI < numCells; ++globI) {
            this->temperature_[globI] = simulator_.problem().initialFluidState(globI).temperature(0);
        }
        // keep a copy of the intensive quantities to simplify the update during
        // the newton iterations
        intQuants_.resize(numCells);

        // find and store the overlap cells
        const auto& elemMapper = simulator_.model().elementMapper();
        detail::findOverlapAndInterior(simulator_.vanguard().grid(), elemMapper, overlapRows_, interiorRows_);

        // set the scaling factor
        scalingFactor_ = getPropValue<TypeTag, Properties::BlackOilEnergyScalingFactor>();

        // find the sparsity pattern of the temperature matrix
        using NeighborSet = std::set<unsigned>;
        std::vector<NeighborSet> neighbors(numCells);
        Stencil stencil(this->gridView_, this->dofMapper_);
        neighborInfo_.reserve(numCells, 6 * numCells);
        std::vector<NeighborInfoCPU> loc_nbinfo;
        for (const auto& elem : elements(this->gridView_)) {
            stencil.update(elem);
            for (unsigned primaryDofIdx = 0; primaryDofIdx < stencil.numPrimaryDof(); ++primaryDofIdx) {
                const unsigned myIdx = stencil.globalSpaceIndex(primaryDofIdx);
                loc_nbinfo.resize(stencil.numDof() - 1); // Do not include the primary dof in neighborInfo_
                for (unsigned dofIdx = 0; dofIdx < stencil.numDof(); ++dofIdx) {
                    const unsigned neighborIdx = stencil.globalSpaceIndex(dofIdx);
                    neighbors[myIdx].insert(neighborIdx);
                    if (dofIdx > 0) {
                        const auto scvfIdx = dofIdx - 1;
                        const auto& scvf = stencil.interiorFace(scvfIdx);
                        const Scalar area = scvf.area();
                        Scalar inAlpha = simulator_.problem().thermalHalfTransmissibility(myIdx, neighborIdx);
                        Scalar outAlpha = simulator_.problem().thermalHalfTransmissibility(neighborIdx, myIdx);
                        ResidualNBInfo nbinfo{area, inAlpha, outAlpha};
                        loc_nbinfo[dofIdx - 1] = NeighborInfoCPU{neighborIdx, nbinfo, nullptr};
                    }
                }
                neighborInfo_.appendRow(loc_nbinfo.begin(), loc_nbinfo.end());
            }
        }
        // allocate matrix for storing the Jacobian of the temperature residual
        energyMatrix_ = std::make_unique<SpareMatrixEnergyAdapter>(simulator_);
        diagMatAddress_.resize(numCells);
        energyMatrix_->reserve(neighbors);
        for (unsigned globI = 0; globI < numCells; globI++) {
            const auto& nbInfos = neighborInfo_[globI];
            diagMatAddress_[globI] = energyMatrix_->blockAddress(globI, globI);
            for (auto& nbInfo : nbInfos) {
                nbInfo.matBlockAddress = energyMatrix_->blockAddress(nbInfo.neighbor, globI);
            }
        }
    }

    void beginTimeStep()
    {
        if (!this->doTemp()) {
            return;
        }

        // We use the specialized intensive quantities here with only the temperature derivative
        const unsigned int numCells = simulator_.model().numTotalDof();
        #ifdef _OPENMP
        #pragma omp parallel for
        #endif
        for (unsigned globI = 0; globI < numCells; ++globI) {
            intQuants_[globI].setFluidState(simulator_.model().intensiveQuantities(globI, /*timeIdx*/ 0).fluidState());
            intQuants_[globI].updateTemperature_(simulator_.problem(), globI, /*timeIdx*/ 0);
            intQuants_[globI].updateEnergyQuantities_(simulator_.problem(), globI, /*timeIdx*/ 0);
        }
        updateStorageCache();

        const int nw = simulator_.problem().wellModel().wellState().numWells();
        this->energy_rates_.resize(nw, 0.0);
    }

    /*!
     * \brief Informs the temperature model that a time step has just been finished.
     */
    void endTimeStep(WellStateType& wellState)
    {
        if (!this->doTemp()) {
            return;
        }

        // We use the specialized intensive quantities here with only the temperature derivative
        const unsigned int numCells = simulator_.model().numTotalDof();
        #ifdef _OPENMP
        #pragma omp parallel for
        #endif
        for (unsigned globI = 0; globI < numCells; ++globI) {
            intQuants_[globI].setFluidState(simulator_.model().intensiveQuantities(globI, /*timeIdx*/ 0).fluidState());
            intQuants_[globI].updateTemperature_(simulator_.problem(), globI, /*timeIdx*/ 0);
            intQuants_[globI].updateEnergyQuantities_(simulator_.problem(), globI, /*timeIdx*/ 0);
        }
        advanceTemperatureFields();

        // update energy_rates
        const int nw = wellState.numWells();
        for (auto wellID = 0*nw; wellID < nw; ++wellID) {
            auto& ws = wellState.well(wellID);
            ws.energy_rate = this->energy_rates_[wellID];
        }

        // Update well temperature
        const auto& wellPtrs = simulator_.problem().wellModel().localNonshutWells();
        for (const auto& wellPtr : wellPtrs) {
            auto& ws = wellState.well(wellPtr->name());
            this->computeWellTemperature(*wellPtr, ws);
        }
    }

    /*!
     * \brief This method writes the complete state of all temperature
     *        to the hard disk.
     */
    template <class Restarter>
    void serialize(Restarter&)
    { /* not implemented */ }

    /*!
     * \brief This method restores the complete state of the temperature
     *        from disk.
     *
     * It is the inverse of the serialize() method.
     */
    template <class Restarter>
    void deserialize(Restarter&)
    { /* not implemented */ }

protected:
    void updateStorageCache()
    {
        // we need the storage term at start of the iteration (timeIdx = 1)
        const unsigned int numCells = simulator_.model().numTotalDof();
        #ifdef _OPENMP
        #pragma omp parallel for
        #endif
        for (unsigned globI = 0; globI < numCells; ++globI) {
            Scalar storage = 0.0;
            computeStorageTerm(globI, storage);
            storage1_[globI] = storage;
        }
    }

    void advanceTemperatureFields()
    {
        const int max_iter = 20;
        const int min_iter = 1;
        bool is_converged = false;
        // solve using Newton
        for (int iter = 0; iter < max_iter; ++iter) {
            assembleEquations();
            if (iter >= min_iter && converged(iter)) {
                is_converged = true;
                break;
            }
            solveAndUpdate();
        }
        if (!is_converged) {
            const auto msg =
                fmt::format(fmt::runtime("Temperature model (TEMP): Newton did not converge after {} iterations. \n"
                                         "The Simulator will continue to the next step with an unconverged solution."),
                            max_iter);
            OpmLog::debug(msg);
        }
    }

    void solveAndUpdate()
    {
        const unsigned int numCells = simulator_.model().numTotalDof();
        EnergyVector dx(numCells);
        bool conv = this->linearSolve_(this->energyMatrix_->istlMatrix(), dx, this->energyVector_);
        if (!conv) {
            if (simulator_.gridView().comm().rank() == 0) {
                OpmLog::warning("Temp model: Linear solver did not converge. Temperature values not updated.");
            }
        }
        else {
            #ifdef _OPENMP
            #pragma omp parallel for
            #endif
            for (unsigned globI = 0; globI < numCells; ++globI) {
                this->temperature_[globI] -= std::clamp(dx[globI][0], -this->maxTempChange_, this->maxTempChange_);
                intQuants_[globI].updateTemperature_(simulator_.problem(), globI, /*timeIdx*/ 0);
                intQuants_[globI].updateEnergyQuantities_(simulator_.problem(), globI, /*timeIdx*/ 0);
            }
        }
    }

    bool converged(const int iter)
    {
        Scalar dt = simulator_.timeStepSize();
        Scalar maxNorm = 0.0;
        Scalar sumNorm = 0.0;
        const auto tolerance_cnv_energy_strict = Parameters::Get<Parameters::ToleranceCnvEnergy<Scalar>>();
        const auto& elemMapper = simulator_.model().elementMapper();
        const IsNumericalAquiferCell isNumericalAquiferCell(simulator_.gridView().grid());
        Scalar sum_pv = 0.0;
        Scalar sum_pv_not_converged = 0.0;
        for (const auto& elem : elements(simulator_.gridView(), Dune::Partitions::interior)) {
            unsigned globI = elemMapper.index(elem);
            const auto pvValue =  simulator_.problem().referencePorosity(globI, /*timeIdx=*/0)
            *  simulator_.model().dofTotalVolume(globI);

            const Scalar scaled_norm = dt * std::abs(this->energyVector_[globI])/ pvValue;
            maxNorm = max(maxNorm, scaled_norm);
            sumNorm += scaled_norm;
            if (!isNumericalAquiferCell(elem)) {
                if (scaled_norm > tolerance_cnv_energy_strict) {
                    sum_pv_not_converged += pvValue;
                }
                sum_pv += pvValue;
            }
        }
        maxNorm = simulator_.gridView().comm().max(maxNorm);
        sumNorm = simulator_.gridView().comm().sum(sumNorm);
        sum_pv = simulator_.gridView().comm().sum(sum_pv);
        sumNorm /= sum_pv;

        // Use relaxed tolerance if the fraction of unconverged cells porevolume is less than relaxed_max_pv_fraction
        sum_pv_not_converged = simulator_.gridView().comm().sum(sum_pv_not_converged);
        Scalar relaxed_max_pv_fraction = Parameters::Get<Parameters::RelaxedMaxPvFraction<Scalar>>();
        const bool relax = (sum_pv_not_converged / sum_pv) <  relaxed_max_pv_fraction;
        const auto tolerance_energy_balance = relax? Parameters::Get<Parameters::ToleranceEnergyBalanceRelaxed<Scalar>>():
                                            Parameters::Get<Parameters::ToleranceEnergyBalance<Scalar>>();
        const bool tolerance_cnv_energy = relax? Parameters::Get<Parameters::ToleranceCnvEnergyRelaxed<Scalar>>():
                                            tolerance_cnv_energy_strict;

        const auto msg = fmt::format(fmt::runtime("Temperature model (TEMP): Newton iter {}: "
                                     "CNV(E): {:.1e}, EB: {:.1e}"),
                                     iter, maxNorm, sumNorm);
        OpmLog::debug(msg);
        if (maxNorm < tolerance_cnv_energy && sumNorm < tolerance_energy_balance) {
            const auto msg2 =
                fmt::format(fmt::runtime("Temperature model (TEMP): Newton converged after {} iterations"),
                                         iter);
            OpmLog::debug(msg2);
            return true;
        }
        return false;
    }

    template< class LhsEval>
    void computeStorageTerm(unsigned globI, LhsEval& storage)
    {
        const auto& intQuants = intQuants_[globI];
        const auto& poro = getValue(simulator_.model().intensiveQuantities(globI, /*timeIdx*/ 0).porosity());
        // accumulate the internal energy of the fluids
        const auto& fs = intQuants.fluidStateTemp();
        for (unsigned phaseIdx = 0; phaseIdx < numPhases; ++ phaseIdx) {
            if (!FluidSystem::phaseIsActive(phaseIdx)) {
                continue;
            }

            const auto& u = decay<LhsEval>(fs.internalEnergy(phaseIdx));
            const auto& S = decay<LhsEval>(fs.saturation(phaseIdx));
            const auto& rho = decay<LhsEval>(fs.density(phaseIdx));

            storage += poro*S*u*rho;
        }

        // add the internal energy of the rock
        const Scalar rockFraction = intQuants.rockFraction();
        const auto& uRock = decay<LhsEval>(intQuants.rockInternalEnergy());
        storage += rockFraction*uRock;
        storage*= scalingFactor_;
    }

    template <class RateVector>
    void computeFluxTerm(const FluidStateTemp& fsIn,
                         const FluidStateTemp& fsEx,
                         const RateVector& darcyFlux,
                         Evaluation& flux)
    {
        for (unsigned phaseIdx = 0; phaseIdx < numPhases; ++ phaseIdx) {
            if (!FluidSystem::phaseIsActive(phaseIdx)) {
                continue;
            }

            const unsigned activeCompIdx =
                FluidSystem::canonicalToActiveCompIdx(FluidSystem::solventComponentIndex(phaseIdx));
            bool inIsUp = darcyFlux[activeCompIdx] > 0;
            const auto& fs = inIsUp ? fsIn : fsEx;
            if (inIsUp) {
                flux += fs.enthalpy(phaseIdx)
                    * fs.density(phaseIdx)
                    * getValue(darcyFlux[activeCompIdx]);
            }
            else {
                flux += getValue(fs.enthalpy(phaseIdx))
                    * getValue(fs.density(phaseIdx))
                    * getValue(darcyFlux[activeCompIdx]);
            }
        }
        flux *= scalingFactor_;
    }

    template < class ResidualNBInfo>
    void computeHeatFluxTerm(const IntensiveQuantitiesTemp& intQuantsIn,
                             const IntensiveQuantitiesTemp& intQuantsEx,
                             const ResidualNBInfo& res_nbinfo,
                             Evaluation& heatFlux)
    {
        short interiorDofIdx = 0; // NB
        short exteriorDofIdx = 1; // NB
        EnergyModule::ExtensiveQuantities::updateEnergy(heatFlux,
                                                        interiorDofIdx, // focusDofIndex,
                                                        interiorDofIdx,
                                                        exteriorDofIdx,
                                                        intQuantsIn,
                                                        intQuantsEx,
                                                        intQuantsIn.fluidStateTemp(),
                                                        intQuantsEx.fluidStateTemp(),
                                                        res_nbinfo.inAlpha,
                                                        res_nbinfo.outAlpha,
                                                        res_nbinfo.faceArea);
        heatFlux *= scalingFactor_*res_nbinfo.faceArea;
    }

    void assembleEquations()
    {
        const unsigned int numCells = simulator_.model().numTotalDof();
        for (unsigned globI = 0; globI < numCells; ++globI) {
            this->energyVector_[globI] = 0.0;
            energyMatrix_->clearRow(globI, 0.0);
        }
        Scalar dt = simulator_.timeStepSize();
        #ifdef _OPENMP
        #pragma omp parallel for
        #endif
        for (unsigned globI = 0; globI < numCells; ++globI) {
            MatrixBlockTemp bMat;
            Scalar volume = simulator_.model().dofTotalVolume(globI);
            Scalar storefac = volume / dt;
            Evaluation storage = 0.0;
            computeStorageTerm(globI, storage);
            this->energyVector_[globI] += storefac * ( getValue(storage) - storage1_[globI] );
            bMat[0][0] = storefac * storage.derivative(temperatureIdx);
            *diagMatAddress_[globI] += bMat;
        }

        const auto& floresInfo = this->simulator_.problem().model().linearizer().getFloresInfo();
        const bool enableDriftCompensation = Parameters::Get<Parameters::EnableDriftCompensationTemp>();
        const auto& problem = simulator_.problem();

        #ifdef _OPENMP
        #pragma omp parallel for
        #endif
        for (unsigned globI = 0; globI < numCells; ++globI) {
            const auto& nbInfos = neighborInfo_[globI];
            const auto& floresInfos = floresInfo[globI];
            int loc = 0;
            const auto& intQuantsIn = intQuants_[globI];
            MatrixBlockTemp bMat;
            for (const auto& nbInfo : nbInfos) {
                unsigned globJ = nbInfo.neighbor;
                const auto& intQuantsEx = intQuants_[globJ];
                assert(globJ != globI);
                const auto& darcyflux = floresInfos[loc].flow;
                // compute convective flux
                Evaluation flux = 0.0;
                computeFluxTerm(intQuantsIn.fluidStateTemp(), intQuantsEx.fluidStateTemp(), darcyflux, flux);
                // compute conductive flux
                Evaluation heatFlux = 0.0;
                computeHeatFluxTerm(intQuantsIn, intQuantsEx, nbInfo.res_nbinfo, heatFlux);
                heatFlux += flux;
                this->energyVector_[globI] += getValue(heatFlux);
                bMat[0][0] = heatFlux.derivative(temperatureIdx);
                *diagMatAddress_[globI] += bMat;
                bMat *= -1.0;
                //SparseAdapter syntax: jacobian_->addToBlock(globJ, globI, bMat);
                *nbInfo.matBlockAddress += bMat;
                loc++;
            }

            if (enableDriftCompensation) {
                auto dofDriftRate = problem.drift()[globI]/dt;
                const auto& fs = intQuantsIn.fluidStateTemp();
                for (unsigned phaseIdx = 0; phaseIdx < numPhases; ++ phaseIdx) {
                   const unsigned activeCompIdx =
                        FluidSystem::canonicalToActiveCompIdx(FluidSystem::solventComponentIndex(phaseIdx));
                   auto drift_hrate = dofDriftRate[activeCompIdx]*getValue(fs.enthalpy(phaseIdx)) * getValue(fs.density(phaseIdx)) /  getValue(fs.invB(phaseIdx));
                   this->energyVector_[globI] -= drift_hrate*scalingFactor_;
                }
            }
        }

        // Well terms
        const auto& wellPtrs = simulator_.problem().wellModel().localNonshutWells();
        for (const auto& wellPtr : wellPtrs) {
            this->assembleEquationWell(*wellPtr);
        }

        // For standard ilu0 we need to set the overlapping cells to identity. But since we use "paroverilu0"
        // here this is not needed.
        //       if (simulator_.gridView().comm().size() > 1) {
        //           // Set dirichlet conditions for overlapping cells
        //            // loop over precalculated overlap rows and columns
        //            for (const auto row : overlapRows_) {
        //                // Zero out row.
        //               (*this->energyMatrix_)[row] = 0.0;
        //
        //        //diagonal block set to diag(1.0).
        //        (*this->energyMatrix_)[row][row][0][0] = 1.0;
        //    }
        //}
    }

    template<class Well>
    void assembleEquationWell(const Well& well)
    {
        const auto& eclWell = well.wellEcl();
        std::size_t well_index = simulator_.problem().wellModel().wellState().index(well.name()).value();
        const auto& ws = simulator_.problem().wellModel().wellState().well(well_index);
        this->energy_rates_[well_index] = 0.0;
        MatrixBlockTemp bMat;
        for (std::size_t i = 0; i < ws.perf_data.size(); ++i) {
            const auto globI = ws.perf_data.cell_index[i];
            auto fs = intQuants_[globI].fluidStateTemp(); //copy to make it possible to change the temp in the injector
            for (unsigned phaseIdx = 0; phaseIdx < numPhases; ++ phaseIdx) {
                if (!FluidSystem::phaseIsActive(phaseIdx)) {
                    continue;
                }

                Evaluation rate = well.volumetricSurfaceRateForConnection(globI, phaseIdx);
                if (rate > 0 && eclWell.isInjector()) {
                    fs.setTemperature(eclWell.inj_temperature());
                    const auto& rho = FluidSystem::density(fs, phaseIdx, fs.pvtRegionIndex());
                    fs.setDensity(phaseIdx, rho);
                    const auto& h = FluidSystem::enthalpy(fs, phaseIdx, fs.pvtRegionIndex());
                    fs.setEnthalpy(phaseIdx, h);
                    rate *= getValue(fs.enthalpy(phaseIdx)) * getValue(fs.density(phaseIdx)) / getValue(fs.invB(phaseIdx));
                } else {
                    const Evaluation d =  1.0 - fs.Rv() * fs.Rs();
                    if (phaseIdx == gasPhaseIdx && d > 0) {
                        const auto& oilrate = well.volumetricSurfaceRateForConnection(globI, oilPhaseIdx);
                        rate -= oilrate * getValue(fs.Rs());
                        rate /= d;
                    }
                    if (phaseIdx == oilPhaseIdx && d > 0) {
                        const auto& gasrate = well.volumetricSurfaceRateForConnection(globI, gasPhaseIdx);
                        rate -= gasrate * getValue(fs.Rv());
                        rate /= d;
                    }
                    rate *= fs.enthalpy(phaseIdx) * getValue(fs.density(phaseIdx)) / getValue(fs.invB(phaseIdx));
                }
                this->energy_rates_[well_index] += getValue(rate);
                rate *= scalingFactor_;
                this->energyVector_[globI] -= getValue(rate);
                bMat[0][0] = -rate.derivative(temperatureIdx);
                *diagMatAddress_[globI] += bMat;
            }
        }
    }

    template<class Well, class SingleWellState>
    void computeWellTemperature(const Well& well, SingleWellState& ws)
    {
        if (well.isInjector()) {
            if (ws.status != WellStatus::STOP) {
                ws.temperature = well.wellEcl().inj_temperature();
                return;
            }
        }
        const int np = simulator_.problem().wellModel().wellState().numPhases();
        std::array<Scalar,2> weighted{0.0,0.0};
        auto& [weighted_temperature, total_weight] = weighted;
        for (std::size_t i = 0; i < ws.perf_data.size(); ++i) {
            const auto globI = ws.perf_data.cell_index[i];
            const auto& fs = intQuants_[globI].fluidStateTemp();
            Scalar weight_factor = simulator_.problem().wellModel().computeTemperatureWeightFactor(i, np, fs, ws);
            total_weight += weight_factor;
            weighted_temperature += weight_factor * fs.temperature(/*phaseIdx*/0).value();
        }
        //simulator_.gridView().comm().sum(weighted.data(), 2);
        ws.temperature = weighted_temperature / total_weight;
    }

    const Simulator& simulator_;
    EnergyVector storage1_;
    std::vector<IntensiveQuantitiesTemp> intQuants_;
    SparseTable<NeighborInfoCPU> neighborInfo_{};
    std::vector<MatrixBlockTemp*> diagMatAddress_{};
    std::unique_ptr<SpareMatrixEnergyAdapter> energyMatrix_;
    std::vector<int> overlapRows_;
    std::vector<int> interiorRows_;
    Scalar scalingFactor_{1.0};
};

// need for the old linearizer
template <class TypeTag>
class TemperatureModel<TypeTag, false>
{
    using Simulator = GetPropType<TypeTag, Properties::Simulator>;
    using Scalar = GetPropType<TypeTag, Properties::Scalar>;
    using FluidSystem = GetPropType<TypeTag, Properties::FluidSystem>;

public:
    TemperatureModel(Simulator&)
    { }

    /*!
     * \brief This method writes the complete state of all temperature
     *        to the hard disk.
     */
    template <class Restarter>
    void serialize(Restarter&)
    { /* not implemented */ }

    /*!
     * \brief This method restores the complete state of the temperature
     *        from disk.
     *
     * It is the inverse of the serialize() method.
     */
    template <class Restarter>
    void deserialize(Restarter&)
    { /* not implemented */ }

    void init() {}
    void beginTimeStep() {}
    const Scalar temperature(size_t /*globalIdx*/) const
    {
        return 273.15; // return 0C to make the compiler happy
    }
};

} // namespace Opm

#endif // OPM_TEMPERATURE_MODEL_HPP
