import warnings
warnings.simplefilter('ignore')  # ignore warnings for nicer output
import numpy as np
from sympy import symbols, log, lambdify, solve, I
import scipy.constants 
from ipywidgets import interact
from bokeh.io import push_notebook, show, output_notebook
from bokeh.plotting import figure
from bokeh.models import Span, layouts
output_notebook()

"""
Explanation: Miscibility Gaps
End of explanation
"""


# define the Gibbs free energy symbolically
ua, ub, S, Hmix, T, xb = symbols('ua ub S Hmix T xb')
xa = 1 - xb
G = ua*xa + ub*xb + Hmix*xa*xb - T*S


# define the ideal mixing entropy and substitute that in our Gibbs energy function
ideal_mixing_entropy = -scipy.constants.R*(xa*log(xa)+xb*log(xb))
G = G.subs(S, ideal_mixing_entropy)

# we'll also take the second derivative to speed up processing time later
d2G_dx2 = G.diff(xb).diff(xb)

"""
Explanation: We define a function G that returns $G$ for a given phase. $G$ is given as a regular solution for a binary A, B system. The mixing enthalpy, $H_\textrm{mix}$, is a changable paramter (Hmix). The entropy, $S$, could concievably be modeled in several ways, but here a simple ideal mixing approximation is substituted.
End of explanation
"""


xb_plot = np.linspace(0, 1, 1000)  # discritized plotting domain

p = figure(title="Free Energy, $G$", plot_height=300, plot_width=600, x_range=(0,1))
p.xaxis.axis_label = 'X_B'
p.yaxis.axis_label = 'G'
r = p.line(xb_plot, lambdify(xb, G.subs({T: 300, Hmix:0, ua: -2000, ub:-1000}), 'numpy')(xb_plot), line_width=3)
s1 = Span(location=None, dimension='height', line_color='red', line_dash='dashed', line_width=3)
s2 = Span(location=None, dimension='height', line_color='red', line_dash='dashed', line_width=3)
p.add_layout(s1)
p.add_layout(s2)

phase_diag = figure(title="Phase diagram", plot_height=300, plot_width=600, x_range=(0,1), y_range=(0,2000))
phase_diag.xaxis.axis_label = 'X_B'
phase_diag.yaxis.axis_label = 'T'
xt = phase_diag.line(np.linspace(0,0,20000), np.tile(np.linspace(0,2000,10000), 2), line_width=3)
ps1 = Span(location=None, dimension='height', line_color='red', line_dash='dashed', line_width=3)
ps2 = Span(location=None, dimension='height', line_color='red', line_dash='dashed', line_width=3)
hspan = Span(location=300, dimension='width', line_color='red', line_dash='dashed', line_width=3)
phase_diag.add_layout(ps1)
phase_diag.add_layout(ps2)
phase_diag.add_layout(hspan)

fig = layouts.Column(p, phase_diag)
show(fig, notebook_handle=True)

@interact(T=(0,2000, 10), Hmix=(-10000,20000, 200), ua=(-10000, 0, 100), ub=(-10000, 0, 100))
def update(T=300, Hmix=0, ua=-2000, ub=-1000):
    # we have to do a little wrangling to keep the labels nice.
    # first alias the args with prefix underscores
    _T, _Hmix, _ua, _ub = T, Hmix, ua, ub
    # redefine as symbols for the substitution and update the curve plot
    ua, ub, Hmix, T = symbols('ua ub Hmix T')
    r.data_source.data['y'] = lambdify(xb, G.subs({T: _T, Hmix: _Hmix, ua: _ua, ub:_ub}), 'numpy')(xb_plot)
    # create a partially substituted second derivative
    p_d2G_dx2 = d2G_dx2.subs({Hmix: _Hmix, ua: _ua, ub:_ub})
    # find the roots of the second derivative and plot those as spans
    roots = solve(p_d2G_dx2.subs({T: _T}), xb) or [I]
    # calculate the Spans
    if all([root.is_real for root in roots]):
        s1.location = float(roots[0])
        s2.location = float(roots[1])
        ps1.location = float(roots[0])
        ps2.location = float(roots[1])
    else:
        s1.location = None
        s2.location = None
        ps1.location = None
        ps2.location = None
        
    # calculate the miscibility gap, leaving T free.
    hspan.location = _T
    misc_roots = solve(p_d2G_dx2, xb)
    # calculate the miscibility gap
    misc_xs = np.ravel(lambdify(T, misc_roots)(np.linspace(0,2000,10000)))
    if (misc_xs.size > 0) and (np.nanmin(misc_xs) >= 0.0):
        xt.data_source.data['x'] = misc_xs
    else:
        xt.data_source.data['x'] = np.linspace(0,0,20000)

    push_notebook()

"""
Explanation: Next we will create a plotting domain for $x_\textrm{B}$ and create a plot with widgets that allow the temperature, mixing enthalpy, and chemical potentials of the pure elements, $\mu_\textrm{A}$ and $\mu_\textrm{B}$, to be changed.
There are two plots below. The first is the free energy curve for the phase that is controlled by the temperature, mixing enthalpy and chemical potential sliders.
The bottom plot is the phase diagram for the phase. It plots the miscibility gap, if there is one. If there is no miscibility gap, the plot will be blank. The horizontal line indicates the temperature of the free energy curve.
End of explanation
"""


