#%% - Imports
import numpy as np
from numpy import *
import matplotlib.pylab as plt
from fipy import *
import pandas as pd

#%% - Constants
rho_b_MH = 589. # kg m^-3
Cp_b = 181.170 #J (mol*K)^-1, Thermal conductivity and specific heat measurements of metal hydrides

# I'm using L in place of lambda
L_eff = 8.4 # W m^-1 K^-1

E_c = 45.*1000 # J/mol
Beta = 13.5 #%
P = 100. #atm
M = 2.02 #g/mol
dH = 28.*1000. #J/mol

h = 150. # W/m^-2 K^-1
T_c = 60.+273.15 #K, Coolant Temperature
To = 60+273.15 # K, Bed Temperature
R = 8.314 # J/mol/K

Peq = ((E_c/dH)/(1+E_c/dH))*P #atm
A = 1.1319e25

#%% - Mesh

L = 1.
nx = 10.
dx = L/nx

# Create the mesh
mesh = Grid1D(nx=nx,dx=dx)
#mesh = CylindricalGrid1D(nx=nx,dx=dx)

##### Sorbtion Setup #####
X = CellVariable(name = "Sorbtion",mesh = mesh, value=0.1, hasOld=True)
CC = mesh.cellCenters[0]

#### PDE Variable Setup #####
#T = CellVariable(name = "Temperature", mesh = mesh,value = 60.+273.15)
T = CellVariable(name = "Temperature", mesh = mesh, value = 60.+273.15, hasOld=True)

# I want to use the robin boundary condition but in either case, the constraint doesn't seem to be working
#T.constrain(60.+273.15,mesh.facesLeft)
constraint_value = FaceVariable(mesh=mesh)
T.faceGrad.constrain(constraint_value,mesh.facesLeft)


# I want to use the robin boundary condition
#T.constrain(110.+273,mesh.facesRight)
T.faceGrad.constrain(0.,mesh.facesRight)
import ipdb; ipdb.set_trace()
k_c = A*T**-9.73625
rc = k_c*(1-X)*(P-Peq)

param = k_c * (P - Peq)
maxX = 0.9
largeValue = 1e9
constraint = largeValue * (X > maxX)

big_value = 1e9

eqX = TransientTerm(var = X) == param - ImplicitSourceTerm(param, var=X) + maxX * constraint - ImplicitSourceTerm(constraint, var=X)
eqT = TransientTerm(rho_b_MH * Cp_b, var=T) == DiffusionTerm(coeff = L_eff, var=T) + rho_b_MH*(dH/M)*rc
eq = eqT & eqX

#%% - Equation Solution
TSD = .005  # This seems to be a stable time step for this function
sweeps = 150 # identify the number of sweeps to run, running 150 for now
ts = np.linspace(0,TSD*sweeps,sweeps+1)  # ts is the number of time slots plotted for the ODE. create an array to plot for the ode

solver = LinearLUSolver() # Trying a different solver

#Prepare an array for storing calculated values of X
Xsolution = np.zeros((sweeps+1,int(nx))) #Create a matrix a row for initial value + 1 row for each sweep. There will be nx columns per row.
#Store the initial value of X at all points in the first row of sweeps
Xsolution[0,:] = X.value

#Prepare an array for storing calculated values of T
Tsolution = np.zeros((sweeps+1,int(nx))) #Create a matrix a row for initial value + 1 row for each sweep. There will be nx columns per row.
#Store the initial value of T at all points in the first row of sweeps
# This matrix has the form with the rows are T|t=0...T|t=t_max, and columns T|r=0....T|r=R
Tsolution[0,:] = T.value

viewer  = Viewer(vars = T, datamin = 330., datamax = 450., title = "Temperature")
viewer2 = Viewer(vars = X, datamin = 0., datamax = 1.05, title = "X")

for sweep in range(sweeps):
    constraint_value[:] = h * ((60. + 273.15) - T.faceValue)
    X.updateOld()
    T.updateOld()
    eq.sweep(dt=TSD,solver=solver)
    # eqX.sweep(X, dt=TSD, solver=solver)
    # eqT.sweep(T, dt=TSD, solver=solver)
    viewer.plot()
    viewer2.plot()
    Xsolution[sweep+1] = X.value #Store the recently computed X value in the current sweep
    Tsolution[sweep+1] = T.value #Store the recently computed T value in the current sweep
    #print("X values",X.value[-1])
    #print("T values",T.value[-1])
    #print(rc[-1])
    df=pd.DataFrame([X.value[-1],T.value[-1],k_c[-1],rc[-1]],index = ["X","T","k_c","rc"],columns=[""])
    print(df.T)

#print(Xsolution[:,0])
#plt.plot(CC.value,X.value,'ro') # plot the points where calculations are made in the fipy solver.
fig,ax = plt.subplots()
ax.plot(ts,Xsolution[:,0],'bo')
ax.set_xlabel("time")
ax.set_ylabel(X.name)
ax.set_title("ODE Solution")

fig,ax = plt.subplots()
ax.plot(ts,Tsolution[:,0],"bo")
ax.set_xlabel("time")
ax.set_ylabel(T.name)
ax.set_title("Time vs. Temp")

"""
fig,ax = plt.subplots()
ax.plot(ts,Tsolution[-1,:],"bo")
ax.set_xlabel("time")
ax.set_ylabel(T.name)
ax.set_title("Time vs. Temp")
"""