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

from notebooks.lib import linear


def cr_n(df, ignore_prefix=0, data_color="tab:blue", fit_color="tab:red"):
    p, pcov = scipy.optimize.curve_fit(linear, np.log(df.cr[df.N > ignore_prefix]), np.log(df.N[df.N > ignore_prefix]))
    linear_extent = np.linspace(0, np.max(np.log(df.cr)))

    plt.scatter(
        np.log(df.cr),
        np.log(df.N),
        s=1, marker='.', color=data_color
    )

    plt.plot(
        linear_extent,
        linear(linear_extent, *p),
        color=fit_color
    )

    plt.xlabel("$\\log r_{max}$")
    plt.ylabel("$\\log N$")
    return p, pcov