'''
Main Model Runs
===============

This notebook demonstrates how to use Peridoxia to reproduce the main model runs in Birner et al. (2024).

***To run this code interactively in your web browser, click the "Launch Binder" button below:***
(Note that the file may take a long time to load, and outputs will not be saved.)

.. image:: https://mybinder.org/badge_logo.svg
  :target: https://mybinder.org/v2/gl/skbirner%2Fperidoxia/main?urlpath=%2Fdoc%2Ftree%2F.%2Fdoc%2Fsource%2Fauto_examples%2Fplot_main_model_runs.ipynb
'''

# %%
# Import Statements
# -----------------

# ---------------------------------------------------------------------------------------------------------------------
# General import statements
# ---------------------------------------------------------------------------------------------------------------------

import numpy as np
import matplotlib.pyplot as plt

# ---------------------------------------------------------------------------------------------------------------------
# These notebooks contain the functions and classes necessary for running the model
# ---------------------------------------------------------------------------------------------------------------------

from peridoxia import input_composition as input_comp
from peridoxia import main
from peridoxia import plotting_and_saving

# Suppress deprecation warnings in Binder script
import warnings
warnings.filterwarnings("ignore", category=DeprecationWarning)

# %%
# Input Parameters
# ----------------

# ---------------------------------------------------------------------------------------------------------------------
# Initial composition
# ---------------------------------------------------------------------------------------------------------------------

initial_assemblage = input_comp.get_WH05_starting_comp(bulk_Fe2O3=0.3, display=False) # Initial oxides/modes from Workman and Hart (2005)

# ---------------------------------------------------------------------------------------------------------------------
# Initial pressure 
# ---------------------------------------------------------------------------------------------------------------------

PGPa_0 = 4.0

# ---------------------------------------------------------------------------------------------------------------------
# Potential temperatures and colors for plotting
# ---------------------------------------------------------------------------------------------------------------------

pot_TC_vals = [1350, 1400, 1450, 1500, 1550]

colors = plt.cm.inferno(np.linspace(0.175, 0.825, 5))
pot_TC_cols = {
    1350: colors[0],
    1400: colors[1],
    1450: colors[2],
    1500: colors[3],
    1550: colors[4]
}

# %% 
# Projection to Garnet Field
# --------------------------

initial_assemblages_gt_field  = {}
for pot_TC in pot_TC_vals:
    new_assemblage = input_comp.garnet_project(initial_assemblage, pot_TC, PGPa_0, spl_final=0.05, display=False)
    initial_assemblages_gt_field[pot_TC] = new_assemblage

# %%
# Model Runs (Liquid and Garnet present)
# --------------------------------------

run_parameters = {
    'P_range_GPa':          [PGPa_0, 0.1],
    'P_step':               -0.01,
    'melting':              True,
    'melting_method':       'fractional',
    'gt_out_P_range_GPa':   [3.0, 2.8-0.001],
    'gt_out_rxn_wts':       {'gt_olv__spl_pyx':1, 'gt__pyx_pyxTs':2},
    'debug':                0,
    'run_model':            True
}            

empirical = {}

for pot_TC in pot_TC_vals:

    empirical[pot_TC] = main.ModelRun(
             label = f'Tp = {pot_TC}°C',
             color = pot_TC_cols[pot_TC],          
             potential_temperature_C = pot_TC,
             initial_assemblage = initial_assemblages_gt_field[pot_TC],
             **run_parameters
            )
    
# %% 
# Results
# -------

fig, ax = plotting_and_saving.plot_model_runs(empirical, title='Evolution of $f_{O2}$ during peridotite \ndecompression and melting', legend='upper_right')