diff --git a/dumux/freeflow/staggeredni/fluxvariables.hh b/dumux/freeflow/staggeredni/fluxvariables.hh
index a594070a2837c78f6db114bb8cb5c63b875b4e0d..ac33aceab492309b4a0c4a293ab7f9a58e81274b 100644
--- a/dumux/freeflow/staggeredni/fluxvariables.hh
+++ b/dumux/freeflow/staggeredni/fluxvariables.hh
@@ -34,7 +34,7 @@ namespace Properties
 {
 // forward declaration
 //NEW_PROP_TAG(EnableComponentTransport);
-NEW_PROP_TAG(EnableEnergyBalance);
+NEW_PROP_TAG(EnableEnergyBalanceStokes);
 NEW_PROP_TAG(EnableInertiaTerms);
 }
 
@@ -93,10 +93,8 @@ public:
                                                         const FluxVariablesCache& fluxVarsCache)
     {
         CellCenterPrimaryVariables flux(0.0);
-
         flux += advectiveFluxForCellCenter_(problem, fvGeometry, elemVolVars, globalFaceVars, scvf);
-        flux += HeatConductionType::diffusiveFluxForCellCenter_(problem, fvGeometry, elemVolVars, scvf);
-        flux *= scvf.area() * sign(scvf.outerNormalScalar());
+        flux += HeatConductionType::diffusiveFluxForCellCenter(problem, element, fvGeometry, elemVolVars, scvf);
         return flux;
     }
 
@@ -119,7 +117,7 @@ private:
         if(scvf.boundary())
         {
             const auto bcTypes = problem.boundaryTypesAtPos(scvf.center());
-                if(bcTypes.isOutflow(momentumBalanceIdx)) // TODO ??
+                if(bcTypes.isOutflow(momentumBalanceIdx))
                     isOutflow = true;
         }
 
@@ -140,6 +138,8 @@ private:
         flux[energyBalanceIdx] = (upWindWeight * upstreamDensity * upstreamEnthalpy
                                  + (1.0 - upWindWeight) * downstreamDensity * downstreamEnthalpy)
                                  * velocity;
+
+        flux *= scvf.area() * sign(scvf.outerNormalScalar());
         return flux;
     }
 };
diff --git a/dumux/freeflow/staggeredni/indices.hh b/dumux/freeflow/staggeredni/indices.hh
index ca488bdbc3d46264acea63791b81405b42acd30c..3e2fc7d6919c036d2863d19da7db898864e5cb2c 100644
--- a/dumux/freeflow/staggeredni/indices.hh
+++ b/dumux/freeflow/staggeredni/indices.hh
@@ -40,8 +40,9 @@ template <class TypeTag, int PVOffset = 0>
 class NavierStokesNIIndices : public NavierStokesCommonIndices<TypeTag, PVOffset>
 {
 public:
-    static const int numEq = GET_PROP_VALUE(TypeTag, NumEq);
-    static constexpr int energyBalanceIdx = PVOffset + numEq - 1;
+    static const int numEqCC = GET_PROP_VALUE(TypeTag, NumEqCellCenter);
+
+    static constexpr int energyBalanceIdx = PVOffset + numEqCC -1;
     static constexpr int temperatureIdx = energyBalanceIdx;
 };
 } // end namespace
diff --git a/dumux/freeflow/staggeredni/localresidual.hh b/dumux/freeflow/staggeredni/localresidual.hh
index 33baa4c0a392db925e8521fb28f26bdacffda566..fb3276b461bcbede296898dc6ee2d3167be979da 100644
--- a/dumux/freeflow/staggeredni/localresidual.hh
+++ b/dumux/freeflow/staggeredni/localresidual.hh
@@ -26,12 +26,12 @@
 #ifndef DUMUX_STAGGERED_NAVIERSTOKES_NI_LOCAL_RESIDUAL_HH
 #define DUMUX_STAGGERED_NAVIERSTOKES_NI_LOCAL_RESIDUAL_HH
 
-#include <dune/istl/matrix.hh> // TODO ?
+#include <dune/istl/matrix.hh>
 
-#include <dumux/common/valgrind.hh> // TODO ?
+#include <dumux/common/valgrind.hh>
 #include <dumux/implicit/staggered/localresidual.hh>
 
-#include "properties.hh" // TODO ?
+#include "properties.hh"
 
 namespace Dumux
 {
@@ -39,17 +39,17 @@ namespace Dumux
 namespace Properties
 {
 // forward declaration
-//NEW_PROP_TAG(EnableComponentTransport);
-NEW_PROP_TAG(EnableEnergyBalance);
+NEW_PROP_TAG(EnableComponentTransport);
+NEW_PROP_TAG(EnableEnergyBalanceStokes);
 NEW_PROP_TAG(EnableInertiaTerms);
-//NEW_PROP_TAG(ReplaceCompEqIdx);
+NEW_PROP_TAG(ReplaceCompEqIdx);
 }
 
 /*!
- * \ingroup CCModel
+ * \ingroup StaggeredModel
  * \ingroup StaggeredLocalResidual
  * \brief Element-wise calculation of the residual for models
- * 			based on the fully implicit cell-centered scheme
+ * 			based on the staggered grid scheme
  *
  * \todo Please doc me more!
  */
@@ -58,24 +58,17 @@ NEW_PROP_TAG(EnableInertiaTerms);
 template<class TypeTag, bool enableComponentTransport, bool enableEnergyBalance>
 class StaggeredNavierStokesResidualImpl;
 
-template<class TypeTag>
-using StaggeredNavierStokesResidual = StaggeredNavierStokesResidualImpl<TypeTag, GET_PROP_VALUE(TypeTag, EnableComponentTransport),
-                                                                 GET_PROP_VALUE(TypeTag, EnableEnergyBalance)>;
-
-
 // specialization for immiscible, non-isothermal flow
 template<class TypeTag>
 class StaggeredNavierStokesResidualImpl<TypeTag, false, true> : public StaggeredNavierStokesResidualImpl<TypeTag, false, false>
 {
     friend class StaggeredLocalResidual<TypeTag>;
-
     using ParentType = StaggeredNavierStokesResidualImpl<TypeTag, false, false>;
 
     using Scalar = typename GET_PROP_TYPE(TypeTag, Scalar);
     using SubControlVolume = typename GET_PROP_TYPE(TypeTag, SubControlVolume);
     using CellCenterPrimaryVariables = typename GET_PROP_TYPE(TypeTag, CellCenterPrimaryVariables);
     using Indices = typename GET_PROP_TYPE(TypeTag, Indices);
-
     using FluxVariables = typename GET_PROP_TYPE(TypeTag, FluxVariables);
 
     using DofTypeIndices = typename GET_PROP(TypeTag, DofTypeIndices);
@@ -83,10 +76,12 @@ class StaggeredNavierStokesResidualImpl<TypeTag, false, true> : public Staggered
     typename DofTypeIndices::FaceIdx faceIdx;
 
     enum {
+        massBalanceIdx = Indices::massBalanceIdx,
         energyBalanceIdx = Indices::energyBalanceIdx
     };
 
     using VolumeVariables = typename GET_PROP_TYPE(TypeTag, VolumeVariables);
+    using GlobalFaceVars = typename GET_PROP_TYPE(TypeTag, GlobalFaceVars);
 
     static constexpr bool useMoles = GET_PROP_VALUE(TypeTag, UseMoles);
 
@@ -98,19 +93,16 @@ public:
      * The result should be averaged over the volume (e.g. phase mass
      * inside a sub control volume divided by the volume)
      */
-    CellCenterPrimaryVariables computeStorageForCellCenter(const SubControlVolume& scv, const VolumeVariables& volVars) const
+    CellCenterPrimaryVariables computeStorageForCellCenter(const SubControlVolume& scv, const VolumeVariables& volVars)
     {
         CellCenterPrimaryVariables storage(0.0);
         const Scalar density = useMoles? volVars.molarDensity() : volVars.density();
 
         // compute storage of mass
-        storage = ParentType::computeStorageForCellCenter(scv, volVars);
+        storage[massBalanceIdx] = volVars.density();
 
         // compute the storage of energy
-        storage[energyBalanceIdx] = volVars.density(0) * volVars.internalEnergy(0);
-
-        std::cout << "** Subcontrolvolume " << scv.index() << ": energy storage = "
-               << storage[energyBalanceIdx] << std::endl;
+        storage[energyBalanceIdx] = density * volVars.internalEnergy();
 
         return storage;
     }
diff --git a/dumux/freeflow/staggeredni/model.hh b/dumux/freeflow/staggeredni/model.hh
index 8ee83383478493fb72b7696cc3e8e3b419491f56..0f8b3a7a3a41751779313e4205cce206d7502df7 100644
--- a/dumux/freeflow/staggeredni/model.hh
+++ b/dumux/freeflow/staggeredni/model.hh
@@ -79,6 +79,8 @@ class NavierStokesNIModel : public NavierStokesModel<TypeTag>
 
     using GlobalPosition = Dune::FieldVector<Scalar, dimWorld>;
 
+    static constexpr int numComponents = GET_PROP_VALUE(TypeTag, NumComponents);
+
     using DofTypeIndices = typename GET_PROP(TypeTag, DofTypeIndices);
     typename DofTypeIndices::CellCenterIdx cellCenterIdx;
     typename DofTypeIndices::FaceIdx faceIdx;
@@ -93,8 +95,8 @@ public:
 
         // add temperature to output
         auto& vtkOutputModule = problem.vtkOutputModule();
-        vtkOutputModule.addSecondaryVariable("temperature", [](const VolumeVariables& v){ return v.temperature();});
-
+//        vtkOutputModule.addPrimaryVariable("temperature", [](const VolumeVariables& v){ return v.temperature();});
+        vtkOutputModule.addPrimaryVariable("temperature", Indices::temperatureIdx);
     }
 };
 }
diff --git a/dumux/freeflow/staggeredni/properties.hh b/dumux/freeflow/staggeredni/properties.hh
index c7198e3e9e9adfc6b683b43b53b41fb83a05d32d..dbe0b517b5d2dc8c804066b3c0b64ffa705323aa 100644
--- a/dumux/freeflow/staggeredni/properties.hh
+++ b/dumux/freeflow/staggeredni/properties.hh
@@ -45,6 +45,8 @@ namespace Properties {
 //! The type tags for the non-isothermal Navier Stokes problems
 NEW_TYPE_TAG(NavierStokesNI, INHERITS_FROM(NavierStokes));
 
+NEW_PROP_TAG(EnableEnergyBalanceStokes);
+
 //! The type tags for the corresponding non-isothermal problems
 // NEW_TYPE_TAG(NavierStokesNI, INHERITS_FROM(NavierStokes, NonIsothermal));
 
@@ -67,6 +69,8 @@ NEW_TYPE_TAG(NavierStokesNI, INHERITS_FROM(NavierStokes));
 //NEW_PROP_TAG(EnableEnergyTransport); //!<  Returns whether to consider energy transport or not
 //NEW_PROP_TAG(FaceVariables); //!<  Returns whether to consider energy transport or not
 //NEW_PROP_TAG(UseMoles); //!< Defines whether molar (true) or mass (false) density is used
+NEW_PROP_TAG(PhaseIdx); //!< Defines the phaseIdx
+
 // \}
 }
 
diff --git a/dumux/freeflow/staggeredni/propertydefaults.hh b/dumux/freeflow/staggeredni/propertydefaults.hh
index 4433539c0fd9fd4ae8fa4e5a1b9d6bc98a7fb780..dc1917b603b7b4c596235217dfb84c05c5321e0b 100644
--- a/dumux/freeflow/staggeredni/propertydefaults.hh
+++ b/dumux/freeflow/staggeredni/propertydefaults.hh
@@ -27,15 +27,15 @@
 #ifndef DUMUX_NAVIER_STOKES_NI_PROPERTY_DEFAULTS_HH
 #define DUMUX_NAVIER_STOKES_NI_PROPERTY_DEFAULTS_HH
 
+// TODO clean-up
 #include "properties.hh"
 
 #include "model.hh"
-#include "../staggered/volumevariables.hh"
+#include "volumevariables.hh"
 #include "indices.hh"
 #include "localresidual.hh"
 #include "fluxvariables.hh"
 #include "../staggered/problem.hh"
-// #include "../staggered/model.hh"
 #include "../staggered/propertydefaults.hh"
 
 #include <dumux/implicit/staggered/localresidual.hh>
@@ -63,13 +63,30 @@ NEW_PROP_TAG(FluxVariablesCache);
 ///////////////////////////////////////////////////////////////////////////
 namespace Properties {
 
-SET_INT_PROP(NavierStokesNI, NumEqCellCenter, 1); // temp. (additional to pressure)
+SET_INT_PROP(NavierStokesNI, NumEqCellCenter, 2);
 
 //! the VolumeVariables property
-SET_TYPE_PROP(NavierStokesNI, VolumeVariables, NavierStokesVolumeVariables<TypeTag>);
+SET_TYPE_PROP(NavierStokesNI, VolumeVariables, NavierStokesNIVolumeVariables<TypeTag>);
 SET_TYPE_PROP(NavierStokesNI, Model, NavierStokesNIModel<TypeTag>);
 SET_TYPE_PROP(NavierStokesNI, Indices, NavierStokesNIIndices<TypeTag>);
 
+SET_BOOL_PROP(NavierStokesNI, EnableEnergyBalanceStokes, true);
+
+SET_BOOL_PROP(NavierStokesNI, UseMoles, true);
+
+SET_TYPE_PROP(NavierStokesNI, HeatConductionType, FouriersLaw<TypeTag>);
+
+SET_INT_PROP(NavierStokesNI, PhaseIdx, 0); //!< Defines the phaseIdx
+
+
+////! average is used as default model to compute the effective thermal heat conductivity
+//SET_PROP(NavierStokesNI, ThermalConductivityModel)
+//{ private :
+//    typedef typename GET_PROP_TYPE(TypeTag, Scalar) Scalar;
+//  public:
+//    typedef ThermalConductivityAverage<Scalar> type;
+//};
+
 
 /*!
  * \brief The fluid state which is used by the volume variables to
@@ -77,14 +94,14 @@ SET_TYPE_PROP(NavierStokesNI, Indices, NavierStokesNIIndices<TypeTag>);
  *        appropriately for the model ((non-)isothermal, equilibrium, ...).
  *        This can be done in the problem.
  */
-SET_PROP(NavierStokesNI, FluidState)
-{
-    private:
-        typedef typename GET_PROP_TYPE(TypeTag, Scalar) Scalar;
-        typedef typename GET_PROP_TYPE(TypeTag, FluidSystem) FluidSystem;
-    public:
-        typedef ImmiscibleFluidState<Scalar, FluidSystem> type;
-};
+//SET_PROP(NavierStokesNI, FluidState)
+//{
+//    private:
+//        typedef typename GET_PROP_TYPE(TypeTag, Scalar) Scalar;
+//        typedef typename GET_PROP_TYPE(TypeTag, FluidSystem) FluidSystem;
+//    public:
+//        typedef ImmiscibleFluidState<Scalar, FluidSystem> type;
+//};
 
 // //! Enable advection
 // SET_BOOL_PROP(NavierStokes, EnableAdvection, true);
@@ -138,19 +155,9 @@ SET_PROP(NavierStokesNI, FluidState)
 //
 // SET_BOOL_PROP(NavierStokes, EnableInertiaTerms, true);
 //
-// SET_BOOL_PROP(NavierStokes, EnableEnergyTransport, false);
+// SET_BOOL_PROP(NavierStokes, EnableEnergyTransport, true);
 //
 
-SET_TYPE_PROP(NavierStokesNI, HeatConductionType, FouriersLaw<TypeTag>);
-
-//! average is used as default model to compute the effective thermal heat conductivity
-// SET_PROP(NavierStokesNI, ThermalConductivityModel)
-// { private :
-//     typedef typename GET_PROP_TYPE(TypeTag, Scalar) Scalar;
-//   public:
-//     typedef ThermalConductivityAverage<Scalar> type;
-// };
-
 //////////////////////////////////////////////////////////////////
 // Property values for isothermal model required for the general non-isothermal model
 //////////////////////////////////////////////////////////////////