#################################################################################
# The Institute for the Design of Advanced Energy Systems Integrated Platform
# Framework (IDAES IP) was produced under the DOE Institute for the
# Design of Advanced Energy Systems (IDAES).
#
# Copyright (c) 2018-2023 by the software owners: The Regents of the
# University of California, through Lawrence Berkeley National Laboratory,
# National Technology & Engineering Solutions of Sandia, LLC, Carnegie Mellon
# University, West Virginia University Research Corporation, et al.
# All rights reserved.  Please see the files COPYRIGHT.md and LICENSE.md
# for full copyright and license information.
#################################################################################
import pytest
import pyomo.environ as pyo
import idaes.core
from idaes.models_extra.power_generation.unit_models.helm import (
    HelmMixer,
    MomentumMixingType,
)
from idaes.models.properties import iapws95
from idaes.models.properties.general_helmholtz import helmholtz_available


@pytest.mark.skipif(not helmholtz_available(), reason="General Helmholtz not available")
@pytest.mark.component
def test_mixer():
    m = pyo.ConcreteModel()
    m.fs = idaes.core.FlowsheetBlock(dynamic=False)
    m.fs.properties = iapws95.Iapws95ParameterBlock()
    m.fs.unit = HelmMixer(
        property_package=m.fs.properties, inlet_list=["i1", "i2", "i3"]
    )

    Fin1 = 1.1e4  # mol/s
    hin1 = 4000  # J/mol
    Pin1 = 1.2e5  # Pa
    Fin2 = 1e4  # mol/s
    hin2 = 5000  # J/mol
    Pin2 = 2e5  # Pa
    Fin3 = 1.3e4  # mol/s
    hin3 = 6000  # J/mol
    Pin3 = 3e5  # Pa
    Pout = 1.5e5  # Pa

    m.fs.unit.i1.flow_mol[0].fix(Fin1)
    m.fs.unit.i1.enth_mol[0].fix(hin1)
    m.fs.unit.i1.pressure[0].fix(Pin1)
    m.fs.unit.i2.flow_mol[0].fix(Fin2)
    m.fs.unit.i2.enth_mol[0].fix(hin2)
    m.fs.unit.i2.pressure[0].fix(Pin2)
    m.fs.unit.i3.flow_mol[0].fix(Fin3)
    m.fs.unit.i3.enth_mol[0].fix(hin3)
    m.fs.unit.i3.pressure[0].fix(Pin3)

    m.fs.unit.initialize()
    Fout = Fin1 + Fin2 + Fin3
    hout = (hin1 * Fin1 + hin2 * Fin2 + hin3 * Fin3) / Fout

    assert pyo.value(m.fs.unit.outlet.flow_mol[0]) == pytest.approx(Fout, rel=1e-7)
    assert pyo.value(m.fs.unit.outlet.enth_mol[0]) == pytest.approx(hout, rel=1e-7)
    assert pyo.value(m.fs.unit.outlet.pressure[0]) == pytest.approx(Pin1, rel=1e-7)


@pytest.mark.skipif(not helmholtz_available(), reason="General Helmholtz not available")
@pytest.mark.component
def test_mixer2():
    m = pyo.ConcreteModel()
    m.fs = idaes.core.FlowsheetBlock(dynamic=False)
    m.fs.properties = iapws95.Iapws95ParameterBlock()
    m.fs.unit = HelmMixer(
        momentum_mixing_type=MomentumMixingType.equality,
        property_package=m.fs.properties,
        inlet_list=["i1", "i2", "i3"],
    )

    Fin1 = 1.1e4  # mol/s
    hin1 = 4000  # J/mol
    Pin1 = 1.2e5  # Pa
    Fin2 = 1e4  # mol/s
    hin2 = 5000  # J/mol
    Pin2 = 2e5  # Pa
    Fin3 = 1.3e4  # mol/s
    hin3 = 6000  # J/mol
    Pin3 = 3e5  # Pa
    Pout = 1.5e5  # Pa

    m.fs.unit.i1.flow_mol[0].fix(Fin1)
    m.fs.unit.i1.enth_mol[0].fix(hin1)
    m.fs.unit.i1.pressure[0] = Pin1
    m.fs.unit.i2.flow_mol[0].fix(Fin2)
    m.fs.unit.i2.enth_mol[0].fix(hin2)
    m.fs.unit.i2.pressure[0] = Pin2
    m.fs.unit.i3.flow_mol[0].fix(Fin3)
    m.fs.unit.i3.enth_mol[0].fix(hin3)
    m.fs.unit.i3.pressure[0] = Pin3
    m.fs.unit.outlet.pressure[0].fix(Pout)

    m.fs.unit.initialize()
    Fout = Fin1 + Fin2 + Fin3
    hout = (hin1 * Fin1 + hin2 * Fin2 + hin3 * Fin3) / Fout

    assert pyo.value(m.fs.unit.outlet.flow_mol[0]) == pytest.approx(Fout, rel=1e-7)
    assert pyo.value(m.fs.unit.outlet.enth_mol[0]) == pytest.approx(hout, rel=1e-7)
    assert pyo.value(m.fs.unit.outlet.pressure[0]) == pytest.approx(Pout, rel=1e-7)
    assert pyo.value(m.fs.unit.i1.pressure[0]) == pytest.approx(Pout, rel=1e-7)
    assert pyo.value(m.fs.unit.i2.pressure[0]) == pytest.approx(Pout, rel=1e-7)
    assert pyo.value(m.fs.unit.i3.pressure[0]) == pytest.approx(Pout, rel=1e-7)


@pytest.mark.skipif(not helmholtz_available(), reason="General Helmholtz not available")
@pytest.mark.unit
def test_get_stream_table_contents():
    m = pyo.ConcreteModel()
    m.fs = idaes.core.FlowsheetBlock(dynamic=False)
    m.fs.properties = iapws95.Iapws95ParameterBlock()
    m.fs.unit = HelmMixer(
        momentum_mixing_type=MomentumMixingType.equality,
        property_package=m.fs.properties,
        inlet_list=["i1", "i2", "i3"],
    )

    stable = m.fs.unit._get_stream_table_contents()

    expected = {
        "Units": {
            "Mass Flow": getattr(pyo.units.pint_registry, "kg/s"),
            "Molar Flow": getattr(pyo.units.pint_registry, "mol/s"),
            "Molar Enthalpy": getattr(pyo.units.pint_registry, "J/mol"),
            "P": getattr(pyo.units.pint_registry, "Pa"),
            "T": getattr(pyo.units.pint_registry, "K"),
            "Vapor Fraction": getattr(pyo.units.pint_registry, "dimensionless"),
        },
        "i1": {
            "Mass Flow": pytest.approx(0.01801527, rel=1e-5),
            "Molar Flow": pytest.approx(1.0, rel=1e-5),
            "Molar Enthalpy": pytest.approx(0.01102139, rel=1e-5),
            "P": pytest.approx(11032300, rel=1e-5),
            "T": pytest.approx(270.4877, rel=1e-5),
            "Vapor Fraction": pytest.approx(0.0, abs=1e-5),
        },
        "i2": {
            "Mass Flow": pytest.approx(0.01801527, rel=1e-5),
            "Molar Flow": pytest.approx(1.0, rel=1e-5),
            "Molar Enthalpy": pytest.approx(0.01102139, rel=1e-5),
            "P": pytest.approx(11032300, rel=1e-5),
            "T": pytest.approx(270.4877, rel=1e-5),
            "Vapor Fraction": pytest.approx(0.0, abs=1e-5),
        },
        "i3": {
            "Mass Flow": pytest.approx(0.01801527, rel=1e-5),
            "Molar Flow": pytest.approx(1.0, rel=1e-5),
            "Molar Enthalpy": pytest.approx(0.01102139, rel=1e-5),
            "P": pytest.approx(11032300, rel=1e-5),
            "T": pytest.approx(270.4877, rel=1e-5),
            "Vapor Fraction": pytest.approx(0.0, abs=1e-5),
        },
        "outlet": {
            "Mass Flow": pytest.approx(0.01801527, rel=1e-5),
            "Molar Flow": pytest.approx(1.0, rel=1e-5),
            "Molar Enthalpy": pytest.approx(0.01102139, rel=1e-5),
            "P": pytest.approx(11032300, rel=1e-5),
            "T": pytest.approx(270.4877, rel=1e-5),
            "Vapor Fraction": pytest.approx(0.0, abs=1e-5),
        },
    }

    assert stable.to_dict() == expected
