'''
MELTS2.0: Bayes sims of morb crystallization
============================================

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_bayes_morb_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

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



# %%
# Load the MELTS2 parameters & draw from uncertainties
# ----------------------------------------------------

NDRAW = 5

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)

model_fit['param_svd_mean']
param_svd_err = model_fit['param_svd_err']
param_svd_mean = model_fit['param_svd_mean']
rnd_draws = np.random.randn(NDRAW,len(param_svd_err))
pdraws_svd = param_svd_mean + param_svd_err*rnd_draws
pdraws = pdraws_svd.dot(model_fit['Vh_param_svd_proj'])
pdraws[:,0]


# %%
database = thermo.Database('MELTS_v1_0')
Liq = database.get_phase('Liq')


# %%


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},)


# %%
# Run sims and store results
# --------------------------

phs_frac_tables = []
liq_comp_tables = []

for ind in np.arange(NDRAW):

    ipdraw = pdraws[ind]

    Liq.set_param_values(param_names=model_fit['param_names'], 
                         param_values=ipdraw)

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

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


    phs_frac_tbl = sys.history.get_phase_frac_table()
    liq_comp_tbl = sys.history.get_phase_comp_table('Liquid')

    phs_frac_tbl = phs_frac_tbl.reset_index().rename(columns={'index':'T_K'})
    liq_comp_tbl = liq_comp_tbl.reset_index().rename(columns={'index':'T_K'})

    #magmaforge.plot.magma_evolution(sys.history)
    phs_frac_tables.append(phs_frac_tbl)
    liq_comp_tables.append(liq_comp_tbl)


# %%
def plot_spaghetti_diagram(prop_tables, col_list, color_list, ylbl, alpha=0.5):
    for i in np.arange(len(prop_tables)):
        ifrac_tbl = prop_tables[i]
        ifrac_tbl['T_C'] = ifrac_tbl['T_K']-273.15

        for iphsnm, icolor in zip(col_list, color_list):
            if(i==0):
                ilabel = iphsnm
            else:
                ilabel = None

            sns.lineplot(ifrac_tbl,x='T_C', y=iphsnm, label=ilabel,
                          color=icolor, alpha=alpha)

    plt.ylabel(ylbl)
    plt.legend()

# %%
# Plot phase fraction evolution
# -----------------------------

col_list = ['Feldspar','Olivine','Clinopyroxene','Orthopyroxene','Liquid']
color_list = ['orange','g','r','c','purple']
plot_spaghetti_diagram(phs_frac_tables, col_list, color_list, 'Phase Fraction')


# %%
# Plot liquid composition evolution
# ---------------------------------

col_list = ['SiO2','MgO','FeO','Fe2O3','Al2O3','K2O','Na2O','H2O']
color_list = ['cyan','orange','green','red','purple','brown','pink']
plot_spaghetti_diagram(liq_comp_tables, col_list, color_list, 'Liquid Comp [wt%]')



# %%