from numpy import zeros,linspace
from simple_PDE import *


def _test_constant_solution():
    """
    Test problem where u=u_const is the exact solution, to be
    reproduced (to machine precision) by any relevant method.
    """
    def exact_solution(t):
        return C_s,C_i,C_r
    
    def lam(t,x):
        return C_s

    def beta(t,x):
        return (C_s*C_i)/float(C_r)

    #Constant values
    C_s = 1.2
    C_i = 0.8
    C_r = 0.6
    
    #lam = C_s
    #beta = (lam*C_i)/float(C_r)
    
    T = 2; Nt = 200
    X = 20; Nx = 40
    S_1 = np.ones(Nx+3)*C_s
    I_1 = np.ones(Nx+3)*C_i
    R_1 = np.ones(Nx+3)*C_r
    
    t,x,S,I,R = simple_PDE(T,Nx,Nt,X,lam,beta,S_1,I_1,R_1)
    
    S_e,I_e,R_e = exact_solution(t)
    difference = abs(S_e - S).max()  # max deviation
    tol = 1E-14
    assert difference < tol

    difference = abs(I_e - I).max()  # max deviation
    tol = 1E-14
    assert difference < tol

    difference = abs(R_e - R).max()  # max deviation
    tol = 1E-14
    assert difference < tol

def _test_manufactured_solution(T,Nt,X,Nx):
    
    def exact_solution_S(t,x):
        return np.cos(np.pi*x)*t

    def exact_solution_I(t,x):
        return np.cos(np.pi*x)*t

    def exact_solution_R(t,x):
        return np.cos(np.pi*x)*t


    def beta(t,x):
        return exact_solution_S(t,x)*exact_solution_I(t,x)/exact_solution_R(t,x)
   
    def f(t,x):
        return (t**2*np.cos(np.pi*x) + (np.pi/float(X))**2*t + 1)*np.cos(np.pi*x) 

    def g(t,x):
        return (lam*t - t**2*np.cos(np.pi*x) + (np.pi/float(X))**2*t + 1)*np.cos(np.pi*x)

    def h(t,x):
        return (-lam*t + (np.pi/float(X))**2*t + 1)*np.cos(np.pi*x)
        

    dx = X/float(Nx)
    dt = T/float(Nt)
    S_1 = exact_solution_S(0,np.linspace(0-dx,X+dx,Nx+3))
    I_1 = exact_solution_I(0,np.linspace(0-dx,X+dx,Nx+3))
    R_1 = exact_solution_R(0,np.linspace(0-dx,X+dx,Nx+3))
     
    t,x,S,I,R = simple_PDE(T,Nx,Nt,X,lam,beta,S_1,I_1,R_1,f,g,h)
    S_e = exact_solution_S(t[-1],x)
    I_e = exact_solution_I(t[-1],x)
    R_e = exact_solution_R(t[-1],x)
    
    difference_S = abs(S_e - S).max()  # max deviation
    difference_I = abs(I_e - I).max()  # max deviation
    difference_R = abs(R_e - R).max()  # max deviation

    return difference_S,difference_I,difference_R
     
def solver(I, a, b, T, dt, theta):
    """
    Solve u'=-a(t)*u + b(t), u(0)=I,
    for t in (0,T] with steps of dt.
    a and b are Python functions of t.
    """
    dt = float(dt)            # avoid integer division
    Nt = int(round(T/dt))     # no of time intervals
    T = Nt*dt                 # adjust T to fit time step dt
    u = zeros(Nt+1)           # array of u[n] values
    t = linspace(0, T, Nt+1)  # time mesh

    u[0] = I                  # assign initial condition
    for n in range(0, Nt):    # n=0,1,...,Nt-1
        u[n+1] = ((1 - dt*(1-theta)*a(t[n]))*u[n] +\
                  dt*(theta*b(t[n+1]) + (1-theta)*b(t[n])))/\
                  (1 + dt*theta*a(t[n+1]))
    return u, t

#test_constant_solution()
a = 15
T = 1; Nt = [200,200*(2**2),200*(4**2),200*(8**2)]
X = 1; Nx = [10,10*2,10*4,10*8]
#Nt = [400,400*(2**2),400*(3**2),400*(4**2)]
#Nt = [800,800*(2**2),800*(3**2),800*(4**2)]
#Nt = [1600,1600*(2**2),1600*(3**2),1600*(4**2)]
for i in range(4):
    d_s,d_i,d_r = test_manufactured_solution(T,Nt[i],X,Nx[0])
    if i > 0:
        #print "-----------------"
        #print "dt", T/float(Nt[i])     
        #print "dt_1", T/float(Nt[i-1])     
        #print "d_s",d_s
        #print "d_s_1",d_s_1
        print "r",np.log(d_s_1/d_s)/np.log((T/float(Nt[i-1]))/(T/float(Nt[i])))
    d_s_1 = d_s

