'''
MagmaForge: Fractional Decompression
=======================================

Demonstrates how to run a fractional (isentropic) decompression routine using MagmaForge.

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_decompression_frac.ipynb
'''

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

from thermoengine import magmaforge

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

# %%
# Define a bulk composition:

DMM_comp_adj = {
    'SiO2': 44.7354,
    'TiO2': 0.130074,
    'Al2O3': 3.98226,
    'Fe2O3': 0.275456,
    'Cr2O3': 0.570323,
    'FeO': 8.01254,
    'MgO': 38.752,
    'NiO': 0.240136,
    'CaO': 3.1718,
    'Na2O': 0.130074
}

P_GPa = 1.45

# %%
# Find the solidus
# ----------------
# Currently, ThermoEngine cannot equilibrate subsolidus, as a different equilibration routine is required when the liquid is not present.
# We plan to get subsolidus equilibration implemented soon.

sys = magmaforge.System(comp=DMM_comp_adj, 
                        T_liquidus=True,
                        P_GPa=P_GPa, 
                        database='MELTS_pMELTS',
                        )

sys.crystallize(method='equil',
                T_step=5,
                fix_fO2=False,
                calc_args={'debug':0})

solidus_T = sys.T_C

print(f"Solidus found at {solidus_T :.1f}˚C")


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

sys = magmaforge.System(comp=DMM_comp_adj,
                        T_C = solidus_T,
                        P_GPa=P_GPa,
                        database='MELTS_pMELTS'
                        )

# %%
# Melt via isentropic decompression:
#
# * For fractional melting, the fraction retained is treated as a fixed porosity. By default, this value is set to 0.01.
# * Fractionation happens at the beginning of each step, prior to equilibration, so that the final state is an equilibrium state.

sys.melt_by_decompression(P_final_GPa=0.25,
                          method='frac', 
                          P_step_GPa=0.01,
                          fix_fO2=False,
                          )

# %% 
# Output
# ------
pd.set_option('display.max_columns', None)
pd.set_option('display.max_row', None)
phase_masses = sys.history.get_phase_mass_table(phases='present', index='P_GPa')
phase_masses['Aggregated Liquid'] = sys.history.get_fractionated_masses()
print(phase_masses)


# %%
# Plot System Evolution
# ---------------------
#
# Automatic plotting functions for decompression melting coming soon!
# The dotted lines show the instantaneous fractional melt compositions.

oxides_to_plot = ['MgO', 'SiO2', 'Al2O3', 'FeO', 'Fe2O3', 'CaO']

pressures = sys.history.get_pressures('GPa')
temperatures = sys.history.get_temperatures('C')

phase_masses = sys.history.get_phase_mass_table(phases='present', index='P_GPa')
mass_agg_melt = sys.history.get_fractionated_masses('aggregated')
phase_masses['Agg Liquid'] = mass_agg_melt

agg_liq_comps = sys.history.get_fractionated_comps(unit='wt_oxides', method='aggregated', norm=True)
instant_liq_comps = sys.history.get_fractionated_comps(unit='wt_oxides', method="incremental", norm=True)

phase_colors = {'Spinel': 'black',
                'Garnet': 'purple',
                'Olivine': 'yellowgreen',
                'Clinopyroxene' : 'green',
                'Orthopyroxene': 'darkolivegreen',}

# Create plot

fig, ax = plt.subplots(1, 3, sharey='row', gridspec_kw={'wspace': 0.1, 'width_ratios': [1,1,2]}, figsize=(8,4))

# Plot phase masses
for ph in phase_masses.columns:
    if (ph != 'Liquid') and (ph != 'Agg Liquid'):
        ax[0].plot(phase_masses[ph], pressures, 
                   color=phase_colors[ph],label=ph)

ax[0].plot(phase_masses['Agg Liquid'], pressures,
           lw=2, c='red', label='Agg. Melt')

# Plot temperature
ax[1].plot(temperatures, pressures, c='k', lw=2)

# Plot liquid composition
i=0
for ox in oxides_to_plot:
  ax[2].plot(agg_liq_comps[ox], pressures, label=ox, c=f"C{i}")
  ax[2].plot(instant_liq_comps[ox], pressures, c=f"C{i}", ls=':')
  i += 1

# Add instantaneous melts to legend
ax[2].set_xlim(ax[2].get_xlim())
ax[2].set_ylim(ax[2].get_ylim())
ax[2].plot([100]*2, [100]*2, label=" ", lw=0)
ax[2].plot([100]*2, [100]*2, c='k', ls=':', label="Instantaneous\nMelts")

# Format axes
ax[0].invert_yaxis()
ax[0].legend(ncol=2, loc='upper left', bbox_to_anchor=(-0.4, 1.3))
ax[2].legend(ncol=3, loc='upper left', bbox_to_anchor=(-0.3, 1.3))
plt.subplots_adjust(top=0.8)

ax[0].set_ylabel('Pressure (GPa)')
ax[0].set_xlabel('Phase Mass')
ax[1].set_xlabel('Temperature (˚C)')
ax[2].set_xlabel('Aggregated Liquid Comp (wt%)')


plt.show()


# %%
# Save Output
# -----------
#
# All of the output can be saved to csv files for future use, for example, uncomment the following:

# phase_masses.to_csv('pMELTS_phase_masses.csv')
# agg_liq_comps.to_csv('pMELTS_aggregated_liquid_comps.csv')


# %%