'''
Teaching Exercise: Chemical Reactions in Earth
==============================================

This notebook is designed to be used for a class activity where students are randomly
assigned a chemical reaction from a list stored in a datafile. They must then find the
chemical formulae for each reactant and product, balance the equation, and then calculate
the Gibbs Free Energy of reaction across a range of pressures and temperatures.

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%2Freaction_teaching_exercise%2Fplot_reaction_teaching_exercise.ipynb


'''

# %% [markdown]
# Coding preliminaries
# --------------------
#
# First we will bring in some useful code for the exercise:

# %%
from thermoengine import model
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import pprint

# %% [markdown]
# Here I have written some functions which will help with the exercise. 
# You can ignore this part of the notebook if you like.

# %%
def check_reaction_is_balanced(phase_dict):
    totalmass = np.zeros((1,106))
    for ph in phase_dict:
        if ph in phobj:
            totalmass += phase_dict[ph] * phobj[ph].endmember_element_comps[:106]
        else:
            raise AttributeError("The mineral wasn't recognised. Check it is spelled correctly.")
    
    if np.any(totalmass != 0):
        return False
    else:
        return True

def pick_a_reaction():
    reactions = pd.read_csv('input_files/reaction_teaching_exercise.csv')
    row_rand = reactions.sample()

    print(row_rand['Reaction'].iloc[0])
    print("")
    description = row_rand['Information'].iloc[0]
    fmt = pprint.pformat(description)
    fmt = fmt.replace("'","")
    fmt = fmt[1:]
    fmt = fmt[:-1]
    print(fmt)

def make_DGr_grid(phases_in_reaction, Pmin=0.1, Pmax=30000, Tmin=0, Tmax=2000, nP=100, nT=100, seefullcolorbar=False):

    if not check_reaction_is_balanced(phases_in_reaction):
        raise AttributeError("Your reaction is not balanced!")
        
    p = np.linspace(Pmin, Pmax, nP)
    t = np.linspace(Tmin, Tmax, nT)

    tt, pp = np.meshgrid(t, p)

    dGr = np.zeros(np.shape(pp))

    for i in range(np.shape(pp)[0]):
        for j in range(np.shape(pp)[1]):
            for ph in phases_in_reaction:
                dGr[i, j] += phobj[ph].gibbs_energy(tt[i,j] + 273.15, pp[i,j]) * phases_in_reaction[ph]
    
    if np.any(dGr < 0):
        mostneg = np.min(dGr)
    else:
        mostneg = 0
    
    if np.any(dGr > 0):
        mostpos = np.max(dGr)
    else:
        mostpos = 0
    
    cscale_ends = 0.0
    if - mostneg > mostpos:
        cscale_ends = - mostneg / 1000
    else:
        cscale_ends = mostpos / 1000
    
    f, a = plt.subplots()

    cf = a.contourf(tt, pp/1000, dGr/1000, vmin=-cscale_ends, vmax=cscale_ends, cmap=plt.cm.seismic)

    a.set_xlim(a.get_xlim())
    a.set_ylim(a.get_ylim())

    if seefullcolorbar is True:
        cf = a.scatter([a.get_xlim()[0] - 10]*2, [a.get_ylim()[0] - 10]*2, c=[-cscale_ends, cscale_ends], vmin=-cscale_ends, vmax=cscale_ends, cmap=plt.cm.seismic)

    a.set_xlabel('Temperature (˚C)')
    a.set_ylabel('Pressure (kbar)')

    cbar = f.colorbar(cf)
    cbar.set_label(r'$\Delta G_r$ (kJ mol$^{-1}$)')

    reactionstr = ''
    for ph in phases_in_reaction:
        if phases_in_reaction[ph] < 0:
            if len(reactionstr) > 0:
                reactionstr += '+ '
            reactionstr += '{0} {1} '.format(-phases_in_reaction[ph], phobj[ph].props['formula'][0])
    reactionstr += '= '
    for ph in phases_in_reaction:
        if phases_in_reaction[ph] > 0:
            if reactionstr[-2:] != '= ':
                reactionstr += '+ '
            reactionstr += '{0} {1} '.format(phases_in_reaction[ph], phobj[ph].props['formula'][0])
    
    a.set_title(reactionstr)

    plt.show()


# %% [markdown]
# Access the properties for all the minerals we might be interested in today:

# %%
db = model.Database('MELTS_v1_0_Berman_PurePhases')

phobj = {
         'albite':      db.get_phase('Ab'),
         'almandine':   db.get_phase('Alm'),
         'andalusite':  db.get_phase('And'),
         'anorthite':   db.get_phase('An'),
         'antigorite':  db.get_phase('Atg'),
         'coesite':     db.get_phase('Coe'),
         'corundum':    db.get_phase('Crn'),
         'diopside':    db.get_phase('Di'),
         'enstatite':   db.get_phase('En'),
         'fayalite':    db.get_phase('Fa'),
         'ferrosilite': db.get_phase('Fs'),
         'forsterite':  db.get_phase('Fo'),
         'H2O':         db.get_phase('H2O'),
         'hematite':    db.get_phase('Hem'),
         'hydrogen':    db.get_phase('H2'),
         'jadeite':     db.get_phase('Jd'),
         'K-feldspar':  db.get_phase('Or'),
         'kyanite':     db.get_phase('Ky'),
         'magnetite':   db.get_phase('Mag'),
         'muscovite':   db.get_phase('Ms'),
         'oxygen':      db.get_phase('O2'),
         'pyrope':      db.get_phase('Prp'),
         'quartz':      db.get_phase('Qz'),
         'sillimanite': db.get_phase('Sil'),
         'spinel':      db.get_phase('Spl'),
         'tremolite':   db.get_phase('Tr'),
         'tridymite':   db.get_phase('Trd'),
         }

# %% [markdown]
# The minerals and their formulae
# -------------------------------
#
# It is important you balance the equations using the formulae given here, otherwise 
# the molar thermodynamic properties will be incorrect

print('The minerals available and their formulae:')
for ph in phobj:
    print('{0: <15} {1}'.format(ph, phobj[ph].props['formula'][0]))


# %%
# Pick a reaction
# ---------------
#
# Running this cell will pick a reaction for you to model:
#

pick_a_reaction()


# %% [markdown]
# Balance the reaction
# --------------------
#
# It might be easier to do this on paper, then transfer it to the computer. 
# Reactants should have negative numbers (per mole), and products should have 
# positive numbers (per mole).
#
# 1. Find the mineral formulae you need from the list above
# 2. Balance it (set the number of moles of each reactant and product)
# 3. Check that you have balanced it correctly
#
# The example corresponds to:
#
# 1 Mg~2~SiO~4~ + 1 SiO~2~ → 2 MgSiO~3~
#

phases_in_reaction = {
                      'forsterite': 1,
                      'enstatite': -2,
                      'quartz': 1
                      }

# %% [markdown]
#
# This cell will check if it is balanced. If it returns `True` then it *is* balanced. 
# If it returns `False` you need to double check your reaction.

check_reaction_is_balanced(phases_in_reaction)


# %%
# Calculate the Gibbs Energy of Reaction
# --------------------------------------


make_DGr_grid(phases_in_reaction)


# %%
# Why is the reaction favourable or unfavorable?
# -----------------------------------------------
#
# You could try calculating volumes, enthalpies and entropies for the phases at particular 
# pressures and temperatures to better understand why the reaction is favourable or 
# unfavourable:



phobj['forsterite'].entropy(
                            100 + 273.15, # Temperature in K
                            1000          # Pressure in bar
                            )

# %%
phobj['forsterite'].volume(
                           100 + 273.15, # Temperature in K
                           1000          # Pressure in bar
                           )

# %%
phobj['enstatite'].enthalpy(
                             100 + 273.15, # Temperature in K
                             1000          # Pressure in bar
                             )
# %%