#def runner(i_ext,tspan):

"""
Author:
    Amy LeBar (20 August 2018)

Li-O2 Battery Model:
    This model examines the reactions taking place within the carbon-based
    cathode of a Li-O2 battery. Electrolyte = 1 M LiTFSI in TEGDME

"""
""" Load any needed modules """
"============================================================================"
import numpy as np
import cantera as ct
import matplotlib.pyplot as plt
from scipy.integrate import solve_ivp

""" BEGIN USER INPUTS """
"============================================================================"
phi_elyte_init = -3.19                  # double layer voltage [V]
eps_elyte_init = 0.5                      # initial electrolyte volume fraction [-]
eps_oxide_init = 1e-12                    # initial oxide volume fraction [-]
eps_binder_init = 0.                      # initial binder volume fraction [-]
eps_carbon = 1. - eps_elyte_init - eps_binder_init - eps_oxide_init      # initial carbon volume fraction [-]

atol = 1e-10
rtol = 2.5e-6

tspan = 7824                            # [s]

i_ext = -1e-3                       # [A/m2]
cap = 1e-3*2.1733333                    # battery capacity

Nx = 1                                  # 1D model
Ny = 5                                  # no. cells in the y-direction
Nvars = 3                               # no. of variables
th_ca = 50e-6                           # cathode thickness [m]
dy = th_ca/Ny                           # [m]
d_part = 10e-6                          # carbon particle diameter [m]
d_oxide = 2e-6                          # oxide particle diameter [m]
th_oxide = 5e-6                         # thickness of oxide ellipsoid [m]
V_part = 2/3 * np.pi * (d_part / 2)**3  # particle volume [m3]
A_part = 4 * np.pi * (d_part / 2)**2    # particle surface area [m2]
A_int = eps_carbon * A_part / V_part      # interface area [m2/m3 total]
A_oxide = np.pi * d_oxide**2 / 4        # oxide area contacting carbon particle
V_oxide = 2/3 * np.pi * (d_oxide/2)**2 * th_oxide   # oxide volume [m3]
C_dl = 1.1e-6                           # double layer capacitance [F/m2]

TP = 300, 101325                        # inital temp, pressure [K, Pa]

ctifile = 'LiAir_mod.cti'

""" END USER INPUTS """
"============================================================================"
# Import necessary phases
gas = ct.Solution(ctifile,'air')
cath_b = ct.Solution(ctifile,'graphite')
elyte = ct.Solution(ctifile,'electrolyte')
oxide = ct.Solution(ctifile,'Li2O2')
inter = ct.Interface(ctifile,'cathode_surf',[elyte,oxide,cath_b])
air_elyte = ct.Interface(ctifile,'air_elyte',[gas,elyte])
Li_b = ct.Solution(ctifile,'Lithium')
Li_s = ct.Interface(ctifile,'Li_surface',[Li_b,elyte])

oxide.TP = TP
elyte.TP = TP
inter.TP = TP
cath_b.TP = TP

# Store these phases in a common 'objs' dict
objs = {}
objs['gas'] = gas
objs['cath_b'] = cath_b
objs['elyte'] = elyte
objs['oxide'] = oxide
objs['inter'] = inter
objs['air_elyte'] = air_elyte
objs['Li_b'] = Li_b
objs['Li_s'] = Li_s

# Store parameters in a common 'params' dict
params = {}
params['i_ext'] = i_ext
params['T'] = TP[0]
params['eps_elyte_0'] = eps_elyte_init
params['eps_oxide_0'] = eps_oxide_init
params['rtol'] = rtol
params['atol'] = atol

# Store pointers in a common 'ptr' dict
ptr = {}
ptr['elec'] = elyte.n_species + oxide.n_species         # electron in the inter net_production_rates vector
ptr['oxide'] = elyte.n_species                          # oxide in the inter net_production_rates vector
ptr['elyte'] = np.arange(0,elyte.n_species)             # electrolyte in the inter net_production_rates vector

# Store solution vector pointers in a common 'SVptr' dict
SVptr = {}
SVptr['phi'] = 0                                        # double layer potential in solution vector SV
SVptr['oxide'] = 1                                      # oxide density in solution vector SV
SVptr['elyte'] = np.arange(2,elyte.n_species + 2)       # electrolyte densities in solution vector SV

# Store plot pointers in a common 'pltptr' dict
pltptr = {}
pltptr['O2'] = 2
pltptr['Li+'] = 3
pltptr['PF6-'] = 4
pltptr['EC'] = 5
pltptr['EMC'] = 6

# Set inital values
rho_oxide_init = oxide.density*params['eps_oxide_0']          # oxide concentraion
rho_elyte_init = elyte.Y*elyte.density*params['eps_elyte_0']  # electrolyte concentrations
SV0 = np.r_[phi_elyte_init,rho_oxide_init,rho_elyte_init]   # store in an array
SV_0 = np.tile(SV0,Ny)                                      # tile SV0 based on discritization

# Define function to solve
def LiO2_func(t,SV,params,objs,ptr,SVptr):

    dSVdt = np.zeros_like(SV * (Ny + 2))

    dPhidt = np.zeros_like(dSVdt)
    dRhoOxidedt = np.zeros_like(dSVdt)
    dRhoElytedt = np.zeros_like(dSVdt)

    # Pull phases out of 'objs' inside function
    gas = objs['gas']
    cath_b = objs['cath_b']
    elyte = objs['elyte']
    oxide = objs['oxide']
    inter = objs['inter']
    air_elyte = objs['air_elyte']
    Li_b = objs['Li_b']
    Li_s = objs['Li_s']
    
    " --- Pre-loop --- "
    # Set potentials for 'next'
    Phi_cathode = SV[SVptr['phi']]
    cath_b.electric_potential = 0
    oxide.electric_potential = 0
    elyte.electric_potential = Phi_cathode
    phi_elyte_next = elyte.electric_potential
    
    # Set concentrations for 'next'
    rho_oxide_next = SV[SVptr['oxide']]
    rho_k_elyte_next = SV[SVptr['elyte']]
    eps_oxide = rho_oxide_next / oxide.density_mass
    eps_elyte = params['eps_elyte_0'] - (eps_oxide - params['eps_oxide_0'])
    rho_elyte = (sum(rho_k_elyte_next)) / eps_elyte
    elyte.TDY = params['T'], rho_elyte, rho_k_elyte_next
    
    # Set transport properties for 'next'
    Xk_next = elyte.mole_fractions
    Dk_next = elyte.binary_diff_coeffs
    XT_next = elyte.density_mole

    # Set transport properties for 'this'
    rho_k_elyte_this = elyte.Y*elyte.density*params['eps_elyte_0']
    elyte.TDY = params['T'], rho_elyte, rho_k_elyte_this
    Xk_this = elyte.mole_fractions
    Dk_this = elyte.binary_diff_coeffs
    XT_this = elyte.density_mole
    
    # Mass transport and ionic current
    Dk_down = (Dk_this + Dk_next) / 2
    XT_down = (XT_this + XT_next) / 2
    phi_elyte_this = elyte.electric_potential
    Jk_down = Dk_down * XT_down * (Xk_this - Xk_next) / dy
    sigma_io = elyte.electrical_conductivity
    i_io_down = sigma_io * (phi_elyte_this - phi_elyte_next) / dy
    

    for j in np.arange(Ny):
        
        # Old 'next' become new 'this' and old 'down' become new 'up'
        phi_elyte_this = phi_elyte_next
        rho_k_elyte_this = rho_k_elyte_next
        rho_oxide_this = rho_oxide_next
        Xk_this = Xk_next
        Dk_this = Dk_next
        XT_this = XT_next
        Jk_up = Jk_down
        i_io_up = i_io_down
        Dk_up = Dk_down
        XT_up = XT_down

        # Set potentials
        Phi_cathode = SV[SVptr['phi']]
        cath_b.electric_potential = 0
        oxide.electric_potential = 0
        elyte.electric_potential = Phi_cathode
        phi_elyte_next = elyte.electric_potential
        
        # Set concentrations
        rho_k_elyte_next = SV[SVptr['elyte']]
        rho_oxide_next = SV[SVptr['oxide']]
        eps_oxide = rho_oxide_next / oxide.density_mass
        eps_elyte = params['eps_elyte_0'] - (eps_oxide - params['eps_oxide_0'])
        rho_elyte = (sum(rho_k_elyte_next)) / eps_elyte
        elyte.TDY = params['T'], rho_elyte, rho_k_elyte_next

        # Mass transport and ionic current
        Xk_next = elyte.mole_fractions
        Dk_next = elyte.binary_diff_coeffs
        XT_next = elyte.density_mole
        Dk_down = (Dk_this - Dk_next) / 2
        XT_down = (XT_this - XT_next) / 2
        Jk_down = Dk_down * XT_down * (Xk_this - Xk_next) / dy
        sigma_io = elyte.electrical_conductivity
        i_io_down = sigma_io * (phi_elyte_this - phi_elyte_next) / dy

        # Calculate net production rates at interface
        sdot = inter.net_production_rates                 # interface production rates

        # Calculate Faradaic current
        i_far = -sdot[ptr['elec']] * ct.faraday            # Faradaic current

        # Calculate change in oxide concentration
        W_oxide = oxide.mean_molecular_weight             # oxide molecular weight
        A_int_avail = A_int - eps_oxide / th_oxide          # available interface area on carbon particle

        dRhoOxidedt = sdot[ptr['oxide']] * A_int_avail * W_oxide

        # Calculate change in double layer potential
        i_dl = (i_io_down - i_io_up) / dy - i_far*A_int_avail    # double layer current
        dPhidt = i_dl / (C_dl*A_int)                           # double layer potential

        # Calculate change in electrolyte concentrations
        W_elyte = elyte.molecular_weights
        dRhoElytedt = (Jk_down - Jk_up) / dy + (sdot[ptr['elyte']] * A_int_avail * W_elyte)

    " --- Post-loop --- "
    # Old 'next' become new 'this' and old 'down' become new 'up'
    phi_elyte_this = phi_elyte_next
    rho_k_elyte_this = rho_k_elyte_next
    rho_oxide_this = rho_oxide_next
    Xk_this = Xk_next
    Dk_this = Dk_next
    XT_this = XT_next
    Jk_up = Jk_down
    i_io_up = i_io_down
    Dk_up = Dk_down
    XT_up = XT_down
    
    # BC's at separator
    i_io_down = params['i_ext']
    Jk_down = np.zeros_like(elyte.mole_fractions)

    # Calculate net production rates at interface
    sdot = inter.net_production_rates                 # interface production rates

    # Calculate Faradaic current
    i_far = -sdot[ptr['elec']] * ct.faraday            # Faradaic current

    # Calculate change in oxide concentration
    W_oxide = oxide.mean_molecular_weight             # oxide molecular weight
    A_int_avail = A_int - eps_oxide / th_oxide          # available interface area on carbon particle

    dRhoOxidedt = sdot[ptr['oxide']] * A_int_avail * W_oxide

    # Calculate change in double layer potential
    i_dl = (i_io_down - i_io_up) / dy - i_far*A_int_avail    # double layer current
    dPhidt = i_dl / (C_dl*A_int)                           # double layer potential

    # Calculate change in electrolyte concentrations
    W_elyte = elyte.molecular_weights
    dRhoElytedt = (Jk_down - Jk_up) / dy + (sdot[ptr['elyte']] * A_int_avail * W_elyte)

    # Load differentials into dSVdt
    dSVdt[SVptr['phi']] = dPhidt                            # double layer potential
    dSVdt[SVptr['oxide']] = dRhoOxidedt                     # oxide concentration
    dSVdt[SVptr['elyte']] = dRhoElytedt                     # electrolyte concentration

    if 0:
        print('-----------------------------------------------------------')
        print('Phi =',Phi_cathode)
        print('eps_oxide =',eps_oxide)
        print('eps_elyte =',eps_elyte)
        print('sdot =',sdot)
        print('i_Far =',i_far)
        print('i_dl =',i_dl)
        elyte()

    return dSVdt

# Solve function using IVP solver
SV = solve_ivp(lambda t, y: LiO2_func(t,y,params,objs,ptr,SVptr), [0, tspan], SV_0, method='BDF',atol=params['atol'],rtol=params['rtol'])

#    Phi_dl = SV.y[SVptr['phi'],-1]

#    return SV

""" Plot solutions to concentrations and potentials """
"============================================================================"
#plt.figure(1)
#plt.plot(SV.t,SV.y[SVptr['phi']])
#plt.xlabel('Time (s)')
#plt.ylabel('Double Layer Potential (V)')

#plt.figure(2)
#plt.plot(SV.t,SV.y[SVptr['oxide']])
#plt.xlabel('Time (s)')
#plt.ylabel('Oxide Concentration (kg/m3)')

eps_oxide = SV.y[SVptr['oxide']] / oxide.density_mass      # oxide volume fraction
eps_elyte = params['eps_elyte_0'] - (eps_oxide - params['eps_oxide_0'])
A_int_avail = A_int - eps_oxide / th_oxide

#plt.figure(3)
#plt.plot(SV.t,eps_elyte)
#plt.xlabel('Time (s)')
#plt.ylabel('Elyte Volume Fraction')
#plt.show()

#plt.figure(4)
#plt.plot(SV.t/3600 * -i_ext,SV.y[SVptr['phi']])
#plt.xlabel('Capacity (Ah/m2)')
#plt.ylabel('Voltage (V)')

plt.figure(5)
plt.plot(SV.t,A_int_avail)
plt.xlabel('Time (s)')
plt.ylabel('Available Area (m2)')

#plt.figure(3)
#plt.plot(SV.t,SV.y[pltptr['O2']],SV.t,SV.y[pltptr['Li+']],SV.t,SV.y[pltptr['PF6-']],SV.t,SV.y[pltptr['EC']],SV.t,SV.y[pltptr['EMC']])
#plt.legend(['O2','Li+','PF6-','EC','EMC'])
#plt.xlabel('Time (s)')
#plt.ylabel('Electrolyte Concentration (kg/m3)')
#plt.show()


#t = SV.t
#dl = SV.y[SVptr['phi']]
#Ck_ox = SV.y[SVptr['oxide']]
#
#df = DataFrame({'Time': t, 'Double Layer': dl, 'Oxide Concentration': Ck_ox})
#
#with ExcelWriter('path_to_file.xlsx') as writer:
#    df.to_excel(writer)
