import numpy as np
import matplotlib.pyplot as plt

# Set up figure font sizes for presentation
import matplotlib.pylab as pylab
from scipy.interpolate import interp1d

params = {'legend.fontsize': 22,
          'figure.figsize': (15, 10),
         'axes.labelsize': 22,
         'axes.titlesize':30,
          'lines.markersize':15,
         'xtick.labelsize':22,
         'ytick.labelsize':22}
pylab.rcParams.update(params)


# Set up the time grid
# We use a general procedure that also works
# when t_final is not an integer multiplum of h.
# In that case, t_final is adjusted afterwards
t_final = 10.0    # end point
h       = 0.1    # time step
t_vec   = np.arange(0, t_final, h) # create evenly spaced points
t_vec   = np.append(t_vec, t_vec[-1] + h)  # add endpoint. may be different from t_final.
n = len(t_vec) - 1  # number of intervals, or Euler steps
t_final = t_vec[-1] # adjust t_final

# Initialize arrays
# that are coputed in the Euler loop.
# Note the different lengths of these arrays.
A = np.zeros(n+1)
M_A = np.zeros(n+1)
Q = np.zeros(n)

# We fake a dataset of times t2 and measurements of
# the water input stream, P2
# We then interpolate this dataset
P2 = [2, 2, 2, 2, 1, 1, 1, 0.5, 0.4, 0.2, 0 ]
t2 = [0, 1, 2, 3, 4, 5, 6, 7,   8,   9, 10]
P2_interp = interp1d(t2, P2, kind = 'zero')
P = P2_interp(t_vec)

# Simlar to P, but with incoming concentration of
# sulphate
C_P2 = [1.2, 1.1, 1.3, 1.4, 1.2, 1, 2, 3, 4, 4, 4]
C_P2_interp = interp1d(t2, C_P2, kind = 'zero')
C_P = C_P2_interp(t_vec)


#
# Initialize Euler's method
#
A[0] = 0.001
M_A[0] = 0.0
K = 1.0
Amin = 0.5

# Euler loop.
for k in range(n):
    Q[k] = K * max(A[k] - Amin, 0.0)
    A[k+1] = A[k] + h * (P[k] - Q[k])
    M_A[k+1] = M_A[k] + h * (P[k]*C_P[k] - Q[k]*M_A[k]/A[k]  )

#Plotting of results
plt.plot(t_vec, A)
plt.plot(t_vec, P)
plt.plot(t_vec, C_P)
plt.plot(t_vec[0:n], Q)
plt.plot(t_vec, M_A/A)
plt.legend(['A','P','C_P','Q', 'C_A'])

plt.show()

# Mass balance check for water A
VS = h * np.sum(P[:-1] - Q)
HS = A[-1] - A[0]
print VS - HS