import numpy as np
import matplotlib.pyplot as plt
import scipy

from notebooks.lib import read_load


def linear(x, a, b):
    return x * a + b


alpha = read_load("../data/alpha")
meaned_by_N = alpha.groupby('N').agg({'fd': ['mean', 'std']}) \
    .reset_index() \
    .replace([np.inf, -np.inf], np.nan)

without_prefix = alpha[alpha.N > 50]

p, pcov = scipy.optimize.curve_fit(linear, np.log(without_prefix.cr), np.log(without_prefix.N))
linear_extent = np.linspace(0, np.max(np.log(alpha.cr)))

plt.scatter(np.log(alpha.cr), np.log(alpha.N), s=1, marker='.', color="tab:blue")
plt.plot(linear_extent, linear(linear_extent, *p), color="tab:red")

plt.xlabel("$\\log r_{max}$")
plt.ylabel("$\\log N$")

plt.savefig('../figures/rmax-n.svg')
plt.savefig('../figures/rmax-n.png')
plt.show()