import numpy as np
# jolymer is my personal library
from jolymer import jocolors as joc
import matplotlib.pyplot as plt
from scipy.signal import savgol_filter

plt.style.use('../jowebstyle.mplstyle')

num=150
np.random.seed(3)
# some parameters for the plots
lw = 3  # linewidth

# set up the x and y data for the plots
# xI vs I* are the intensity traces
# Is is for small, Ib for big and Isb for s+b
xI = np.linspace(0, 10, num=num)
xg2 = np.logspace(-3, 4, num=num)
Is = 1 - np.random.rand(num)
Is = 4.5 + Is*0.3
Ib = 1 - np.random.rand(num)
Ib = 2 + Ib*2.5
# smooth the Intensity for the big particle:
Ib = savgol_filter(Ib, 20, 2)
Isb = Is - Ib
g2s = np.exp(-xg2)
g2b = np.exp(-xg2 / 1000)
g2bs = 0.5*g2s + 0.5*g2b

# now create the plots:
fig, (axtrace, axg2) = plt.subplots(ncols=2,
                                    figsize=(6, 3))

axtrace.plot(xI, Is, color=joc.Magenta)
axtrace.plot(xI, Ib, color=joc.Cyan)
axtrace.plot(xI, Isb, color=joc.GreenYellow)
axtrace.set_ylim(0.5, 5.5)
axtrace.set_xticks([])
axtrace.set_yticks([])
axtrace.set_ylabel('intensity $I(t)$')
axtrace.set_xlabel('time $t$')
axtrace.annotate('small', xy=(1, 5),
                 color=joc.Magenta)
axtrace.annotate('big', xy=(1, 3.5),
                 color=joc.Cyan)
axtrace.annotate('big + small', xy=(1, 2),
                 color=joc.GreenYellow)

axg2.plot(xg2, g2s,
          color=joc.Magenta)
axg2.plot(xg2, g2b,
          color=joc.Cyan)
axg2.plot(xg2, g2bs,
          color=joc.GreenYellow)
axg2.set_xscale('log')
axg2.set_xticks([])
axg2.set_yticks([])
axg2.set_ylabel('auto correlation $g_2(\\tau)$')
axg2.set_xlabel('delay time $\\tau$')

plt.savefig('/home/johannes/jorepos/jowebsite/figures/dls/trace-g2.svg')
plt.show()