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, (axpr, axpo) = plt.subplots(ncols=2, figsize=(5,3))
axpr.set_prop_cycle(color=[joc.Magenta, joc.Blue, joc.GreenYellow])
axpo.set_prop_cycle(color=[joc.Red, joc.Orange, joc.GreenYellow])

xdat = np.linspace(0,5, num=600)
ydat16 = gupo(xdat, 1, 0.06, 1.9) * xdat**2
ydato1 = gupo(xdat, 1, 0.06, 5/3) * xdat**2
ydato2 = gupo(xdat, 1, 0.06, 2) * xdat**2
ydato3 = gupo(xdat, 1, 0.06, 3) * xdat**2
ydat2 = gupo(xdat, 1.5, 0.15, 3) * xdat**2 # part
ydat25 = gupo(xdat, 3, 1, 8) * xdat**2 # fold

axpr.plot(xdat, ydat16, label='unfolded', lw=4)
axpr.plot(xdat, ydat2,  label='partially unfolded', lw=4)
axpr.plot(xdat, ydat25, label='folded', lw=4)

axpo.plot(xdat, ydato1,  label='good solvent', lw=4)
axpo.plot(xdat, ydato2,  label='$\\Theta$ solvent', lw=4)
axpo.plot(xdat, ydato3,   label='bad solvent', lw=4)


axpr.set_ylabel('$q^2 I$')
axpo.set_ylabel('$q^2 I$')
axpr.set_xlabel('$q$')
axpo.set_xlabel('$q$')
axpr.set_yticks([])
axpr.set_xticks([])
axpo.set_yticks([])
axpo.set_xticks([])
axpo.set_ylim(bottom=0)
axpr.legend()
axpo.legend()
axpr.annotate('(a)', xy=(-0.1, 1), xycoords='axes fraction')
axpo.annotate('(b)', xy=(-0.1, 1), xycoords='axes fraction')

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