#!/usr/bin/env python3

import numpy as np
from matplotlib import pyplot as plt
import mpmath
import sys

def fixpoly(s,x):
    # The mpmath polylog sometimes returns a tiny spurious imaginary part, and
    # it throws exceptions for weird cases. Clean this up.
    try:
        return np.real(mpmath.fp.polylog(s,x))
    except (OverflowError, ValueError):
        return np.nan
polylog = np.vectorize(fixpoly, otypes=[float])

# assumed density of states G(E) = V * g_0 * (E - E_0) ^ (alpha - 1) / Gamma(alpha)
# as appropriate for 3D, nonrelativistic, massive particles.
# The prefactor g_0 includes things like spin degeneracy, mass, planck constant, factors of pi, etc.
# V is volume of the system (where applicable - for particles in harmonic well,
# just make sure V*g_0 is the right value).
# We will only plot things that are independent of V and g_0, and we assume E_0 = 0.

# The key parameter in this density of states is alpha.
# for massive particle in a box, alpha = dimensionality/2
# for particle in a harmonic well, alpha = dimensionality
# for hyperrelativistic particles in a box, alpha = dimensionality
if len(sys.argv) > 1:
    # allow massive-particle-in-box dimensionality to be provided as command line arg
    alpha = float(sys.argv[1]) / 2
else:
    # default to 3D massive case:
    alpha = 1.5

# For the above density of states function, our three ideal gases are easily calculated
# in grand canonical ensemble, yielding the following grand potential:
#   Omega = - V * g_0 * (kT)^(alpha + 1) * I(alpha, mu/kT)
# where I(s,t) is an integral depending on the statistics.
#
# Fermi gas has I = fdint(s, t), the complete fermi-dirac integral,
#   I(s,t) = fdint(s, t) = -polylog(s + 1, -exp(t))
# Bose gas (only valid for mu <=0) has I = beint(s, t), the complete bose-einstein integral,
#   I(s,t) = beint(s, t) = +polylog(s + 1, +exp(t))
# The classical ideal gas has simply:
#   I(s,t) = exp(t)
#
# Note that our generic I(s,t) obeys in all cases:
#     d/dt I(s, t) = I(s-1, t)
# so:
#     d/d(mu) I(s, mu/kT) = (1/kT) I(s-1, mu/kT)
#     d/d(T)  I(s, mu/kT) = -(mu/kT)^2 I(s-1, mu/kT)
# which will make it very easy to compute the various derivative quantities.

# Quantities:
# Pressure        P = -d(Omega)/d(V) = -Omega/V
# Particle number N = -d(Omega)/d(mu)
# Entropy         S = -d(Omega)/d(T)
# Compressibility uses d^2(Omega)/d(mu)^2, and we will represent via bulk modulus (see below)
#
# Using the I(s,t) derivation these derivatives are easy to compute:
#   N = V * g_0 * (kT)^(alpha) * I(alpha - 1 , mu/kT)
#   S = k * N * ( (alpha + 1) * I(alpha, mu/kT) / I(alpha-1, mu/kT) - mu/kT )
#
# Unfortunately, there is no analytic formula for mu in terms of N, only the other way.
# This is unfortunate since we want to plot temperature-dependence of a system
# with fixed N and fixed V.
# However we can see that both P and N are functions of kT and mu/kT. So, we can
# sweep the value of mu/kT, and then for each point, calculate the T value for the
# N/V we have chosen. That means turning the above N relation into a definition
# for T:
#    kT = [(N/V) / (I(alpha - 1, mu/kT) * g_0)]^(1/alpha)

# Finally, we will construct some reference quantities just to nondimensionalize
# the plots. We will define the following convenient reference temperature T' based
# on our fixed density N/V:
#    kT' = [(N/V) / g_0]^(-1/alpha)
# (Very roughly speaking, this is the temperature at which the thermal de broglie
#  wavelength is comparable to the distance between identical particles.)
#
# So our plotted x axis quantity will be:
#    T/T' = kT/kT' = (I(alpha - 1, mu/kT))^(1/alpha)
# The Fermi temperature and Bose temperature are constant multiples of T',
# so they are fixed reference points on the x axis (though they move around
# depending on alpha).
#
# Likewise, a reference pressure to match this temperature and density:
#    P' = N k T' / V
#       = g_0 (kT')^(alpha + 1)
# So our pressure plot y axis will be:
#    P/P' = (T/T')^(alpha + 1) * I(alpha, mu/kT)

# Bulk modulus (isothermal)
#    K = V (\partial P/\partial V)_{N,T}
#      = (N/V)^2 / (d^2(P)/d(mu)^2)_{V,T})
#      = g_0 * (kT)^(alpha+1) * I(alpha - 1 , mu/kT)^2/I(alpha - 2, mu/kT)
#        (for fermi gas, similar for bose)
# and it has pressure units so we plot dimensionless:
#    K/P' = (T/T')^(alpha + 1) * I(alpha - 1 , mu/kT)^2/I(alpha - 2, mu/kT)



# Program variables used below:
# z = exp(mu/kT)
# T refers to T/T', where T' defined above.
# P refers to P/P', where P' defined above.
# S refers to S/(N.k)
# mu refers to mu/(kT') = (mu/kT) * (T/T')
# K refers to K/P'

# Pressure vs temperature graph
fig1 = plt.figure()
fig1.set_size_inches(3,2.3)
ax1 = plt.axes((0.09, 0.17, 0.90, 0.82))

# Entropy vs temperature graph
fig2 = plt.figure()
fig2.set_size_inches(3,2.3)
ax2 = plt.axes((0.18, 0.17, 0.81, 0.82))

# Chemical potential vs temperature graph
fig3 = plt.figure()
fig3.set_size_inches(3,2.3)
ax3 = plt.axes((0.15, 0.17, 0.84, 0.82))

# Bulk modulus vs temperature graph
fig4 = plt.figure()
fig4.set_size_inches(3,2.3)
ax4 = plt.axes((0.09, 0.17, 0.90, 0.82))

#####
# Fermi gas
# sweep z; make the last point to be almost T=0
z = np.exp(np.linspace(-2, 20, 201))
z[-1] = 1e100
I_m0 = -polylog(alpha + 1, -z)
I_m1 = -polylog(alpha, -z)
I_m2 = -polylog(alpha - 1, -z)
T = I_m1 ** (-1./alpha)
P = (T)**(alpha + 1) * I_m0
S = (alpha+1) * (I_m0 / I_m1) - np.log(z)
K = (T)**(alpha + 1) * I_m1**2 / I_m2
mu = np.log(z) * T
# Define 0-temperature and extend traces to exactly T=0
T_fermi = mpmath.fp.gamma(alpha+1) ** (1./alpha)
P_fermi = T_fermi / (alpha+1)
K_fermi = T_fermi / (alpha)
T = np.concatenate((T, [0.]))
P = np.concatenate((P, [P_fermi]))
S = np.concatenate((S, [0.]))
K = np.concatenate((K, [K_fermi]))
mu = np.concatenate((mu, [T_fermi]))

color_fermi = '#1f77b4'
for ax,y in [(ax1,P),(ax2,S),(ax3,mu),(ax4,K)]:
    ax.plot(T,y, label="Fermi", color=color_fermi)
    # Draw fermi temperature line by drawing dashed line up to the trace
    y_at_T_fermi = np.interp(T_fermi, T[::-1], y[::-1])
    ax.plot([T_fermi, T_fermi], [y_at_T_fermi, -100 ], color=color_fermi, lw=1, ls=(0, (1,5)))


#####
# Ideal gas
# This is trivial, but for consistency, calculate
# it similarly to the bose and fermi cases.
z = np.exp(np.linspace(-2, 20, 200))
I_m0 = I_m1 = I_m2 = z
T = I_m1 ** (-1./alpha)
P = (T)**(alpha + 1) * I_m0
S = (alpha+1) * (I_m0 / I_m1) - np.log(z)
K = (T)**(alpha + 1) * I_m1**2 / I_m2
mu = np.log(z) * T
# extend traces to exactly T=0
T = np.concatenate((T, [0.]))
P = np.concatenate((P, [0.]))
S = np.concatenate((S, [-np.inf]))
mu = np.concatenate((mu, [0.]))
K = np.concatenate((K, [0]))

color_classical = '#ff7f0e'
for ax,y in [(ax1,P),(ax2,S),(ax3,mu),(ax4,K)]:
    ax.plot(T,y, label="classical", color=color_classical)


#####
# Bose gas
# Approach mu=0 from below, making sure to include the last floating point
# number smaller than 1.
z = np.concatenate((np.linspace(0.01, 0.99, 99), [0.999, 0.9999, 1 - 1e-16]))
#z = 1 - np.exp(np.linspace(-0.1, -34, 200))
#z = np.exp(np.linspace(-2, -1e-15, 100))

I_m0 = polylog(alpha + 1, +z)
I_m1 = polylog(alpha, +z)
I_m2 = polylog(alpha - 1, +z)
T = I_m1 ** (-1./alpha)
P = (T)**(alpha + 1) * I_m0
S = (alpha+1) * (I_m0 / I_m1) - np.log(z)
K = (T)**(alpha + 1) * I_m1**2 / I_m2
mu = np.log(z) * T

if alpha > 1:
    # In >2 dimensions, the bose gas starts condensing at nonzero temperature,
    # right at the point when mu=0.
    Tcrit = (polylog(alpha, 1)) ** (-1./alpha)
    Pcrit = (Tcrit)**(alpha + 1) * polylog(alpha + 1, 1)
    Scrit = (alpha+1) * polylog(alpha + 1, 1)/polylog(alpha, 1)
    # What about T < Tcrit?
    # Basically now mu is pinned to 0. It cannot go any higher because that would
    # mean infinite particles in every state with energy below mu.
    # Instead as temperature lowers, mu stays at 0 and the particle number in the
    # continuum of states with energy above mu will drop accordingly.
    # The continuum is called the 'noncondensed phase'; the particles that
    # have disappeared have all necessarily gone into some lowest-energy state
    # that is infinitesimally above mu, the 'condensed phase'.
    # So, let's set mu=0, and drop our constraint on N, and calculate the pressure
    # from the continuum phase just as before. (the condensed phase contributes
    # no pressure, in macroscopic ideal gas).

    z2 = 1.
    I_m0 = polylog(alpha + 1, z2)
    T2 = np.linspace(Tcrit, 0, 101)
    P2 = T2 ** (alpha + 1) * I_m0
    # In the case of entropy "S/Nk", we have to be careful since now the total N' is
    # distinct from our continuum N (excited).
    S2 = T2 ** alpha * (alpha+1) * I_m0
    K2 = 0*T2
    mu2 = 0*T2
else:
    # In <= 2 dimensions, just tack on a single data point for 0 temperature.
    T2=np.array([0.])
    P2=0*T2
    S2=0*T2
    K2 = 0*T2
    mu2 = 0*T2

# concatenate to existing traces
T = np.concatenate((T, T2))
P = np.concatenate((P, P2))
S = np.concatenate((S, S2))
mu = np.concatenate((mu, mu2))
K = np.concatenate((K, K2))

color_bose = '#2ca02c'
for ax,y,ycrit in [(ax1,P,P2[0]),(ax2,S,S2[0]),(ax3,mu,mu2[0]),(ax4,K,K2[0])]:
    ax.plot(T,y, label="Bose", color=color_bose)
    if alpha > 1 and ycrit is not None:
        # Mark the critical temperature with a * and a vertical line.
        ax.plot([Tcrit],[ycrit], '*', color='k', ms=12, mew=0, alpha=0.6, zorder=0)
        ax.plot([Tcrit, Tcrit], [ycrit, -100 ], color=color_bose, lw=1, ls=(0, (4,2)))

#####
# set view ranges
Tmin = -0.03
Tmax = 1.8
ax1.set_ylim(Tmin,Tmax)
yrange = 0.2 + 1 + (1 + np.log(1.8))*alpha   # range to fit the right side
ax2.set_ylim(-0.4*yrange, +yrange)
yrange = 1.5 + 0.7*alpha
ax3.set_ylim(-yrange, +0.6*yrange)
ax4.set_ylim(Tmin,Tmax)

# format temperature axis nicely
for ax in [ax1, ax2, ax3, ax4]:
    ax.set_xlim(Tmin,Tmax)
    ax.set_xlabel('temperature', labelpad=0)
    tl = []
    tl.append((0, "0", 'k'))
    if alpha > 1:
        tl.append((Tcrit, r"$T_{\rm B}$", color_bose))
    tl.append((T_fermi, r"$T_{\rm F}$", color_fermi))
    ticks, labels, colors = zip(*tl)
    ax.set_xticks(ticks)
    ax.set_xticklabels(labels)
    for label,color in zip(ax.xaxis.get_ticklabels(), colors):
        label.set_color(color)

ax1.set_yticks([0])
ax1.set_ylabel('pressure', labelpad=-5)
ax1.legend(loc='lower right')
fig1.savefig('quantum ideal gas pressure %gd.svg'%(alpha*2,))

ax2.set_ylabel('entropy per particle $S/Nk$')
ax2.legend(loc='lower right')
fig2.savefig('quantum ideal gas entropy %gd.svg'%(alpha*2,))

ax3.set_yticks([0])
ax3.set_yticklabels([r'$\varepsilon_0$'])
ax3.set_ylabel('chemical potential $\\mu$')
ax3.legend(loc='lower left')
fig3.savefig('quantum ideal gas chemical potential %gd.svg'%(alpha*2,))

ax4.set_yticks([0])
ax4.set_ylabel(r'bulk modulus $K_{T}$', labelpad=-5)
ax4.legend(loc='lower right')
fig4.savefig('quantum ideal gas bulk modulus %gd.svg'%(alpha*2,))