import numpy
from matplotlib import pyplot
import scipy
from scipy import stats
import seaborn

seaborn.set()
%matplotlib inline

print('numpy: {}'.format(numpy.version.full_version))
print('scipy: {}'.format(scipy.version.full_version))

numpy.random.seed(0)
N = 37000
wikiLoc = 2.25
wikiScale = 0.875

# _n denotes numpy
X_n = numpy.random.lognormal(mean=wikiLoc, sigma=wikiScale, size=N)

bins = numpy.logspace(-1, 3, num=250)

fig, ax = pyplot.subplots(figsize=(8, 5))
_ = ax.hist(X_n, bins=bins, cumulative=False, histtype='step', 
            normed=True, lw=1.25, label='Hist of random numpy sample')

ax.set_xscale('log')
ax.set_xlabel(r'$X$')
ax.set_ylabel('PDF')
ax.legend()

spShape, spLoc, spScale = stats.lognorm.fit(X_n, loc=0)
ln_dist = stats.lognorm(spShape, spLoc, spScale)
print("""
spShape = {:.3f}
spLoc = {:.3f}
spScale = {:.3f}
""".format(spShape, spLoc, spScale))

fig, ax = pyplot.subplots(figsize=(8, 5))
_ = ax.hist(X_n, bins=bins, cumulative=False, histtype='step', 
            normed=True, lw=1.25, label='Numpy random number')

# theoretical PDF
x_hat = numpy.logspace(-1, 2, num=1000)
y_hat = ln_dist.pdf(x_hat)

ax.plot(x_hat, y_hat, '-', lw=1.25, label='Fit with scipy')
ax.set_xscale('log')
ax.set_xlabel(r'$X$')
ax.set_ylabel('PDF')
ax.legend()

wikiM = numpy.exp(wikiLoc + 0.5*wikiScale**2)
wikiM_est = X_n.mean()

wikiS = numpy.sqrt((numpy.exp(wikiScale**2) - 1)  * numpy.exp(2*wikiLoc + wikiScale**2))
wikiS_est = X_n.std()

summary = """
\tTheory\tSample
wikiM\t{0:.3f}\t{1:.3f}
wikiS\t{2:.3f}\t{3:.3f}
"""
print(summary.format(wikiM, wikiM_est, wikiS, wikiS_est))

sampleLoc = numpy.log((wikiM_est**2) / numpy.sqrt(wikiS_est**2 + wikiM_est**2))
sampleScale = numpy.sqrt(numpy.log(1 + (wikiS_est/wikiM_est)**2))

summary = """
         \tTheory\tSample
wikiLoc  \t{0:.3f}\t{1:.3f}
wikiScale\t{2:.3f}\t{3:.3f}
"""

print(summary.format(wikiLoc, sampleLoc, wikiScale, sampleScale))

print("""
spShape = {:.3f}
spLoc = {:.3f}
spScale = {:.3f}
""".format(spShape, spLoc, spScale))

sp_wikiLoc = numpy.log(spScale)
sp_wikiScale = spShape

summary = """
         \tTheory\tSample\tScipy
wikiLoc  \t{0:.3f}\t{1:.3f}\t{2:.3f}
wikiScale\t{3:.3f}\t{4:.3f}\t{5:.3f}
"""
print(summary.format(wikiLoc, sampleLoc, sp_wikiLoc, wikiScale, sampleScale, sp_wikiScale))

frozen_dist = stats.lognorm(wikiScale, loc=0, scale=numpy.exp(wikiLoc))

fig, ax = pyplot.subplots(figsize=(8, 5))
_ = ax.hist(X_n, bins=bins, cumulative=False, histtype='step', 
            normed=True, lw=1.25, alpha=0.5,
            label='Numpy random number')

# theoretical PDF
x_hat = numpy.logspace(-1, 2, num=1000)
y_hat_1 = ln_dist.pdf(x_hat)
y_hat_2 = frozen_dist.pdf(x_hat)

ax.plot(x_hat, y_hat_1, '-', lw=1.25, label='Fit with scipy')
ax.plot(x_hat, y_hat_2, '--', lw=1.25, label='Theoretical')
ax.set_xscale('log')
ax.set_xlabel(r'$X$')
ax.set_ylabel('PDF')
ax.legend()