'''
Model Runs Varying Initial Ferric Iron
======================================

This notebook demonstrates the effect of varying initial ferric iron content.

***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%2Fempiricalredoxmodel/main?urlpath=%2Fdoc%2Ftree%2F.%2Fdoc%2Fsource%2Fauto_examples%2Fplot_vary_ferric.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 empirical_redox_model import input_composition as input_comp
from empirical_redox_model import main
from empirical_redox_model import plotting_and_saving

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

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

### -----------------------------------------------------------------------------------------------
### Define parameters that vary across runs
### -----------------------------------------------------------------------------------------------

values = [1350, 1550]
potential_temperatures = {}
colors = {}
for val in values:
    key = f'pot T = {val:.0f} °C'
    potential_temperatures[key] = val

    # Define colors-- change c_range values to adjust bounds of color map
    c_range = [0.175,0.825]
    scale_0_1 = (val-min(values))/(max(values)-min(values))
    scale_c_range = scale_0_1*(c_range[1]-c_range[0]) + c_range[0]
    colors[key] = plt.cm.inferno(scale_c_range)

values = [0.25,0.30,0.35]
Fe2O3_values = {}
color_shades = {}
for val in values:
    key = f'bulk Fe2O3 = {val:.2f} wt%'
    Fe2O3_values[key] = val

    # Define color shades-- change c_range values to adjust bounds of color darkness
    c_range = [0.25,1]
    scale_0_1 = (val-min(values))/(max(values)-min(values))
    scale_c_range = scale_0_1*(c_range[1]-c_range[0]) + c_range[0]
    color_shades[key] = scale_c_range

### -----------------------------------------------------------------------------------------------
### Define parameters that are constant across runs
### -----------------------------------------------------------------------------------------------

# Pressure
PGPa_0 = 4.0        # initial pressure, in GPa
P_step = -0.01      # pressure step, in GPa

# Melting
melting = True      # 'True' allows melt to form, 'False' forces subsolidus

# %%
# Iterate over potential temperature and Fe2O3 values
# ---------------------------------------------------

model_runs = {}

# Iterate over potential temperatures
for key_T, val_T in potential_temperatures.items():

    # Define potential temperature and color (for plotting)
    pot_TC = potential_temperatures[key_T]
    color_val = colors[key_T]

    # Iterate over bulk Fe2O3 values
    for key_A, val_A in Fe2O3_values.items():
        
        combined_key = key_T+', '+key_A

        # Get gt-field assemblage at Fe2O3 of interest (phase compositions adapted from Workman and Hart, 2005)
        initial_assemblage_spl_field = input_comp.get_WH05_starting_comp(bulk_Fe2O3=val_A)
        initial_assemblage_gt_field = input_comp.garnet_project(initial_assemblage_spl_field, pot_TC, PGPa_0, spl_final=0.05)

        # Change shade in RGBA value
        color = (color_val[0], color_val[1], color_val[2], color_shades[key_A])

        model_runs[combined_key] = main.ModelRun(
                label = combined_key,
                color = color,          
                melting = melting,
                potential_temperature_C = pot_TC,
                P_range_GPa = [PGPa_0, 0.1], # ends automatically at cpx-out
                P_step = P_step,
                initial_assemblage = initial_assemblage_gt_field,
                debug = 0,
                )
        
# %%
# Results
# -------

# Without legend
fig, ax = plotting_and_saving.plot_model_runs(model_runs, legend=None)

# %%

# With legend
fig, ax = plotting_and_saving.plot_model_runs(model_runs, legend='below_plot')

# %%