from math import gcd
import numpy as np
import matplotlib.pyplot as plt

## Cantor staircase

def cantor(n):
    return [0.0] + cant(0.0, 1.0, n) + [1.0]


def cant(x, y, n):
    if n == 0:
        return []

    new_pts = [2.0 * x / 3.0 + y / 3.0, x / 3.0 + 2.0 * y / 3.0]
    return cant(x, new_pts[0], n - 1) + new_pts + cant(new_pts[1], y, n - 1)


xcantor = np.array(cantor(5))
ycantor = np.cumsum(np.ones(len(xcantor)) / (len(xcantor) - 2)) - 1.0 / (len(xcantor) - 2)
ycantor[-1] = 1

# Unit circle
tcircle = np.linspace(0, 2 * np.pi, 100)

# ## Brownian motion
# nbrow = 300
# dt = 1.0 / (nbrow - 1)
# brow = np.cumsum(np.sqrt(dt) * np.random.randn(nbrow))

## Weierstrass function
def weirstrass(x, a, b):
    n = 200  # number of terms in the series
    result = 0.0
    for k in range(n):
        result += a**k * np.cos(b**k * np.pi * x)
    return result

## Thomae's function
xthom = []
ythom = []
for q in range(1, 101):
    for p in range(1, q):
        xthom.append(p * 1.0 / q * 1.0)
        ythom.append(1 / (q / gcd(p, q)))

fig, axs = plt.subplots(2, 2, figsize=(10, 10))
# plot a filed circle
axs[0,0].fill(np.cos(tcircle), np.sin(tcircle), color="k")
axs[0,0].set_title("Zeros of $x^2 + y^2 - 1$")
axs[0,1].plot(xcantor, ycantor, color="k")
axs[0,1].set_title("Cantor staircase graph")
# axs[0,1].plot(np.arange(nbrow)/nbrow, brow, color="k")
# axs[0,1].set_title("Realization of 1D brownian motion")
axs[1,0].plot(np.linspace(0, 1, 300), weirstrass(np.linspace(-1, 1, 300), 0.5, 3.0), color="k")
axs[1,0].set_title("Graph of Weierstrass function")
axs[1,1].scatter(xthom, ythom, color="k", marker='x')
axs[1,1].set_title("Graph of Thomae's function")

plt.savefig("smooth_function_on_examples.png")