import numpy as np
import matplotlib.pyplot as plt



def riemann_zeta_fn(n : int, s : complex) -> np.ndarray:
    """

    Args_
        n - number of  points to plot
        s - exponential

    Returns:
        numpy array of length n; values for the first n
            terms in the series defining the riemann zeta function

    """

    n_values = np.arange(1, n+1)

    sequence_values = np.array([1/(n**s) for n in n_values])

    return n_values, np.cumsum(sequence_values)


n = 30
n_values, riemann1 = riemann_zeta_fn(n, 1)
n_values, riemann2 = riemann_zeta_fn(n, 2)

label1 = r"$\sum_{n = 1}^{\infty} \frac{1}{n}$"
label2 = r"$\sum_{n = 1}^{\infty} \frac{1}{n^2}$"

plt.plot(n_values, riemann1, '*', label=label1)
plt.plot(n_values, riemann2, 'xb', label=label2)

plt.xlabel("n")
plt.ylabel("Partial sums")

plt.legend(fontsize=15)
plt.savefig("riemann_zeta_s1and2.svg")