import numpy as np
import matplotlib.pyplot as plt
# sample parameter §§ Gold
la0 = 429 # conductivity in W/mK
gma = 62.8 # thermal constant J/m^3K^2
Cl = 2.6*10**6 # phonon specific heat  in J/m^3K
Gel = 3.5*10**16 # lelectron phonon coupling constant
tau_e = 0.04 # e relaxation t const
tau_l = 0.6 # l relaxation t const
Tm = 300
L = 100 # sample thickness in nm
# Laser pulse parameters
T = 6

R = 0.93 # reflectivity
I0 = 100 # fulence in J/m^2
tp = 0.1 # in ps
zz0 = 15.3 # penetration depth in nm


#F = 0.5

def las(x,t):
    pt = 0.93*((1-R)/(zz0*tp))*I0*np.exp(-(x/zz0)-2.772*((t-2*tp)**2/tp**2))
    return pt

def I(x):
    tin = Tm
    return tin

a = la0



def solver_FE(I, a, las, L, T):
    """

    """
    import time;  t0 = time.clock()  # For measuring the CPU time

    dt = 0.001
    Nt = int(round(T/float(dt)))
    t = np.linspace(0, Nt*dt, Nt+1)   # Mesh points in time

    dx = 1
    Nx = int(round(L/dx))
    x = np.linspace(0, L, Nx+1)       # Mesh points in space


    F = a*(dt/dx**2)
    u   = np.zeros(Nx+1)
    u_n = np.zeros(Nx+1)
    v   = np.zeros(Nx+1)
    v_n = np.zeros(Nx+1)



    # Set initial condition u(x,0) = I(x)
    for i in range(0, Nx+1):
        u_n[i] = I(x[i])
        v_n[i] = I(x[i])



    for n in range(0, Nt):
        # Compute u, v  at inner mesh points
        for i in range(1, Nx):
            u[i] = u_n[i] + F*(u_n[i-1] - 2*u_n[i] + u_n[i+1]) + dt*(-Gel*(u_n[i] - v_n[i]) + las(x[i], t[n]))
            v[i] = v_n[i] + F*(v_n[i-1] - 2*v_n[i] + v_n[i+1]) + dt*(Gel*(u_n[i] - v_n[i]))

        # Insert boundary conditions
        u[0] = 0;  u[Nx] = 0
        v[0] = 0;  v[Nx] = 0

       # this step is to save each value of v and u in each time step before updating it, such that I can plot u and v with respect to time.
         ut = np.array([])
         vt = np.array([])
        for m in range(0, Nt):
            ut = np.append(ut,u)
            vt = np.append(vt,v)


        # Switch variables before next step

        u_n, u = u, u_n
        v_n, v = v, v_n



    t1 = time.clock()
    return u_n, v_n, x, t, t1-t0  # u_n holds latest u


u, v, x, t, cpu = solver_FE(I, a, las, L, T)

#fig = plt.figure(1)
##t = np.linspace(0, T, h)
#plt.plot(t,u,'r',label=r'$T_e$')
#plt.plot(t,v,'b',label=r'$T_l$')
##plt.plot(t,sol0[:,2],'g',label=r'$T_s$')
#plt.ylabel('$Temperature$',fontsize=20)
#plt.xlabel('Delay',fontsize=20)
#plt.legend(loc='best')
#plt.show()
##
fig = plt.figure(2)
#t = np.linspace(0, T, h)
plt.plot(x,u,'r',label=r'$T_e$')
plt.plot(x,v,'g',label=r'$T_l$')
#plt.plot(t,sol0[:,2],'g',label=r'$T_s$')
plt.ylabel('$Temperature$',fontsize=20)
plt.xlabel('length',fontsize=20)
plt.legend(loc='best')
plt.show()

fig = plt.figure(3)
#t = np.linspace(0, T, h)
plt.plot(t,las(0,t),'r',label=r'$T_e$')
#plt.plot(x,v,'g',label=r'$T_l$')
#plt.plot(t,sol0[:,2],'g',label=r'$T_s$')
plt.ylabel('$Intensity$',fontsize=20)
plt.xlabel('Delay',fontsize=20)
plt.legend(loc='best')
plt.show()

RuntimeWarning: overflow encountered in double_scalars
  u[i] = u_n[i] + F*(u_n[i-1] - 2*u_n[i] + u_n[i+1]) + dt*(-Gel*(u_n[i] - v_n[i]) + las(x[i], t[n]))
C:/Users/jayas/OneDrive/Documents/Python Scripts/llg/diff_ttm_v_test.py:90: RuntimeWarning: overflow encountered in double_scalars
  v[i] = v_n[i] + F*(v_n[i-1] - 2*v_n[i] + v_n[i+1]) + dt*(Gel*(u_n[i] - v_n[i]))
C:/Users/jayas/OneDrive/Documents/Python Scripts/llg/diff_ttm_v_test.py:89: RuntimeWarning: invalid value encountered in double_scalars
  u[i] = u_n[i] + F*(u_n[i-1] - 2*u_n[i] + u_n[i+1]) + dt*(-Gel*(u_n[i] - v_n[i]) + las(x[i], t[n]))
C:/Users/jayas/OneDrive/Documents/Python Scripts/llg/diff_ttm_v_test.py:90: RuntimeWarning: invalid value encountered in double_scalars
  v[i] = v_n[i] + F*(v_n[i-1] - 2*v_n[i] + v_n[i+1]) + dt*(Gel*(u_n[i] - v_n[i]))