#!/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

# Both Fermi gas and Bose gas are easily calculated in grand canonical ensemble.
# Fermi gas has grand potential given by:
#   Omega = - V * g_0 * (kT)^(alpha + 1) * fdint(alpha, mu/kT)
# where fdint(s, t) is the complete fermi-dirac integral,
#   fdint(s, t) = -polylog(s + 1, -exp(t))
# Bose gas has grand potential:
#   Omega = - V * g_0 * (kT)^(alpha + 1) * beint(alpha, mu/kT)
# (only valid for mu <= 0) where beint(s, t) is
#   beint(s, t) = +polylog(s + 1, +exp(t))
# The classical ideal gas is in-between these, with:
#   Omega = - V * g_0 * (kT)^(alpha + 1) * exp(mu/kT)

# Quantities:
# Pressure        P = -d(Omega)/d(V) = -Omega/V
# Particle number N = -d(Omega)/d(mu)
# Entropy         S = -d(Omega)/d(T)
# Note that:
#    d/dt polylog(s, exp(t)) = polylog(s-1, exp(t))
#
# So these derivatives are easy to compute, e.g., for Fermi:
#   N = -V * g_0 * (kT)^(alpha)  * fdint(alpha - 1 , mu/kT)
#   S = k * N * ( (alpha + 1) * fdint(alpha, mu/kT) / fdint(alpha-1, mu/kT) - mu/kT )
#  ( .. likewise substituting beint or exp for the bose/classical cases .. )
#
# 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 at each point, choose T such that N has the
# desired value.

# To nondimensionalize, we will use a characteristic temperature scale T' based on
# the fixed number N', using the classical case for mu=0:
#    g_0 (kT')^(alpha) = N'/V
# Roughly speaking, this is the temperature at which the thermal de broglie
# wavelength is comparable to the distance between identical particles.
# Likewise, a characteristic pressure to match this temperature and density:
#    P' = N' k T' / V
#       = g_0 (kT')^(alpha + 1)
# With T' and P' in hand, we can rescale all our formulas to be independent of
# g_0, N', and volume, and also do our mu-scanning trick.
# e.g., in the fermi case we get:
#    (T/T')^(alpha) = 1/fdint(alpha - 1, mu/kT)
#    P/P' = (T/T')^(alpha + 1) * fdint(alpha, mu/kT)

# 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')

# 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))

# Fermi gas
color_fermi = '#1f77b4'
T_fermi = mpmath.fp.gamma(alpha+1) ** (1./alpha)
P_fermi = T_fermi / (alpha+1)
# sweep z; make the last point to be basically T=0
z = np.exp(np.linspace(-2, 20, 201))
z[-1] = 1e100
T = (-polylog(alpha, -z)) ** (-1./alpha)
P = (T)**(alpha + 1) * -polylog(alpha + 1, -z)
S = (alpha+1) * polylog(alpha + 1, -z)/polylog(alpha, -z) - np.log(z)
mu = np.log(z) * T
# extend traces to exactly T=0
T = np.concatenate((T, [0.]))
P = np.concatenate((P, [P_fermi]))
S = np.concatenate((S, [0.]))
mu = np.concatenate((mu, [T_fermi]))

ax1.plot(T,P, label="Fermi", color=color_fermi)
ax2.plot(T,S, label="Fermi", color=color_fermi)
ax3.plot(T,mu, label="Fermi", color=color_fermi)

# Indicate fermi temperature
P_at_T_fermi = np.interp(T_fermi, T[::-1], P[::-1])
S_at_T_fermi = np.interp(T_fermi, T[::-1], S[::-1])
mu_at_T_fermi = np.interp(T_fermi, T[::-1], mu[::-1])
ax1.plot([T_fermi, T_fermi], [P_at_T_fermi, -100 ], color=color_fermi, lw=1, ls=(0, (1,5)))
ax2.plot([T_fermi, T_fermi], [S_at_T_fermi, -100 ], color=color_fermi, lw=1, ls=(0, (1,5)))
ax3.plot([T_fermi, T_fermi], [mu_at_T_fermi, -100 ], color=color_fermi, lw=1, ls=(0, (1,5)))


# Ideal gas -- this is just a straight line T=P but for consistency, calculate
# it similarly to the bose and fermi cases.
color_classical = '#ff7f0e'
z = np.exp(np.linspace(-2, 20, 200))
T = (z) ** (-1./alpha)
P = (T)**(alpha + 1) * z
S = (alpha+1) - np.log(z)
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.]))

ax1.plot(T,P, label="classical", color=color_classical)
ax2.plot(T,S, label="classical", color=color_classical)
ax3.plot(T,mu, label="classical", color=color_classical)


# Bose gas
color_bose = '#2ca02c'
# 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))
T = (polylog(alpha, z)) ** (-1./alpha)
P = (T)**(alpha + 1) * polylog(alpha + 1, z)
S = (alpha+1) * polylog(alpha + 1, z)/polylog(alpha, z) - np.log(z)
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).
    T2 = np.linspace(Tcrit, 0, 101)
    P2 = T2 ** (alpha + 1) * polylog(alpha + 1, 1)
    # 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) * polylog(alpha + 1, 1)
    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))

    # Mark the critical temperature with a *
    ax1.plot([Tcrit],[P2[0]], '*', color='k', ms=12, mew=0, alpha=0.6, zorder=2.5)
    ax2.plot([Tcrit],[S2[0]], '*', color='k', ms=12, mew=0, alpha=0.6, zorder=2.5)
    ax3.plot([Tcrit],[0], '*', color='k', ms=12, mew=0, alpha=0.6, zorder=2.5)

    ax1.plot([Tcrit, Tcrit], [P2[0], -100 ], color=color_bose, lw=1, ls=(0, (4,2)))
    ax2.plot([Tcrit, Tcrit], [S2[0], -100 ], color=color_bose, lw=1, ls=(0, (4,2)))
    ax3.plot([Tcrit, Tcrit], [0, -100 ], color=color_bose, lw=1, ls=(0, (4,2)))

else:
    # extend traces to 0
    T = np.concatenate((T, [0.]))
    P = np.concatenate((P, [0.]))
    S = np.concatenate((S, [0.]))
    mu = np.concatenate((mu, [0.]))

ax1.plot(T,P, label="Bose", color=color_bose)
ax2.plot(T,S, label="Bose", color=color_bose)
ax3.plot(T,mu, label="Bose", color=color_bose)


# format temperature axis nicely
for ax in [ax1, ax2, ax3]:
    ax.set_xlim(0,1.8)
    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_ylim(0,1.8)
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,))

yrange = 0.2 + 1 + (1 + np.log(1.8))*alpha   # range to fit the right side
ax2.set_ylim(-0.4*yrange, +yrange)
ax2.set_ylabel('entropy per particle $S/Nk$')
ax2.legend(loc='lower right')
fig2.savefig('quantum ideal gas entropy %gd.svg'%(alpha*2,))

yrange = 1.5 + 0.7*alpha
ax3.set_ylim(-yrange, +0.6*yrange)
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,))