'''
MELTS2.0: Simulations with a calibrated model
=============================================

Load the best-fit calibrated MELTS2.0 model and use it 
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%2Fmanual_examples%2Fmelts2_sims%2Fplot_melts2_sims.ipynb


'''

# %%
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import thermoengine as thermo
from thermoengine import rockychem, redox
import json

from thermoengine import magmaforge
import pandas as pd # a useful data analysis package


# %%
# Load the MELTS2 calibration parameters
# --------------------------------------

filenm = "model_fit.json"
# filenm = "model_fit_med.json"
# filenm = "model_fit_unconstr.json"
# filenm = "model_fit_uniform.json"
with open(f"input_files/{filenm}") as f:
    model_fit = json.load(f)

param_best_fit = np.array(model_fit['param_svd_mean']).dot(np.array(model_fit['Vh_param_svd_proj']))
param_best_fit

database = thermo.Database('MELTS_v1_0')

Liq = database.get_phase('Liq')
Liq.set_param_values(param_names=model_fit['param_names'], param_values=param_best_fit )


# %%


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},) # Values in grams (extensive units)




# %%
# Equil crystallization
# ---------------------

sys = magmaforge.System(comp = morb_oxides,
                        P_bar = 1000.0, # bars
                        T_C = 1300, # C
                        logfO2 = ('NNO',-1), # setting fO2 equal to 1 log unit below the QFM buffer
                        database=database, # liquid model name
                        )

sys.crystallize(
      method='equil', # equilibrium crystallization
      T_step=5, # decrease in temperature at each step
      fix_fO2=True, # hold fO2 constant at initial fO2 value, relative to indicated buffer
              )

# magmaforge.plot.phase_fractions(sys.history)
magmaforge.plot.magma_evolution(sys.history)


# %%
# Frac crystallization
# --------------------

sys = magmaforge.System(comp = morb_oxides,
                        P_bar = 1000.0, # bars
                        T_C = 1300, # C
                        logfO2 = ('NNO',0), 
                        database=database, 
                        # database="MELTS_v1_0", # liquid model name
                        )

sys.crystallize(
      method='frac',
      T_step=5, # decrease in temperature at each step
      fix_fO2=True, # hold fO2 constant at initial fO2 value, relative to indicated buffer
      T_final_C=800
    )

# magmaforge.plot.phase_fractions(sys.history)
magmaforge.plot.magma_evolution(sys.history)