#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
@author: jmw418
"""
###Packages###
import numpy as np
import matplotlib.pyplot as plt
import scipy.linalg
import scipy.special

### Initial condition function###
def jump(x, alpha=0.1, beta=0.3):
    """   
SUMMARY: 
     This is a function that produces a square wave height 1 in the range between
     alpha and beta it is 0 elsewhere
     
FLEXIBILITY:
     you can vary alpha and beta
     
INPUT:
    x     = the input of the jump function
    alpha = the left constant
    beta  = the right constant
    
OUTPUT: 
    1 if x belongs between alpha and beta and 0 else """
    one = lambda x: 1
    return np.where((x<beta) & (x>=alpha), one(x), 0.)

### Initial condition function###
def cosBell(x, alpha=0.1, beta=0.3):
    """   
SUMMARY: 
     This is a function that produces a cos wave height 1 in the range between
     alpha and beta it is 0 elsewhere
     
FLEXIBILITY:
     you can vary alpha and beta
     
INPUT:
    x     = the input of the jump function
    alpha = the left constant
    beta  = the right constant
    
OUTPUT: 
    0.5*(1 - np.cos(2*np.pi*(x-alpha)/width)) if x belongs between alpha 
    and beta and 0 else """
    
    width = beta - alpha
    bell = lambda x: 0.5*(1 - np.cos(2*np.pi*(x-alpha)/width))
    return np.where((x<beta) & (x>=alpha), bell(x), 0.)


###Timestepping methods###
def picard(X, nx, nt, Tfinal, mu, theta, form):
    """
SUMMARY:
    This is a timestepping algorithm that solves the burgers equation \
    when it is discretised by the picard method.

FLEXIBILITY:
    The flexibility is not in this function, it is in the SolveBurger\
    this function is called upon by the solve burger and shouldnt be ran
    
INPUT: X generated by the initial condition
    
OUTPUT: X generated by the timestepping method, in this case picards method
    """
    ### Parameters ###
    dx = (1-0)/(nx-1) 
    dt = (Tfinal-0)/(nt-1)
    C = 1*dt/(2*dx)
    D = mu*dt/(dx**2)
    
    ### Initialise structure ###
    beta = np.zeros([nx])
    A = np.zeros([nx,nx])
    
    if form == "non conservative":   ### solve the conservative form
        if theta > 0:                ### we need to invert a matrix 
            for i in range(0,nt-1):  ### time stepping###
                                     
                                     ### X[] this is now the old timestep###
                                     
                for j in range(0,nx):### create vector on RHS
                    beta[j] = X[j]\
                    - C*(1-theta)*X[j]*(X[(j+1)%nx] - X[(j-1)%nx])\
                    + (D)*(1-theta)*(X[(j-1)%nx]-2*X[j]+X[(j+1)%nx])
        
                for p in range(0,nx): ### Create a matrix on LHS
                    A[p,(p+1)%nx] =  theta*C*X[p] - theta*D  
                    A[p,(p-1)%nx] = -theta*C*X[p] - theta*D 
                    A[p,p] = 1 + 2*theta*D
            
                X[:] = scipy.linalg.solve(A, beta)### Solving for next timestep 
                                                  ### 
                
        if theta == 0:  ### to avoid inverting a identity matrix pointlessly.
                        ### we consider the theta = 0 case seperately to increase speed
                        ### we could remove the below and the code would still work
                        
            for i in range(0,nt-1):###time stepping###
                for j in range(0,nx): ### create vector on RHS###
                    beta[j] = X[j]\
                    - C*X[j]*(X[(j+1)%nx] - X[(j-1)%nx])\
                    + D*(X[(j-1)%nx]-2*X[j]+X[(j+1)%nx])
    
                X[:] = beta ### update to the next timestep ###

    if form == "conservative":
        if theta >0:
            for i in range(0,nt-1):   ###    time stepping     ###
                for j in range(0,nx): ### create vector on RHS###
                    beta[j] = X[j]\
                    - C*0.5*(1)*(X[(j+1)%nx]**2 - X[(j-1)%nx]**2)\
                    + (D)*(1)*(X[(j-1)%nx]-2*X[j]+X[(j+1)%nx])
           
                for p in range(0,nx): ### Create Matrix on LHS ###
                    A[p,(p+1)%nx] =  0.5*theta*C*X[p] - theta*D # b 
                    A[p,(p-1)%nx] = -0.5*theta*C*X[p] - theta*D  #c  
                    A[p,p] = 1 + 2*theta*D + C*theta*0.5*(X[(p+1)%nx] - X[(p-1)%nx])
        
                X[:] = scipy.linalg.solve(A, beta)### Solving for next timestep ###
        ###
        if theta == 0: ## to avoid inverting a identity matrix pointlessly.
            for i in range(0,nt-1):###time stepping###
                for j in range(0,nx): ### create vector on RHS###
                    beta[j] = X[j]\
                    - C*0.5*(X[(j+1)%nx]**2 - X[(j-1)%nx]**2)\
                    + D*(X[(j-1)%nx]-2*X[j]+X[(j+1)%nx])
    
                X[:] = beta ### Solving for next timestep ###        
 
    return X

   
def newton(X, nx, nt, Tfinal, mu, theta, form):
    """
SUMMARY:
    This is a timestepping algorithm that solves the burgers equation \
    when it is discretised by the picard method.

FLEXIBILITY:
    The flexibility is not in this function, it is in the SolveBurger, \
    this function is called upon by the solve burger and shouldnt be ran
    
INPUT: X generated by the initial condition
    
OUTPUT: X generated by the timestepping method, in this case newtons method
    """
    ### derived parameters ###
    dx = (1-0)/(nx-1) 
    dt = (Tfinal-0)/(nt-1)
    C = 1*dt/(2*dx)
    D = mu*dt/(dx**2)
    
    ##Creation of newton method structure 
    beta = np.zeros([nx])
    w = np.zeros([nx])
    dw = np.zeros([nx])
    
    if form == "non conservative":
        for i in range(0,nt-1):### time loop ###
                     ### first initialise w^0 = phi^n as initial guess###
            w[:]=X[:]### this is the starting guess for newton method ###
        
            ### Construct the newton loop
            tol = 10**(-13)### extreeme accuracy can be imposed
            err  = 2*tol ### set error larger than the tolerance
            while (err>tol): ### until convergence  newton loop###
            
                ### Create beta^k_j ###
                for q in range(0,nx): 
                    beta[q] = X[q] - C*(1-theta)*X[q]*(X[(q+1)%nx] - X[(q-1)%nx]) \
                    + (D)*(1-theta)*(X[(q-1)%nx]-2*X[q]+X[(q+1)%nx]) \
                    -w[q] - C*theta*w[q]*(w[(q+1)%nx]-w[(q-1)%nx]) + D*theta*(w[(q+1)%nx]-2*w[q] +w[(q-1)%nx])
            
                ##Create A^k_j ###
                A = np.zeros([nx,nx])
                for p in range(0,nx):
                    A[p,(p+1)%nx] = ((theta*C*w[p])-(theta*D)) # b
                    A[p,(p-1)%nx] = ((-theta*C*w[(p+1)%nx])-(theta*D)) # c
                    A[p,p] = (1+ 2*theta*D + C*theta*(w[(p+1)%nx]-w[(p-1)%nx]))
            
            
                ## solving for dw
                dw = scipy.linalg.solve(A, beta)
                err = np.linalg.norm(dw,2)
                ## improve the newton loop
                w = w + dw 
                ## we have w^k
                if err > 10**6:
                    print("error = %a"%err)
                    
            X[:] = w[:]
            
    if form == "conservative":
        for i in range(0,nt-1):### time loop 
            ## first initialise w^0 = phi^n
            w[:]=X[:]

            ### Construct the newton loop
            tol = 10**(-13) ### accuracy can be imposed ###
            err  = 2*tol### set error larger than the tolerance
            while (err>tol): ## eventually replace with while loop and stopping criterion

                ### Create beta^k ###
                for q in range(0,nx): 
                    beta[q] = X[q] - C*0.5*(1-theta)*(X[(q+1)%nx]**2 - X[(q-1)%nx]**2) \
                    + (D)*(1-theta)*(X[(q-1)%nx]-2*X[q]+X[(q+1)%nx]) \
                    -w[q] - C*0.5*theta*(w[(q+1)%nx]**2-w[(q-1)%nx]**2) \
                    + D*theta*(w[(q+1)%nx]-2*w[q] +w[(q-1)%nx])
            
                ### Create A^k_j ###
                A = np.zeros([nx,nx])
                for p in range(0,nx):
                    A[p,(p+1)%nx] = ((theta*C*w[(p+1)%nx])-(theta*D)) 
                    A[p,(p-1)%nx] = ((-theta*C*w[(p-1)%nx])-(theta*D)) 
                    A[p,p] = (1+ 2*theta*D + C*theta*(w[(p+1)%nx]-w[(p-1)%nx]))
        
                ## solving for dw
                dw = scipy.linalg.solve(A, beta)
                err = np.linalg.norm(dw,2)
                ## improve the newton loop
                w = w + dw 
                ## we have w^k
                if err > 10**6:
                    print("error = %a"%err)
                
            X[:] = w[:]

    return X 

def SolveBurger(nx ,nt ,Tfinal ,mu, method, theta, form, IC):
    """
SUMMARY:
    This function takes an "initial_condition" and solves
    the "conservative" or "non conservative" burgers equation 
    numerically for some later time "Tfinal".
                    
FLEXIBILITY:
    You can choose:  the spatial discretisation used "nx", the number of timesteps "nt",
    the viscosity term "mu", the "method" used, the implicitness of the scheme
    "theta", the initial conditions "IC"
      
INPUT:      
    nx      = number of space points
    nt      = number of time points 
    Tfinal  = the final time  
    mu      = the viscosity in the burgers equation 
    method  = "picard" SOLVES APPROXIMATION TO BURGERS EQUATION
              "newton" SOLVES EXACT BURGERS EQUATION
    theta   = 1 is implicit,
              0 is explicit,
              You can put in other real values, 
    form    = "conservative"     -->    u_t + (u^2/2)_x - mu u_{xx} = 0 
              "non conservative" -->    u_t + u(u_x) - mu u_{xx} = 0 
    IC      = "jump"             -->    a square wave initial condition defined by a function
            = "cosBell"          -->    a cos wave initial condition defined by a function
Output:
    A vector representing the numerical solution 
                """
    
    ##**Creating the structure**## 
    X = np.zeros([nx])
    x = np.linspace(0,1,nx)
    ### Initialisation of Initial conditions, as a d vector###
    if IC == "jump":
        X[:] = jump(x, 0.1, 0.3)
    if IC == "cosBell":
        X[:] = cosBell(x,0.1,0.3)
    
    ### Timestepping methods ###
    if method == "newton":
        newton(X, nx, nt, Tfinal, mu, theta, form)
    if method == "picard":
        picard(X, nx, nt, Tfinal, mu, theta, form)
    return X


def __main__():
###main is the function that runs some scripts, 
###the majority of the code is actually above this;
###this is just some examples of running the code
    ###-----  -Tests-  -----###
    
    ###Test0: high accuracy solutions, these results are interesting, 
    ###however take too long to compute longer 2minuits###
#    nx = 401
#    nt = 1001
#    x = np.linspace(0,1,nx) 
#    plt.figure(figsize= (9,5)) 
#    plt.plot(x, X, 'r',linewidth=0.5,label = "t = 1, ftcs")
#    X = SolveBurger(nx,nt,1,0.001,"picard",0,"non conservative","jump")
#    nx = 401
#    x = np.linspace(0,1,nx) 
#   nt = 1001
#    X = SolveBurger(nx,nt,1,0.001,"newton",1,"non conservative","jump")
#    plt.plot(x, X, color = 'b' ,linewidth=0.5,label = "newt")
#    nx = 401
#    nt = 1001
#    x = np.linspace(0,1,nx) 
#    plt.plot(x, X, 'k',linewidth=0.5,label = "t = 1, ftcs-cons")
#    X = SolveBurger(nx,nt,1,0.001,"picard",0,"conservative","jump")
#    nx = 401
#    nt = 1001
#    x = np.linspace(0,1,nx) 
#    X = SolveBurger(nx,nt,1,0.001,"newton",1,"conservative","jump")
#    plt.plot(x, X, color = 'g' ,linewidth=0.5,label = "newt-cons")
#    plt.ylabel('$\phi$')
#    plt.xlabel('x')
#    plt.legend()
#    plt.show()
    
    ### The solutions of most methods resolve the jump at a spatial discretisation level of 
    ### 200, so we will pay close attention to the case when nx = 100 where we see which methods
    ### deal with the shocks
    
    
    
    


    
    ### Test1: Problems and limitations with ftcs method ###
    ###aim: demonstrate naive scheme
    
        
        ### advection diffusion implies problems on the time step:
        #1)nt has to be large if nx is very large and mu is big
        #2)nt has to be large if mu is small and u is large fortunately u travels as a wave.
        #3)for the burges equation we get problems if the solution grows in time,
        
        
    nx = 1001 
    nt = 2001
    x = np.linspace(0,1,nx) 
    plt.figure(1,figsize= (9,5)) 
    EA = SolveBurger(nx,nt,1,0.001,"picard",0,"conservative","jump") ###this is ftcs on conservative burgers mu = 0.001
    plt.plot(x, EA, color = 'r' ,linewidth=0.5 ,label = "mu = 0.001, nx = %a ,nt = %a " %(nx,nt)) ###plot solution
    plt.ylabel('$\phi$')
    plt.xlabel('x')
    plt.legend()
    plt.title("ftcs-conservative, extenuating circumatances parameters ")

    print("ftcs-conservative, where in extenuating circumatances parameters are chosen to make methodology valueable ")

    nx = 101
    nt = 20001
    x = np.linspace(0,1,nx) 
    plt.figure(2,figsize= (9,5)) 
    X = SolveBurger(nx,nt,1,1,"picard",0,"conservative","jump")
    plt.plot(x, X, color = 'r' ,linewidth=0.5 ,label = "mu = 1, nx = %a ,nt = %a " %(nx,nt))
    plt.ylabel('$\phi$')
    plt.xlabel('x')
    plt.legend()
    plt.title("ftcs-conservative failure of convergence")

    print("instability even with 20000 timesteps, for large mu we have problems")    
    


    nx = 101
    nt = 20001
    x = np.linspace(0,1,nx) 
    plt.figure(3,figsize= (9,5)) 
    X = SolveBurger(nx,nt,1,0.000001,"picard",0,"conservative","jump")
    plt.plot(x, X, color = 'r' ,linewidth=0.5 ,label = "mu = 0.0001, nx = %a ,nt = %a " %(nx,nt))
    plt.ylabel('$\phi$')
    plt.xlabel('x')
    plt.legend()
    plt.title("ftcs-conservative: not resolving the shock")

    print("nx needs to be large, to deal with the shock, and as mu tends to zero \
          the shock becomes harsh. The timestep needs increasing greatly \
          hence the method becomes absolutely unuseable as mu tends to zero")   
    
    
    
    
    ###----- Test3: Demonstration of limitation of picards method -----###
    ### we suppose that we have had some space discretisation limitation nx =101
    ### we have theta = 1 and now the method is unconditionally stable ###

    nx = 101
    x = np.linspace(0,1,nx) 
    plt.figure(4,figsize= (9,5))
    for i in range(0,4):
        nt = int(10**(i+1) + 1)
        X = SolveBurger(nx,nt,1,0.001,"picard",0.5,"non conservative","jump")
        plt.plot(x, X, color = (0,0,0.2+0.2*i),linewidth=0.5,label = " nt = %a " %(nt))
        plt.ylabel('$\phi$')
        plt.xlabel('x')
        plt.legend()
        plt.title("picard, $theta$ = 0.5")

    print("Conclusion picards method will always converge, however the equation\
          you are solving is only accurate for small timestepping, and the \
          shock isnt resolved for small nx")
    
    
    
###----- Section: small spacial discretisation schemes
###-----Test: Convergence results, we consider how spatial resolution changes the convergence of 4 methods

    
### Test3: forward time centered difference method as spatial discretisation increases ###
    #a) for "non conservative" form
    ntf = np.zeros([4])
    ntf[:] = [1001,1001,1001,1001] ### we test with about 1000 timepoints, to ensure solution converges
    plt.figure(5,figsize= (9,5))
    for j in range(0,4):
        nx = 25*(2**j) + 1  ### we change the spatial discretisation as about 25,50,100,200
        nt = int(ntf[j])
        x = np.linspace(0,1,nx) 
        X = SolveBurger(nx,nt,1,0.001,"picard",0,"non conservative","jump")
        plt.plot(x, X, color = (1-0.2*j,0.2+0.2*j,0.0) ,linewidth=0.5,label = "nx = %a ,nt = %a " %(nx,nt))
        plt.ylabel('$\phi$')
        plt.xlabel('x')
        plt.legend()
        plt.title("ftcs")
    #b) for "conservative" form
    ntf = np.zeros([4])
    ntf[:] = [1001,1001,1001,1001] 
    plt.figure(6,figsize= (9,5))
    for j in range(0,4):
        nx = 25*(2**j) + 1
        nt = int(ntf[j])
        x = np.linspace(0,1,nx) 
        X = SolveBurger(nx,nt,1,0.001,"picard",0,"conservative","jump")
        plt.plot(x, X, color = (1-0.2*j,0.2+0.2*j,0.0) ,linewidth=0.5,label = "nx = %a ,nt = %a " %(nx,nt))
        plt.ylabel('$\phi$')
        plt.xlabel('x')
        plt.legend()
        plt.title("ftcs-conservative")
    #c) Timestep restrictions with newton method advective form 
    ntf = np.zeros([4])
    ntf[:] = [25,42,70,74] ### I have found the smallest nt available ie largest timesteps possible
    plt.figure(7,figsize= (9,5))
    for j in range(0,4):
        nx = 25*(2**j) + 1
        nt = int(ntf[j])
        x = np.linspace(0,1,nx) 
        X = SolveBurger(nx,nt,1,0.001,"newton",1,"non conservative","jump")
        plt.plot(x, X, color = (1-0.2*j,0.2,1-0.1*j) ,linewidth=0.5,label = "nx = %a ,nt = %a " %(nx,nt))
        plt.ylabel('$\phi$')
        plt.xlabel('x')
        plt.legend()
        plt.title("Fully Implicit Newton Method on Advective Form of Burgers")
    #d)Timestep restrictions with newton method conservative form 
    ntf = np.zeros([4])
    ntf[:] = [34,70,127,237] ### I have found the smallest nt available ie largest timesteps possible
    plt.figure(8,figsize= (9,5))
    for j in range(0,4):
        nx = 25*(2**j) + 1
        nt = int(ntf[j])
        x = np.linspace(0,1,nx) 
        X = SolveBurger(nx,nt,1,0.001,"newton",1,"conservative","jump")
        plt.plot(x, X, color = (1-0.2*j,0.2,1-0.1*j) ,linewidth=0.5,label = "nx = %a ,nt = %a " %(nx,nt))
        plt.ylabel('$\phi$')
        plt.xlabel('x')
        plt.legend()
        plt.title("Fully Implicit Newton Method on Conservative Form of Burgers")

    
    ###----- Test: nx = 101 comparisons of all different methods as mu varies -----####
    ### test for 1 unit time later after the jump initial condition, using implicit newton and theta 
    nx = 101
    nt = 201
    x = np.linspace(0,1,nx) 
    plt.figure(9,figsize= (9,5))
    X = SolveBurger(nx,nt,1,0.001,"newton",1,"conservative","jump")
    plt.plot(x, X, color = 'k' ,linewidth=0.5,label = "newton conservative" )
    X = SolveBurger(nx,nt,1,0.001,"newton",1,"non conservative","jump")
    plt.plot(x, X, color = 'g' ,linewidth=0.5,label = "newton non conservative" )
    X = SolveBurger(nx,nt,1,0.001,"picard",0,"conservative","jump")
    plt.plot(x, X, color = 'b' ,linewidth=0.5,label = "ftcs conservative" )
    X = SolveBurger(nx,nt,1,0.001,"picard",0,"non conservative","jump")
    plt.plot(x, X, color = 'r' ,linewidth=0.5,label = "ftcs non conservative" )

    plt.ylabel('$\phi$')
    plt.xlabel('x')
    plt.legend()
    plt.title("nx =101, nt = 501")
    plt.show()
    
    

    

    
    
    
    

__main__()