'''
MELTS2.0: Bayes sims of mantle melting
======================================

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_mantle_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')


# %%

MM3_oxides = pd.Series({
        "SiO2":45.47,
        "TiO2":0.11,
        "Al2O3":4.00,
        "Cr2O3":0.68,
        "Fe2O3":1.0,
        "FeO":6.22,
        "MnO": 0.0,
        "MgO":38.53,
        "CaO":3.59,
        "Na2O":0.31,
        "K2O":0.0,
        "P2O5": 0.0,
        "H2O":  0.0})




# %%
# 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 = MM3_oxides,
                            P_bar = 10e3, # bars
                            T_K = 2200, # 
                            database=database, # liquid model name
                            )

    sys.crystallize(
          method='equil', # equilibrium crystallization
          T_step=15, # decrease in temperature at each step
          T_final_K = 1550
                  )


    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%]')



# %%