# LUCAS.jl
#
# Version 4b, 8 Nov 2015
#
# (c) 2015 Andrew J. Campbell
# 
# Simulates growth of planet and consequent core formation, with metal/silicate reactions
# 
# 0. Initialize element data and model parameters
# 1. Read composition of initial body.
# 2. Equilibrate metal/silicate in this body at 1 bar and specified high T and fO2.
# 3. For each impactor:
# 3a. Read composition of impactor.
# 3b. Equilibrate metal/silicate in impactor at 1 bar and specified T, fO2.
# 3c. Calculate P,T at specified depth of impactor/target mixing.
# 3d. Mix impactor mass with fraction of target mantle.
# 3e. Equilibrate metal/silicate of mixed impactor/target mass.
# 3f. Add resulting metal to core; silicate to mantle. Mix core and mantle.
# 


# --------------------------------------------------------------------------

function impactors(numsteps)

# Produces array of compositions of impactors
# Edit this for new impact history
	time0 = time()

	# Reduced (IW-3.5) composition from Fischer et al. EPSL 2015
	elem = ["O", "Fe", "Mg", "Al", "Si", "Ca", "Co", "Ni"]
	oxide = [0, 2.24, 37.5, 4.62, 51.41, 3.75, 0.00051, 0.00101]
	metal = [0.04, 91.1, 0, 0, 2.4, 0, 0.26, 5.55]
	metalfrac = 0.313
	
	# Oxidized (IW-1.5) composition from Fischer et al. EPSL 2015
	elem = ["O", "Fe", "Mg", "Al", "Si", "Ca", "Co", "Ni"]
	oxide = [0, 21.13, 29.4, 3.63, 42.19, 2.95, 0.0083, 0.0174]
	metal = [0.4, 89.07, 0, 0, 0.0205, 0, 0.44, 10.0]
	metalfrac = 0.165
	
	earthmass = 5.972e27   # grams
	impmass = earthmass / numsteps
	
	maw = getmaws(elem)
	valence = getvalences(elem)
	oxidewt = maw + maw[1] * valence/2.
	oxidewt[1] = 16.   # to avoid div/0 error
	
	Moxide = impmass * (1-metalfrac) * oxide / sum(oxide)
	Mmetal = impmass * metalfrac * metal / sum(metal)
	Noxide = Moxide ./ oxidewt
	Noxide[1] = 0   # just to make sure
	Noxide[1] = sum(Noxide .* valence/2.)
	Nmetal = Mmetal ./ maw
	Nbulk = Nmetal + Noxide
	
	println("impactors took ",time()-time0," s")
		
	return elem, Nmetal, Noxide

end  # function impactors

# --------------------------------------------------------------------------

function getmaws(elem)

# Produces array of mean atomic weights
# Edit this when adding new elements

	num = length(elem)
	maw = zeros(num)
	for i in 1:num
		if elem[i] == "O"
			maw[i] = 15.999
		elseif elem[i] == "Fe"
			maw[i] = 55.845
		elseif elem[i] == "Mg"
			maw[i] = 24.305
		elseif elem[i] == "Al"
			maw[i] = 26.982
		elseif elem[i] == "Si"
			maw[i] = 28.085
		elseif elem[i] == "Ca"
			maw[i] = 40.078
		elseif elem[i] == "V"
			maw[i] = 50.942
		elseif elem[i] == "Cr"
			maw[i] = 51.996
		elseif elem[i] == "Co"
			maw[i] = 58.933
		elseif elem[i] == "Ni"
			maw[i] = 58.693
		elseif elem[i] == "Nb"
			maw[i] = 92.906
		elseif elem[i] == "Ta"
			maw[i] = 180.948
		elseif elem[i] == "W"
			maw[i] = 183.84
		elseif elem[i] == "Th"
			maw[i] = 232.04
		elseif elem[i] == "U"
			maw[i] = 238.03
		end 
	end
		
	return maw

end  # function getmaws

# --------------------------------------------------------------------------

function getvalences(elem)

# Produces array of valences
# Edit this when adding new elements

	num = length(elem)
	val = zeros(num)
	for i in 1:num
		if elem[i] == "O"
			val[i] = -2
		elseif elem[i] == "Fe"
			val[i] = 2
		elseif elem[i] == "Mg"
			val[i] = 2
		elseif elem[i] == "Al"
			val[i] = 3
		elseif elem[i] == "Si"
			val[i] = 4
		elseif elem[i] == "Ca"
			val[i] = 2
		elseif elem[i] == "V"
			val[i] = 3
		elseif elem[i] == "Cr"
			val[i] = 2
		elseif elem[i] == "Co"
			val[i] = 2
		elseif elem[i] == "Ni"
			val[i] = 2
		elseif elem[i] == "Nb"
			val[i] = 5
		elseif elem[i] == "Ta"
			val[i] = 5
		elseif elem[i] == "W"
			val[i] = 6
		elseif elem[i] == "Th"
			val[i] = 4
		elseif elem[i] == "U"
			val[i] = 4
		end 
	end
		
	return val

end  # function valences

# --------------------------------------------------------------------------

function kd0(p,t)

# Calculates KD values for each element, with no compositional dependence included
# From Fischer et al. GCA 2015 mainly
# Edit this when adding new elements

	# Arguments p,t are pressure, temperature

	kd = zeros(numelements)
	a = zeros(numelements)
	b = zeros(numelements)
	c = zeros(numelements)
	
	for i in 1:numelements
		if element[i] == "O"
			a[i] = 0.6
			b[i] = -3800
			c[i] = 22
		elseif element[i] == "Fe"
			a[i] = 0
		elseif element[i] == "Mg"
			a[i] = -10
		elseif element[i] == "Al"
			a[i] = -10
		elseif element[i] == "Si"
			a[i] = 1.3
			b[i] = -13500
		elseif element[i] == "Ca"
			a[i] = -10
		elseif element[i] == "V"
			a[i] = -0.3
			b[i] = -5400
			c[i] = 19
		elseif element[i] == "Cr"
			a[i] = 0.0
			b[i] = -2900
			c[i] = 9
		elseif element[i] == "Co"
			a[i] = 0.36
			b[i] = 1500
			c[i] = -33
		elseif element[i] == "Ni"
			a[i] = 0.46
			b[i] = 2700
			c[i] = -61
		elseif element[i] == "Nb"  # Mann 2009
			a[i] = 2.66
			b[i] = -14032
			c[i] = -199
		elseif element[i] == "Ta"  # Mann 2009
			a[i] = 0.84
			b[i] = -13806
			c[i] = -105
		elseif element[i] == "W"
			a[i] = 0
		elseif element[i] == "Th"
			a[i] = -10
		elseif element[i] == "U"
			a[i] = -10
		end 
		kd[i] = 10^(a[i] + b[i]/t + c[i]*p/t)
	end
	
	return kd

end  # function kd0

# --------------------------------------------------------------------------

function kdX(Xmet, Xsil, p, t)

# Calculates KD values for each element, including compositional dependence
# From Fischer et al. GCA 2015
# Edit this when adding new elements
	
	# Xmet, Xsil are metal, silicate compositions. p,t are pressure, temperature

	kd = zeros(numelements)
	e = zeros(numelements, numelements)
	eps = zeros(numelements, numelements)
	epsterm1 = zeros(numelements)
	epsterm2 = zeros(numelements)
	epsterm3 = zeros(numelements)
	
	e[5,1] = -0.06   # e[Si,O] e - superscript Si - subscript O
	e[1,1] = -0.12   # e[O,O]
	e[1,5] = -0.11   # e[O,Si]
	
	for i = 1:numelements, k = 1:numelements
		eps[i,k] = e[i,k] * maw[i] * 1873 / 0.242 / t - maw[i]/maw[2] + 1
	end
	for i = 1:numelements
		epsterm1[i] = eps[i,i] * log(1-Xmet[i]) / log(10)
		for k = 1:numelements
			if (k != i && k != 2)
				epsterm2[i] = epsterm2[i]+eps[i,k]/log(10)*Xmet[k]*(1+log(1-Xmet[k])/Xmet[k]-1/(1-Xmet[i]))
				epsterm3[i] = epsterm3[i]-eps[i,k]/log(10)*Xmet[k]^2*Xmet[i]*(1/(1-Xmet[i])+1/(1-Xmet[k])+Xmet[i]/2/(1-Xmet[i])^2-1)
			end
		end
	end
	kd[1] = 10^(0.1 -2200 / t + 5 * p/t + epsterm1[1] + epsterm2[1] + epsterm3[1]) # O
	kd[2] = 1  # Fe
	kd[3] = 1.e-10  # Mg
	kd[4] = 1.e-10  # Al
	kd[5] = 10^(0.6 -11700 / t + epsterm1[5] + epsterm2[5] + epsterm3[5])# Si
	kd[6] = 1.e-10  # Ca
	kd[7] = 10^(0.36 + 1500 / t -33 * p/t)  # Co
	kd[8] = 10^(0.46 + 2700 / t -61 * p/t)  # Ni
	
	return kd

end  # function kdX

# --------------------------------------------------------------------------

function equilibrate(Nbulk, p, t, iwtry)

# Equilibrates bulk composition

	# Argument Nbulk is a list of molar compositions. p,t are pressure, temperature
	# Element 1 is Oxygen; element 2 is Fe
	
	# Initally, put everything in the silicate except Fe
	Nsil = Nbulk+0
	Nsil[2] = 0
	Nmet = Nbulk - Nsil
	# Ntotmet, Ntotsil are total number of moles in each phase
	Ntotmet = sum(Nmet)
	Ntotsil = sum(Nsil)
	
	# Set up temporary arrays
	newNmet = Nmet+0
	newNsil = Nsil+0
	
	# Initialize IW, KD, and D terms
	iwbegin = iwtry
	iwend = iwbegin + 0.001
	kd = kd0(p, t)
	d = ones(numelements)
	
	while abs(iwend-iwbegin) > 0.0005
	# Scan to find convergence. Doesn't interate well.
		if iwend < iwbegin
			iwbegin = iwbegin - 0.0002
		else
			iwbegin = iwbegin + 0.0002
		end
		# First is O budget
		newNmet[1] = Ntotmet * kd[1] * 10^(iwbegin/2.)
		newNsil[1] = Nbulk[1] - Nmet[1]
		# Now cations except Fe
		for i = 3:numelements
			d[i] = kd[i] * 10^(-iwbegin*valence[i]/4)
			newNsil[i] = Nbulk[i] / (1 + d[i] * Ntotmet/Ntotsil)
			newNmet[i] = Nbulk[i] - Nsil[i]
		end
		# Balance Fe:  FeO accounts for rest of oxygen
		Osum = newNmet[1] + 0
		for i in 3:numelements
			Osum = Osum + newNsil[i]*(valence[i]/2)
		end
		newNsil[2] = Nbulk[1] - Osum
		if newNsil[2] <= 0.0
			newNsil[2] = Nbulk[2] * 1.e-8
		end
		newNmet[2] = Nbulk[2] - newNsil[2]
		# Recalculate new compositions, dIW, KD for next iteration
		Nmet = Nmet/2 + newNmet/2
		Nsil = Nsil/2 + newNsil/2
		Ntotmet = sum(Nmet)
		Ntotsil = sum(Nsil[2:numelements])
		Xsil = Nsil / Ntotsil
		Xmet = Nmet / Ntotmet
		iwend = 2*log10(Xsil[2]/Xmet[2])
#		kd = kdX(Xmet, Xsil, p, t)       # Ignore this to make direct comparison to FischerEPSL2015
	end # convergence loop
	
	return Nmet, Nsil
	
end # function equilibrate


# --------------------------------------------------------------------------

function equilibratefixedfO2(Nbulk, p, t, dIW)

# Equilibrates bulk composition at fixed delta-IW

	# Argument Nbulk is a list of molar compositions
	# Argument dIW is the fO2 (delta-IW) at which to equilibrate
	# Element 1 is Oxygen; element 2 is Fe
	
	# Initally, put everything in the silicate except Fe
	Nbulk[1] = 0. # Oxygen will float to fix fO2
	Nsil = Nbulk+0
	Nsil[2] = 0
	Nmet = Nbulk - Nsil # Hence zero except for Fe
	for i in 2:numelements
		Nsil[1] = Nsil[1] + Nsil[i]*valence[i]/2.
	end
	Xsil = Nsil / sum(Nsil[2:numelements])
	Xmet = Nmet / sum(Nmet)
	kd = kd0(p, t)
	
	iwtest = 0
	d = ones(numelements)
	
	while abs(iwtest-dIW) > 0.0001
	# Iterate to converge on solution
		for i = 1:numelements
			d[i] = kd[i] * 10^(-dIW*valence[i]/4)
		end
		# First Fe
		Nsil[2] = sum(Nsil[3:numelements]) / (d[2]/Xmet[2] - 1.)
		Nmet[2] = Nbulk[2] - Nsil[2]
		Ntotmet = sum(Nmet)
		Ntotsil = sum(Nsil[2:numelements])
		# Then other cations
		for i = 3:numelements
			Nsil[i] = Nbulk[i] / (1 + d[i] * Ntotmet/Ntotsil)
			Nmet[i] = Nbulk[i] - Nsil[i]
			Ntotmet = sum(Nmet)
			Ntotsil = sum(Nsil[2:numelements])
		end
		# Then O
		Nmet[1] = Ntotmet * kd[1] * 10^(dIW/2.)
		Nsil[1] = 0
		for i in 2:numelements
			Nsil[1] = Nsil[1] + Nsil[i]*valence[i]/2.
		end
		Ntotmet = sum(Nmet)
		Ntotsil = sum(Nsil[2:numelements])
				
		# Calculate new dIW to test convergence
		Xsil = Nsil / Ntotsil
		Xmet = Nmet / Ntotmet
#		kd = kdX(Xmet, Xsil, p, t)
		iwtest = 2*log10(Xsil[2]/Xmet[2])
	end # iteration loop
	
	return Nmet, Nsil
	
end # function equilibratefixedfO2


# --------------------------------------------------------------------------

function cmbpressure(element, Nmet, Nsil)

# Calculates mass, density, radius, gravity, pressure distribution inside model planet
# Returns CMB pressure

	# Parameters
	numsteps = 1000  # numsteps each in core and mantle
	
	# Calculate mass of metal and rock
	maw = getmaws(element)
	totalmass_metal = sum(maw .* Nmet) / 1000   # convert to kg
	totalmass_rock = sum(maw .* Nsil) / 1000   # convert to kg
	
	# Initialize
	step_metal = totalmass_metal / numsteps
	step_rock = totalmass_rock / numsteps
	steps = [1:numsteps]
	mass_metal = step_metal * steps[:]  # kg
	mass_rock = step_rock * steps[:]  # kg
	density_metal = steps[:] * 0.0  # kg/m3
	density_rock = steps[:] * 0.0  # kg/m3
	volume_metal = steps[:] * 0.0  # m3
	volume_rock = steps[:] * 0.0  # m3
	radius_metal = steps[:] * 0.0  # m
	radius_rock = steps[:] * 0.0  # m
	gravity_metal = steps[:] * 0.0  # m/s2
	gravity_rock = steps[:] * 0.0  # m/s2
	pressure_metal = steps[:] * 0.0  # GPa
	pressure_rock = steps[:] * 0.0  # GPa
	rms_pressure = 100.  # large initial misfit
	gravconst = 6.673e-11  # m^3 kg^-1 s^-2
	iterations = 0
	
	while rms_pressure > 0.01  # iterate until P input is close to P output
	
	  # Calculate core
	  for i in 1:numsteps
	    # Equation of state for metal parameterized against PREM outer core
	    density_metal[i] = 5851.3 + pressure_metal[i]^0.5 * 348.24
	    if i == 1
	    	volume_metal[i] = step_metal/density_metal[i]
	    else
	    	volume_metal[i] = volume_metal[i-1] + step_metal/density_metal[i]
	    end
	    radius_metal[i] = (volume_metal[i] *3/4/pi)^(1./3.)
	    gravity_metal[i] = gravconst * mass_metal[i] / radius_metal[i]^2.
	  end
	
	  # Calculate mantle
	  for i in 1:numsteps
	    # Equation of state for rock parameterized against PREM lower mantle
	    density_rock[i] = 3479.2 + pressure_rock[i]^0.5 * 176.944
	    if i == 1
	    	volume_rock[i] = volume_metal[numsteps] + step_rock/density_rock[i]
	    else
	    	volume_rock[i] = volume_rock[i-1] + step_rock/density_rock[i]
	    end
	    radius_rock[i] = (volume_rock[i] *3/4/pi)^(1./3.)
	    gravity_rock[i] = gravconst * (mass_metal[numsteps]+mass_rock[i]) / radius_rock[i]^2.
	  end
	
	  # Calculate pressure
	  old_pressure_rock = pressure_rock[:]
	  old_pressure_metal = pressure_metal[:]
	
	  pressure_rock[numsteps] = 0.0
	  # Integrate from surface inward
	  for i in (numsteps-1):-1:1
	    pressure_rock[i] = pressure_rock[i+1]+density_rock[i+1]*gravity_rock[i+1]*(radius_rock[i+1]-radius_rock[i])
	  end
	  pressure_metal[numsteps] = pressure_rock[1]+density_rock[1]*gravity_rock[1]*(radius_rock[1]-radius_metal[numsteps])
	  for i in (numsteps-1):-1:1
	    pressure_metal[i] = pressure_metal[i+1]+density_metal[i+1]*gravity_metal[i+1]*(radius_metal[i+1]-radius_metal[i])
	  end
	  pressure_rock = pressure_rock * 1.e-9  # convert to GPa
	  pressure_metal = pressure_metal * 1.e-9  # convert to GPa
	
	  rms_pressure = sum((pressure_rock - old_pressure_rock).^2) + sum((pressure_metal - old_pressure_metal).^2)
	  rms_pressure = (rms_pressure / (2.*numsteps))^0.5
	  iterations = iterations + 1
	  if iterations > 100
	    println("no convergence")
	    break
	  end
	end
	
	# CMB pressure
	return pressure_metal[numsteps]
	
end


# --------------------------------------------------------------------------

# LUCAS

# Main body of program. v4 has series of impactors, all same composition, and calculated CMB pressure
# Impactor composition from FischerEPSL2015 'reduced'

# Initialize elements and other data. 
time0 = time()
numsteps = 1000
element, Nmetinit, Nsilinit = impactors(numsteps)
println("impactors time = ",time()-time0," seconds")

time0 = time()
numelements = length(element)
maw = getmaws(element)
valence = getvalences(element)
oxidewt = maw + maw[1] .* valence / 2
Nbulkinit = Nsilinit + Nmetinit
println("elementdata time = ",time()-time0," seconds")

time0 = time()
# Initialize arrays
Nsil = zeros(numelements, numsteps)
Nmet = zeros(numelements, numsteps)
Nsilplanet = zeros(numelements)
Nmetplanet = zeros(numelements)
pressure = zeros(numsteps)
temperature = zeros(numsteps)
iw = zeros(numsteps)
iwtry = -3.5
println("array initialization time = ",time()-time0," seconds")

time0 = time()
# At each step, equilibrate bulk composition at new P,T and add to growing planet
for i in 1:numsteps
	if i == 1
		pressure[i] = 0
	else
		pressure[i] = 0.6 * cmbpressure(element, Nmetplanet, Nsilplanet)
	end
	# Liquidus T from Andrault EPSL 2011, parameterized in Fischer PhD thesis
	temperature[i] = 2066.6 + 25.939 * pressure[i] - 0.0453 * pressure[i]^2
	Nmeti, Nsili = equilibrate(Nbulkinit, pressure[i], temperature[i], iwtry)
	Nsil[:,i] = Nsili
	Nmet[:,i] = Nmeti
	Xsil = Nsili/sum(Nsili[2:numelements])
	Xmet = Nmeti/sum(Nmeti)
	iw[i] = 2*log10(Xsil[2]/Xmet[2])
	iwtry = iw[i]
	Nsilplanet = Nsilplanet + Nsili
	Nmetplanet = Nmetplanet + Nmeti
	Xsilplanet = Nsilplanet / sum(Nsilplanet[2:numelements])
	Xmetplanet = Nmetplanet / sum(Nmetplanet)
	iwplanet = 2*log10(Xsilplanet[2]/Xmetplanet[2])
#	println(i,"  ",iw[i],"  ",iwplanet,"  ",pressure[i])
end
time1 = time()
println("loop time: ",time1-time0," seconds")

# --------------------------------------------------------------------------


# Need to do:

# Make kdX calculation more flexible w.r.t. list of input elements and their order

# Include epsilon data for V, Cr, others

# Output results