/*
  Copyright 2014 SINTEF ICT, Applied Mathematics.
  Copyright 2014 STATOIL ASA.

  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 3 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/>.
*/

#include "config.h"

#include <opm/polymer/fullyimplicit/FullyImplicitCompressiblePolymerSolver.hpp>


#include <opm/autodiff/AutoDiffBlock.hpp>
#include <opm/autodiff/AutoDiffHelpers.hpp>
#include <opm/autodiff/BlackoilPropsAdFromDeck.hpp>
#include <opm/autodiff/GeoProps.hpp>
#include <opm/autodiff/WellStateFullyImplicitBlackoil.hpp>

#include <opm/core/grid.h>
#include <opm/core/linalg/LinearSolverInterface.hpp>
#include <opm/core/props/rock/RockCompressibility.hpp>
#include <opm/simulators/timestepping/SimulatorTimerInterface.hpp>
#include <opm/common/OpmLog/OpmLog.hpp>
#include <opm/polymer/PolymerBlackoilState.hpp>
#include <opm/common/ErrorMacros.hpp>
#include <opm/core/well_controls.h>
#include <cassert>
#include <cmath>
#include <iostream>
#include <iomanip>

// A debugging utility.
#define DUMP(foo)                                                       \
    do {                                                                \
        std::cout << "==========================================\n"     \
                  << #foo ":\n"                                         \
                  << collapseJacs(foo) << std::endl;                    \
    } while (0)



namespace Opm {

typedef AutoDiffBlock<double> ADB;
typedef ADB::V V;
typedef ADB::M M;
typedef Eigen::Array<double,
                     Eigen::Dynamic,
                     Eigen::Dynamic,
                     Eigen::RowMajor> DataBlock;





namespace {

    std::vector<int>
    buildAllCells(const int nc)
    {
        std::vector<int> all_cells(nc);

        for (int c = 0; c < nc; ++c) { all_cells[c] = c; }

        return all_cells;
    }





    template <class GeoProps>
    Eigen::SparseMatrix<double>
    gravityOperator(const UnstructuredGrid& grid,
                    const HelperOps&        ops ,
                    const GeoProps&         geo )
    {
        const int nc = grid.number_of_cells;

        std::vector<int> f2hf(2 * grid.number_of_faces, -1);
        for (int c = 0, i = 0; c < nc; ++c) {
            for (; i < grid.cell_facepos[c + 1]; ++i) {
                const int f = grid.cell_faces[ i ];
                const int p = 0 + (grid.face_cells[2*f + 0] != c);

                f2hf[2*f + p] = i;
            }
        }

        typedef AutoDiffBlock<double>::V V;

        const V& gpot  = geo.gravityPotential();
        const V& trans = geo.transmissibility();

        const HelperOps::IFaces::Index ni = ops.internal_faces.size();

        typedef Eigen::Triplet<double> Tri;
        std::vector<Tri> grav;  grav.reserve(2 * ni);
        for (HelperOps::IFaces::Index i = 0; i < ni; ++i) {
            const int f  = ops.internal_faces[ i ];
            const int c1 = grid.face_cells[2*f + 0];
            const int c2 = grid.face_cells[2*f + 1];

            assert ((c1 >= 0) && (c2 >= 0));

            const double dG1 = gpot[ f2hf[2*f + 0] ];
            const double dG2 = gpot[ f2hf[2*f + 1] ];
            const double t   = trans[ f ];

            grav.push_back(Tri(i, c1,   t * dG1));
            grav.push_back(Tri(i, c2, - t * dG2));
        }

        Eigen::SparseMatrix<double> G(ni, nc);
        G.setFromTriplets(grav.begin(), grav.end());

        return G;
    }





    V computePerfPress(const UnstructuredGrid& grid, const Wells& wells, const V& rho, const double grav)
    {
        const int nw = wells.number_of_wells;
        const int nperf = wells.well_connpos[nw];
        const int dim = grid.dimensions;
        V wdp = V::Zero(nperf,1);
        assert(wdp.size() == rho.size());

        // Main loop, iterate over all perforations,
        // using the following formula:
        //    wdp(perf) = g*(perf_z - well_ref_z)*rho(perf)
        // where the total density rho(perf) is taken to be
        //    sum_p (rho_p*saturation_p) in the perforation cell.
        // [although this is computed on the outside of this function].
        for (int w = 0; w < nw; ++w) {
            const double ref_depth = wells.depth_ref[w];
            for (int j = wells.well_connpos[w]; j < wells.well_connpos[w + 1]; ++j) {
                const int cell = wells.well_cells[j];
                const double cell_depth = grid.cell_centroids[dim * cell + dim - 1];
                wdp[j] = rho[j]*grav*(cell_depth - ref_depth);
            }
        }
        return wdp;
    }

} // Anonymous namespace





    FullyImplicitCompressiblePolymerSolver::
    FullyImplicitCompressiblePolymerSolver(const UnstructuredGrid&  		grid,
                                		   const BlackoilPropsAdFromDeck&  fluid,
                                		   const DerivedGeology&           	geo  ,
                                		   const RockCompressibility*      	rock_comp_props,
                                		   const PolymerPropsAd&           	polymer_props_ad,
                                		   const Wells&                    	wells,
                                		   const NewtonIterationBlackoilInterface&    	linsolver)
        : grid_  (grid)
        , fluid_ (fluid)
        , geo_   (geo)
        , rock_comp_props_(rock_comp_props)
        , polymer_props_ad_(polymer_props_ad)
        , wells_ (wells)
        , linsolver_ (linsolver)
        , cells_ (buildAllCells(grid.number_of_cells))
        , ops_   (grid)
        , wops_  (wells)
        , grav_  (gravityOperator(grid_, ops_, geo_))
		, cmax_(V::Zero(grid.number_of_cells))
        , phaseCondition_ (grid.number_of_cells)
        , sd_    (fluid.numPhases() + 1)
        , residual_ ( { std::vector<ADB>(fluid.numPhases() + 1, ADB::null()),
                        ADB::null(),
                        ADB::null(),
                        { 1.1169, 1.0031, 0.0031, 1.0 },
                        false } ) // default scaling
        , linearizations_(0)
    {
    }





    int
    FullyImplicitCompressiblePolymerSolver::
    step(const SimulatorTimerInterface& timer,
         PolymerBlackoilState& x ,
         WellStateFullyImplicitBlackoilPolymer& xw)
    {
        const std::vector<double>& polymer_inflow = xw.polymerInflow();
        const double dt = timer.currentStepLength();

        // Initial max concentration of this time step from PolymerBlackoilState.
        cmax_ = Eigen::Map<V>(&x.getCellData( x.CMAX )[0], Opm::AutoDiffGrid::numCells(grid_));

        const SolutionState state = constantState(x, xw);
        computeAccum(state, 0);

        const double atol  = 1.0e-12;
        const double rtol  = 5.0e-8;
        const int    maxit = 15;
        assemble(dt, x, xw, polymer_inflow);

        const double r0  = residualNorm();
        const double r_polymer = residual_.material_balance_eq[2].value().matrix().lpNorm<Eigen::Infinity>();
        int it  = 0;
        // this solver does not solve well equations first.
        wellIterations_ = 0;
        std::cout << "\nIteration         Residual     Polymer Res\n"
                  << std::setw(9) << it << std::setprecision(9)
                  << std::setw(18) << r0 << std::setprecision(9)
                  << std::setw(18) << r_polymer << std::endl;
        bool resTooLarge = r0 > atol;
        while (resTooLarge && (it < maxit)) {
            const V dx = solveJacobianSystem();

            // update the number of linear iterations used.
            linearIterations_ += linsolver_.iterations();

            updateState(dx, x, xw);
            assemble(dt, x, xw, polymer_inflow);

            const double r = residualNorm();

        	const double rr_polymer = residual_.material_balance_eq[2].value().matrix().lpNorm<Eigen::Infinity>();
            resTooLarge = (r > atol) && (r > rtol*r0);

            it += 1;
            linearizations_ += 1;
            newtonIterations_ += 1;
            std::cout << std::setw(9) << it << std::setprecision(9)
                      << std::setw(18) << r << std::setprecision(9)
                      << std::setw(18) << rr_polymer << std::endl;
        }

        if (resTooLarge) {
            std::cerr << "Failed to compute converged solution in " << it << " iterations. Ignoring!\n";
            return -1;
            // OPM_THROW(std::runtime_error, "Failed to compute converged solution in " << it << " iterations.");
        }

        // Update max concentration.
        computeCmax(x);

        return it;
    }

    int FullyImplicitCompressiblePolymerSolver::linearizations() const
    {
        return linearizations_;
    }

    int FullyImplicitCompressiblePolymerSolver::nonlinearIterations() const
    {
        return newtonIterations_;
    }

    int FullyImplicitCompressiblePolymerSolver::linearIterations() const
    {
        return linearIterations_;
    }

    int FullyImplicitCompressiblePolymerSolver::wellIterations() const
    {
        return wellIterations_;
    }

    FullyImplicitCompressiblePolymerSolver::ReservoirResidualQuant::ReservoirResidualQuant()
        : accum(2, ADB::null())
        , mflux(   ADB::null())
        , b    (   ADB::null())
        , mu   (   ADB::null())
        , rho  (   ADB::null())
        , kr   (   ADB::null())
        , head (   ADB::null())
        , mob  (   ADB::null())
        , ads  (2, ADB::null())
    {
    }





    FullyImplicitCompressiblePolymerSolver::SolutionState::SolutionState(const int np)
        : pressure  (    ADB::null())
        , temperature(   ADB::null())
        , saturation(np, ADB::null())
        , concentration( ADB::null())
        , qs        (    ADB::null())
        , bhp       (    ADB::null())
    {
    }





    FullyImplicitCompressiblePolymerSolver::
    WellOps::WellOps(const Wells& wells)
        : w2p(wells.well_connpos[ wells.number_of_wells ],
              wells.number_of_wells)
        , p2w(wells.number_of_wells,
              wells.well_connpos[ wells.number_of_wells ])
    {
        const int        nw   = wells.number_of_wells;
        const int* const wpos = wells.well_connpos;

        typedef Eigen::Triplet<double> Tri;

        std::vector<Tri> scatter, gather;
        scatter.reserve(wpos[nw]);
        gather .reserve(wpos[nw]);

        for (int w = 0, i = 0; w < nw; ++w) {
            for (; i < wpos[ w + 1 ]; ++i) {
                scatter.push_back(Tri(i, w, 1.0));
                gather .push_back(Tri(w, i, 1.0));
            }
        }

        w2p.setFromTriplets(scatter.begin(), scatter.end());
        p2w.setFromTriplets(gather .begin(), gather .end());
    }





    FullyImplicitCompressiblePolymerSolver::SolutionState
    FullyImplicitCompressiblePolymerSolver::constantState(const PolymerBlackoilState& x,
                                               			  const WellStateFullyImplicitBlackoil&     xw)
    {
        const int nc = grid_.number_of_cells;
        const int np = x.numPhases();

        // The block pattern assumes the following primary variables:
        //    pressure
        //    water saturation (if water present)
        //    polymer concentration
        //    well rates per active phase and well
        //    well bottom-hole pressure
        // Note that oil is assumed to always be present, but is never
        // a primary variable.
        std::vector<int> bpat(np + 1, nc);
        bpat.push_back(xw.bhp().size() * np);
        bpat.push_back(xw.bhp().size());

        SolutionState state(np);

        // Pressure.
        assert (not x.pressure().empty());
        const V p = Eigen::Map<const V>(& x.pressure()[0], nc, 1);
        state.pressure = ADB::constant(p, bpat);

        // Temperature.
        const V T = Eigen::Map<const V>(& x.temperature()[0], nc, 1);
        state.temperature = ADB::constant(T);

        // Saturation.
        assert (not x.saturation().empty());
        const DataBlock s = Eigen::Map<const DataBlock>(& x.saturation()[0], nc, np);
        V so = V::Ones(nc, 1);
        const V sw  = s.col(0);
        so -= sw;
        state.saturation[0] = ADB::constant(sw, bpat);
        state.saturation[1] = ADB::constant(so, bpat);

        // Concentration
        {
            auto& concentration = x.getCellData( x.CONCENTRATION );
            assert(concentration.empty());
            const V c = Eigen::Map<const V>(concentration.data(), nc);
            // Do not understand:
            //concentration = ADB::constant(c);
            // Old code based on concentraton() method had the statement:
            //
            //   state.concentration = ADB::constant(c)
            //
            // This looks like it was a method assignment - how did it even compile?
        }

        // Well rates.
        assert (not xw.wellRates().empty());
        // Need to reshuffle well rates, from ordered by wells, then phase,
        // to ordered by phase, then wells.
        const int nw = wells_.number_of_wells;
        // The transpose() below switches the ordering.
        const DataBlock wrates = Eigen::Map<const DataBlock>(& xw.wellRates()[0], nw, np).transpose();
        const V qs = Eigen::Map<const V>(wrates.data(), nw * np);
        state.qs = ADB::constant(qs, bpat);

        // Well bottom-hole pressure.
        assert (not xw.bhp().empty());
        const V bhp = Eigen::Map<const V>(& xw.bhp()[0], xw.bhp().size());
        state.bhp = ADB::constant(bhp, bpat);

        return state;
    }





    FullyImplicitCompressiblePolymerSolver::SolutionState
    FullyImplicitCompressiblePolymerSolver::variableState(const PolymerBlackoilState& x,
                                               			  const WellStateFullyImplicitBlackoil&     xw)
    {
        const int nc = grid_.number_of_cells;
        const int np = x.numPhases();

        std::vector<V> vars0;

        // Initial pressure.
        assert (not x.pressure().empty());
        const V p = Eigen::Map<const V>(& x.pressure()[0], nc, 1);
        vars0.push_back(p);

        // Initial saturation.
        assert (not x.saturation().empty());
        const DataBlock s = Eigen::Map<const DataBlock>(& x.saturation()[0], nc, np);
        const V sw0 = s.col(0);
        vars0.push_back(sw0);

        // Initial concentration.
        {
            auto& concentration = x.getCellData( x.CONCENTRATION );
            assert (not concentration.empty());
            const V c = Eigen::Map<const V>(concentration.data() , nc);
            vars0.push_back(c);
        }

        // Initial well rates.
        assert (not xw.wellRates().empty());
        // Need to reshuffle well rates, from ordered by wells, then phase,
        // to ordered by phase, then wells.
        const int nw = wells_.number_of_wells;
        // The transpose() below switches the ordering.
        const DataBlock wrates = Eigen::Map<const DataBlock>(& xw.wellRates()[0], nw, np).transpose();
        const V qs = Eigen::Map<const V>(wrates.data(), nw*np);
        vars0.push_back(qs);

        // Initial well bottom-hole pressure.
        assert (not xw.bhp().empty());
        const V bhp = Eigen::Map<const V>(& xw.bhp()[0], xw.bhp().size());
        vars0.push_back(bhp);

        std::vector<ADB> vars = ADB::variables(vars0);

        SolutionState state(np);

        // Pressure.
        int nextvar = 0;
        state.pressure = vars[ nextvar++ ];

        // Temperature. (this is always a constant so far)
        const V T = Eigen::Map<const V>(& x.temperature()[0], nc, 1);
        state.temperature = ADB::constant(T);

        // Saturation.
        const std::vector<int>& bpat = vars[0].blockPattern();
        {
            ADB so  = ADB::constant(V::Ones(nc, 1), bpat);
            ADB& sw = vars[ nextvar++ ];
            state.saturation[0] = sw;
            so = so - sw;
            state.saturation[1] = so;
         }

        // Concentration.
        state.concentration = vars[nextvar++];

        // Qs.
        state.qs = vars[ nextvar++ ];

        // Bhp.
        state.bhp = vars[ nextvar++ ];

        assert(nextvar == int(vars.size()));

        return state;
    }





    void
    FullyImplicitCompressiblePolymerSolver::computeAccum(const SolutionState& state,
                                              		     const int            aix  )
    {

        const ADB&              press = state.pressure;
        const ADB&              temp  = state.temperature;
        const std::vector<ADB>& sat   = state.saturation;
        const ADB&              c     = state.concentration;

        const std::vector<PhasePresence> cond = phaseCondition();
		std::vector<ADB> pressure = computePressures(state);

        const ADB pv_mult = poroMult(press);

        for (int phase = 0; phase < 2; ++phase) {
            sd_.rq[phase].b = fluidReciprocFVF(phase, pressure[phase], temp, cond, cells_);
        }
        sd_.rq[0].accum[aix] = pv_mult * sd_.rq[0].b * sat[0];
        sd_.rq[1].accum[aix] = pv_mult * sd_.rq[1].b * sat[1];
	const ADB cmax = ADB::constant(cmax_, state.concentration.blockPattern());
        const ADB ads = polymer_props_ad_.adsorption(state.concentration, cmax);
        const double rho_rock = polymer_props_ad_.rockDensity();
        const V phi = Eigen::Map<const V>(&fluid_.porosity()[0], grid_.number_of_cells, 1);

        const double dead_pore_vol = polymer_props_ad_.deadPoreVol();
        sd_.rq[2].accum[aix] = pv_mult * sd_.rq[0].b * sat[0] * c * (1. - dead_pore_vol) + pv_mult *  rho_rock * (1. - phi) / phi * ads;
    }




    std::vector<std::vector<double> >
    FullyImplicitCompressiblePolymerSolver::computeFluidInPlace(const PolymerBlackoilState& x,
                                                                const std::vector<int>& fipnum)
    {
        const int np = x.numPhases();
        const int nc = grid_.number_of_cells;

        SolutionState state(np);
        state.pressure = ADB::constant(Eigen::Map<const V>(& x.pressure()[0], nc, 1));
        state.temperature = ADB::constant(Eigen::Map<const V>(& x.temperature()[0], nc, 1));
        const DataBlock s = Eigen::Map<const DataBlock>(& x.saturation()[0], nc, np);
        for (int phase = 0; phase < np; ++phase) {
            state.saturation[phase] = ADB::constant(s.col(phase));
        }

        const ADB&              press = state.pressure;
        const ADB&              temp  = state.temperature;
        const std::vector<ADB>& sat   = state.saturation;

        const std::vector<PhasePresence> cond = phaseCondition();
        std::vector<ADB> pressure = computePressures(state);

        const ADB pv_mult = poroMult(press);
        const V& pv = geo_.poreVolume();
        std::vector<V> fip(5, V::Zero(nc));
        for (int phase = 0; phase < 2; ++phase) {
            const ADB& b = fluidReciprocFVF(phase, pressure[phase], temp, cond, cells_);
            fip[phase] = (pv_mult * b * sat[phase] * pv).value();
        }


        const int dims = *std::max_element(fipnum.begin(), fipnum.end());
        std::vector<std::vector<double> > values(dims);
        for (int i=0; i < dims; ++i) {
            values[i].resize(7, 0.0);
        }

        V hcpv = V::Zero(nc);
        V pres = V::Zero(nc);
        for (int i = 0; i < 5; ++i) {
            for (int c = 0; c < nc; ++c) {
                if (fipnum[c] != 0) {
                    values[fipnum[c]-1][i] += fip[i][c];
                }
            }
        }

        // compute PAV and PORV or every regions.
        for (int c = 0; c < nc; ++c) {
            if (fipnum[c] != 0) {
                hcpv[fipnum[c]-1] += pv[c] * s.col(Oil)[c];
                pres[fipnum[c]-1] += pv[c] * state.pressure.value()[c];
                values[fipnum[c]-1][5] += pv[c];
                values[fipnum[c]-1][6] += pv[c] * state.pressure.value()[c] * s.col(Oil)[c];
            }
        }

        for (int reg = 0; reg < dims; ++reg) {
             if (hcpv[reg] != 0) {
                 values[reg][6] /= hcpv[reg];
             } else {
                 values[reg][6] = pres[reg] / values[reg][5];
             }
        }

        // Fluid in place is not implemented in this class.
        // See BlackoilModelBase::computeFluidInPlace(...) for how it's implemented there
        // FIXME: This following code has not been tested.
        if (sd_.fip[0].size() == 0) {
            OpmLog::warning("NOT_COMPUTING_FIP",
                    "Computing fluid in place is not implemented for summary files.");
            for (int i = 0; i < 7; ++i) {
                sd_.fip[i] = V::Zero(nc);
            }
        }


        return values;
    }





    void
    FullyImplicitCompressiblePolymerSolver::
    computeCmax(PolymerBlackoilState& state)
    {
        const int nc = grid_.number_of_cells;
        V tmp = V::Zero(nc);
        const auto& concentration = state.getCellData( state.CONCENTRATION );
        auto& cmax = state.getCellData( state.CMAX );

        for (int i = 0; i < nc; ++i) {
            tmp[i] = std::max(cmax[i], concentration[i]);
        }

        std::copy(&tmp[0], &tmp[0] + nc, cmax.begin());
    }





    void
    FullyImplicitCompressiblePolymerSolver::
    assemble(const double             dt,
             const PolymerBlackoilState& x,
             const WellStateFullyImplicitBlackoil& xw,
             const std::vector<double>& polymer_inflow)
    {
        // Create the primary variables.
        //
        const SolutionState state = variableState(x, xw);
		const V pvdt = geo_.poreVolume() / dt;
        // -------- Mass balance equations --------

        // Compute b_p and the accumulation term b_p*s_p for each phase,
        // except gas. For gas, we compute b_g*s_g + Rs*b_o*s_o.
        // These quantities are stored in sd_.rq[phase].accum[1].
        // The corresponding accumulation terms from the start of
        // the timestep (b^0_p*s^0_p etc.) were already computed
        // in step() and stored in sd_.rq[phase].accum[0].
        computeAccum(state, 1);

        // Set up the common parts of the mass balance equations
        // for each active phase.
        const V trans = subset(geo_.transmissibility(), ops_.internal_faces);
        {
            const std::vector<ADB> kr = computeRelPerm(state);
            const auto& pu = fluid_.phaseUsage();
            for (int phaseIdx=0; phaseIdx < fluid_.numPhases(); ++phaseIdx) {
                sd_.rq[phaseIdx].kr = kr[pu.phase_pos[phaseIdx]];
            }
        }
		const ADB cmax = ADB::constant(cmax_, state.concentration.blockPattern());
        const ADB krw_eff = polymer_props_ad_.effectiveRelPerm(state.concentration, cmax, sd_.rq[0].kr);
        const ADB mc = computeMc(state);
        computeMassFlux(trans, mc, sd_.rq[1].kr, krw_eff, state);
        residual_.material_balance_eq[0] = pvdt*(sd_.rq[0].accum[1] - sd_.rq[0].accum[0])
                                    + ops_.div*sd_.rq[0].mflux;
        residual_.material_balance_eq[1] = pvdt*(sd_.rq[1].accum[1] - sd_.rq[1].accum[0])
                                    + ops_.div*sd_.rq[1].mflux;
        residual_.material_balance_eq[2] = pvdt*(sd_.rq[2].accum[1] - sd_.rq[2].accum[0]) //+ cell / dt * (sd_.rq[2].ads[1] - sd_.rq[2].ads[0])
                                    + ops_.div*sd_.rq[2].mflux;

        // -------- Extra (optional) sg or rs equation, and rs contributions to the mass balance equations --------

        // Add the extra (flux) terms to the gas mass balance equations
        // from gas dissolved in the oil phase.
        // The extra terms in the accumulation part of the equation are already handled.
        // -------- Well equation, and well contributions to the mass balance equations --------

        // Contribution to mass balance will have to wait.

        const int nc = grid_.number_of_cells;
        const int np = wells_.number_of_phases;
        const int nw = wells_.number_of_wells;
        const int nperf = wells_.well_connpos[nw];

        const std::vector<int> well_cells(wells_.well_cells, wells_.well_cells + nperf);
        const V transw = Eigen::Map<const V>(wells_.WI, nperf);

        const ADB& bhp = state.bhp;

        // Extract variables for perforation cell pressures
        // and corresponding perforation well pressures.
        const ADB p_perfcell = subset(state.pressure, well_cells);
        // Finally construct well perforation pressures and well flows.

        // Compute well pressure differentials.
        // Construct pressure difference vector for wells.
        const int dim = grid_.dimensions;
        const double* g = geo_.gravity();
        if (g) {
            // Guard against gravity in anything but last dimension.
            for (int dd = 0; dd < dim - 1; ++dd) {
                assert(g[dd] == 0.0);
            }
        }
        ADB cell_rho_total = ADB::constant(V::Zero(nc), state.pressure.blockPattern());
		std::vector<ADB> press = computePressures(state);
        const ADB& temp = state.temperature;
        const std::vector<PhasePresence> cond = phaseCondition();
        for (int phase = 0; phase < 2; ++phase) {
            const ADB cell_rho = fluidDensity(phase, press[phase], temp, cond, cells_);
            cell_rho_total += state.saturation[phase] * cell_rho;
        }
        ADB inj_rho_total = ADB::constant(V::Zero(nperf), state.pressure.blockPattern());
        assert(np == wells_.number_of_phases);
        const DataBlock compi = Eigen::Map<const DataBlock>(wells_.comp_frac, nw, np);
        for (int phase = 0; phase < 2; ++phase) {
            const ADB cell_rho = fluidDensity(phase, press[phase], temp, cond, cells_);
            const V fraction = compi.col(phase);
            inj_rho_total += (wops_.w2p * fraction.matrix()).array() * subset(cell_rho, well_cells);
        }
        const V rho_perf_cell = subset(cell_rho_total, well_cells).value();
        const V rho_perf_well = inj_rho_total.value();
        V prodperfs = V::Constant(nperf, -1.0);
        for (int w = 0; w < nw; ++w) {
            if (wells_.type[w] == PRODUCER) {
                std::fill(prodperfs.data() + wells_.well_connpos[w],
                          prodperfs.data() + wells_.well_connpos[w+1], 1.0);
            }
        }
        const Selector<double> producer(prodperfs);
        const V rho_perf = producer.select(rho_perf_cell, rho_perf_well);
        const V well_perf_dp = computePerfPress(grid_, wells_, rho_perf, g ? g[dim-1] : 0.0);

        const ADB p_perfwell = wops_.w2p * bhp + well_perf_dp;
        const ADB nkgradp_well = transw * (p_perfcell - p_perfwell);
        // DUMP(nkgradp_well);
        const Selector<double> cell_to_well_selector(nkgradp_well.value());
        ADB well_rates_all = ADB::constant(V::Zero(nw*np), state.bhp.blockPattern());
        ADB perf_total_mob = subset(sd_.rq[0].mob, well_cells) + subset(sd_.rq[1].mob, well_cells);
        std::vector<ADB> well_contribs(np, ADB::null());
        std::vector<ADB> well_perf_rates(np, ADB::null());
        for (int phase = 0; phase < np; ++phase) {
            const ADB& cell_b = sd_.rq[phase].b;
            const ADB perf_b = subset(cell_b, well_cells);
            const ADB& cell_mob = sd_.rq[phase].mob;
            const V well_fraction = compi.col(phase);
            // Using total mobilities for all phases for injection.
            const ADB perf_mob_injector = (wops_.w2p * well_fraction.matrix()).array() * perf_total_mob;
            const ADB perf_mob = producer.select(subset(cell_mob, well_cells),
                                                 perf_mob_injector);
            const ADB perf_flux = perf_mob * (nkgradp_well); // No gravity term for perforations.
            well_perf_rates[phase] = (perf_flux * perf_b);
            const ADB well_rates = wops_.p2w * well_perf_rates[phase];
            well_rates_all += superset(well_rates, Span(nw, 1, phase*nw), nw*np);

            // const ADB well_contrib = superset(perf_flux*perf_b, well_cells, nc);
            well_contribs[phase] = superset(perf_flux*perf_b, well_cells, nc);
            // DUMP(well_contribs[phase]);
            residual_.material_balance_eq[phase] += well_contribs[phase];
        }

        // well rates contribs to polymer mass balance eqn.
        // for injection wells.
        const V polyin = Eigen::Map<const V>(& polymer_inflow[0], nc);
        const V poly_in_perf = subset(polyin, well_cells);
		const V poly_mc_cell = subset(mc, well_cells).value();
		const V poly_in_c = poly_in_perf;// * poly_mc_cell;
        const V poly_mc = producer.select(poly_mc_cell, poly_in_c);

		residual_.material_balance_eq[2] += superset(well_perf_rates[0] * poly_mc, well_cells, nc);
        // Set the well flux equation
        residual_.well_flux_eq = state.qs + well_rates_all;
        // DUMP(residual_.well_flux_eq);

        // Handling BHP and SURFACE_RATE wells.
        V bhp_targets(nw);
        V rate_targets(nw);
        Eigen::SparseMatrix<double> rate_distr(nw, np*nw);
        for (int w = 0; w < nw; ++w) {
            const WellControls* wc = wells_.ctrls[w];
            if (well_controls_get_current_type(wc) == BHP) {
                bhp_targets[w] = well_controls_get_current_target(wc);
                rate_targets[w] = -1e100;
            } else if (well_controls_get_current_type(wc) == SURFACE_RATE) {
                bhp_targets[w] = -1e100;
                rate_targets[w] = well_controls_get_current_target(wc);
                {
                    const double* distr = well_controls_get_current_distr(wc);
            	    for (int phase = 0; phase < np; ++phase) {
                        rate_distr.insert(w, phase*nw + w) = distr[phase];
                    }
                }
            } else {
                OPM_THROW(std::runtime_error, "Can only handle BHP and SURFACE_RATE type controls.");
            }
        }
        const ADB bhp_residual = bhp - bhp_targets;
        const ADB rate_residual = rate_distr * state.qs - rate_targets;
        // Choose bhp residual for positive bhp targets.
        Selector<double> bhp_selector(bhp_targets);
        residual_.well_eq = bhp_selector.select(bhp_residual, rate_residual);
        // DUMP(residual_.well_eq);
    }





    V FullyImplicitCompressiblePolymerSolver::solveJacobianSystem() const
    {
        return linsolver_.computeNewtonIncrement(residual_);
    }





    namespace {
        struct Chop01 {
            double operator()(double x) const { return std::max(std::min(x, 1.0), 0.0); }
        };
    }





    void FullyImplicitCompressiblePolymerSolver::
	updateState(const V& 				dx,
                PolymerBlackoilState& 	state,
                WellStateFullyImplicitBlackoil& 				well_state) const
    {
        const int np = fluid_.numPhases();
        const int nc = grid_.number_of_cells;
        const int nw = wells_.number_of_wells;
        const V one = V::Constant(nc, 1.0);
        const V zero = V::Zero(nc);
        // Extract parts of dx corresponding to each part.
        const V dp = subset(dx, Span(nc));
        int varstart = nc;
        const V dsw =subset(dx, Span(nc, 1, varstart));
        varstart += dsw.size();
        const V dc = subset(dx, Span(nc, 1, varstart));
        varstart += dc.size();
        const V dqs = subset(dx, Span(np*nw, 1, varstart));
        varstart += dqs.size();
        const V dbhp = subset(dx, Span(nw, 1, varstart));
        varstart += dbhp.size();
        assert(varstart == dx.size());

        // Pressure update.
        const double dpmaxrel = 0.8;
        const V p_old = Eigen::Map<const V>(&state.pressure()[0], nc, 1);
        const V absdpmax = dpmaxrel*p_old.abs();
        const V dp_limited = sign(dp) * dp.abs().min(absdpmax);
        const V p = (p_old - dp_limited).max(zero);
        std::copy(&p[0], &p[0] + nc, state.pressure().begin());

        // Saturation updates.
        const double dsmax = 0.3;
        const DataBlock s_old = Eigen::Map<const DataBlock>(&state.saturation()[0], nc, np);
        V so = one;
        const V sw_old = s_old.col(0);
        const V dsw_limited = sign(dsw) * dsw.abs().min(dsmax);
        const V sw = (sw_old - dsw_limited).unaryExpr(Chop01());
        so -= sw;
        for (int c = 0; c < nc; ++c) {
            state.saturation()[c*np] = sw[c];
            state.saturation()[c*np + 1] = so[c];
        }

        // Concentration updates.
//        const double dcmax = 0.3 * polymer_props_ad_.cMax();
//		std::cout << "\n the max concentration: " << dcmax / 0.3 << std::endl;
        {
            auto& concentration = state.getCellData( state.CONCENTRATION );
            const V c_old = Eigen::Map<const V>(concentration.data() , nc, 1);
            //        const V dc_limited = sign(dc) * dc.abs().min(dcmax);
            //        const V c = (c_old - dc_limited).max(zero);//unaryExpr(Chop02());
            const V c = (c_old - dc).max(zero);
            std::copy(&c[0], &c[0] + nc, concentration.begin());
        }

        // Qs update.
        // Since we need to update the wellrates, that are ordered by wells,
        // from dqs which are ordered by phase, the simplest is to compute
        // dwr, which is the data from dqs but ordered by wells.
        const DataBlock wwr = Eigen::Map<const DataBlock>(dqs.data(), np, nw).transpose();
        const V dwr = Eigen::Map<const V>(wwr.data(), nw*np);
        const V wr_old = Eigen::Map<const V>(&well_state.wellRates()[0], nw*np);
        const V wr = wr_old - dwr;
        std::copy(&wr[0], &wr[0] + wr.size(), well_state.wellRates().begin());

        // Bhp update.
        const V bhp_old = Eigen::Map<const V>(&well_state.bhp()[0], nw, 1);
        const V bhp = bhp_old - dbhp;
        std::copy(&bhp[0], &bhp[0] + bhp.size(), well_state.bhp().begin());

    }





    std::vector<ADB>
    FullyImplicitCompressiblePolymerSolver::computeRelPerm(const SolutionState& state) const
    {
        const int               nc   = grid_.number_of_cells;
        const std::vector<int>& bpat = state.pressure.blockPattern();

        const ADB null = ADB::constant(V::Zero(nc, 1), bpat);

        const ADB sw = state.saturation[0];
        const ADB so = state.saturation[1];
        const ADB sg = null;
        return fluid_.relperm(sw, so, sg, cells_);
    }





    std::vector<ADB>
    FullyImplicitCompressiblePolymerSolver::
	computePressures(const SolutionState& state) const
    {
        const int               nc   = grid_.number_of_cells;
        const std::vector<int>& bpat = state.pressure.blockPattern();

        const ADB null = ADB::constant(V::Zero(nc, 1), bpat);

        const ADB sw = state.saturation[0];

        const ADB so = state.saturation[1];

        const ADB sg = null;

        // convert the pressure offsets to the capillary pressures
        std::vector<ADB> pressure = fluid_.capPress(sw, so, sg, cells_);
		pressure[0] = pressure[0] - pressure[1];

        // add the total pressure to the capillary pressures
        pressure[0] = state.pressure - pressure[0];
        pressure[1] = state.pressure + pressure[1];

        return pressure;
    }





    void
    FullyImplicitCompressiblePolymerSolver::computeMassFlux(
                                                 const V&                transi,
                                                 const ADB&              mc,
                                                 const ADB&              kro,
                                                 const ADB&              krw_eff,
                                                 const SolutionState&    state )
    {
        const ADB tr_mult = transMult(state.pressure);
        const std::vector<PhasePresence> cond = phaseCondition();
        std::vector<ADB> press = computePressures(state);
        const ADB& temp = state.temperature;

        {
            const ADB mu_w = fluidViscosity(0, press[0], temp, cond, cells_);
            ADB inv_wat_eff_vis = polymer_props_ad_.effectiveInvWaterVisc(state.concentration, mu_w.value());
            sd_.rq[0].mob = tr_mult * krw_eff * inv_wat_eff_vis;
            sd_.rq[2].mob = tr_mult * mc * krw_eff * inv_wat_eff_vis;
            const ADB mu_o = fluidViscosity(1, press[1], temp, cond, cells_);
            sd_.rq[1].mob = tr_mult * kro / mu_o;

            sd_.rq[0].mu = mu_w;
            sd_.rq[1].mu = mu_o;
        }

        for (int phase = 0; phase < 2; ++phase) {
            sd_.rq[phase].rho = fluidDensity(phase, press[phase], temp, cond, cells_);
            ADB& head = sd_.rq[phase].head;
            // compute gravity potensial using the face average as in eclipse and MRST
            const ADB rhoavg = ops_.caver * sd_.rq[phase].rho;
            const ADB dp = ops_.ngrad * press[phase] - geo_.gravity()[2] * (rhoavg * (ops_.ngrad * geo_.z().matrix()));
            head = transi*dp;
            UpwindSelector<double> upwind(grid_, ops_, head.value());
            const ADB& b       = sd_.rq[phase].b;
            const ADB& mob     = sd_.rq[phase].mob;
            sd_.rq[phase].mflux = upwind.select(b * mob) * head;
        }
        sd_.rq[2].b = sd_.rq[0].b;
        sd_.rq[2].head = sd_.rq[0].head;
        UpwindSelector<double> upwind(grid_, ops_, sd_.rq[2].head.value());
        sd_.rq[2].mflux = upwind.select(sd_.rq[2].b * sd_.rq[2].mob) * sd_.rq[2].head;
    }





    double
    FullyImplicitCompressiblePolymerSolver::residualNorm() const
    {
        double r = 0;
        for (std::vector<ADB>::const_iterator
                 b = residual_.material_balance_eq.begin(),
                 e = residual_.material_balance_eq.end();
             b != e; ++b)
        {
            r = std::max(r, (*b).value().matrix().lpNorm<Eigen::Infinity>());
        }
        r = std::max(r, residual_.well_flux_eq.value().matrix().lpNorm<Eigen::Infinity>());
        r = std::max(r, residual_.well_eq.value().matrix().lpNorm<Eigen::Infinity>());

        return r;
    }





    ADB
    FullyImplicitCompressiblePolymerSolver::fluidViscosity(const int                         phase,
		                                                   const ADB&                        p    ,
		                                                   const ADB&                        T    ,
                                                           const std::vector<PhasePresence>& cond,
           			                                       const std::vector<int>&           cells) const
    {
        const ADB null = ADB::constant(V::Zero(grid_.number_of_cells, 1), p.blockPattern());
        switch (phase) {
        case Water:
            return fluid_.muWat(p, T, cells);
        case Oil: {
            return fluid_.muOil(p, T, null, cond, cells);
        }
        default:
            OPM_THROW(std::runtime_error, "Unknown phase index " << phase);
        }
    }





    ADB
    FullyImplicitCompressiblePolymerSolver::fluidReciprocFVF(const int                         phase,
                                                  		     const ADB&                        p    ,
                                                             const ADB&                        T    ,
                                                             const std::vector<PhasePresence>& cond,
                                                  		     const std::vector<int>&           cells) const
    {
        const ADB null = ADB::constant(V::Zero(grid_.number_of_cells, 1), p.blockPattern());
        switch (phase) {
        case Water:
            return fluid_.bWat(p, T, cells);
        case Oil: {
            return fluid_.bOil(p, T, null, cond, cells);
        }
        default:
            OPM_THROW(std::runtime_error, "Unknown phase index " << phase);
        }
    }





    ADB
    FullyImplicitCompressiblePolymerSolver::fluidDensity(const int                         phase,
                                              			 const ADB&                        p    ,
                                                         const ADB&                        T    ,
                                                         const std::vector<PhasePresence>& cond,
                                              		     const std::vector<int>&           cells) const
    {
        const V rhos = fluid_.surfaceDensity(phase, cells);
        const ADB b = fluidReciprocFVF(phase, p, T, cond, cells);
        const ADB rho = rhos * b;
        return rho;
    }





    // here mc means m(c) * c.
    ADB
    FullyImplicitCompressiblePolymerSolver::computeMc(const SolutionState& state) const
    {
        ADB c = state.concentration;
        return polymer_props_ad_.polymerWaterVelocityRatio(c);
    }





    ADB
    FullyImplicitCompressiblePolymerSolver::poroMult(const ADB& p) const
    {
        const int n = p.size();
        if (rock_comp_props_ && rock_comp_props_->isActive()) {
            V pm(n);
            V dpm(n);
            for (int i = 0; i < n; ++i) {
                pm[i] = rock_comp_props_->poroMult(p.value()[i]);
                dpm[i] = rock_comp_props_->poroMultDeriv(p.value()[i]);
            }
            ADB::M dpm_diag(dpm.matrix().asDiagonal());
            const int num_blocks = p.numBlocks();
            std::vector<ADB::M> jacs(num_blocks);
            for (int block = 0; block < num_blocks; ++block) {
                jacs[block] = dpm_diag * p.derivative()[block];
            }
            return ADB::function(std::move(pm), std::move(jacs));
        } else {
            return ADB::constant(V::Constant(n, 1.0), p.blockPattern());
        }
    }





    ADB
    FullyImplicitCompressiblePolymerSolver::transMult(const ADB& p) const
    {
        const int n = p.size();
        if (rock_comp_props_ && rock_comp_props_->isActive()) {
            V tm(n);
            V dtm(n);
            for (int i = 0; i < n; ++i) {
                tm[i] = rock_comp_props_->transMult(p.value()[i]);
                dtm[i] = rock_comp_props_->transMultDeriv(p.value()[i]);
            }
            ADB::M dtm_diag(dtm.matrix().asDiagonal());
            const int num_blocks = p.numBlocks();
            std::vector<ADB::M> jacs(num_blocks);
            for (int block = 0; block < num_blocks; ++block) {
                jacs[block] = dtm_diag * p.derivative()[block];
            }
            return ADB::function(std::move(tm), std::move(jacs));
        } else {
            return ADB::constant(V::Constant(n, 1.0), p.blockPattern());
        }
    }


    void
    FullyImplicitCompressiblePolymerSolver::classifyCondition(const PolymerBlackoilState& state)
    {
        const int nc = grid_.number_of_cells;
        const DataBlock s = Eigen::Map<const DataBlock>(& state.saturation()[0], nc, 2);

        const V so = s.col(1);
        for (V::Index c = 0, e = so.size(); c != e; ++c) {
            phaseCondition_[c].setFreeWater();
            if (so[c] > 0) {
                phaseCondition_[c].setFreeOil();
            }
        }
    }





    const FullyImplicitCompressiblePolymerSolver&
    FullyImplicitCompressiblePolymerSolver::model() const
    {
        return *this;
    }





    namespace detail {
        template <class RAIter>
        double euclidianNormSquared(RAIter it, const RAIter end)
        {
            double product = 0.0 ;
            for( ; it != end; ++it ) {
                product += ( *it * *it );
            }
            return product;
        }
    } // namespace detail





    double
    FullyImplicitCompressiblePolymerSolver::
    relativeChange(const PolymerBlackoilState& previous,
                   const PolymerBlackoilState& current ) const
    {
        std::vector< double > p0  ( previous.pressure() );
        std::vector< double > sat0( previous.saturation() );

        const std::size_t pSize = p0.size();
        const std::size_t satSize = sat0.size();

        // compute u^n - u^n+1
        for( std::size_t i=0; i<pSize; ++i ) {
            p0[ i ] -= current.pressure()[ i ];
        }

        for( std::size_t i=0; i<satSize; ++i ) {
            sat0[ i ] -= current.saturation()[ i ];
        }

        // compute || u^n - u^n+1 ||
        const double normDiff  = detail::euclidianNormSquared( p0.begin(),   p0.end())
                               + detail::euclidianNormSquared( sat0.begin(), sat0.end());

        // compute || u^n+1 ||
        const double normNewState  = detail::euclidianNormSquared( current.pressure().begin(),   current.pressure().end())
                                   + detail::euclidianNormSquared( current.saturation().begin(), current.saturation().end());

        if ( normNewState > 0.0 ) {
            return normDiff / normNewState;
        } else {
            return 0.0;
        }
    }

} //namespace Opm