import numpy as np
import matplotlib.pyplot as plt
import burnman
from burnman import equilibrate
from burnman.minerals import SLB_2011

# Set the pressure, temperature and composition
pressure = 3.e9
temperature = 1500.
composition = {'Na': 0.02, 'Fe': 0.2, 'Mg': 2.0, 'Si': 1.9,
               'Ca': 0.2, 'Al': 0.4, 'O': 6.81}

# Create the assemblage
gt = SLB_2011.garnet()
ol = SLB_2011.mg_fe_olivine()
opx = SLB_2011.orthopyroxene()
assemblage = burnman.Composite(phases=[ol, opx, gt],
                               fractions=[0.7, 0.1, 0.2],
                               name='NCFMAS ol-opx-gt assemblage')

# The solver uses the current compositions of each solution as a starting guess,
# so we have to set them here
ol.set_composition([0.93, 0.07])
opx.set_composition([0.8, 0.1, 0.05, 0.05])
gt.set_composition([0.8, 0.1, 0.05, 0.03, 0.02])

equality_constraints = [('P', 10.e9), ('T', 1500.)]

sol, prm = equilibrate(composition, assemblage, equality_constraints)

print(f'It is {sol.success} that equilibrate was successful')
print(sol.assemblage)

# The total entropy of the assemblage is the molar entropy
# multiplied by the number of moles in the assemblage
entropy = sol.assemblage.S*sol.assemblage.n_moles

"""
Explanation: <h1>The BurnMan Tutorial</h1>

Part 5: Equilibrium problems
This file is part of BurnMan - a thermoelastic and thermodynamic toolkit
for the Earth and Planetary Sciences
Copyright (C) 2012 - 2021 by the BurnMan team,
released under the GNU GPL v2 or later.
Introduction
This ipython notebook is the fifth in a series designed to introduce new users to the code structure and functionalities present in BurnMan.
<b>Demonstrates</b>

burnman.equilibrate, an experimental function that determines the bulk elemental composition, pressure, temperature, phase proportions and compositions of an assemblage subject to user-defined constraints.

Everything in BurnMan and in this tutorial is defined in SI units. 
Phase equilibria
What BurnMan does and doesn't do
Members of the BurnMan Team are often asked whether BurnMan does Gibbs energy minimization. The short answer to that is no, for three reasons:
1) Python is ill-suited to such computationally intensive problems. 
2) There are many pieces of software already in the community that do Gibbs energy minimization, including but not limited to: PerpleX, HeFESTo, Theriak Domino, MELTS, ENKI, FactSAGE (proprietary), and MMA-EoS.
3) Gibbs minimization is a hard problem. The brute-force pseudocompound/simplex technique employed by Perple_X is the only globally robust method, but clever techniques have to be used to make the computations tractable, and the solution found is generally only a (very close) approximation to the true minimum assemblage. More refined Newton / higher order schemes (e.g. HeFESTo, MELTS, ENKI) provide an exact solution, but can get stuck in local minima or even fail to find a solution.
So, with those things in mind, what does BurnMan do? Well, because BurnMan can compute the Gibbs energy and analytical derivatives of composite materials, it is well suited to solving the equilibrium relations for fixed assemblages. This is done using the burnman.equilibrate function, which acts in a similar (but slightly more general) way to the THERMOCALC software developed by Tim Holland, Roger Powell and coworkers. Essentially, one chooses an assemblage (e.g. olivine + garnet + orthopyroxene) and some equality constraints (typically related to bulk composition, pressure, temperature, entropy, volume, phase proportions or phase compositions) and the equilibrate function attempts to find the remaining unknowns that satisfy those constraints.
In a sense, then, the equilibrate function is simultaneously more powerful and more limited than Gibbs minimization techniques. It allows the user to investigate and plot metastable reactions, and quickly obtain answers to questions like "at what pressure does wadsleyite first become stable along a given isentrope?". However, it is not designed to create P-T tables of equilibrium assemblages. If a user wishes to do this for a complex problem, we refer them to other existing codes. BurnMan also contains a useful utility material called burnman.PerplexMaterial that is specifically designed to read in and interrogate P-T data from PerpleX. 
There are a couple more caveats to bear in mind. Firstly, the equilibrate function is experimental and can certainly be improved. Equilibrium problems are highly nonlinear, and sometimes solvers struggle to find a solution. If you have a better, more robust way of solving these problems, we would love to hear from you! Secondly, the equilibrate function is not completely free from the curse of multiple roots - sometimes there is more than one solution to the equilibrium problem, and BurnMan (and indeed any equilibrium software) may find one a metastable root.
Equilibrating at fixed bulk composition
Fixed bulk composition problems are most similar to those asked by Gibbs minimization software like HeFESTo. Essentially, the only difference is that rather than allowing the assemblage to change to minimize the Gibbs energy, the assemblage is instead fixed.
In the following code block, we calculate the equilibrium assemblage of olivine, orthopyroxene and garnet for a mantle composition in the system NCFMAS at 10 GPa and 1500 K.
End of explanation
"""


equality_constraints = [('P', np.linspace(3.e9, 13.e9, 21)),
                        ('S', entropy)]

sols, prm = equilibrate(composition, assemblage, equality_constraints)

"""
Explanation: Each equality constraint can be a list of constraints, in which case equilibrate will loop over them. In the next code block we change the equality constraints to be a series of pressures which correspond to the total entropy obtained from the previous solve.
End of explanation
"""


data = np.array([[sol.assemblage.pressure,
                  sol.assemblage.temperature,
                  sol.assemblage.p_wave_velocity,
                  sol.assemblage.shear_wave_velocity,
                  sol.assemblage.molar_fractions[0],
                  sol.assemblage.molar_fractions[1],
                  sol.assemblage.molar_fractions[2]]
                 for sol in sols if sol.success])

"""
Explanation: The object sols is now a 1D list of solution objects. Each one of these contains an equilibrium assemblage object that can be interrogated for any properties:
End of explanation
"""


fig = plt.figure(figsize=(12, 4))
ax = [fig.add_subplot(1, 3, i) for i in range(1, 4)]

P, T, V_p, V_s = data.T[:4]
phase_proportions = data.T[4:]
ax[0].plot(P/1.e9, T)
ax[1].plot(P/1.e9, V_p/1.e3)
ax[1].plot(P/1.e9, V_s/1.e3)

for i in range(3):
    ax[2].plot(P/1.e9, phase_proportions[i], label=sol.assemblage.phases[i].name)

for i in range(3):
    ax[i].set_xlabel('Pressure (GPa)')
ax[0].set_ylabel('Temperature (K)')
ax[1].set_ylabel('Seismic velocities (km/s)')
ax[2].set_ylabel('Molar phase proportions')
ax[2].legend()

plt.show()

"""
Explanation: The next code block plots these properties.
End of explanation
"""


equality_constraints = [('phase_fraction', [opx, 0.]),
                        ('S', entropy)]
sol, prm = equilibrate(composition, assemblage, equality_constraints)

print(f'Orthopyroxene is exhausted from the assemblage at {sol.assemblage.pressure/1.e9:.2f} GPa, {sol.assemblage.temperature:.2f} K.')

"""
Explanation: From the above figure, we can see that the proportion of orthopyroxene is decreasing rapidly and is exhausted near 13 GPa. In the next code block, we determine the exact pressure at which orthopyroxene is exhausted.
End of explanation
"""


# Initialize the minerals we will use in this example.
ol = SLB_2011.mg_fe_olivine()
wad = SLB_2011.mg_fe_wadsleyite()
rw = SLB_2011.mg_fe_ringwoodite()

# Set the starting guess compositions for each of the solutions
ol.set_composition([0.90, 0.10])
wad.set_composition([0.90, 0.10])
rw.set_composition([0.80, 0.20])

"""
Explanation: Equilibrating while allowing bulk composition to vary
End of explanation
"""


T = 1600.
composition = {'Fe': 0.2, 'Mg': 1.8, 'Si': 1.0, 'O': 4.0}
assemblage = burnman.Composite([ol, wad, rw], [1., 0., 0.])
equality_constraints = [('T', T),
                        ('phase_fraction', (ol, 0.0)),
                        ('phase_fraction', (rw, 0.0))]
free_compositional_vectors = [{'Mg': 1., 'Fe': -1.}]

sol, prm = equilibrate(composition, assemblage, equality_constraints,
                        free_compositional_vectors,
                        verbose=False)
if not sol.success:
    raise Exception('Could not find solution for the univariant using '
                    'provided starting guesses.')

P_univariant = sol.assemblage.pressure
phase_names = [sol.assemblage.phases[i].name for i in range(3)]
x_fe_mbr = [sol.assemblage.phases[i].molar_fractions[1] for i in range(3)]

print(f'Univariant pressure at {T:.0f} K: {P_univariant/1.e9:.3f} GPa')
print('Fe2SiO4 concentrations at the univariant:')
for i in range(3):
    print(f'{phase_names[i]}: {x_fe_mbr[i]:.2f}')

"""
Explanation: First, we find the compositions of the three phases at the univariant.
End of explanation
"""


output = []
for (m1, m2, x_fe_m1) in [[ol, wad, np.linspace(x_fe_mbr[0], 0.001, 20)],
                          [ol, rw, np.linspace(x_fe_mbr[0], 0.999, 20)],
                          [wad, rw, np.linspace(x_fe_mbr[1], 0.001, 20)]]:

    assemblage = burnman.Composite([m1, m2], [1., 0.])

    # Reset the compositions of the two phases to have compositions
    # close to those at the univariant point
    m1.set_composition([1.-x_fe_mbr[1], x_fe_mbr[1]])
    m2.set_composition([1.-x_fe_mbr[1], x_fe_mbr[1]])

    # Also set the pressure and temperature
    assemblage.set_state(P_univariant, T)

    # Here our equality constraints are temperature,
    # the phase fraction of the second phase,
    # and we loop over the composition of the first phase.
    equality_constraints = [('T', T),
                            ('phase_composition',
                             (m1, [['Mg_A', 'Fe_A'],
                                   [0., 1.], [1., 1.], x_fe_m1])),
                            ('phase_fraction', (m2, 0.0))]

    sols, prm = equilibrate(composition, assemblage,
                            equality_constraints,
                            free_compositional_vectors,
                            verbose=False)

    # Process the solutions
    out = np.array([[sol.assemblage.pressure,
                     sol.assemblage.phases[0].molar_fractions[1],
                     sol.assemblage.phases[1].molar_fractions[1]]
                    for sol in sols if sol.success])
    output.append(out)

output = np.array(output)

"""
Explanation: Now we solve for the stable sections of the three binary loops
End of explanation
"""


fig = plt.figure()
ax = [fig.add_subplot(1, 1, 1)]

color='purple'
# Plot the line connecting the three phases
ax[0].plot([x_fe_mbr[0], x_fe_mbr[2]],
            [P_univariant/1.e9, P_univariant/1.e9], color=color)

for i in range(3):
    if i == 0:
        ax[0].plot(output[i,:,1], output[i,:,0]/1.e9, color=color, label=f'{T} K')
    else:
        ax[0].plot(output[i,:,1], output[i,:,0]/1.e9, color=color)

    ax[0].plot(output[i,:,2], output[i,:,0]/1.e9, color=color)
    ax[0].fill_betweenx(output[i,:,0]/1.e9, output[i,:,1], output[i,:,2],
                        color=color, alpha=0.2)

ax[0].text(0.1, 6., 'olivine', horizontalalignment='left')
ax[0].text(0.015, 14.2, 'wadsleyite', horizontalalignment='left',
        bbox=dict(facecolor='white',
                    edgecolor='white',
                    boxstyle='round,pad=0.2'))
ax[0].text(0.9, 15., 'ringwoodite', horizontalalignment='right')

ax[0].set_xlim(0., 1.)
ax[0].set_ylim(0.,20.)
ax[0].set_xlabel('p(Fe$_2$SiO$_4$)')
ax[0].set_ylabel('Pressure (GPa)')
ax[0].legend()
plt.show()


"""
Explanation: Finally, we do some plotting
End of explanation
"""


