import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit
from matplotlib import rcParams

rcParams['font.size'] = 14

# Data
data = np.array([[
            1952,     1952,     1957,     1961,     1961,     1963,
            1963,     1971,     1978,     1979,     1979,     1982,
            1983,     1985,     1989,     1992,     1994,     1996,
            1996,     1997,     1998,     1999,     2001,     2003,
            2004,     2005,     2005,     2006,     2008,     2013,
            2016,     2017,     2018,     2024],
       [     157,      687,      969,     1281,     1332,     2917,
            3376,     6002,     6533,     6987,    13395,    25962,
           39751,    65050,    65087,   227832,   258716,   378632,
          420921,   895932,   909526,  2098960,  4053946,  6320430,
         7235733,  7816230,  9152052,  9808358, 12978189, 17425170,
        22338618, 23249425, 24862048, 41024320]])
year = data[0, :]
primes = data[1, :]

year_of_lin_growth = 1999

plt.plot(year, primes, color="#377eb8", marker=".", drawstyle='steps-post')
plt.yscale('log')
# Fit function
def lin_fun(x, a, b):
    return a*x + b

def exp_fun(x, a, b):
    return np.exp(a*x + b)
    
## Plot the older part which follows exponential growth of digits with time
popt, pcov = curve_fit(lin_fun, year[year <= year_of_lin_growth], np.log(primes[year <= year_of_lin_growth]))
x_vals = np.linspace(np.min(year), year_of_lin_growth, 20)
y_vals = exp_fun(x_vals, *popt)
plt.plot(x_vals, y_vals, color="#1ae41c", linestyle="dashed", label='digit number doubling every {:.1f} y'.format(np.log(2)/popt[0]))
print('Exponential fit (green curve): y = exp({:.4f} * t  +  {:.4f})'.format(*popt))



year_of_lin_growth = 1998
## Plot the newer part which follows linear growth of digits with time
popt, pcov = curve_fit(lin_fun, year[year >= year_of_lin_growth], primes[year >= year_of_lin_growth])
x_vals = np.linspace(year_of_lin_growth, np.max(year), 100)
y_vals = lin_fun(x_vals, *popt)
plt.plot(x_vals, y_vals, color="#e41a1c", linestyle="dashed", label='digit number growing by $10^6$ every {:.2f} y'.format(1e6/popt[0]))
print('Linear fit (red curve): y = {:.4g} * t  +  {:.4g}'.format(*popt))

# Plot details
plt.legend(prop={'size':10})
plt.xlabel("Year")
plt.ylabel("Number of digits in largest known prime")
plt.savefig("Digits_in_largest_found_prime_as_a_function_of_time.svg", bbox_inches="tight")
plt.show()