from jolymer.sas.beaucage import beaucage_function
from jolymer import jocolors as joc

import matplotlib.pyplot as plt
import numpy as np
plt.style.use('../jowebstyle.mplstyle')

def gupo(q, Rg, G, dd, s=0):
    q = np.atleast_1d(q)
    Q1 = (1 / Rg) * ((dd - s) * (3 - s) / 2) ** 0.5
    D = G * np.exp(-Q1 ** 2 * Rg ** 2 / (3 - s)) * Q1 ** (dd - s)
    # define functions in different regions
    def _I1_2regions(q):
        res = G / q ** s * np.exp(-q ** 2 * Rg ** 2 / (3 - s))
        return res
    def _I2_2regions(q):
        res = D / q ** dd
        return res
    I = np.piecewise(q, [q < Q1, q >= Q1], [_I1_2regions, _I2_2regions])
    return I

fig, (axrg, axexp) = plt.subplots(ncols=2, figsize=(5,3))
axrg.set_prop_cycle(color=[joc.Magenta, joc.Cyan, joc.GreenYellow])
axexp.set_prop_cycle(color=[joc.Red, joc.Blue, joc.Brown])


xdata = np.logspace(-3, 1, num=10000)
b4 = beaucage_function(xdata, 1., 1., 4.)
b3 = beaucage_function(xdata, 100., 1., 3.)
b2 = beaucage_function(xdata, 100., 1., 2.)

b4b = beaucage_function(xdata, 10., 1., 4.)
b4c = beaucage_function(xdata, 100., 1., 4.)

axrg.loglog(xdata, b4c,  label='$R_G=100$ nm', lw=4)
axrg.loglog(xdata, b4b,  label='$R_G=10$ nm', lw=4)
axrg.loglog(xdata, b4,  label='$R_G=1$ nm', lw=4)
axrg.set_ylabel('$P(q)$')
axrg.set_xlabel('$q\\,\\mathrm{[nm^{-1}]}$')
# axrg.annotate('(a)', xy=(-0.3, 1), xycoords='axes fraction')
# axexp.annotate('(b)', xy=(-0.3, 1), xycoords='axes fraction')
axrg.legend()

axexp.loglog(xdata, b4c,  label='$n_P=4$', lw=4)
axexp.loglog(xdata, b3,  label='$n_P=3$', lw=4)
axexp.loglog(xdata, b2,  label='$n_P=2$', lw=4)
axexp.set_xlabel('$q\\,\\mathrm{[nm^{-1}]}$')
axexp.legend()

plt.savefig('/home/johannes/jorepos/jowebsite/figures/saxs/beaucage.svg')
plt.show()