import numpy as np
%matplotlib inline
import matplotlib.pyplot as plt

n = 10
a, b = 1, 10
warp = np.log
unwarp = np.exp
linlin = np.linspace(a, b, num=n)

"""
Explanation: Warping Demonstration
This notebook explains how to obtain points which are evenly spaced in the latent space. I think that the warping function must be monotonic (and so, order-preserving) and continuous in order for this to work.
End of explanation
"""


plt.figure(figsize=(6, 1.5))
plt.plot(linlin, [0]*n, 'o')
plt.margins(0.2)

"""
Explanation: Linearly spaced points in linear space
End of explanation
"""


xs = np.linspace(a, b, num=100)
plt.plot(xs, warp(xs))
for x in linlin:
    l = plt.plot([x, x, 0], [0, warp(x), warp(x)], '--', linewidth=0.5)
    plt.plot([0, x], [warp(x), 0], 'o', color=l[0].get_color())
plt.xlabel('linear space')
plt.ylabel('warped space')
#plt.xscale('log')
plt.margins(0.1)

"""
Explanation: Warping the linearly spaced points
End of explanation
"""


wa, wb = warp(a), warp(b)
xs = np.linspace(wa, wb, num=100)
plt.plot(xs, unwarp(xs))
for x in warp(linlin):
    l = plt.plot([x, x, 0], [0, unwarp(x), unwarp(x)], '--', linewidth=0.5)
    plt.plot([0, x], [unwarp(x), 0], 'o', color=l[0].get_color())
plt.xlabel('warped space')
plt.ylabel('linear space')
plt.margins(0.1)

"""
Explanation: Unwarping the warped points
End of explanation
"""


wa, wb = warp(a), warp(b)
linwarp = np.linspace(wa, wb, num=n)

xs = np.linspace(wa, wb, num=100)
plt.plot(xs, unwarp(xs))
for x in linwarp:
    l = plt.plot([x, x, 0], [0, unwarp(x), unwarp(x)], '--', linewidth=0.5)
    plt.plot([0, x], [unwarp(x), 0], 'o', color=l[0].get_color())
plt.xlabel('warped space')
plt.ylabel('linear space')
plt.margins(0.1)

plt.figure(figsize=(6, 1.5))
plt.plot(unwarp(linwarp), [0]*n, 'o')
plt.margins(0.1)

"""
Explanation: To obtain points in the linear space which are distributed evenly in the warped space:
space points evenly between the warped endpoints, then unwarp.
End of explanation
"""


xs = np.logspace(np.log(a), np.log(b), num=n, base=np.e)
print(np.allclose(xs, unwarp(linwarp)))

"""
Explanation: When the warping function is the natural logarithm, the evenly spaced points can be generated using the np.logspace() function:
End of explanation
"""


