'''
MagmaForge: Fluid Saturated Crystallization
===========================================

Demonstrates how to run a model for crystallising magma when H2O is at saturation.

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_2_magmaforge%2Fplot_magmaforge_crystallization_fluid_present.ipynb

'''

# %% 
# Initialization
# --------------
#
# Import necessary packages:

import thermoengine
from thermoengine import magmaforge, model

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

# %%
# Define a bulk composition:

morb_oxides = pd.Series({ 
            'SiO2':  48.68, 
            'TiO2':   1.01, 
            'Al2O3': 17.64, 
            'Fe2O3':  0.89,
            'Cr2O3':  0.0425, 
            'FeO':    7.59, 
            'MgO':    9.10, 
            'CaO':   12.45, 
            'Na2O':   2.65, 
            'K2O':    0.03, 
            'P2O5':   0.08, 
            'H2O':    0.2,
            },)

# %%
# Set up fluid-saturated condition for a H2O fluid using
# rhyolite-MELTS v1.0

db = 'MELTS_v1_0'
fluid = model.Database(db).get_phase('H2O')

def mu_H2O(t, p, state):
    return fluid.gibbs_energy(t, p)

chempot_constraints = [({'H':2.0,'O':1.0},mu_H2O)]


# %%
# Set up fluid-saturated conditions for a H2O-CO2 fluid
# with rhyolite-MELTS v1.2 (uncomment below code and run)

# X_H2O=0.4
# morb_oxides['CO2'] = 0.1
# db = 'MELTS_v1_2'
# fluid = model.Database(db).get_phase('Fl')

# def mu_H2O(t, p, state):
#     return fluid.gibbs_energy(t, p, 
#                               mol=np.array([X_H2O, 1-X_H2O]),
#                               deriv={'dmol':1})[0][0]

# def mu_CO2(t, p, state):
#     return fluid.gibbs_energy(t, p, 
#                               mol=np.array([X_H2O, 1-X_H2O]),
#                               deriv={'dmol':1})[0][1]

# chempot_constraints = [({'H':2.0,'O':1.0},mu_H2O),
#                        ({'C':1.0,'O':2.0},mu_CO2)]

# db = model.Database(db)
# db.remove_phases(['Fl'])
# db.activate()

# %% 
# System and Calculations
# -----------------------
#
# Define system using magmaforge and run the calculation:

sys = magmaforge.System(comp=morb_oxides, 
                        # T_C=1350,
                        T_liquidus=True,
                        P_bar=1000.0, 
                        logfO2 = ('QFM', -1.0),
                        database=db,
                        chemical_potential_constraints=chempot_constraints,
                        )


sys.crystallize(method='equil', 
                T_step=10,
                fix_fO2=True, 
                )

# %% 
# Plot results
# --------------

magmaforge.plot.magma_evolution(sys.history)

# %%