'''
Calculator: Liquid Properties
===========================================

This notebook calculates the activity of endmembers within a liquid.
It replaces the functionality previously implemented by the MELTS Supplemental Calculator (https://melts.ofm-research.org/CalcForms/)

This calculator is a work in progress -- please let us know if you find any errors!

Open this code in an executable MyBinder instance (MyBinder links may be slow to load-- please be patient!):

.. image:: https://mybinder.org/badge_logo.svg
  :target: https://mybinder.org/v2/gl/swmatthews-research%2FThermoEngineLite/main?urlpath=%2Fdoc%2Ftree%2F.%2Fdoc%2Fsource%2Fauto_examples%2F_4_unorganized%2Fplot_supplemental_calculator.ipynb

'''

# %%
# User Input
# ----------
#
# 1. Input liquid composition in wt% oxides, temperature, pressure, and log(fO2) below
# 2. Once you have input your data in the cell above, click "Run" --> "Run All Cells" at the top of the notebook. This will run the entire notebook and produce the output table at the bottom of the file.

oxide_wt_dict = {
    'SiO2':    48.4012,	
    'TiO2':     0.9208,	
    'Al2O3':   14.9349, 
    'FeO':      9.0825,
    'MnO':      0.1587, 
    'MgO':      8.1285, 
    'CaO':      14.1368, 
    'Na2O':      1.3124, 
    'K2O':       0.1462, 
    'P2O5':      0.001
    }

TK = 1500 + 273.15 # temperature value in K
Pbar = 15000 # pressure value in bar
logfO2 = None # can replace this with a value for absolute logfO2 (not relative to a buffer). Otherwise, calculates with given Fe2O3 value.
logfO2_buffer = ('QFM', 0.0) # or, specify a buffer and offset, e.g., QFM, NNO, IW

# Alternatively, you can explicitly identify the redox state and run this unbuffered by uncommenting lines below
# oxide_wt_dict['FeO'] = 8.807
# oxide_wt_dict['Fe2O3'] = 1.389
# logfO2_buffer = None

# Liquid model (uncomment the one you wish to use):
database = 'MELTS_v1_2' # rhyolite-MELTS v1.2
# database = 'MELTS_v1_0' # rhyolite-MELTS v1.0
# database = 'MELTS_pMELTS' # pMELTS

# %%
# Calculations
# ------------
#
# The user does not need to modify this code-- it will run automatically when "Run All Cells" is triggered.

###-----------------------------------------------------------------------------------------
### 0. Import modules
###-----------------------------------------------------------------------------------------

import numpy as np
import pandas as pd
from thermoengine import model, rockychem, chemistry
from thermoengine.core import chem
from thermoengine import redox
from IPython.display import display, HTML

###-----------------------------------------------------------------------------------------
### 1. Calculate oxide moles
###-----------------------------------------------------------------------------------------

oxide_moles = rockychem.convert_bulk_comp(wt_oxides=oxide_wt_dict, to='mol_oxides')
db = model.Database(database_name=database)
oxide_moles_ser = pd.Series(oxide_moles, index=db.oxides)
oxide_moles_ser = oxide_moles_ser.fillna(0)


###-----------------------------------------------------------------------------------------
### 2. Adjust composition to be consistent with log fO2 input
###-----------------------------------------------------------------------------------------

if logfO2_buffer is not None:
    logfO2 = redox.redox_buffer(TK, Pbar, logfO2_buffer[0]) + logfO2_buffer[1]
if logfO2 is not None:
    ln_Fe_oxide_ratio = redox._redox_state_Kress91(TK, Pbar, oxide_comp=[oxide_moles_ser], logfO2=logfO2)
    Fe2O3_FeO = np.exp(ln_Fe_oxide_ratio)
    oxide_moles_ser = redox.impose_redox_ratio(Fe2O3_FeO, oxide_moles_ser)
elif (('Fe2O3' in oxide_wt_dict and oxide_wt_dict['Fe2O3'] > 0) 
      and ('FeO' in oxide_wt_dict and oxide_wt_dict['FeO'] > 0)):
    logfO2 = redox.Kress91.calc_logfO2(TK, Pbar, oxide_moles_ser)

###-----------------------------------------------------------------------------------------
### 3. Create liquid phase instance
###-----------------------------------------------------------------------------------------

liquid = db.get_phase('Liq')
liq_end = liquid.calc_endmember_comp(oxide_moles_ser.values, method='least_squares',
                            output_residual=False, normalize=False,
                            decimals=10)

###-----------------------------------------------------------------------------------------
### 4. Calculate Gibbs Free Energies of Formation
###-----------------------------------------------------------------------------------------

end_array = np.identity(liquid.endmember_num)
G_array = np.zeros(liquid.endmember_num)
for i in range(liquid.endmember_num):
    G = liquid.gibbs_energy(TK,Pbar,mol=end_array[i,:])
    G_array[i] = G

###-----------------------------------------------------------------------------------------
### 5. Calculate Activities
###-----------------------------------------------------------------------------------------

activities = liquid.activity(TK, Pbar, mol=liq_end)

###-----------------------------------------------------------------------------------------
### 6. Calculate other properties
###-----------------------------------------------------------------------------------------

properties = {}

mass_scaling = 100 / liq_end.dot(liquid.endmember_molwts)
volume = liquid.volume(TK, Pbar, mol=liq_end)

properties['Volume of 100g (cm3)'] = (volume * mass_scaling * 10)

properties['Density (g/cm3)'] = 10.0 / volume / mass_scaling

properties['Heat Capacity of 100g (J/K)'] = (liquid.heat_capacity(TK, Pbar, mol=liq_end)
                                             * mass_scaling)

if not (logfO2 is None or np.isinf(logfO2)):
    properties['log10(fO2)'] = logfO2

properties['Bulk Modulus (GPa)'] = (- volume * 1e-4
                                    / liquid.gibbs_energy(TK, Pbar, mol=liq_end, deriv={'dP':2})) 

properties_df = pd.DataFrame(pd.Series(properties), columns=[''])

oxide_wts_after_fO2 = rockychem.convert_bulk_comp(mol_oxides=oxide_moles_ser, to='wt_oxides')
oxide_moles_ser /= oxide_moles_ser.sum()
oxide_moles_ser['Total'] = 1.0
oxide_wts_after_fO2['Total'] = oxide_wts_after_fO2.sum()
comp_table = pd.DataFrame([oxide_moles_ser, oxide_wts_after_fO2], 
                          index=(['mol frac', 'grams'])).fillna(0.0).T

###-----------------------------------------------------------------------------------------
### 6. Combine and display data
###-----------------------------------------------------------------------------------------

data = {'Endmember': liquid.endmember_names, 'Gibbs Energy (kJ)': G_array/1000, 'Activity': activities}
endmember_df = pd.DataFrame(data)
pd.set_option('display.max_rows', None)

# These lines will provide nicer printing, but only work in mybinder:
# display(HTML(properties_df.to_html()))
# display(HTML(comp_table.to_html()))

print("\n-------------------------------------------")
for key in properties:
    print(f" {key :30s} {properties[key]:>10.5f}")

print("\n Oxide        mol frac    grams")
print("--------------------------------")
for ox in oxide_wts_after_fO2.index:
    print(f" {ox :8s} {oxide_moles_ser[ox]/oxide_moles_ser.sum():>10.3f} {oxide_wts_after_fO2[ox]:>10.3f}")
print("--------------------------------")
print(f" Total         1.000 {oxide_wts_after_fO2.sum():>10.3f}")
print("--------------------------------\n")

endmember_df

# %%