#!/usr/bin/python3.5

#if __name__ == "__main__":

import sys
import scipy
import numpy as np
import matplotlib
matplotlib.use('TkAgg') 
import matplotlib.pyplot as plt
import pylab
import time as tim

    #grid and array
nx = 100
u  = np.zeros (nx)
total_length = 1
width = 0.1
#width = 0.01
dx = total_length / float (nx)
x  = np.arange (0, total_length, dx)
kappa = 1e5
diffeff = 1.0 / (3.0 * kappa)

tmax = 1.0
time = 0.0

#time step
cs = np.sqrt(1.0/3.0)
dt = dx / cs * 0.4

print ('Time step limit for two-moment system:',dx/cs)
print ('Time step limit for diffusion equation:',0.5*dx**2/diffeff)

    # initial conditions
j_rad = np.exp (- (x-0.5) ** 2 / (2 * width **2))
h_rad = np.zeros (len(x))

j0 = j_rad
u_exact = j_rad

plt.plot (x,j_rad)
plt.ylim ( (-2,2) )
plt.xlabel ('x')
plt.ylabel ('J,H')
plt.draw()
#plt.show()
plt.show(block=False)
plt.ioff() 
    
while (time < tmax):


    print("time:",time)
    j_left = np.roll (j_rad,1)
    h_left = np.roll (h_rad,1)
    j_rght = np.roll (j_rad,-1)
    h_rght = np.roll (h_rad,-1)

# Boundary conditions: no incoming flux
    j_left [ 0] = j_rad [ 0]
    j_rght [-1] = j_rad [-1]
    h_left [ 0] = max ([h_rad [ 0], 0.0])
    h_rght [-1] = min ([h_rad [-1], 0.0])

# slope limiter (MC)
    djdxl = (j_rad - j_left) / dx
    djdxr = (j_rght - j_rad) / dx
    djdxav = 0.5 * (djdxl + djdxr)
    ii = np.where (djdxl*djdxr <=0)
    djdx = np.minimum(np.abs(djdxl), np.abs(djdxr))
    djdx = np.minimum(2*djdx, np.abs(djdxav))
    djdx = djdx * np.sign(djdxl)
    djdx [ii] = 0.0

    dhdxl = (h_rad - h_left) / dx
    dhdxr = (h_rght - h_rad) / dx
    dhdxav = 0.5 * (dhdxl + dhdxr)
    ii = np.where (dhdxl*dhdxr <=0)
    dhdx = np.minimum(np.abs(dhdxl), np.abs(dhdxr))
    dhdx = np.minimum(2*dhdx, np.abs(dhdxav))
    dhdx = dhdx * np.sign(dhdxl)
    dhdx [ii] = 0.0

    j_lp = j_rad - 0.5 * djdx * dx
    h_lp = h_rad - 0.5 * dhdx * dx
    j_rm = j_rad + 0.5 * djdx * dx
    h_rm = h_rad + 0.5 * dhdx * dx
    j_lm = np.roll (j_rm, 1)
    h_lm = np.roll (h_rm, 1)
    j_rp = np.roll (j_lp, -1)
    h_rp = np.roll (h_lp, -1)

# Boundary conditions: no incoming flux
    j_lm [0] = j_rad [0]
    h_lm [0] = min ([h_rad [0],0.0])
    j_rp [0] = j_rad [-1]
    h_rp [0] = max ([h_rad [-1],0.0])

    f_j_l = 0.5 * (h_lm + h_lp - cs * (j_lp - j_lm))
    f_h_l = 0.5 * (j_lm/3.0 + j_lp/3.0 - cs * (h_lp - h_lm)) 

    f_j_r = 0.5 * (h_rm + h_rp - cs * (j_rp - j_rm))
    f_h_r = 0.5 * (j_rm/3.0 + j_rp/3.0 - cs * (h_rp - h_rm)) 

# Update: hyperbolic part
    j_rad = j_rad + dt / dx * (f_j_l - f_j_r)
    h_rad = h_rad + dt / dx * (f_h_l - f_h_r)

# Update: momentum source term (implicit)
    h_rad = h_rad / (1.0 + kappa * dt)

    time = time + dt

    # analytic solution
    j_exact = np.exp (- (x-0.5) ** 2 / (2 * (width **2 + 2*time*diffeff))) * np.sqrt(width**2/(width**2+2*time*diffeff))

    plt.clf()
    plt.ylim ( (-2,2) )
    plt.xlabel ('x')
    plt.ylabel ('J,H')
    c0, =plt.plot(x,j0,label='initial condition (J)')
    c1, =plt.plot(x,j_rad,label='numerical solution (J)')
    c2, =plt.plot(x,j_exact,ls='dashed',label='diffusion approximation')
    c3, =plt.plot(x,h_rad,label='numerical solution (H)')
    plt.legend(handles=[c0,c1,c2,c3])
    plt.draw()
    plt.show(block = False)
    plt.pause(0.05)

plt.ion()
plt.clf()
plt.ylim ( (-2,2) )
plt.xlabel ('x')
plt.ylabel ('J,H')
c0, =plt.plot(x,j0,label='initial condition (J)')
c1, =plt.plot(x,j_rad,label='numerical solution (J)')
c2, =plt.plot(x,j_exact,ls='dashed',label='diffusion approximation')
c3, =plt.plot(x,h_rad,label='numerical solution (H)')
plt.legend(handles=[c0,c1,c2,c3])
plt.draw()
plt.show(block = True)

        
        
