import pymc3 as pm 
import theano.tensor as tt 

with pm.Model() as model: 

  U_wall = pm.Uniform(name="U_wall", lower=0.1, upper=0.3)
  U_out= pm.Uniform(name="U_out", lower=0.01, upper=0.1)
  C_in = pm.Uniform(name="C_in", lower=2, upper=5)
  C_in=C_in*1e3
  C_wall = pm.Uniform(name="C_wall", lower=2, upper= 5)
  C_wall=C_wall*1e4 
  sd=pm.Normal(name="sd",mu=0.1, sigma=0.05)
 
  y_mean=tt.zeros((len(y_observed),2))
  y_mean = tt.inc_subtensor(y_mean[0], np.array([15,12]))

  # Define f as tensor 
  def f_rk(y,t_now):
    step=int(t_now/dt)
    f=tt.zeros(2)
    f=tt.inc_subtensor(f[0], U_wall/C_in*(y[1]-y[0]))
    if step % 2 ==0:
      f=tt.inc_subtensor(f[1], U_wall/C_wall*(y[0]-y[1])+U_out/C_wall*(y_out[step,0]-y[1]))
    else: 
      f=tt.inc_subtensor(f[1], U_wall/C_wall*(y[0]-y[1])+U_out/C_wall*( (y_out[int(step+0.5),0]+ y_out[int(step-0.5),0])/2-y[1]))
    return f

  # compute y_sol    
  for i in range(len(y_observed)-1):
    h = t[i+1] - t[i]
    k1 = f_rk(y_mean[i],t[i])
    k2 = f_rk(y_mean[i] + k1 * h / 2, t[i] + h / 2)
    k3 = f_rk(y_mean[i] + k2 * h / 2, t[i] + h / 2)
    k4 = f_rk(y_mean[i] + k3 * h, t[i] + h )
    y_mean = tt.inc_subtensor(y_mean[i+1], y_mean[i] + (h / 6.) * (k1 + 2*k2 + 2*k3 + k4))
  
  y_temperature=pm.Normal(name="y_temperature", mu=y_mean,sigma=sd, observed=y_observed)