/*
  Copyright 2017 SINTEF Digital, Mathematics and Cybernetics.
  Copyright 2017 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/simulators/wells/MultisegmentWellSegments.hpp>

#include <opm/common/ErrorMacros.hpp>

#include <opm/input/eclipse/Schedule/MSW/AICD.hpp>
#include <opm/input/eclipse/Schedule/MSW/WellSegments.hpp>
#include <opm/input/eclipse/Schedule/Well/WellConnections.hpp>

#include <opm/material/densead/EvaluationFormat.hpp>

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

#include <opm/models/blackoil/blackoilvariableandequationindices.hh>
#include <opm/models/blackoil/blackoilonephaseindices.hh>
#include <opm/models/blackoil/blackoiltwophaseindices.hh>

#include <opm/simulators/utils/DeferredLogger.hpp>
#include <opm/simulators/wells/MSWellHelpers.hpp>
#include <opm/simulators/wells/SegmentState.hpp>
#include <opm/simulators/wells/WellInterfaceGeneric.hpp>

#include <fmt/format.h>

#include <algorithm>
#include <array>
#include <cmath>
#include <cstddef>
#include <functional>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>

namespace {

template <typename FluidSystem>
std::vector<typename FluidSystem::Scalar> surfaceDensities(const int pvt_region, const std::size_t num_quantities)
{
    using Scalar = typename FluidSystem::Scalar;
    auto surf_dens = std::vector<Scalar>(num_quantities);

    for (auto phIdx = 0*FluidSystem::numPhases; phIdx < FluidSystem::numPhases; ++phIdx) {
        if (!FluidSystem::phaseIsActive(phIdx)) {
            continue;
        }

        const auto compIdx = FluidSystem::canonicalToActiveCompIdx(FluidSystem::solventComponentIndex(phIdx));
        surf_dens[compIdx] = FluidSystem::referenceDensity(phIdx, pvt_region);
    }

    return surf_dens;
}

} // anonymous namespace

namespace Opm
{

template<class FluidSystem, class Indices>
MultisegmentWellSegments<FluidSystem,Indices>::
MultisegmentWellSegments(const int numSegments,
                         const ParallelWellInfo<Scalar>& parallel_well_info,
                         WellInterfaceGeneric<Scalar, IndexTraits>& well)
    : perforations_(numSegments)
    , local_perforation_depth_diffs_(well.numLocalPerfs(), 0.0)
    // Generally, the info stored with the class MultisegmentWellSegments is global, i.e., the same across all
    // processes, since segments do not have a 1:1 correspondence to the grid. All members of this class store
    // global information, except for local_perforation_depth_diffs_, which is the only one that contains only
    // local information. This is an exception and intentionally, since here, we only need the local entries.
    , inlets_(well.wellEcl().getSegments().size())
    , depth_diffs_(numSegments, 0.0)
    , surface_densities_(surfaceDensities<FluidSystem>(well.pvtRegionIdx(), well.numConservationQuantities()))
    , densities_(numSegments, 0.0)
    , mass_rates_(numSegments, 0.0)
    , viscosities_(numSegments, 0.0)
    , upwinding_segments_(numSegments, 0)
    , phase_densities_(numSegments, std::vector<EvalWell>(well.numConservationQuantities(), 0.0)) // number of phase here?
    , phase_fractions_(numSegments, std::vector<EvalWell>(well.numConservationQuantities(), 0.0)) // number of phase here?
    , phase_viscosities_(numSegments, std::vector<EvalWell>(well.numConservationQuantities(), 0.0)) // number of phase here?
    , well_(well)
{
    // since we decide to use the WellSegments from the well parser. we can reuse a lot from it.
    // for other facilities needed but not available from parser, we need to process them here

    // initialize the segment_perforations_ and update perforation_segment_depth_diffs_
    const WellConnections& completion_set = well_.wellEcl().getConnections();
    // index of the perforation within wells struct
    // there might be some perforations not active, which causes the number of the perforations in
    // well_ecl_ and wells struct different
    // the current implementation is a temporary solution for now, it should be corrected from the parser
    // side
    int i_perf_wells = 0;
    // The perfDepth vector will contain the depths of all perforations across all processes of this well!
    int num_perfs_whole_mswell = parallel_well_info.communication().sum(well.numLocalPerfs());
    well.perfDepth().resize(num_perfs_whole_mswell, 0.0);
    const auto& segment_set = well_.wellEcl().getSegments();
    for (std::size_t perf = 0; perf < completion_set.size(); ++perf) {
        const Connection& connection = completion_set.get(perf);
        if (connection.state() == Connection::State::OPEN) {
            const int segment_index = segment_set.segmentNumberToIndex(connection.segment());
            if (segment_index == -1) {
                OPM_THROW(std::logic_error,
                          fmt::format("COMPSEGS: Well {} has connection in cell {}, {}, {} "
                                      "without associated segment.", well_.wellEcl().name(),
                                      connection.getI() + 1, connection.getJ() + 1,
                                      connection.getK() + 1));
            }
            perforations_[segment_index].push_back(i_perf_wells);
            well.perfDepth()[i_perf_wells] = connection.depth();
            const Scalar segment_depth = segment_set[segment_index].depth();
            int local_perf_index = parallel_well_info.activePerfToLocalPerf(i_perf_wells);
            if (local_perf_index > -1) // If local_perf_index == -1, then the perforation is not on this process
                local_perforation_depth_diffs_[local_perf_index] = well_.perfDepth()[i_perf_wells] - segment_depth;
            i_perf_wells++;
        }
    }

    // initialize the segment_inlets_
    for (const Segment& segment : segment_set) {
        const int segment_number = segment.segmentNumber();
        const int outlet_segment_number = segment.outletSegment();
        if (outlet_segment_number > 0) {
            const int segment_index = segment_set.segmentNumberToIndex(segment_number);
            const int outlet_segment_index = segment_set.segmentNumberToIndex(outlet_segment_number);
            inlets_[outlet_segment_index].push_back(segment_index);
        }
    }

    // calculating the depth difference between the segment and its oulet_segments
    // for the top segment, we will make its zero unless we find other purpose to use this value
    for (int seg = 1; seg < numSegments; ++seg) {
        const Scalar segment_depth = segment_set[seg].depth();
        const int outlet_segment_number = segment_set[seg].outletSegment();
        const Segment& outlet_segment = segment_set[segment_set.segmentNumberToIndex(outlet_segment_number)];
        const Scalar outlet_depth = outlet_segment.depth();
        depth_diffs_[seg] = segment_depth - outlet_depth;
    }
}

template<class FluidSystem, class Indices>
void MultisegmentWellSegments<FluidSystem,Indices>::
computeFluidProperties(const EvalWell& temperature,
                       const EvalWell& saltConcentration,
                       const PrimaryVariables& primary_variables,
                       DeferredLogger& deferred_logger)
{
    const int num_quantities = well_.numConservationQuantities();
    PhaseCalcResult result(static_cast<size_t>(num_quantities));

    for (std::size_t seg = 0; seg < perforations_.size(); ++seg) {
        calculatePhaseProperties(result, temperature, saltConcentration,
                                 primary_variables, seg, true, deferred_logger);

        phase_densities_[seg] = result.phase_densities;
        phase_viscosities_[seg] = result.phase_viscosities;

        const auto& mix = result.mix;
        const auto& mix_s = result.mix_s;
        const auto& b = result.b;
        const auto& volrat = result.vol_ratio;

        viscosities_[seg] = 0.;
        // calculate the average viscosity
        for (int comp_idx = 0; comp_idx < num_quantities; ++comp_idx) {
            const EvalWell fraction =  mix[comp_idx] / b[comp_idx] / volrat;
            // TODO: a little more work needs to be done to handle the negative fractions here
            phase_fractions_[seg][comp_idx] = fraction; // >= 0.0 ? fraction : 0.0;
            viscosities_[seg] += phase_viscosities_[seg][comp_idx] * phase_fractions_[seg][comp_idx];
        }

        EvalWell density(0.0);
        for (int comp_idx = 0; comp_idx < num_quantities; ++comp_idx) {
            density += surface_densities_[comp_idx] * mix_s[comp_idx];
        }
        densities_[seg] = density / volrat;

        // calculate the mass rates
        mass_rates_[seg] = 0.;
        for (int comp_idx = 0; comp_idx < well_.numConservationQuantities(); ++comp_idx) {
            const int upwind_seg = upwinding_segments_[seg];
            const EvalWell rate = primary_variables.getSegmentRateUpwinding(seg,
                                                                            upwind_seg,
                                                                            comp_idx);
            mass_rates_[seg] += rate * surface_densities_[comp_idx];
        }
    }
}

template<class FluidSystem, class Indices>
void MultisegmentWellSegments<FluidSystem,Indices>::
updateUpwindingSegments(const PrimaryVariables& primary_variables)
{
    for (std::size_t seg = 0; seg < perforations_.size(); ++seg) {
        // special treatment is needed for segment 0
        if (seg == 0) {
            // we are not supposed to have injecting producers and producing injectors
            assert(!(well_.isProducer() && primary_variables.eval(seg)[primary_variables.WQTotal] > 0.));
            assert(!(well_.isInjector() && primary_variables.eval(seg)[primary_variables.WQTotal] < 0.));
            upwinding_segments_[seg] = seg;
            continue;
        }

        // for other normal segments
        if (primary_variables.eval(seg)[primary_variables.WQTotal] <= 0.) {
            upwinding_segments_[seg] = seg;
        } else {
            const auto& segment_set = well_.wellEcl().getSegments();
            const int outlet_segment_index = segment_set.segmentNumberToIndex(segment_set[seg].outletSegment());
            upwinding_segments_[seg] = outlet_segment_index;
        }
    }
}

template<class FluidSystem, class Indices>
typename MultisegmentWellSegments<FluidSystem,Indices>::EvalWell
MultisegmentWellSegments<FluidSystem,Indices>::
getHydroPressureLoss(const int seg,
                     const int seg_density) const
{
    return densities_[seg_density] * well_.gravity() * depth_diffs_[seg];
}

template<class FluidSystem, class Indices>
typename MultisegmentWellSegments<FluidSystem,Indices>::Scalar
MultisegmentWellSegments<FluidSystem,Indices>::
getPressureDiffSegLocalPerf(const int seg,
                            const int local_perf_index) const
{
    return well_.gravity() * densities_[seg].value() * local_perforation_depth_diffs_[local_perf_index];
}

template<class FluidSystem, class Indices>
typename MultisegmentWellSegments<FluidSystem,Indices>::EvalWell
MultisegmentWellSegments<FluidSystem,Indices>::
getSurfaceVolume(const EvalWell& temperature,
                 const EvalWell& saltConcentration,
                 const PrimaryVariables& primary_variables,
                 const int seg_idx,
                 DeferredLogger& deferred_logger) const
{
    PhaseCalcResult result(static_cast<size_t>(well_.numConservationQuantities()));
    calculatePhaseProperties(result, temperature, saltConcentration,
                             primary_variables, seg_idx, false, deferred_logger);

    const EvalWell& vol_ratio = result.vol_ratio;

    // We increase the segment volume with a factor 10 to stabilize the system.
    const Scalar volume = well_.wellEcl().getSegments()[seg_idx].volume();

    return volume / vol_ratio;
}

template<class FluidSystem, class Indices>
typename MultisegmentWellSegments<FluidSystem,Indices>::EvalWell
MultisegmentWellSegments<FluidSystem,Indices>::
getFrictionPressureLoss(const int seg,
                        const bool extra_reverse_flow_derivatives /*false*/) const
{
    EvalWell mass_rate = mass_rates_[seg];
    const int seg_upwind = upwinding_segments_[seg];
    EvalWell density = densities_[seg_upwind];
    EvalWell visc = viscosities_[seg_upwind];
    // In the reverse flow case, we don't have enough slots for all derivatives, e.g.,
    // upwind pressure and flow. We amend this by a second function call optioin, where
    // only these remaining derivatives are considered.
    // For reference: the pressure equation assumes pressure/flow derivatives are given
    // at segment node while fraction derivatives are given at upwind node.

    if (seg != seg_upwind) {
        if (!extra_reverse_flow_derivatives) {
            constexpr int WQTotal = Indices::numEq + PrimaryVariables::WQTotal;
            constexpr int SPres = Indices::numEq + PrimaryVariables::SPres;
            density.setDerivative(WQTotal, 0.0);
            density.setDerivative(SPres, 0.0);
            visc.setDerivative(WQTotal, 0.0);
            visc.setDerivative(SPres, 0.0);
        } else {
            if constexpr (PrimaryVariables::has_wfrac_variable) {
                constexpr int WFrac = Indices::numEq + PrimaryVariables::WFrac;
                density.setDerivative(WFrac, 0.0);
                visc.setDerivative(WFrac, 0.0);
            }
            if constexpr (PrimaryVariables::has_gfrac_variable) {
                constexpr int GFrac = Indices::numEq + PrimaryVariables::GFrac;
                density.setDerivative(GFrac, 0.0);
                visc.setDerivative(GFrac, 0.0);
            }
            mass_rate.clearDerivatives();
        }
    }

    const auto& segment_set = well_.wellEcl().getSegments();
    const int outlet_segment_index = segment_set.segmentNumberToIndex(segment_set[seg].outletSegment());
    const Scalar length = segment_set[seg].totalLength() - segment_set[outlet_segment_index].totalLength();
    assert(length > 0.);
    const Scalar roughness = segment_set[seg].roughness();
    const Scalar area = segment_set[seg].crossArea();
    const Scalar diameter = segment_set[seg].internalDiameter();

    const Scalar sign = mass_rate < 0. ? 1.0 : - 1.0;

    return sign * mswellhelpers::frictionPressureLoss(length, diameter, area, roughness, density, mass_rate, visc);
}

template<class FluidSystem, class Indices>
typename MultisegmentWellSegments<FluidSystem,Indices>::EvalWell
MultisegmentWellSegments<FluidSystem,Indices>::
pressureDropSpiralICD(const int seg,
                      const bool extra_reverse_flow_derivatives /*false*/) const
{
    const auto& segment_set = well_.wellEcl().getSegments();
    const SICD& sicd = segment_set[seg].spiralICD();

    const int seg_upwind = upwinding_segments_[seg];
    const std::vector<EvalWell>& phase_fractions = phase_fractions_[seg_upwind];
    const std::vector<EvalWell>& phase_viscosities = phase_viscosities_[seg_upwind];

    EvalWell water_fraction = 0.;
    EvalWell water_viscosity = 0.;
    if (FluidSystem::phaseIsActive(FluidSystem::waterPhaseIdx)) {
        const int water_pos = FluidSystem::canonicalToActiveCompIdx(FluidSystem::waterCompIdx);
        water_fraction = phase_fractions[water_pos];
        water_viscosity = phase_viscosities[water_pos];
    }

    EvalWell oil_fraction = 0.;
    EvalWell oil_viscosity = 0.;
    if (FluidSystem::phaseIsActive(FluidSystem::oilPhaseIdx)) {
        const int oil_pos = FluidSystem::canonicalToActiveCompIdx(FluidSystem::oilCompIdx);
        oil_fraction = phase_fractions[oil_pos];
        oil_viscosity = phase_viscosities[oil_pos];
    }

    EvalWell gas_fraction = 0.;
    EvalWell gas_viscosity = 0.;
    if (FluidSystem::phaseIsActive(FluidSystem::gasPhaseIdx)) {
        const int gas_pos = FluidSystem::canonicalToActiveCompIdx(FluidSystem::gasCompIdx);
        gas_fraction = phase_fractions[gas_pos];
        gas_viscosity = phase_viscosities[gas_pos];
    }

    EvalWell density = densities_[seg_upwind];
    EvalWell mass_rate = mass_rates_[seg];
    // In the reverse flow case, we don't have enough slots for all derivatives, e.g.,
    // upwind pressure and flow. We amend this by a second function call option, where
    // only these remaining derivatives are considered.
    // For reference: the pressure equation assumes pressure/flow derivatives are given
    // at segment node while fraction derivatives are given at upwind node.
    if (seg != seg_upwind) {
        constexpr int nvar = PrimaryVariables::numWellEq;
        std::vector<bool> zero_mask(nvar, false);
        if (!extra_reverse_flow_derivatives){
            zero_mask[PrimaryVariables::WQTotal] = true;
            zero_mask[PrimaryVariables::SPres] = true;
        } else {
            if constexpr (PrimaryVariables::has_wfrac_variable) {
                zero_mask[PrimaryVariables::WFrac] = true;
            }
            if constexpr (PrimaryVariables::has_gfrac_variable) {
                zero_mask[PrimaryVariables::GFrac] = true;
            }
            // mass_rate has no extra derivatives (they are organized as in equations)
            mass_rate.clearDerivatives();
        }
        for (int ii = 0; ii < nvar; ++ii) {
            if (zero_mask[ii]) {
                water_fraction.setDerivative(Indices::numEq + ii, 0.0);
                water_viscosity.setDerivative(Indices::numEq + ii, 0.0);
                oil_fraction.setDerivative(Indices::numEq + ii, 0.0);
                oil_viscosity.setDerivative(Indices::numEq + ii, 0.0);
                gas_fraction.setDerivative(Indices::numEq + ii, 0.0);
                gas_viscosity.setDerivative(Indices::numEq + ii, 0.0);
                density.setDerivative(Indices::numEq + ii, 0.0);
            }
        }
    }

    const EvalWell liquid_fraction = water_fraction + oil_fraction;

    // viscosity contribution from the liquid
    const EvalWell liquid_viscosity_fraction = liquid_fraction < 1.e-30 ? oil_fraction * oil_viscosity + water_fraction * water_viscosity :
            liquid_fraction * mswellhelpers::emulsionViscosity<EvalWell,Scalar>(water_fraction, water_viscosity, oil_fraction, oil_viscosity, sicd);

    const EvalWell mixture_viscosity = liquid_viscosity_fraction + gas_fraction * gas_viscosity;

    const EvalWell reservoir_rate = mass_rate / density;

    const EvalWell reservoir_rate_icd = reservoir_rate * sicd.scalingFactor();

    const Scalar viscosity_cali = sicd.viscosityCalibration();

    using MathTool = MathToolbox<EvalWell>;

    const Scalar density_cali = sicd.densityCalibration();
    // make sure we don't pass negative base to powers
    const EvalWell temp_value1 = density > 0.0 ? MathTool::pow(density / density_cali, 0.75) : 0.0;
    const EvalWell temp_value2 = mixture_viscosity > 0.0 ? MathTool::pow(mixture_viscosity / viscosity_cali, 0.25) : 0.0;

    // formulation before 2016, base_strength is used
    // const double base_strength = sicd.strength() / density_cali;
    // formulation since 2016, strength is used instead
    const Scalar strength = sicd.strength();

    const Scalar sign = reservoir_rate_icd <= 0. ? 1.0 : -1.0;

    return sign * temp_value1 * temp_value2 * strength * reservoir_rate_icd * reservoir_rate_icd;
}

template<class FluidSystem, class Indices>
typename MultisegmentWellSegments<FluidSystem,Indices>::EvalWell
MultisegmentWellSegments<FluidSystem,Indices>::
pressureDropAutoICD(const int seg,
                    const UnitSystem& unit_system,
                    const bool extra_reverse_flow_derivatives /*false*/) const
{
    const auto& segment_set = well_.wellEcl().getSegments();
    const AutoICD& aicd = segment_set[seg].autoICD();

    const int seg_upwind = upwinding_segments_[seg];
    const std::vector<EvalWell>& phase_fractions = phase_fractions_[seg_upwind];
    const std::vector<EvalWell>& phase_viscosities = phase_viscosities_[seg_upwind];
    const std::vector<EvalWell>& phase_densities = phase_densities_[seg_upwind];

    EvalWell water_fraction = 0.;
    EvalWell water_viscosity = 0.;
    EvalWell water_density = 0.;
    if (FluidSystem::phaseIsActive(FluidSystem::waterPhaseIdx)) {
        const int water_pos = FluidSystem::canonicalToActiveCompIdx(FluidSystem::waterCompIdx);
        water_fraction = phase_fractions[water_pos];
        water_viscosity = phase_viscosities[water_pos];
        water_density = phase_densities[water_pos];
    }

    EvalWell oil_fraction = 0.;
    EvalWell oil_viscosity = 0.;
    EvalWell oil_density = 0.;
    if (FluidSystem::phaseIsActive(FluidSystem::oilPhaseIdx)) {
        const int oil_pos = FluidSystem::canonicalToActiveCompIdx(FluidSystem::oilCompIdx);
        oil_fraction = phase_fractions[oil_pos];
        oil_viscosity = phase_viscosities[oil_pos];
        oil_density = phase_densities[oil_pos];
    }

    EvalWell gas_fraction = 0.;
    EvalWell gas_viscosity = 0.;
    EvalWell gas_density = 0.;
    if (FluidSystem::phaseIsActive(FluidSystem::gasPhaseIdx)) {
        const int gas_pos = FluidSystem::canonicalToActiveCompIdx(FluidSystem::gasCompIdx);
        gas_fraction = phase_fractions[gas_pos];
        gas_viscosity = phase_viscosities[gas_pos];
        gas_density = phase_densities[gas_pos];
    }

    EvalWell density = densities_[seg_upwind];
    EvalWell mass_rate = mass_rates_[seg];
    // In the reverse flow case, we don't have enough slots for all derivatives, e.g.,
    // upwind pressure and flow. We amend this by a second function call option, where
    // only these remaining derivatives are considered.
    // For reference: the pressure equation assumes pressure/flow derivatives are given
    // at segment node while fraction derivatives are given at upwind node.
    if (seg != seg_upwind) {
        constexpr int nvar = PrimaryVariables::numWellEq;
        std::vector<bool> zero_mask(nvar, false);
        if (!extra_reverse_flow_derivatives){
            zero_mask[PrimaryVariables::WQTotal] = true;
            zero_mask[PrimaryVariables::SPres] = true;
        } else {
            if constexpr (PrimaryVariables::has_wfrac_variable) {
                zero_mask[PrimaryVariables::WFrac] = true;
            }
            if constexpr (PrimaryVariables::has_gfrac_variable) {
                zero_mask[PrimaryVariables::GFrac] = true;
            }
            // mass_rate has no extra derivatives (they are organized as in equations)
            mass_rate.clearDerivatives();
        }
        for (int ii = 0; ii < nvar; ++ii) {
            if (zero_mask[ii]) {
                water_fraction.setDerivative(Indices::numEq + ii, 0.0);
                water_viscosity.setDerivative(Indices::numEq + ii, 0.0);
                water_density.setDerivative(Indices::numEq + ii, 0.0);
                oil_fraction.setDerivative(Indices::numEq + ii, 0.0);
                oil_viscosity.setDerivative(Indices::numEq + ii, 0.0);
                oil_density.setDerivative(Indices::numEq + ii, 0.0);
                gas_fraction.setDerivative(Indices::numEq + ii, 0.0);
                gas_viscosity.setDerivative(Indices::numEq + ii, 0.0);
                gas_density.setDerivative(Indices::numEq + ii, 0.0);
                density.setDerivative(Indices::numEq + ii, 0.0);
            }
        }
    }

    using MathTool = MathToolbox<EvalWell>;
    // make sure we don't pass negative base to powers
    auto safe_pow = [](const auto& a, const Scalar b) {
        return a > 0.0 ? MathTool::pow(a,b) : 0.0;
    };

    const EvalWell mixture_viscosity = safe_pow(water_fraction, aicd.waterViscExponent()) * water_viscosity
                                     + safe_pow(oil_fraction, aicd.oilViscExponent()) * oil_viscosity
                                     + safe_pow(gas_fraction, aicd.gasViscExponent()) * gas_viscosity;

    const EvalWell mixture_density = safe_pow(water_fraction, aicd.waterDensityExponent()) * water_density
                                   + safe_pow(oil_fraction, aicd.oilDensityExponent()) * oil_density
                                   + safe_pow(gas_fraction, aicd.gasDensityExponent()) * gas_density;

    const Scalar rho_reference = aicd.densityCalibration();
    const Scalar visc_reference = aicd.viscosityCalibration();
    const auto volume_rate_icd = mass_rate * aicd.scalingFactor() / mixture_density;
    const Scalar sign = volume_rate_icd <= 0. ? 1.0 : -1.0;
    // convert 1 unit volume rate
    using M  = UnitSystem::measure;
    const Scalar unit_volume_rate = unit_system.to_si(M::geometric_volume_rate, 1.);

    // TODO: we did not consider the maximum allowed rate here
    const auto result = sign * safe_pow(mixture_density / rho_reference, aicd.densityExponent())
                      * safe_pow(visc_reference/mixture_viscosity, aicd.viscExponent()) * mixture_density
                      * aicd.strength() * safe_pow( -sign * volume_rate_icd, aicd.flowRateExponent())
                      * std::pow(unit_volume_rate, (2. - aicd.flowRateExponent())) ;
    return result;
}

template<class FluidSystem, class Indices>
typename MultisegmentWellSegments<FluidSystem,Indices>::EvalWell
MultisegmentWellSegments<FluidSystem,Indices>::
pressureDropValve(const int seg,
                  const SummaryState& summary_state,
                  const bool extra_reverse_flow_derivatives /*false*/) const
{
    const Opm::WellSegments& segment_set = well_.wellEcl().getSegments();
    const Valve& valve = segment_set[seg].valve();

    EvalWell mass_rate = mass_rates_[seg];
    const int seg_upwind = upwinding_segments_[seg];
    EvalWell visc = viscosities_[seg_upwind];
    EvalWell density = densities_[seg_upwind];
    // In the reverse flow case, we don't have enough slots for all derivatives, e.g.,
    // upwind pressure and flow. We amend this by a second function call optioin, where
    // only these remaining derivatives are considered.
    // For reference: the pressure equation assumes pressure/flow derivatives are given
    // at segment node while fraction derivatives are given at upwind node.
    if (seg != seg_upwind) {
        if (!extra_reverse_flow_derivatives) {
            constexpr int WQTotal = Indices::numEq + PrimaryVariables::WQTotal;
            constexpr int SPres = Indices::numEq + PrimaryVariables::SPres;
            density.setDerivative(WQTotal, 0.0);
            density.setDerivative(SPres, 0.0);
            visc.setDerivative(WQTotal, 0.0);
            visc.setDerivative(SPres, 0.0);
        } else {
            if constexpr (PrimaryVariables::has_wfrac_variable) {
                constexpr int WFrac = Indices::numEq + PrimaryVariables::WFrac;
                density.setDerivative(WFrac, 0.0);
                visc.setDerivative(WFrac, 0.0);
            }
            if constexpr (PrimaryVariables::has_gfrac_variable) {
                constexpr int GFrac = Indices::numEq + PrimaryVariables::GFrac;
                density.setDerivative(GFrac, 0.0);
                visc.setDerivative(GFrac, 0.0);
            }
            mass_rate.clearDerivatives();
        }
    }

    const Scalar additional_length = valve.pipeAdditionalLength();
    const Scalar roughness = valve.pipeRoughness();
    const Scalar diameter = valve.pipeDiameter();
    const Scalar area = valve.pipeCrossArea();

    const EvalWell friction_pressure_loss =
        mswellhelpers::frictionPressureLoss(additional_length, diameter, area, roughness, density, mass_rate, visc);

    const ValveUDAEval uda_eval {summary_state, this->well_.name(), static_cast<std::size_t>(segment_set[seg].segmentNumber())};
    const Scalar area_con = valve.conCrossArea(uda_eval);
    const Scalar cv = valve.conFlowCoefficient();

    const EvalWell constriction_pressure_loss =
        mswellhelpers::valveContrictionPressureLoss(mass_rate, density, area_con, cv);

    const Scalar sign = mass_rate <= 0. ? 1.0 : -1.0;
    return sign * (friction_pressure_loss + constriction_pressure_loss);
}

template<class FluidSystem, class Indices>
typename MultisegmentWellSegments<FluidSystem,Indices>::EvalWell
MultisegmentWellSegments<FluidSystem,Indices>::
accelerationPressureLossContribution(const int seg,
                                     const Scalar area,
                                     const bool extra_reverse_flow_derivatives /*false*/) const
{
    // Compute the *signed* velocity head for given segment (sign is positive for flow towards surface, i.e., negative rate)
    // Optionally return derivatives for reversed flow case
    EvalWell mass_rate = mass_rates_[seg];
    const int seg_upwind = upwinding_segments_[seg];
    EvalWell density = densities_[seg_upwind];
    if (seg != seg_upwind) {
        if (!extra_reverse_flow_derivatives) {
            constexpr int WQTotal = Indices::numEq + PrimaryVariables::WQTotal;
            constexpr int SPres = Indices::numEq + PrimaryVariables::SPres;
            density.setDerivative(WQTotal, 0.0);
            density.setDerivative(SPres, 0.0);
        } else {
            if constexpr (PrimaryVariables::has_wfrac_variable) {
                constexpr int WFrac = Indices::numEq + PrimaryVariables::WFrac;
                density.setDerivative(WFrac, 0.0);
            }
            if constexpr (PrimaryVariables::has_gfrac_variable) {
                constexpr int GFrac = Indices::numEq + PrimaryVariables::GFrac;
                density.setDerivative(GFrac, 0.0);
            }
            mass_rate.clearDerivatives();
        }
    }
    const Scalar sign = mass_rate > 0 ? -1.0 : 1.0;
    return sign * mswellhelpers::velocityHead(area, mass_rate, density);
}

template <class FluidSystem, class Indices>
void
MultisegmentWellSegments<FluidSystem,Indices>::
copyPhaseDensities(SegmentState<Scalar>& segSol) const
{
    // TODO: probably we should change the the name phaseMap
    // this part can be wrong and needs to be checked
    auto* rho = segSol.phase_density.data();

    const auto phaseMap = std::vector {
        FluidSystem::oilPhaseIdx,
        FluidSystem::gasPhaseIdx,
        FluidSystem::waterPhaseIdx
    };

    // Densities stored in 'rho' as
    // [{ p0, p1, ..., (np - 1), mixture, mixture_with_exponents },
    //  { p0, p1, ..., (np - 1), mixture, mixture_with_exponents },
    //  ...
    //  { p0, p1, ..., (np - 1), mixture, mixture_with_exponents }]
    // Stride is np + 2.
    // here it has to be number of active phases
    constexpr int num_phases = Indices::numPhases;
    for (const auto& fsPhaseIdx : phaseMap) {
        if (FluidSystem::phaseIsActive(fsPhaseIdx)) {
            const auto phase_active_index = FluidSystem::canonicalToActivePhaseIdx(fsPhaseIdx);
            this->copyPhaseDensities(fsPhaseIdx, num_phases + 2, rho + phase_active_index);
        }
    }

    // Mixture densities.
    for (auto seg = 0*this->densities_.size(); seg < this->densities_.size(); ++seg) {
        const auto mixOffset = seg*(num_phases + 2) + num_phases;

        rho[mixOffset + 0] = this->mixtureDensity(seg);
        rho[mixOffset + 1] = this->mixtureDensityWithExponents(seg);
    }
}

template <class FluidSystem, class Indices>
void
MultisegmentWellSegments<FluidSystem,Indices>::
copyPhaseDensities(const unsigned    phaseIdx,
                   const std::size_t stride,
                   Scalar*           dens) const
{
    const auto compIdx = FluidSystem::canonicalToActiveCompIdx
        (FluidSystem::solventComponentIndex(phaseIdx));

    for (const auto& phase_density : this->phase_densities_) {
        *dens = phase_density[compIdx].value();
        dens += stride;
    }
}

template <class FluidSystem, class Indices>
typename MultisegmentWellSegments<FluidSystem,Indices>::Scalar
MultisegmentWellSegments<FluidSystem,Indices>::
mixtureDensity(const int seg) const
{
    auto mixDens = 0.0;

    const auto& rho = this->phase_densities_[seg];
    const auto& q   = this->phase_fractions_[seg];

    for (const auto& phIdx : {
            FluidSystem::oilPhaseIdx,
            FluidSystem::gasPhaseIdx,
            FluidSystem::waterPhaseIdx
        })
    {
        if (! FluidSystem::phaseIsActive(phIdx)) {
            continue;
        }

        const auto active_comp_index = FluidSystem::canonicalToActiveCompIdx(FluidSystem::solventComponentIndex(phIdx));

        mixDens += q[active_comp_index].value() * rho[active_comp_index].value();
    }

    return mixDens;
}

template <class FluidSystem, class Indices>
typename MultisegmentWellSegments<FluidSystem,Indices>::Scalar
MultisegmentWellSegments<FluidSystem,Indices>::
mixtureDensityWithExponents(const int seg) const
{
    if (const auto& segment = this->well_.wellEcl().getSegments()[seg];
        segment.isAICD())
    {
        return this->mixtureDensityWithExponents(segment.autoICD(), seg);
    }

    // No other segment type includes exponents of flowing fractions when
    // calculating the mixture/emulsion density.
    return this->mixtureDensity(seg);
}

template <class FluidSystem, class Indices>
typename MultisegmentWellSegments<FluidSystem,Indices>::Scalar
MultisegmentWellSegments<FluidSystem,Indices>::
mixtureDensityWithExponents(const AutoICD& aicd, const int seg) const
{
    auto mixDens = 0.0;

    const auto& rho = this->phase_densities_[seg];
    const auto& q   = this->phase_fractions_[seg];

    constexpr auto densityExponents = std::array {
        std::pair { FluidSystem::oilPhaseIdx  , &AutoICD::oilDensityExponent   },
        std::pair { FluidSystem::gasPhaseIdx  , &AutoICD::gasDensityExponent   },
        std::pair { FluidSystem::waterPhaseIdx, &AutoICD::waterDensityExponent },
    };

    for (const auto& [fsPhaseIdx, densityExponent] : densityExponents) {
        if (FluidSystem::phaseIsActive(fsPhaseIdx)) {
            const auto compIdx = FluidSystem::canonicalToActiveCompIdx(FluidSystem::solventComponentIndex(fsPhaseIdx));

            // exp = (aicd.*densityExponent)() in native syntax.
            const auto exp = std::invoke(densityExponent, aicd);

            mixDens += std::pow(q[compIdx].value(), exp) * rho[compIdx].value();
        }
    }

    return mixDens;
}

template<class FluidSystem, class Indices>
void
MultisegmentWellSegments<FluidSystem, Indices>::
calculatePhaseProperties(PhaseCalcResult& result,
                         const EvalWell& temperature,
                         const EvalWell& saltConcentration,
                         const PrimaryVariables& primary_variables,
                         const int seg,
                         const bool update_visc_and_den,
                         DeferredLogger& deferred_logger) const
{
    result.clear();

    auto& b = result.b;
    auto& mix_s = result.mix_s;
    auto& mix = result.mix;
    auto& vol_ratio = result.vol_ratio;
    auto& phase_viscosities = result.phase_viscosities;
    auto& phase_densities = result.phase_densities;

    // we might use reference here
    const EvalWell seg_pressure = primary_variables.getSegmentPressure(seg);

    const int num_quantities = well_.numConservationQuantities();
    for (int comp_idx = 0; comp_idx < num_quantities; ++comp_idx) {
        mix_s[comp_idx] = primary_variables.surfaceVolumeFraction(seg, comp_idx);
    }

    const bool waterActive = FluidSystem::phaseIsActive(FluidSystem::waterPhaseIdx);
    const bool gasActive = FluidSystem::phaseIsActive(FluidSystem::gasPhaseIdx);
    const bool oilActive = FluidSystem::phaseIsActive(FluidSystem::oilPhaseIdx);

    const int waterActiveCompIdx = waterActive ? FluidSystem::canonicalToActiveCompIdx(FluidSystem::waterCompIdx) : -1;
    const int gasActiveCompIdx = gasActive ? FluidSystem::canonicalToActiveCompIdx(FluidSystem::gasCompIdx) : -1;
    const int oilActiveCompIdx = oilActive ? FluidSystem::canonicalToActiveCompIdx(FluidSystem::oilCompIdx) : -1;

    const int pvt_region_index = well_.pvtRegionIdx();

    // water phase
    if (waterActive) {
        // rsw is only for interface usage
        const EvalWell rsw{0.};
        b[waterActiveCompIdx] = FluidSystem::waterPvt().inverseFormationVolumeFactor(
                                             pvt_region_index, temperature, seg_pressure, rsw, saltConcentration);
        if (update_visc_and_den) {
            // TODO: should not we use phaseIndex here?
            phase_viscosities[waterActiveCompIdx] = FluidSystem::waterPvt().viscosity(
                                             pvt_region_index, temperature, seg_pressure, rsw, saltConcentration);
            phase_densities[waterActiveCompIdx] = b[waterActiveCompIdx] * surface_densities_[waterActiveCompIdx];
        }
    }

    EvalWell rv {0.};
    // gas phase
    if (gasActive) {
        // rvw is only for interface usage
        const EvalWell rvw{0.};
        const bool oil_exist = oilActive && mix_s[oilActiveCompIdx] > 0.0;
        if (oil_exist) {
            const EvalWell rvmax = max(
                    FluidSystem::gasPvt().saturatedOilVaporizationFactor(pvt_region_index, temperature, seg_pressure),
                    0.);
            if (mix_s[gasActiveCompIdx] > 0.0) {
                rv = std::clamp(mix_s[oilActiveCompIdx] / mix_s[gasActiveCompIdx], EvalWell{0.}, rvmax);
            }
            b[gasActiveCompIdx] = FluidSystem::gasPvt().inverseFormationVolumeFactor(
                                               pvt_region_index, temperature, seg_pressure, rv, rvw);
            if (update_visc_and_den) {
                phase_viscosities[gasActiveCompIdx] = FluidSystem::gasPvt().viscosity(
                                                 pvt_region_index, temperature, seg_pressure, rv, rvw);
                phase_densities[gasActiveCompIdx] = b[gasActiveCompIdx] * surface_densities_[gasActiveCompIdx]
                                                    + rv * b[gasActiveCompIdx] * surface_densities_[oilActiveCompIdx];
            }
        } else { // no oil here
            b[gasActiveCompIdx] = FluidSystem::gasPvt().saturatedInverseFormationVolumeFactor(
                                                  pvt_region_index, temperature, seg_pressure);
            if (update_visc_and_den) {
                phase_viscosities[gasActiveCompIdx] = FluidSystem::gasPvt().saturatedViscosity(
                                                       pvt_region_index, temperature, seg_pressure);
                phase_densities[gasActiveCompIdx] = b[gasActiveCompIdx] * surface_densities_[gasActiveCompIdx];
            }
        }
    }

    EvalWell rs {0.};
    // oil phase
    if (oilActive) {
        const bool gas_exist = gasActive && mix_s[gasActiveCompIdx] > 0.0;
        if (gas_exist) {
            const EvalWell rsmax = max(
                    FluidSystem::oilPvt().saturatedGasDissolutionFactor(pvt_region_index, temperature, seg_pressure),
                    0.);
            if (mix_s[oilActiveCompIdx] > 0.0) {
                rs = std::clamp(mix_s[gasActiveCompIdx] / mix_s[oilActiveCompIdx], EvalWell{0.}, rsmax);
            }
            b[oilActiveCompIdx] = FluidSystem::oilPvt().inverseFormationVolumeFactor(
                                                        pvt_region_index, temperature, seg_pressure, rs);
            if (update_visc_and_den) {
                phase_viscosities[oilActiveCompIdx] = FluidSystem::oilPvt().viscosity(pvt_region_index, temperature,
                                                                                      seg_pressure, rs);
                phase_densities[oilActiveCompIdx] = b[oilActiveCompIdx] * surface_densities_[oilActiveCompIdx]
                                              + rs * b[oilActiveCompIdx] * surface_densities_[gasActiveCompIdx];
            }
        } else { // no gas phase
            b[oilActiveCompIdx] = FluidSystem::oilPvt().saturatedInverseFormationVolumeFactor(pvt_region_index,
                                                                                              temperature, seg_pressure);
            if (update_visc_and_den) {
                phase_viscosities[oilActiveCompIdx] = FluidSystem::oilPvt().saturatedViscosity(pvt_region_index,
                                                                                               temperature, seg_pressure);
                phase_densities[oilActiveCompIdx] = b[oilActiveCompIdx] * surface_densities_[oilActiveCompIdx];
            }
        }
    }

    mix = mix_s;
    if (oilActive && gasActive) {
        const EvalWell d = 1.0 - rs * rv;
        if (d <= 0.0) {
            const std::string str =
                fmt::format("Problematic d value {} obtained for well {} during segment density calculations "
                            "with rs {}, rv {} and pressure {}. Continue as if no dissolution (rs = 0) and "
                            "vaporization (rv = 0) for this connection.",
                            d, well_.name(), rs, rv, seg_pressure);
            deferred_logger.debug(str);
        } else {
            if (rs > 0.0) {
                mix[gasActiveCompIdx] = (mix_s[gasActiveCompIdx] - mix_s[oilActiveCompIdx] * rs) / d;
            }
            if (rv > 0.0) {
                mix[oilActiveCompIdx] = (mix_s[oilActiveCompIdx] - mix_s[gasActiveCompIdx] * rv) / d;
            }
        }
    }

    for (int comp_idx = 0; comp_idx < num_quantities; ++comp_idx) {
        vol_ratio += mix[comp_idx] / b[comp_idx];
    }
}

template<typename FluidSystem, typename Indices>
void
MultisegmentWellSegments<FluidSystem, Indices>::
PhaseCalcResult::clear()
{
    std::ranges::fill(b, 0.0);
    std::ranges::fill(mix, 0.0);
    std::ranges::fill(mix_s, 0.0);
    std::ranges::fill(phase_viscosities, 0.0);
    std::ranges::fill(phase_densities, 0.0);
    vol_ratio = 0.0;
}

#include <opm/simulators/utils/InstantiationIndicesMacros.hpp>

INSTANTIATE_TYPE_INDICES(MultisegmentWellSegments, double)

#if FLOW_INSTANTIATE_FLOAT
INSTANTIATE_TYPE_INDICES(MultisegmentWellSegments, float)
#endif

} // namespace Opm
