import matplotlib.pyplot as plt
import numpy as np
import skued
from scipy.optimize import curve_fit

# Load data from file first
# 2 x N array, first row is time-delay, second row is diffracted intensity
block = np.load('data/tseries1.npy')
time, intensity = block[0, :], block[1, :]

# Compute initial guesses for this curve
initial_guesses = (0,                                   # time-zero
                   intensity.max() - intensity.min(),   # amplitude
                   1,                                   # time-constant
                   intensity.min())                     # offset

params, pcov = curve_fit(skued.exponential, time, intensity,
                                                     p0 = initial_guesses)

tzero, amplitude, tconst, offset = params
best_fit_curve = skued.exponential(time, *params)
# Equivalent:
#   best_fit_curve = skued.exponential(time, tzero, amplitude, tconst, offset)

fig = plt.figure()
ax = fig.add_subplot(111)
ax.axvline(x = tzero, linestyle = 'dashed', color = 'grey', label = 'Time-zero')
ax.plot(time, intensity, '.k', label = 'Time-series')
ax.plot(time, best_fit_curve, '-r', label = 'Best fit')
ax.set_xlabel('Time-delay [ps]')
ax.set_ylabel('Diffracted Intensity [a.u.]')
ax.set_xlim([time.min(), 40])
ax.legend()
plt.show()