from feos.eos import *
from feos.epcsaft import *
import si_units as si
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns

sns.set_context('notebook')
sns.set_palette('Dark2')
sns.set_style('ticks')

def molality_to_mole_fraction(
    molality_solute: si.SIObject, 
    molarweight_solvent: si.SIObject
) -> float:
    """Calculate mole fraction of solute from solute molality given molar weight of solvent.

    Args:
        molality_solute: Molality of solute (salt).
        molarweight_solvent: Molar weight of solvent.

    Returns:
        Mole fraction of solute (salt)
    """
    return molality_solute * molarweight_solvent / (1.0 + molality_solute * molarweight_solvent)

def mean_ionic_activity_coefficient(
    eos: EquationOfState,
    temperature: si.SIObject, 
    pressure: si.SIObject, 
    molality_solute: si.SIObject,
    molarweight_solvent: si.SIObject,
    stoichiometric_coefficients: np.ndarray
) -> float:
    """Calculate mean ionic activity coefficient (MIAC). 

    1. calculates fugacity coefficients for given molality and inf. dilution
    2. uses eq. 9 to calculate activity coefficients from fugacities
    3. uses eq. 8 to get MIAC
    
    Args:
        eos: Equation of state to use
        temperature: Temperature
        pressure: Pressure
        molality_solute: Molality of solute (salt).
        molarweight_solvent: Molar weight of solvent.
        stoichiometric_coefficients: 
            Stoichiometric coefficients of salt ions.
            Order has to match parameter order.

    Returns:
        Mean ionic activity coefficient.
    """
    # mole fraction of salt from molality
    x_ions = molality_to_mole_fraction(molality_solute=molality_solute, molarweight_solvent=molarweight_solvent)
    # mole fractions from closure
    molefracs = np.array([1 - x_ions, x_ions / 2, x_ions / 2])

    # thermodynamic condition
    state = State(eos, temperature=temperature, pressure=pressure, molefracs=molefracs)
    # reference: infinite dilution (same for each ion so only calculated once)
    infinite_dilution = State(eos, temperature=temperature, pressure=pressure, molefracs=np.array([1.0, 0.0, 0.0]))

    # Equation 9
    # only use ion activity coefficients (component 0 is solvent)
    ac_ions = (np.exp(state.ln_phi()) / np.exp(infinite_dilution.ln_phi()))[1:]
    # Equation 8
    miac = np.prod(ac_ions**stoichiometric_coefficients)**(1 / stoichiometric_coefficients.sum())
    return miac

# parameters
parameters = ElectrolytePcSaftParameters.from_json(
    substances=['water', 'sodium ion', 'chloride ion'],
    pure_path='../parameters/epcsaft/held2014_w_permittivity_added.json',
    binary_path='../parameters/epcsaft/held2014_binary.json'
)
molarweight_water = parameters.pure_records[0].molarweight * si.GRAM / si.MOL

# equation of state
eos = EquationOfState.epcsaft(parameters, epcsaft_variant=ElectrolytePcSaftVariants.Advanced)

# thermodynamic conditions
temperature = 298.15 * si.KELVIN
pressure = si.BAR
# stoichiometric coefficients for [Na+] and [Cl-]
nu = np.array([1.0, 1.0])

# mean ionic activity coefficients
molality_salt = si.linspace(0 * si.MOL / si.KILOGRAM, 6.0 * si.MOL / si.KILOGRAM, 100) 
gammas = np.array([
    mean_ionic_activity_coefficient(eos, temperature, pressure, b, molarweight_solvent=molarweight_water, stoichiometric_coefficients=nu) 
    for b in molality_salt
])

# plot results
plt.plot(molality_salt / si.MOL * si.KILOGRAM, gammas, clip_on=False)
plt.xlabel(r'Molality $m_\text{salt}$ / mol / kg')
plt.ylabel(r'$\gamma^{*, x}_\pm$')
plt.xlim(0, 6)
plt.ylim(0.7, 1.2)
sns.despine(offset=10);