#!/usr/bin

import numpy as np
from matplotlib import pyplot as plt
from matplotlib import animation
from utils import *
import ODE
import riemann_solver as RS
import pdb
import time as module_time
import sys
print ("**************** BEGINING OF THE COMPUTATION ****************")
start_time=module_time.time()
"""
Implementation of the first-order finite volume method to solve an elastic dynamic bar problem in a solid following a gradient-damage model.
The damage model gives rise to a two transport equations system as well as a source term that is treated using a Godunov splitting approach.
"""

## Length of the solid domain
length = 0.005
## Loading
vd=-0.1

#######################################
## Material parameters: elastic part ##
#######################################
# Young modulus
E = 27.e9
# Mass density
rho = 2500.0
######################################
## Material parameters: damage part ##
######################################
# Critical energy release rate
Gc=1e-1
# Characteristic length
l0=5.e-5
# Power in the degradation function
weightPower=2.0
# Damage power in free-energy (1: AT1 model; 2: AT2 model)
NRGPower=1
if NRGPower==1:
    ## AT1 model
    cw=2.0/3.0
elif NRGPower==2:
    ## AT2 model
    cw=1./2.
else:
    print ("NRG power in damage energy can only be 1 or 2!!")
    pdb.set_trace()
# Coefficient in the exponential change of variable
expoPower=4
# Damage viscosity
eta_d=1.e-4
# Parameters of the Fremond-like formulation
wl=Gc/(4.0*l0*cw);
lc=np.sqrt(Gc*l0/(2.0*cw));
############################################
## Material parameters: micromorphic part ##
############################################
# Micro inertia coefficient
Kbeta=1.e0
beta=Kbeta*rho*lc**2/E
penalty_coefficient=1.e6

## Damage-dependent functions in the free energy and its derivatives
#####################################################################################
## A(d) function and A'(d)
damagedStiffness= lambda d,eps: (0.5*(1+np.sign(eps))*E*np.exp(-weightPower*expoPower*d) + 0.5*(1-np.sign(eps))*E)
damagedStiffness_slope=lambda d,eps: -0.5*weightPower*expoPower*(1+np.sign(eps))*E*np.exp(-weightPower*expoPower*d)
damagedStiffness_curvature=lambda d,eps: 0.5*(1+np.sign(eps))*E*np.exp(-weightPower*expoPower*d)*(weightPower*expoPower)**2
    
## W(d) function and W'(d)
damageEnergyDensity=lambda d:wl*(1.0-np.exp(-expoPower*d))**NRGPower
damageEnergyDensity_slope=lambda d:wl*NRGPower*expoPower*np.exp(-expoPower*d)*(1-np.exp(-expoPower*d))**(NRGPower-1)
if (NRGPower>1):
    damageEnergyDensity_curvature=lambda d:wl*NRGPower*expoPower*( -expoPower*np.exp(-expoPower*d)*(1-np.exp(-expoPower*d))**(NRGPower-1) + (NRGPower-1)*np.exp(-expoPower*d)*np.exp(-expoPower*d)*(1-np.exp(-expoPower*d))**(NRGPower-2))
else :
    damageEnergyDensity_curvature=lambda d:-wl*expoPower*expoPower*np.exp(-expoPower*d)

penalty= lambda d,d1: 0.5*penalty_coefficient*(d1-1+np.exp(-expoPower*d))**2
dpenaltydD= lambda d,d1: -penalty_coefficient*expoPower*(d1-1+np.exp(-expoPower*d))*np.exp(-expoPower*d)
dpenaltydD1= lambda d,d1: penalty_coefficient*(d1-1+np.exp(-expoPower*d))
d2penaltydD2= lambda d,d1: penalty_coefficient*expoPower*expoPower*np.exp(-expoPower*d)*( (d1-1+np.exp(-expoPower*d)) + np.exp(-expoPower*d))
d2penaltydD12= lambda d,d1: penalty_coefficient
d2penaltydDdD1= lambda d,d1: -penalty_coefficient*expoPower*np.exp(-expoPower*d)
#####################################################################################


# Initial wave speeds
c_meca=np.sqrt(E/rho)
c_damage=np.sqrt(lc**2/beta)


# # Elastic waves speed
# c=computeMaxSpeed(np.zeros((1,6)),damagedStiffness,rho,c_damage)
# c_elas=c
    

####################
## Discretization ##
####################
## Number of cells
nCells = 100
## Cells size   
dx=length/nCells
## Finite volume centroids
y=np.linspace(dx/2.,length-dx/2.,nCells)
## Courant number
CFL=0.5
## Physical time
timeOut=12.*length/c_meca
## Initial time step   
dt=CFL*dx/c_meca
## Total number of time steps
NTmaxi=int(timeOut/dt)
print ("Time step set to dt=",'{0:.2E}'.format(dt),"s.")

#####################
## Output settings ##
#####################
## Show or not the animation
anim=True
## Number of time steps recorded
n_stored =200
## Output storage frequency 
save_freq=(int(NTmaxi/n_stored))*(n_stored<=NTmaxi)+(n_stored>NTmaxi)
## Number of time steps already recorded
i_print=0


#############################
## Backward Euler settings ##
#############################
## ODE solver parameters
NiterMax_NewtonRaphson=250
Nsubsteps_BackwardEuler=1
ODE_tol=1.e-6

##########################
##  Mermory allocation  ##
##########################
## Mechanical fields
v=np.zeros((nCells+2,int(NTmaxi/save_freq)+1))
eps= np.zeros((nCells+2,int(NTmaxi/save_freq)+1))
sig=np.zeros((nCells+2,int(NTmaxi/save_freq)+1))
D=np.zeros((nCells+2,int(NTmaxi/save_freq)+1))
D1=np.zeros((nCells+2,int(NTmaxi/save_freq)+1))
## Micromorphic fields
q=np.zeros((nCells+2,int(NTmaxi/save_freq)+1)) # time derivative of D1
p=np.zeros((nCells+2,int(NTmaxi/save_freq)+1))
## time
t=np.zeros(NTmaxi)
## Energies
elast_energy=np.zeros((nCells+2,int(NTmaxi/save_freq)+1))
damage_energy=np.zeros((nCells+2,int(NTmaxi/save_freq)+1))
kinetic_energy=np.zeros((nCells+2,int(NTmaxi/save_freq)+1))
total_energy=np.zeros((nCells+2,int(NTmaxi/save_freq)+1))
time=np.zeros(int(NTmaxi/save_freq)+1)

# Conserved quantities
# U = [epsilon, rho v, D, D1, beta q, p]
U=np.zeros((nCells+2,6))

########################################################
# RIEMANN TYPE INITIAL CONDITIONS IN TERMS OF VELOCITY #
########################################################
U[1:-1,1]=[rho*vd*(nCells-i-0.5)*dx/length for i in range(nCells)]

## Yield surface f(eps,D,D1)
yieldFunc =lambda eps,d,d1: -0.5*damagedStiffness_slope(d,eps)*positivePart(eps)**2 - damageEnergyDensity_slope(d) -dpenaltydD(d,d1) 


ODE_time=0.
print ("Computing first-order FVM solution with AT damage model selected (NRG power="+str(int(NRGPower))+").")

percentage_step=10
progression=10

for i in range(NTmaxi)[1:]:
    #Apply Courant condition
    cmax=computeMaxSpeed(U[1:-1,:],damagedStiffness,rho,c_damage)
    dt = CFL*dx/cmax

    # Compute time at the end of the current step
    if ((t[i-1]+dt)>timeOut):
        dt = timeOut - t[i-1]
    t[i]=t[i-1]+dt

    ### Handmade BCs
    U[0,:]=U[1,:] ; U[-1,:]=U[-2,:]
    ## Velocity
    U[0,1]=2.*rho*vd-U[1,1]
    U[-1,1]=0.-U[-1,1]
    ## Damage gradient
    U[0,5]=-U[1,5]
    U[-1,5]=0.-U[-1,5]
    fR,fL = RS.computeflux(U,rho,beta,lc,damagedStiffness,dt,dx,nCells,vd)

    U=U-dt/dx*(fL-fR)
    
    ODE_time_start=module_time.time()
    U = ODE.integrateODE(U,dt,beta,1./eta_d,damagedStiffness_curvature,damageEnergyDensity_curvature,yieldFunc,dpenaltydD1,d2penaltydD2,d2penaltydD12,d2penaltydDdD1,NiterMax_NewtonRaphson,Nsubsteps_BackwardEuler,ODE_tol)
    ODE_time+= module_time.time()-ODE_time_start

    #Update of fields and Store results
    if (i%save_freq==0) :
        v[:,i_print]=U[:,1]/rho
        eps[:,i_print]=U[:,0]
        A=damagedStiffness(U[:,2],U[:,0])
        sig[:,i_print]=A*eps[:,i_print]
        D[:,i_print]=U[:,2]
        D1[:,i_print]=U[:,3]
        q[:,i_print]=U[:,4]/beta
        p[:,i_print]=U[:,5]
        time[i_print]=t[i]

        elast_energy[:,i_print]=0.5*damagedStiffness(U[:,2],U[:,0])*eps[:,i_print]**2
        damage_energy[:,i_print]=damageEnergyDensity(U[:,2]) + 0.5*(lc*p[:,i_print])**2 +penalty(D[:,i_print],D1[:,i_print])
        kinetic_energy[:,i_print]=0.5*rho*v[:,i_print]**2 +0.5*beta*(q[:,i_print])**2
        total_energy[:,i_print]=elast_energy[:,i_print]+damage_energy[:,i_print]+kinetic_energy[:,i_print]
        i_print+=1

    A=damagedStiffness(U[:,2],U[:,0])
    beta_q_prev = U[:,4]
    
    if (100.*t[i]/timeOut > progression):
        ## Print information every percentage_step% of time steps
        print( "Time step", i," finished ", round((t[i]/timeOut)*100,0),"% (time step dt=",'{0:.2E}'.format(dt),"s)")
        progression+=percentage_step
    
    
        
    if (t[i]==timeOut):
        Nfinal=i
        break


# Delete the columns that won't be filled
sig=np.delete(sig,np.arange(i_print,sig.shape[1]-1,1),1)
v=np.delete(v,np.arange(i_print,v.shape[1]-1,1),1)
eps=np.delete(eps,np.arange(i_print,eps.shape[1]-1,1),1)
D=np.delete(D,np.arange(i_print,D.shape[1]-1,1),1)
D1=np.delete(D1,np.arange(i_print,D1.shape[1]-1,1),1)
q=np.delete(q,np.arange(i_print,q.shape[1]-1,1),1)
p=np.delete(p,np.arange(i_print,p.shape[1]-1,1),1)
elast_energy=np.delete(elast_energy,np.arange(i_print,elast_energy.shape[1]-1,1),1)
damage_energy=np.delete(damage_energy,np.arange(i_print,damage_energy.shape[1]-1,1),1)
kinetic_energy=np.delete(kinetic_energy,np.arange(i_print,kinetic_energy.shape[1]-1,1),1)
total_energy=np.delete(total_energy,np.arange(i_print,total_energy.shape[1]-1,1),1)
time=np.delete(time,np.arange(i_print,len(time)-1,1),0)


print ("final time:",t[i])
## Make sure that the last time step is saved in the tables
v[:,-1]=U[:,1]/rho
eps[:,-1]=U[:,0]
A=damagedStiffness(U[:,2],U[:,0])
sig[:,-1]=A*eps[:,-1]
D[:,-1]=U[:,2]
D1[:,-1]=U[:,3]
q[:,-1]=U[:,4]/beta
p[:,-1]=U[:,5]
damage_energy[:,-1]=damageEnergyDensity(U[:,2]) + 0.5*(lc*U[:,5])**2 +penalty(U[:,2],U[:,3])
time[-1]=t[i]



stored_energy=damageEnergyDensity(D[1:-1,-1]) + 0.5*(lc*p[1:-1,-1])**2 +penalty(D[1:-1,-1],D1[1:-1,-1])

D=1.0-np.exp(-expoPower*D)

## Fragments length
thresh=0.98
lengths=computeFragmentsLength(dx,y,D[1:-1,-1],thresh)

nFrag=len(lengths)
mean_fragment_size=np.sum(lengths)/nFrag

print ("Number of fragments: ",nFrag)
print ("Mean fragement size=",format(mean_fragment_size,'.5e')," m")
print ("Final damage energy (J/m2): ",format(dx*np.sum(damage_energy[1:-1,-1]),'.5e'),"J/m")
print ("Stored energy (J/m2): ",format(dx*np.sum(stored_energy),'.5e'),"J/m")

if anim:
    ## energy plot
    plt.plot(time[:i_print],dx*np.sum(total_energy[1:-1,:i_print],axis=0),label='Tot')
    plt.plot(time[:i_print],dx*np.sum(kinetic_energy[1:-1,:i_print],axis=0),label='Kin')
    plt.plot(time[:i_print],dx*np.sum(elast_energy[1:-1,:i_print]+damage_energy[1:-1,:i_print],axis=0),label='Internal')
    plt.plot(time[:i_print],dx*np.sum(damage_energy[1:-1,:i_print],axis=0),label='Damage')
    plt.legend()
    plt.grid()
    plt.show()
    

T_tot = module_time.time()-start_time
print ("--- %s seconds ---" % T_tot)
print (" Ratio ODE time/ total time  =",ODE_time/T_tot)
print ("Number of time steps recorded: ", np.shape(sig)[1])


if anim:
    #### 4 plots
    # initialization function: plot the background of each frame
    def init():
        line[0].set_data([], [])
        line[1].set_data([], [])
        line[2].set_data([], [])
        line[3].set_data([], [])
        line[4].set_data([], [])
        return line

    # animation function.  This is called sequentially
    def animate(i):
        line[0].set_data(y,sig[1:-1,i])
        line[1].set_data(y,p[1:-1,i])
        line[2].set_data(y,D[1:-1,i])
        line[3].set_data(y,D1[1:-1,i])
        line[4].set_data(y,q[1:-1,i])
        return line

    print( '... building animation ...')
    print( 'Animation finished !')

    ####Animated plot ###########################################################
    # First set up the figure, the axis, and the plot element we want to animate
    fig, (ax1, ax2,ax3,ax4) = plt.subplots(4,1)
    ax1.grid() ; ax2.grid(); ax3.grid();ax4.grid()
    plt.title('Stress and Strain variable')
    line1, = ax1.plot([], [],'ro', lw=2,label=r'$\sigma$')
    line2, = ax3.plot([], [],'b+' , lw=2,label=r'$\nabla D_1$')
    line3, = ax2.plot([], [],'k-' , lw=2,label='D')
    line4, = ax2.plot([], [],'gx' , lw=2,label=r'$D_1$')
    line5, = ax4.plot([], [],'y--' , lw=2,label=r'$\dot{D_1}$')
    line = [line1, line2,line3,line4,line5]
    
    ax1.set_ylim(np.min(sig[1:-1,:]),np.max(sig[1:-1,:]));
    ax3.set_ylim(np.min(p[1:-1,:]),np.max(p[1:-1,:]))
    ax2.set_ylim(0,1.2)
    ax4.set_ylim(np.min(q[1:-1,:]),np.max(q[1:-1,:]))
    ax2.plot([y[0],y[-1]],[1.,1.],'k')
    ax4.set_xlabel('x (m)', fontsize=18)
    ax3.set_ylabel(r'$\nabla D_1$', fontsize=18)
    ax1.set_ylabel(r'$\sigma$', fontsize=18)
    ax2.set_ylabel(r'D', fontsize=18)
    ax4.set_ylabel(r'$\dot{D_1}$', fontsize=18)
    
    ax1.set_xlim(y[0],y[-1]);
    ax2.set_xlim(y[0],y[-1])
    ax3.set_xlim(y[0],y[-1])
    ax4.set_xlim(y[0],y[-1])
    
    # call the animator.  blit=True means only re-draw the parts that have changed.
    n_frames=20
    anim = animation.FuncAnimation(fig, animate, init_func=init,frames=i_print, interval=10, blit=True)


    writervideo = animation.FFMpegWriter(fps=1)
    plt.show()
    plt.close()



print ("")
