"""Demo: practical use of numpy (mexhat-numpy.py)"""
import math
import time

import matplotlib.pyplot as plt
import numpy as np


N = 100000


def mexhat_py(t, sigma=1):
    """Computes Mexican hat shape, see
    http://en.wikipedia.org/wiki/Mexican_hat_wavelet for
    equation (13 Dec 2011)"""
    c = 2.0 / math.sqrt(3 * sigma) * math.pi**0.25
    return c * (1 - t**2 / sigma**2) * math.exp(
        -(t**2) / (2 * sigma**2))


def mexhat_np(t, sigma=1):
    """Computes Mexican hat shape using numpy, see
    http://en.wikipedia.org/wiki/Mexican_hat_wavelet for
    equation (13 Dec 2011)"""
    c = 2.0 / math.sqrt(3 * sigma) * math.pi**0.25
    return c * (1 - t**2 / sigma**2) * np.exp(
        -(t**2) / (2 * sigma**2))


def test_is_really_the_same():
    """Checking whether mexhat_np and mexhat_py produce
    the same results."""
    xs1, ys1 = loop1()
    xs2, ys2 = loop2()
    deviation = math.sqrt(sum((ys1 - ys2) ** 2))
    print("error:", deviation)
    assert deviation < 1e-14


def loop1():
    """Compute arrays xs and ys with mexican hat function
    in ys(xs), returns tuple (xs,ys)"""
    xs = np.linspace(-5, 5, N)
    ys = []
    for x in xs:
        ys.append(mexhat_py(x))
    return xs, ys


def loop2():
    """As loop1, but uses numpy to be faster."""
    xs = np.linspace(-5, 5, N)
    return xs, mexhat_np(xs)


def time_this(f):
    """Call f, measure and return number of seconds
    execution of f() takes"""
    starttime = time.time()
    f()
    stoptime = time.time()
    return stoptime - starttime


def make_plot(filenameroot):
    fig, ax = plt.subplots()
    xs, ys = loop2()
    ax.plot(xs, ys, label="Mexican hat function")
    ax.legend()
    fig.savefig(filenameroot + ".png")
    fig.savefig(filenameroot + ".pdf")


def main():
    test_is_really_the_same()
    make_plot("mexhat1d")
    time1 = time_this(loop1)
    time2 = time_this(loop2)
    print(f"Numpy version is {time1 / time2:.1f} times faster")


if __name__ == "__main__":
    main()