/*
  Copyright 2018 Equinor 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>

#define BOOST_TEST_MODULE WellStateFIBOTest

#include <boost/test/unit_test.hpp>

#include "MpiFixture.hpp"

#include <opm/common/ErrorMacros.hpp>
#include <opm/common/utility/TimeService.hpp>

#include <opm/input/eclipse/EclipseState/EclipseState.hpp>

#include <opm/input/eclipse/Python/Python.hpp>

#include <opm/input/eclipse/Schedule/MSW/WellSegments.hpp>
#include <opm/input/eclipse/Schedule/Schedule.hpp>
#include <opm/input/eclipse/Schedule/SummaryState.hpp>
#include <opm/input/eclipse/Schedule/UDQ/UDQConfig.hpp>
#include <opm/input/eclipse/Schedule/UDQ/UDQParams.hpp>
#include <opm/input/eclipse/Schedule/Well/Well.hpp>
#include <opm/input/eclipse/Schedule/Well/WellConnections.hpp>

#include <opm/material/fluidsystems/PhaseUsageInfo.hpp>
#include <opm/material/fluidsystems/BlackOilDefaultFluidSystemIndices.hpp>

#include <opm/simulators/wells/GlobalWellInfo.hpp>
#include <opm/simulators/wells/ParallelWellInfo.hpp>
#include <opm/simulators/wells/PerfData.hpp>
#include <opm/simulators/wells/PerforationData.hpp>
#include <opm/simulators/wells/SegmentState.hpp>
#include <opm/simulators/wells/SingleWellState.hpp>
#include <opm/simulators/wells/WellContainer.hpp>
#include <opm/simulators/wells/WellState.hpp>

#include <opm/input/eclipse/Units/Units.hpp>

#include <opm/grid/GridHelpers.hpp>
#include <opm/grid/GridManager.hpp>

#include <opm/input/eclipse/Deck/Deck.hpp>

#include <opm/input/eclipse/Parser/Parser.hpp>
#include <opm/input/eclipse/Parser/ParseContext.hpp>

#include <algorithm>
#include <chrono>
#include <cstddef>
#include <ctime>
#include <functional>
#include <memory>
#include <stdexcept>
#include <string>
#include <vector>

BOOST_GLOBAL_FIXTURE(MPIFixture);

struct Setup
{
    using IndexTraits = Opm::BlackOilDefaultFluidSystemIndices;

    Setup(const std::string& filename)
        : Setup(Opm::Parser{}.parseFile(filename))
    {}

    Setup(const Opm::Deck& deck)
        : es   (deck)
        , pu   ()
        , grid (es.getInputGrid())
        , sched(deck, es, std::make_shared<Opm::Python>())
        , st   { Opm::TimeService::from_time_t(sched.getStartTime()),
                 es.runspec().udqParams().undefinedValue() }
    {
        pu.initFromState(es);
        initWellPerfData();
    }

    void initWellPerfData()
    {
        const auto& wells = sched.getWells(0);
        const auto& cartDims = Opm::UgGridHelpers::cartDims(*grid.c_grid());
        const int* compressed_to_cartesian = Opm::UgGridHelpers::globalCell(*grid.c_grid());
        std::vector<int> cartesian_to_compressed(cartDims[0] * cartDims[1] * cartDims[2], -1);
        for (int ii = 0; ii < Opm::UgGridHelpers::numCells(*grid.c_grid()); ++ii) {
            cartesian_to_compressed[compressed_to_cartesian[ii]] = ii;
        }
        well_perf_data.resize(wells.size());
        int well_index = 0;
        for (const auto& well : wells) {
            well_perf_data[well_index].clear();
            well_perf_data[well_index].reserve(well.getConnections().size());
            for (const auto& completion : well.getConnections()) {
                if (completion.state() == Opm::Connection::State::OPEN) {
                    const int i = completion.getI();
                    const int j = completion.getJ();
                    const int k = completion.getK();
                    const int cart_grid_indx = i + cartDims[0] * (j + cartDims[1] * k);
                    const int active_index = cartesian_to_compressed[cart_grid_indx];
                    if (active_index < 0) {
                        const std::string msg
                            = ("Cell with i,j,k indices " + std::to_string(i) + " " + std::to_string(j) + " "
                               + std::to_string(k) + " not found in grid (well = " + well.name() + ").");
                        OPM_THROW(std::runtime_error, msg);
                    } else {
                        Opm::PerforationData<double> pd;
                        pd.cell_index = active_index;
                        pd.connection_transmissibility_factor = completion.CF();
                        pd.connection_d_factor = completion.dFactor();
                        pd.satnum_id = completion.satTableId();
                        well_perf_data[well_index].push_back(pd);
                    }
                } else {
                    if (completion.state() != Opm::Connection::State::SHUT) {
                        OPM_THROW(std::runtime_error,
                                  "Completion state: " +
                                  Opm::Connection::State2String(completion.state()) +
                                  " not handled");
                    }
                }
            }
            ++well_index;
        }
    }

    Opm::EclipseState es;
    Opm::PhaseUsageInfo<IndexTraits>   pu;
    Opm::GridManager  grid;
    std::shared_ptr<Opm::Python> python;
    Opm::Schedule     sched;
    Opm::SummaryState st;
    std::vector<std::vector<Opm::PerforationData<double>>> well_perf_data;
};

namespace {
    using IndexTraits = Opm::BlackOilDefaultFluidSystemIndices;
    using WellState = Opm::WellState<double, IndexTraits>;
    using PhaseUsage = Opm::PhaseUsageInfo<IndexTraits>;

    Opm::WellState<double, IndexTraits>
    buildWellState(const Setup& setup, const std::size_t timeStep,
                   std::vector<Opm::ParallelWellInfo<double>>& pinfos)
    {
        auto state  = WellState{setup.pu};

        const auto cpress =
            std::vector<double>(setup.grid.c_grid()->number_of_cells,
                                100.0*Opm::unit::barsa);

        const auto& unit_sytem = setup.es.getDeckUnitSystem();
        const double temp = unit_sytem.to_si(Opm::UnitSystem::measure::temperature, 25);
        const auto ctemp =
            std::vector<double>(setup.grid.c_grid()->number_of_cells,
                                temp);

        auto wells = setup.sched.getWells(timeStep);
        pinfos.resize(wells.size());
        std::vector<std::reference_wrapper<Opm::ParallelWellInfo<double>>> ppinfos;
        auto pw = pinfos.begin();

        for (const auto& well : wells)
        {
            *pw = Opm::ParallelWellInfo<double>{well.name()};
            ppinfos.push_back(std::ref(*pw));
            pw->communicateFirstPerforation(true);
            ++pw;
        }

        state.init(cpress, ctemp, setup.sched,
                   wells, ppinfos,
                   timeStep, nullptr, setup.well_perf_data, setup.st,
                   false /*enableDistributedWells*/);

        state.initWellStateMSWell(setup.sched.getWells(timeStep),
                                  nullptr);

        return state;
    }

    Opm::Deck makeOilWaterRsconstWellDeck()
    {
        return Opm::Parser{}.parseString(R"(
RUNSPEC
OIL
WATER
METRIC
DIMENS
    2 1 1 /

START
  1 JAN 2000 /

GRID
DXV
  2*100 /
DYV
  1*100 /
DZV
  1*10 /
DEPTHZ
  6*1000 /
PORO
  2*0.20 /
PERMX
  2*100 /
PERMY
  2*100 /
PERMZ
  2*20 /

SCHEDULE
WELSPECS
  'PROD' 'G' 1 1 1000 'OIL' /
  'INJ'  'G' 2 1 1000 'WATER' /
/
COMPDAT
  'PROD' 1 1 1 1 'OPEN' 1* 100 0.1 /
  'INJ'  2 1 1 1 'OPEN' 1* 100 0.1 /
/
WCONPROD
  'PROD' 'OPEN' 'ORAT' 100 4* 100 /
/
WCONINJE
  'INJ' 'WATER' 'OPEN' 'RATE' 100 1* 200 /
/

END
)");
    }


    void setSegPress(const std::vector<Opm::Well>& wells,
                     WellState&                    wstate)
    {
        const auto nWell = wells.size();

        for (auto wellID = 0*nWell; wellID < nWell; ++wellID) {
            const auto& well     = wells[wellID];
            if (! well.isMultiSegment()) {
                continue;
            }

            const auto  pressTop = 100.0 * wellID;
            auto& ws = wstate.well(wellID);
            auto& segments = ws.segments;
            segments.pressure[0] = pressTop;

            const auto& segSet = well.getSegments();
            const auto  nSeg   = segSet.size();

            for (auto segID = 0*nSeg + 1; segID < nSeg; ++segID) {
                // One-based numbering scheme for segments.
                const auto segNo = segSet[segID].segmentNumber();
                segments.pressure[segNo - 1] = pressTop + 1.0*(segNo - 1);
            }
        }
    }


    void setSegRates(const std::vector<Opm::Well>& wells,
                     const PhaseUsage&             pu,
                     WellState&                    wstate)
    {
        const auto wat = pu.phaseIsActive(IndexTraits::waterPhaseIdx);
        const auto iw  = wat ? pu.canonicalToActivePhaseIdx(IndexTraits::waterPhaseIdx) : -1;

        const auto oil = pu.phaseIsActive(IndexTraits::oilPhaseIdx);
        const auto io  = oil ? pu.canonicalToActivePhaseIdx(IndexTraits::oilPhaseIdx) : -1;

        const auto gas = pu.phaseIsActive(IndexTraits::gasPhaseIdx);
        const auto ig  = gas ? pu.canonicalToActivePhaseIdx(IndexTraits::gasPhaseIdx) : -1;

        const auto np = wstate.numPhases();

        const auto nWell = wells.size();

        for (auto wellID = 0*nWell; wellID < nWell; ++wellID) {
            const auto& well     = wells[wellID];
            if (! well.isMultiSegment()) {
                continue;
            }

            const auto  rateTop  = 1000.0 * wellID;
            auto& ws = wstate.well(wellID);
            auto& segments = ws.segments;
            auto& segRates = segments.rates;

            if (wat) { segRates[iw] = rateTop; }
            if (oil) { segRates[io] = rateTop; }
            if (gas) { segRates[ig] = rateTop; }

            const auto& segSet = well.getSegments();
            const auto  nSeg   = segSet.size();

            for (auto segID = 0*nSeg + 1; segID < nSeg; ++segID) {
                // One-based numbering scheme for segments.
                const auto segNo = segSet[segID].segmentNumber();

                auto* rates = &segRates[(segNo - 1) * np];

                if (wat) { rates[iw] = rateTop + 100.0*(segNo - 1); }
                if (oil) { rates[io] = rateTop + 200.0*(segNo - 1); }
                if (gas) { rates[ig] = rateTop + 400.0*(segNo - 1); }
            }
        }
    }
} // Anonymous

BOOST_AUTO_TEST_SUITE(State_Operations)

// ---------------------------------------------------------------------

BOOST_AUTO_TEST_CASE(Linearisation)
{
    const Setup setup{ "msw.data" };
    const auto tstep = std::size_t{0};

    std::vector<Opm::ParallelWellInfo<double>> pinfos;
    const auto wstate = buildWellState(setup, tstep, pinfos);
    const auto& ws = wstate.well("PROD01");

    BOOST_CHECK_EQUAL(ws.segments.size(), 6);

    const auto& wells = setup.sched.getWellsatEnd();
    BOOST_CHECK_EQUAL(wells.size(), 2);
}

// ---------------------------------------------------------------------

BOOST_AUTO_TEST_CASE(Pressure)
{
    const Setup setup{ "msw.data" };
    const auto tstep = std::size_t{0};

    std::vector<Opm::ParallelWellInfo<double>> pinfos;
    auto wstate = buildWellState(setup, tstep, pinfos);

    const auto& wells = setup.sched.getWells(tstep);
    const auto prod01_first = wells[0].name() == "PROD01";

    setSegPress(wells, wstate);

    const auto rpt = wstate.report(setup.grid.c_grid()->global_cell, [](const int){return false;});

    {
        const auto expect_nSeg = 6;
        const auto& xw = rpt.at("PROD01");

        BOOST_CHECK_EQUAL(xw.segments.size(), expect_nSeg);

        const auto pressTop = prod01_first ? 0.0 : 100.0;

        for (auto segID = 0; segID < expect_nSeg; ++segID) {
            const auto& xseg = xw.segments.at(segID + 1);

            BOOST_CHECK_EQUAL(xseg.segNumber, segID + 1);
            const auto pres_idx = Opm::data::SegmentPressures::Value::Pressure;
            BOOST_CHECK_CLOSE(xseg.pressures[pres_idx], pressTop + 1.0*segID, 1.0e-10);
        }
    }

    const auto& ws = wstate.well("PROD01");
    const auto& perf_data = ws.perf_data;
    (void) perf_data;
}

// ---------------------------------------------------------------------

BOOST_AUTO_TEST_CASE(Rates)
{
    const Setup setup{ "msw.data" };
    const auto tstep = std::size_t{0};

    std::vector<Opm::ParallelWellInfo<double>> pinfos;
    auto wstate = buildWellState(setup, tstep, pinfos);

    const auto wells = setup.sched.getWells(tstep);
    const auto prod01_first = wells[0].name() == "PROD01";

    const auto& pu = setup.pu;

    setSegRates(wells, pu,  wstate);

    const auto rpt = wstate.report(setup.grid.c_grid()->global_cell, [](const int){return false;});

    const auto wat = pu.phaseIsActive(IndexTraits::waterPhaseIdx);
    const auto oil = pu.phaseIsActive(IndexTraits::oilPhaseIdx);
    const auto gas = pu.phaseIsActive(IndexTraits::gasPhaseIdx);

    BOOST_CHECK(wat && oil && gas);

    {
        const auto expect_nSeg = 6;
        const auto& xw = rpt.at("PROD01");

        BOOST_CHECK_EQUAL(xw.segments.size(), expect_nSeg);

        const auto rateTop = prod01_first ? 0.0 : 1000.0;

        for (auto segNum = 1; segNum <= expect_nSeg; ++segNum) {
            const auto& xseg = xw.segments.at(segNum);

            BOOST_CHECK_EQUAL(xseg.segNumber, segNum);

            BOOST_CHECK_CLOSE(xseg.rates.get(Opm::data::Rates::opt::wat),
                              rateTop + 100.0*(segNum - 1), 1.0e-10);

            BOOST_CHECK_CLOSE(xseg.rates.get(Opm::data::Rates::opt::oil),
                              rateTop + 200.0*(segNum - 1), 1.0e-10);

            BOOST_CHECK_CLOSE(xseg.rates.get(Opm::data::Rates::opt::gas),
                              rateTop + 400.0*(segNum - 1), 1.0e-10);
        }
    }
}

// ---------------------------------------------------------------------

BOOST_AUTO_TEST_CASE(STOP_well)
{
    /*
      This test verifies that the perforation pressures is correctly initialized
      also for wells in the STOP state.
    */
    const Setup setup{ "wells_manager_data_wellSTOP.data" };

    std::vector<Opm::ParallelWellInfo<double>> pinfos;
    auto wstate = buildWellState(setup, 0, pinfos);
    for (std::size_t well_index = 0; well_index < setup.sched.numWells(0); well_index++) {
        const auto& ws = wstate.well(well_index);
        const auto& perf_data = ws.perf_data;
        for (const auto& p : perf_data.pressure)
            BOOST_CHECK(p > 0);
    }
}

// ---------------------------------------------------------------------

BOOST_AUTO_TEST_CASE(RSCONST_ProducerGasSynthesizedFromOilRate)
{
    const auto setup = Setup { makeOilWaterRsconstWellDeck() };
    auto pinfos = std::vector<Opm::ParallelWellInfo<double>>{};

    const auto wstate = [&setup, &pinfos]() {
        const auto tstep = std::size_t{0};
        auto ws = buildWellState(setup, tstep, pinfos);

        const auto oilPos = setup.pu.canonicalToActivePhaseIdx(IndexTraits::oilPhaseIdx);
        ws.well("PROD").surface_rates[oilPos] = 10.0;

        return ws;
    }();

    const auto rpt = wstate
        .report(setup.grid.c_grid()->global_cell, [](const int) { return false; },
                Opm::WellState<double, IndexTraits>::RsConstInfo { true, 0.37 });

    const auto& rates = rpt.at("PROD").rates;
    BOOST_CHECK_CLOSE(rates.get(Opm::data::Rates::opt::oil), 10.0, 1e-12);
    BOOST_CHECK_CLOSE(rates.get(Opm::data::Rates::opt::gas), 3.7, 1e-12);
    BOOST_CHECK_CLOSE(rates.get(Opm::data::Rates::opt::dissolved_gas), 3.7, 1e-12);
}

// ---------------------------------------------------------------------

BOOST_AUTO_TEST_CASE(RSCONST_InjectionWellDoesNotSynthesizeGas)
{
    const auto setup = Setup { makeOilWaterRsconstWellDeck() };
    auto pinfos = std::vector<Opm::ParallelWellInfo<double>>{};

    const auto wstate = [&setup, &pinfos]() {
        const auto tstep = std::size_t{0};
        return buildWellState(setup, tstep, pinfos);
    }();

    const auto rpt = wstate
        .report(setup.grid.c_grid()->global_cell, [](const int) { return false; },
                Opm::WellState<double, IndexTraits>::RsConstInfo { true, 0.37 });

    const auto& rates = rpt.at("INJ").rates;
    BOOST_CHECK(!rates.has(Opm::data::Rates::opt::gas));
    BOOST_CHECK(!rates.has(Opm::data::Rates::opt::dissolved_gas));
}

// ---------------------------------------------------------------------

//BOOST_AUTO_TEST_CASE(GlobalWellInfo_TEST) {
//    const Setup setup{ "msw.data" };
//    std::vector<Opm::Well> local_wells = { setup.sched.getWell("PROD01", 1) };
//    Opm::GlobalWellInfo gwi(setup.sched, 1, local_wells);
//    Opm::WellContainer<Opm::Well::Status> status({{"PROD01", Opm::Well::Status::OPEN}});
//
//    BOOST_CHECK(!gwi.in_injecting_group("INJE01"));
//    BOOST_CHECK(!gwi.in_injecting_group("PROD01"));
//    BOOST_CHECK(!gwi.in_producing_group("INJE01"));
//    BOOST_CHECK(!gwi.in_producing_group("PROD01"));
//
//    BOOST_CHECK_EQUAL( gwi.well_name(0), "INJE01");
//    BOOST_CHECK_EQUAL( gwi.well_name(1), "PROD01");
//    BOOST_CHECK_EQUAL( gwi.well_index("PROD01"), 1);
//
//    BOOST_CHECK_THROW( gwi.update_group( {}, {}, {} ), std::exception);
//
//
//    Opm::WellContainer<Opm::Well::InjectorCMode> inj_cmode({{"PROD01", Opm::Well::InjectorCMode::CMODE_UNDEFINED}});
//    {
//        Opm::WellContainer<Opm::Well::ProducerCMode> prod_cmode({{"PROD01", Opm::Well::ProducerCMode::GRUP}});
//        gwi.update_group(status, inj_cmode, prod_cmode);
//    }
//    BOOST_CHECK(!gwi.in_producing_group("INJE01"));
//    BOOST_CHECK(gwi.in_producing_group("PROD01"));
//
//    {
//        Opm::WellContainer<Opm::Well::ProducerCMode> prod_cmode(
//            {{"PROD01", Opm::Well::ProducerCMode::CMODE_UNDEFINED}});
//        gwi.update_group(status, inj_cmode, prod_cmode);
//    }
//
//    {
//        Opm::WellContainer<Opm::Well::ProducerCMode> prod_cmode({{"PROD01", Opm::Well::ProducerCMode::GRUP}});
//        gwi.update_group(status, inj_cmode, prod_cmode);
//    }
//    BOOST_CHECK(!gwi.in_producing_group("INJE01"));
//    BOOST_CHECK(gwi.in_producing_group("PROD01"));
//
//    {
//        Opm::WellContainer<Opm::Well::ProducerCMode> prod_cmode({{"PROD01", Opm::Well::ProducerCMode::NONE}});
//        gwi.update_group(status, inj_cmode, prod_cmode);
//    }
//    BOOST_CHECK(!gwi.in_producing_group("INJE01"));
//    BOOST_CHECK(!gwi.in_producing_group("PROD01"));
//}

// ---------------------------------------------------------------------

BOOST_AUTO_TEST_CASE(TESTWellContainer) {
    Opm::WellContainer<int> wc;
    BOOST_CHECK_EQUAL(wc.size(), 0);

    wc.add("W1", 1);
    wc.add("W2", 2);

    BOOST_CHECK_EQUAL(wc.size(), 2);

    BOOST_CHECK_THROW(wc.add("W1", 1), std::exception);

    BOOST_CHECK_THROW(wc[10], std::exception);
    BOOST_CHECK_EQUAL(wc[0], 1);
    BOOST_CHECK_EQUAL(wc[1], 2);

    BOOST_CHECK_THROW(wc["INVALID_WELL"], std::exception);
    BOOST_CHECK_EQUAL(wc["W1"], 1);
    BOOST_CHECK_EQUAL(wc["W2"], 2);
    BOOST_CHECK_EQUAL(wc.well_name(0), "W1");
    BOOST_CHECK_EQUAL(wc.well_name(1), "W2");
    BOOST_CHECK_THROW(wc.well_name(10), std::exception);

    const auto& wells = wc.wells();
    std::vector<std::string> expected = {"W1", "W2"};
    BOOST_CHECK( std::is_permutation( wells.begin(), wells.end(), expected.begin(), expected.end()) );

    Opm::WellContainer<int> wc2;
    wc2.copy_welldata(wc);
    BOOST_CHECK_EQUAL(wc2.size() , 0);

    wc2.add("W1", 100);
    BOOST_CHECK_EQUAL(wc2["W1"], 100);
    wc2.copy_welldata(wc);
    BOOST_CHECK_EQUAL(wc2["W1"], 1);

    Opm::WellContainer<int> wc3;
    wc3.add("W2", 100);
    wc3.copy_welldata(wc);
    BOOST_CHECK_EQUAL(wc3["W2"], 2);
    BOOST_CHECK_EQUAL(wc3[0], 2);

    wc3["W2"] = 200;
    wc3.add("W3", 300);
    wc3.copy_welldata(wc);
    BOOST_CHECK_EQUAL(wc3["W2"], 2);
    BOOST_CHECK_EQUAL(wc3[0], 2);
    BOOST_CHECK_EQUAL(wc3["W3"], 300);
    BOOST_CHECK_EQUAL(wc3[1], 300);

    BOOST_CHECK_THROW(wc3.copy_welldata(wc, "W1"), std::exception);
    BOOST_CHECK_THROW(wc3.copy_welldata(wc, "W3"), std::exception);

    wc.clear();
    BOOST_CHECK_EQUAL(wc.size(), 0);


    BOOST_CHECK(wc3.has("W2"));
    BOOST_CHECK(!wc3.has("NO_SUCH_WELL"));


    std::vector<int> vec_copy(wc3.begin(), wc3.end());
    BOOST_CHECK_EQUAL(vec_copy.size(), wc3.size());
    for (std::size_t i = 0; i < wc3.size(); i++)
        BOOST_CHECK_EQUAL(vec_copy[i], wc3[i]);


    Opm::WellContainer<int> wci({{"W1", 1}, {"W2", 2}, {"W3", 3}});
    BOOST_CHECK_EQUAL(wci.size(), 3);
    BOOST_CHECK(wci.has("W1"));
    BOOST_CHECK_EQUAL(wci[1], 2);
    BOOST_CHECK_EQUAL(wci["W3"], 3);

    auto w3 = wci.well_index("W3");
    BOOST_CHECK(w3.has_value());
    BOOST_CHECK_EQUAL(w3.value(), 2);

    auto wx = wci.well_index("WX");
    BOOST_CHECK(!wx.has_value());
}

// ---------------------------------------------------------------------

BOOST_AUTO_TEST_CASE(TESTSegmentState) {
    const Setup setup{ "msw.data" };
    const auto& well = setup.sched.getWell("PROD01", 0);
    const auto& segments = well.getSegments();
    Opm::SegmentState<double> ss1(3, segments);
    Opm::SegmentState<double> ss2;


    ss1.pressure_drop_hydrostatic[0] = 1;
    ss1.pressure_drop_friction[0] = 2;
    ss1.pressure_drop_accel[0] = 3;

    BOOST_CHECK_EQUAL(ss1.pressure_drop(0), 6);
}

// ---------------------------------------------------------------------

BOOST_AUTO_TEST_CASE(TESTSegmentState2) {
    const Setup setup{ "msw.data" };
    std::vector<Opm::ParallelWellInfo<double>> pinfo;
    const auto wstate = buildWellState(setup, 0, pinfo);
    const auto& well = setup.sched.getWell("PROD01", 0);
    const auto& ws = wstate.well("PROD01");

    auto segments = ws.segments;
    BOOST_CHECK_EQUAL(segments.pressure.size(), well.getSegments().size());

    segments.pressure_drop_friction[0] = 1;
    segments.pressure_drop_accel[0] = 2;
    segments.pressure_drop_hydrostatic[0] = 4;

    BOOST_CHECK_EQUAL(segments.pressure_drop(0), 7);


    for (std::size_t i=0; i < segments.pressure.size(); i++)
        segments.pressure[i] = (i + 1);

    const double bhp = 2.0;
    segments.scale_pressure(bhp);

    for (std::size_t i=0; i < segments.pressure.size(); i++)
        BOOST_CHECK_EQUAL(segments.pressure[i], 1.+(i+1));

    BOOST_CHECK_EQUAL( segments.size(), well.getSegments().size() );
}

// ---------------------------------------------------------------------

BOOST_AUTO_TEST_CASE(TESTPerfData) {
    Opm::PerfData<double> pd1(3, true, 3);
    Opm::PerfData<double> pd2(3, true, 3);
    Opm::PerfData<double> pd3(2, true, 3);
    Opm::PerfData<double> pd4(3, false, 3);


    for (std::size_t i = 0; i < 3; i++) {
        pd1.pressure[i] = i+1;
        pd2.pressure[i] = 10*(i+1);
    }

    BOOST_CHECK(pd1.try_assign(pd2));
    for (std::size_t i = 0; i < 3; i++) {
        BOOST_CHECK(pd2.pressure[i] == 10*(i+1));
        BOOST_CHECK(pd1.pressure[i] == 10*(i+1));
    }

    BOOST_CHECK(!pd1.try_assign(pd3));
    for (std::size_t i = 0; i < 3; i++) {
        BOOST_CHECK(pd1.pressure[i] == 10*(i+1));
    }

    BOOST_CHECK(!pd1.try_assign(pd4));
}

// ---------------------------------------------------------------------

BOOST_AUTO_TEST_CASE(TestSingleWellState) {
    Opm::ParallelWellInfo<double> pinfo;
    std::vector<Opm::PerforationData<double>> connections = {{0,1,1,0,0},{1,1,1,0,1},{2,1,1,0,2}};
    using SingleWellState = Opm::SingleWellState<double, Opm::BlackOilDefaultFluidSystemIndices>;
    using IndexTraits = Opm::BlackOilDefaultFluidSystemIndices;
    using PhaseUsage = Opm::PhaseUsageInfo<IndexTraits>;

    PhaseUsage pu;
    pu.initFromPhases({Opm::Phases{true, true, true}});

    SingleWellState ws1("W1", pinfo, pu, true,  100.0, connections, 1.0);
    SingleWellState ws2("W2", pinfo, pu, true,  100.0, connections, 2.0);
    SingleWellState ws3("W3", pinfo, pu, false, 100.0, connections, 3.0);

    ws1.bhp = 100;
    ws1.thp = 200;

    ws2.init_timestep(ws1);
    BOOST_CHECK_EQUAL(ws2.bhp, ws1.bhp);
    BOOST_CHECK_EQUAL(ws2.thp, ws1.thp);

    ws3.bhp = ws1.bhp * 2;
    ws3.thp = ws1.thp * 2;
    ws3.init_timestep(ws1);
    BOOST_CHECK(ws3.bhp != ws1.bhp);
    BOOST_CHECK(ws3.thp != ws1.thp);
}

BOOST_AUTO_TEST_SUITE_END()
