using JuMP
using Gurobi
using StatsBase
using DataFrames
using CSV
using Bootstrap #External packages, must be added
include("SupportFunction.jl")
include("DataLoad.jl")
println("Leeeeroooy Jenkins")

#Esben's path
#cd("$(homedir())/Documents/GitHub/Thesis/Data")
#path = "$(homedir())/Documents/GitHub/Thesis/Data"

#Skipper's path
path = "/Users/SkipperAfRosenborg/Google Drive/DTU/10. Semester/Thesis/GitHubCode/Thesis/Data"
#HPC path
#path = "/zhome/9f/d/88706/SpecialeCode/Thesis/Data"
mainData = loadIndexDataNoDur(path)
fileName = path*"/Results/IndexData/IndexData"

"""
LASSO WITH WEIGHTED ERRORS
"""
using JuMP
using Gurobi
using StatsBase
using DataFrames
using CSV
using Bootstrap #External packages, must be added
include("SupportFunction.jl")
include("DataLoad.jl")
println("Leeeeroooy Jenkins")

#Esben's path
cd("$(homedir())/Documents/GitHub/Thesis/Data")
path = "$(homedir())/Documents/GitHub/Thesis/Data"
mainData = loadIndexDataNoDur(path)
fileName = path*"/Results/IndexData/IndexData"
mainDataArr = Array(mainData)

colNames = names(mainData)

nRows = size(mainDataArr)[1]
nCols = size(mainDataArr)[2]

mainXarr = mainDataArr[:,1:nCols-1]

# Transform #
mainXarr = expandWithTransformations(mainXarr)

# Standardize
standX = zScoreByColumn(mainXarr)
standY = zScoreByColumn(mainDataArr[:, nCols:nCols])
allData = hcat(standX, standY)
bCols = size(standX)[2]
nRows = size(standX)[1]
println(" \n \n \n solving Convex \n \n \n")

function solveLasso(Xtrain, Ytrain, Xpred, Ypred, gamma, errorWeights)
	bCols = size(Xtrain)[2]
	M = JuMP.Model(solver = GurobiSolver())
	@variables M begin
	        b[1:bCols]
	        t[1:bCols]
	        w
	end
	@objective(M,Min,0.5*w+gamma*ones(bCols)'*t)
	@constraint(M, soc, norm( [1-w;2*(errorWeights.*(Xtrain*b-Ytrain))] ) <= 1+w)
	@constraint(M,  b .<= t)
	@constraint(M, -t .<= b)

	solve(M)
	bSolved = getvalue(b)

	#In-Sample R-squared value
	errors = (Ytrain-Xtrain*bSolved)
	errorTotal = sum(errors[i]^2 for i=1:length(errors))
	errorsMean = Ytrain-mean(Ytrain)
	errorMeanTotal = sum(errorsMean[i]^2 for i=1:length(errorsMean))
	ISRsquared = 1-(errorTotal/errorMeanTotal)

	#OOS R-squared value
	oosErrors = Ypred - Xpred*bSolved
	oosErrorTotal = sum(oosErrors[i]^2 for i=1:length(oosErrors))
	oosErrorsMean = Ypred - mean(Ypred)
	oosErrorsMeanTotal = sum(oosErrorsMean[i]^2 for i=1:length(oosErrorsMean))
	OOSRsquared = 1-(oosErrorTotal/oosErrorsMeanTotal)

	return ISRsquared, OOSRsquared
end
trainingSize = 10
errorWeights = linspace(1,0, trainingSize)
testRuns = Int64(floor(1000/trainingSize))
nGammas = 500
predictions = 5
ISR = zeros(nGammas, testRuns)
OOSR = zeros(nGammas, testRuns)
gammaArray = logspace(0, 7, nGammas)
for r = 1:testRuns
	Xtrain = allData[(1+(r-1)*trainingSize):(trainingSize+(r-1)*trainingSize), 1:bCols]
	Ytrain = allData[(1+(r-1)*trainingSize):(trainingSize+(r-1)*trainingSize), bCols+1]
	Xpred  = allData[(trainingSize+(r-1)*trainingSize):(trainingSize+(r-1)*trainingSize+(predictions-1)), 1:bCols]
	Ypred  = allData[(trainingSize+(r-1)*trainingSize):(trainingSize+(r-1)*trainingSize+(predictions-1)), bCols+1]
	for g = 1:nGammas
		gamma = gammaArray[g]
		ISR[g, r], OOSR[g, r] = solveLasso(Xtrain, Ytrain, Xpred, Ypred, gamma, errorWeights)
	end
end

combinedArray = hcat(round.(gammaArray,3), ISR)
runCounter = collect(0:testRuns)
ISRcomb = vcat(runCounter', combinedArray)
writedlm("ISRsquaredweighted105.CSV", ISRcomb,",")
combinedArray = hcat(round.(gammaArray,3), OOSR)
OOSRcomb = vcat(runCounter', combinedArray)
writedlm("OOSRsquaredweighted105.CSV", OOSRcomb,",")

"""
LASSO
"""
using JuMP
using Gurobi
using StatsBase
using DataFrames
using CSV
using Bootstrap #External packages, must be added
include("SupportFunction.jl")
include("DataLoad.jl")
println("Leeeeroooy Jenkins")

#Esben's path
#cd("$(homedir())/Documents/GitHub/Thesis/Data")
#path = "$(homedir())/Documents/GitHub/Thesis/Data"

#Skipper's path
path = "/Users/SkipperAfRosenborg/Google Drive/DTU/10. Semester/Thesis/GitHubCode/Thesis/Data"

mainData = loadIndexDataNoDur(path)
<<<<<<< HEAD
=======
cd("$(homedir())/Documents/GitHub/Thesis/Data")
path = "$(homedir())/Documents/GitHub/Thesis/Data"
mainData = loadIndexDataNoDur(path)
fileName = path*"/Results/IndexData/IndexData"
>>>>>>> 4fe6210679f2fdd18244c00e7397df164bf0740d
=======
>>>>>>> parent of be6f44e... Merge branch 'master' of https://github.com/skipperafrosenborg/Thesis
mainDataArr = Array(mainData)

colNames = names(mainData)

nRows = size(mainDataArr)[1]
nCols = size(mainDataArr)[2]

mainXarr = mainDataArr[:,1:nCols-1]
mainYarr = mainDataArr[:, nCols:nCols]

# Transform with time elements
mainXarr = expandWithTime3612(mainXarr)

# Transform #
mainXarr = expandWithTransformations(mainXarr)

mainXarr = expandWithMAandMomentum(mainXarr, mainYarr, (nCols-1))

# Standardize
standX = zScoreByColumn(mainXarr)
standY = mainDataArr[:, nCols:nCols]
#standY = zScoreByColumn(mainDataArr[:, nCols:nCols])
allData = hcat(standX, standY)
bCols = size(standX)[2]
nRows = size(standX)[1]
println(" \n \n \n solving Convex \n \n \n")

function solveLasso(Xtrain, Ytrain, Xpred, Ypred, gamma)
	bCols = size(Xtrain)[2]
	M = JuMP.Model(solver = GurobiSolver(OutputFlag = 0))
	@variables M begin
	        b[1:bCols]
	        t[1:bCols]
	        w
	end
	@objective(M,Min,0.5*w+gamma*ones(bCols)'*t)
	@constraint(M, soc, norm( [1-w;2*(Xtrain*b-Ytrain)] ) <= 1+w)
	@constraint(M,  b .<= t)
	@constraint(M, -t .<= b)

	solve(M)
	bSolved = getvalue(b)

	#In-Sample R-squared value
	errors = (Ytrain-Xtrain*bSolved)
	errorTotal = sum(errors[i]^2 for i=1:length(errors))
	errorsMean = Ytrain-mean(Ytrain)
	errorMeanTotal = sum(errorsMean[i]^2 for i=1:length(errorsMean))
	ISRsquared = 1-(errorTotal/errorMeanTotal)
	#=
	#OOS R-squared value (for more predictions)
	oosErrors = Ypred - Xpred*bSolved
	oosErrorTotal = sum(oosErrors[i]^2 for i=1:length(oosErrors))
	oosErrorsMean = Ypred - mean(Ypred)
	oosErrorsMeanTotal = sum(oosErrorsMean[i]^2 for i=1:length(oosErrorsMean))
	OOSRsquared = 1-(oosErrorTotal/oosErrorsMeanTotal)
	=#

	#OOS R-squared value (for single prediction)
	oosErrors = Ypred - Xpred*bSolved
	oosErrorTotal = sum(oosErrors[i]^2 for i=1:length(oosErrors))
	oosSize = sum(Ypred[i]^2 for i=1:length(Ypred))
	OOSRsquared = sqrt(oosErrorTotal) #/ sqrt(oosSize)

	#Indicator Results
	YpredValue = Ypred[1]
	Yestimate = Xpred*bSolved
	YestimateValue = Yestimate[1]
	if YpredValue >= 0 && YestimateValue >= 0
		Indicator = 1
	elseif YpredValue < 0 && YestimateValue < 0
		Indicator = 1
	else
		Indicator = 0
	end

	if YestimateValue >= 0
		Indicator2 = 1
	elseif YestimateValue < 0
		Indicator2 = 0
	else
		Indicator2 = 0
	end

	return ISRsquared, OOSRsquared, Indicator, Indicator2
end

<<<<<<< HEAD
<<<<<<< HEAD
<<<<<<< HEAD
<<<<<<< HEAD
=======
>>>>>>> parent of be6f44e... Merge branch 'master' of https://github.com/skipperafrosenborg/Thesis
nGammas = 10
trainingSize = 12
predictions = 1
testRuns = nRows-trainingSize-predictions
ISR = zeros(testRuns, nGammas)
OOSR = zeros(testRuns, nGammas)
Indi = zeros(testRuns, nGammas)
gammaArray = logspace(0, 3, nGammas)
<<<<<<< HEAD
=======
nGammas = 5
trainingSize = 48
=======
=======
>>>>>>> bee04ad7b02cee3642894bc88ca8f415908cae9a
nGammas = 2
trainingSize = 12
>>>>>>> bee04ad7b02cee3642894bc88ca8f415908cae9a
predictions = 1
testRuns = nRows-trainingSize-predictions
ISR = zeros(nGammas, testRuns)
OOSR = zeros(nGammas, testRuns)
Indi = zeros(nGammas, testRuns)
Indi2 = zeros(nGammas, testRuns)
gammaArray = logspace(0, 1, nGammas)
<<<<<<< HEAD
>>>>>>> 4fe6210679f2fdd18244c00e7397df164bf0740d
=======
>>>>>>> parent of be6f44e... Merge branch 'master' of https://github.com/skipperafrosenborg/Thesis
=======
>>>>>>> bee04ad7b02cee3642894bc88ca8f415908cae9a

for r = 1:(nRows-trainingSize-predictions)
	Xtrain = allData[r:(trainingSize+(r-1)), 1:bCols]
	Ytrain = allData[r:(trainingSize+(r-1)), bCols+1]
	Xpred  = allData[(trainingSize+(r-1)+1):(trainingSize+(r-1)+predictions), 1:bCols]
	Ypred  = allData[(trainingSize+(r-1)+1):(trainingSize+(r-1)+predictions), bCols+1]
	for g = 1:nGammas
		gamma = gammaArray[g]
<<<<<<< HEAD
<<<<<<< HEAD
		ISR[r, g], OOSR[r, g], Indi[r, g] = solveLasso(Xtrain, Ytrain, Xpred, Ypred, gamma)
=======
		ISR[g, r], OOSR[g, r], Indi[g, r], Indi2[g, r] = solveLasso(Xtrain, Ytrain, Xpred, Ypred, gamma)
>>>>>>> bee04ad7b02cee3642894bc88ca8f415908cae9a
=======
		ISR[g, r], OOSR[g, r], Indi[g, r], Indi2[g, r] = solveLasso(Xtrain, Ytrain, Xpred, Ypred, gamma)
>>>>>>> bee04ad7b02cee3642894bc88ca8f415908cae9a
	end
	println("Row $r/",(nRows-trainingSize-predictions))
end

#=
for r = 1:testRuns
	Xtrain = allData[(1+(r-1)*trainingSize):(trainingSize+(r-1)*trainingSize), 1:bCols]
	Ytrain = allData[(1+(r-1)*trainingSize):(trainingSize+(r-1)*trainingSize), bCols+1]
	Xpred  = allData[(trainingSize+(r-1)*trainingSize):(trainingSize+(r-1)*trainingSize+(predictions-1)), 1:bCols]
	Ypred  = allData[(trainingSize+(r-1)*trainingSize):(trainingSize+(r-1)*trainingSize+(predictions-1)), bCols+1]
	for g = 1:nGammas
		gamma = gammaArray[g]
		ISR[g, r], OOSR[g, r] = solveLasso(Xtrain, Ytrain, Xpred, Ypred, gamma)
	end
end
=#
fileName = path*"/Results/IndexData/LassoTests/121TimeExpTA"
combinedArray = vcat((round.(gammaArray,3))', ISR)

runCounter = collect(0:testRuns)
ISRcomb = hcat(runCounter, combinedArray)
writedlm(fileName*"ISRsquared.CSV", ISRcomb,",")
combinedArray = vcat((round.(gammaArray,3))', OOSR)
OOSRcomb = hcat(runCounter, combinedArray)
writedlm(fileName*"OSSRsquared.CSV", OOSRcomb,",")

combinedArray = vcat((round.(gammaArray,3))', Indi)
Indicomb = hcat(runCounter, combinedArray)
writedlm(fileName*"Indi.CSV", Indicomb,",")
<<<<<<< HEAD
=======
combinedArray = hcat(round.(gammaArray,3), ISR)
runCounter = collect(0:testRuns)'
ISRcomb = vcat(runCounter, combinedArray)
writedlm("ISRTestNoDur121IndiPortfolio.CSV", ISRcomb,",")
combinedArray = hcat(round.(gammaArray,3), OOSR)
OOSRcomb = vcat(runCounter, combinedArray)
writedlm("OOSRTestNoDur121Portfolio.CSV", OOSRcomb,",")

combinedArray = hcat(round.(gammaArray,3), Indi)
Indicomb = vcat(runCounter, combinedArray)
<<<<<<< HEAD
<<<<<<< HEAD
writedlm("Indicator481TimeTAOtherNotStandY.CSV", Indicomb,",")
>>>>>>> 4fe6210679f2fdd18244c00e7397df164bf0740d
=======
writedlm("IndicatorTestNoDur121IndiPortfolio.CSV", Indicomb,",")

>>>>>>> bee04ad7b02cee3642894bc88ca8f415908cae9a
=======
>>>>>>> parent of be6f44e... Merge branch 'master' of https://github.com/skipperafrosenborg/Thesis

=======
writedlm("IndicatorTestNoDur121IndiPortfolio.CSV", Indicomb,",")

>>>>>>> bee04ad7b02cee3642894bc88ca8f415908cae9a
combinedArray = hcat(round.(gammaArray,3), Indi2)
Indicomb = vcat(runCounter, combinedArray)
writedlm("Indicator2TestNoDur121IndiPortfolio.CSV", Indicomb,",")

"""
LASSO with daily returns
"""
using JuMP
using Gurobi
using StatsBase
using DataFrames
using CSV
using Bootstrap #External packages, must be added
include("SupportFunction.jl")
include("DataLoad.jl")
println("Leeeeroooy Jenkins")

#Esben's path
cd("$(homedir())/Documents/GitHub/Thesis/Data")
path = "$(homedir())/Documents/GitHub/Thesis/Data/DailyReturns"

"""
REMEMBER: Select index to predict. Then change dataloading. Then decide
observations and predictions to include - change fileName accordingly.
"""

mainData = loadIndexDataDailyOther(path)
fileName = path*"/Results/IndexData/DailyReturnsXobservationsYpredictions"
mainDataArr = Array(mainData)

colNames = names(mainData)

nRows = size(mainDataArr)[1]
nCols = size(mainDataArr)[2]

mainXarr = mainDataArr[:,1:nCols-1]

# Transform #
mainXarr = expandWithTransformations(mainXarr)

# Standardize
standX = zScoreByColumn(mainXarr)
standY = zScoreByColumn(mainDataArr[:, nCols:nCols])
allData = hcat(standX, standY)
bCols = size(standX)[2]
nRows = size(standX)[1]
println(" \n \n \n solving Convex \n \n \n")

function solveLasso(Xtrain, Ytrain, Xpred, Ypred, gamma)
	bCols = size(Xtrain)[2]
	M = JuMP.Model(solver = GurobiSolver())
	@variables M begin
	        b[1:bCols]
	        t[1:bCols]
	        w
	end
	@objective(M,Min,0.5*w+gamma*ones(bCols)'*t)
	@constraint(M, soc, norm( [1-w;2*(Xtrain*b-Ytrain)] ) <= 1+w)
	@constraint(M,  b .<= t)
	@constraint(M, -t .<= b)

	solve(M)
	bSolved = getvalue(b)

	#In-Sample R-squared value
	errors = (Ytrain-Xtrain*bSolved)
	errorTotal = sum(errors[i]^2 for i=1:length(errors))
	isValue = sum(Ytrain[i]^2 for i=1:length(Ytrain))
	isErrorPercentage = (errorTotal / isValue) * 100
	errorsMean = Ytrain-mean(Ytrain)
	errorMeanTotal = sum(errorsMean[i]^2 for i=1:length(errorsMean))
	ISRsquared = 1-(errorTotal/errorMeanTotal)

	#OOS R-squared value
	oosErrors = Ypred - Xpred*bSolved
	oosErrorTotal = sum(oosErrors[i]^2 for i=1:length(oosErrors))
	oosValue = sum(Ypred[i]^2 for i=1:length(Ypred))
	oosErrorPercentage = (oosErrorTotal / oosValue)*100
	oosErrorsMean = Ypred - mean(Ypred)
	oosErrorsMeanTotal = sum(oosErrorsMean[i]^2 for i=1:length(oosErrorsMean))
	OOSRsquared = 1-(oosErrorTotal/oosErrorsMeanTotal)

	return isErrorPercentage, ISRsquared, oosErrorPercentage, OOSRsquared
end

testRuns = 300
nGammas = 50
trainingSize = 50
predictions = 5
ISR = zeros(testRuns, nGammas)
ISP = zeros(testRuns, nGammas)
OOSR = zeros(testRuns, nGammas)
OOSP = zeros(testRuns, nGammas)
varianceTrainArray = zeros(testRuns)
gammaArray = logspace(0, 2, nGammas)
for r = 1:testRuns
	Xtrain = allData[(1+(r-1)*trainingSize):(trainingSize+(r-1)*trainingSize), 1:bCols]
	Ytrain = allData[(1+(r-1)*trainingSize):(trainingSize+(r-1)*trainingSize), bCols+1]
	varianceTrainArray[r] = var(Ytrain)
	Xpred  = allData[(trainingSize+(r-1)*trainingSize):(trainingSize+(r-1)*trainingSize+(predictions-1)), 1:bCols]
	Ypred  = allData[(trainingSize+(r-1)*trainingSize):(trainingSize+(r-1)*trainingSize+(predictions-1)), bCols+1]
	for g = 1:nGammas
		gamma = gammaArray[g]
		ISP[r, g], ISR[r, g], OOSP[r, g], OOSR[r, g] = solveLasso(Xtrain, Ytrain, Xpred, Ypred, gamma)
	end
end

runCounter = collect(0:testRuns)
varArr = vcat("variance",varianceTrainArray)
ISR
gammaArray'
combinedArray = vcat(round.(gammaArray',3), ISR)
ISRcomb = hcat(runCounter, varArr, combinedArray)
writedlm("ISRsquaredDailyOther505.CSV", ISRcomb,",")

combinedArray = vcat(round.(gammaArray',3), ISP)
ISPcomb = hcat(runCounter, varArr, combinedArray)
writedlm("ISPDailyOther505.CSV", ISPcomb,",")

combinedArray = vcat(round.(gammaArray',3), OOSR)
OOSRcomb = hcat(runCounter, varArr, combinedArray)
writedlm("OOSRsquaredDailyOther505.CSV", OOSRcomb,",")

combinedArray = vcat(round.(gammaArray',3), OOSP)
OOSPcomb = hcat(runCounter, varArr, combinedArray)
writedlm("OOSPDailyOther505.CSV", OOSPcomb,",")

"""
LASSO with daily returns and weighted
"""
using JuMP
using Gurobi
using StatsBase
using DataFrames
using CSV
using Bootstrap #External packages, must be added
include("SupportFunction.jl")
include("DataLoad.jl")
println("Leeeeroooy Jenkins")

#Esben's path
cd("$(homedir())/Documents/GitHub/Thesis/Data")
path = "$(homedir())/Documents/GitHub/Thesis/Data/DailyReturns"

"""
REMEMBER: Select index to predict. Then change dataloading. Then decide
observations and predictions to include - change fileName accordingly.
"""

mainData = loadIndexDataDailyOther(path)
fileName = path*"/Results/IndexData/DailyReturnsXobservationsYpredictions"
mainDataArr = Array(mainData)

colNames = names(mainData)

nRows = size(mainDataArr)[1]
nCols = size(mainDataArr)[2]

mainXarr = mainDataArr[:,1:nCols-1]

# Transform #
mainXarr = expandWithTransformations(mainXarr)

# Standardize
standX = zScoreByColumn(mainXarr)
standY = zScoreByColumn(mainDataArr[:, nCols:nCols])
allData = hcat(standX, standY)
bCols = size(standX)[2]
nRows = size(standX)[1]
println(" \n \n \n solving Convex \n \n \n")

function solveLasso(Xtrain, Ytrain, Xpred, Ypred, gamma, errorWeights)
	bCols = size(Xtrain)[2]
	M = JuMP.Model(solver = GurobiSolver())
	@variables M begin
	        b[1:bCols]
	        t[1:bCols]
	        w
	end
	@objective(M,Min,0.5*w+gamma*ones(bCols)'*t)
	@constraint(M, soc, norm( [1-w;2*(errorWeights.*(Xtrain*b-Ytrain))] ) <= 1+w)
	@constraint(M,  b .<= t)
	@constraint(M, -t .<= b)

	solve(M)
	bSolved = getvalue(b)

	#In-Sample R-squared value
	errors = (Ytrain-Xtrain*bSolved)
	errorTotal = sum(errors[i]^2 for i=1:length(errors))
	isValue = sum(Ytrain[i]^2 for i=1:length(Ytrain))
	isErrorPercentage = (errorTotal / isValue) * 100
	errorsMean = Ytrain-mean(Ytrain)
	errorMeanTotal = sum(errorsMean[i]^2 for i=1:length(errorsMean))
	ISRsquared = 1-(errorTotal/errorMeanTotal)

	#OOS R-squared value
	oosErrors = Ypred - Xpred*bSolved
	oosErrorTotal = sum(oosErrors[i]^2 for i=1:length(oosErrors))
	oosValue = sum(Ypred[i]^2 for i=1:length(Ypred))
	oosErrorPercentage = (oosErrorTotal / oosValue)*100
	oosErrorsMean = Ypred - mean(Ypred)
	oosErrorsMeanTotal = sum(oosErrorsMean[i]^2 for i=1:length(oosErrorsMean))
	OOSRsquared = 1-(oosErrorTotal/oosErrorsMeanTotal)

	return isErrorPercentage, ISRsquared, oosErrorPercentage, OOSRsquared
end

testRuns = 300
nGammas = 50
trainingSize = 50
errorWeights = linspace(0, 1, trainingSize)
predictions = 5
ISR = zeros(testRuns, nGammas)
ISP = zeros(testRuns, nGammas)
OOSR = zeros(testRuns, nGammas)
OOSP = zeros(testRuns, nGammas)
varianceTrainArray = zeros(testRuns)
gammaArray = logspace(0, 2, nGammas)
for r = 1:testRuns
	Xtrain = allData[(1+(r-1)*trainingSize):(trainingSize+(r-1)*trainingSize), 1:bCols]
	Ytrain = allData[(1+(r-1)*trainingSize):(trainingSize+(r-1)*trainingSize), bCols+1]
	varianceTrainArray[r] = var(Ytrain)
	Xpred  = allData[(trainingSize+(r-1)*trainingSize):(trainingSize+(r-1)*trainingSize+(predictions-1)), 1:bCols]
	Ypred  = allData[(trainingSize+(r-1)*trainingSize):(trainingSize+(r-1)*trainingSize+(predictions-1)), bCols+1]
	for g = 1:nGammas
		gamma = gammaArray[g]
		ISP[r, g], ISR[r, g], OOSP[r, g], OOSR[r, g] = solveLasso(Xtrain, Ytrain, Xpred, Ypred, gamma, errorWeights)
	end
end

runCounter = collect(0:testRuns)
varArr = vcat("variance",varianceTrainArray)
ISR
gammaArray'
combinedArray = vcat(round.(gammaArray',3), ISR)
ISRcomb = hcat(runCounter, varArr, combinedArray)
writedlm("ISRsquaredDailyOtherWeighted2505.CSV", ISRcomb,",")

combinedArray = vcat(round.(gammaArray',3), ISP)
ISPcomb = hcat(runCounter, varArr, combinedArray)
writedlm("ISPDailyOtherWeighted2505.CSV", ISPcomb,",")

combinedArray = vcat(round.(gammaArray',3), OOSR)
OOSRcomb = hcat(runCounter, varArr, combinedArray)
writedlm("OOSRsquaredDailyOtherWeighted2505.CSV", OOSRcomb,",")

combinedArray = vcat(round.(gammaArray',3), OOSP)
OOSPcomb = hcat(runCounter, varArr, combinedArray)
writedlm("OOSPDailyOtherWeighted2505.CSV", OOSPcomb,",")

"""
LASSO with monthly returns
"""
using JuMP
using Gurobi
using StatsBase
using DataFrames
using CSV
using Bootstrap #External packages, must be added
include("SupportFunction.jl")
include("DataLoad.jl")
println("Leeeeroooy Jenkins")

#Esben's path
cd("$(homedir())/Documents/GitHub/Thesis/Data")
path = "$(homedir())/Documents/GitHub/Thesis/Data"

mainData =loadIndexDataNoDur(path)
mainDataArr = Array(mainData)

colNames = names(mainData)

nRows = size(mainDataArr)[1]
nCols = size(mainDataArr)[2]

mainXarr = mainDataArr[:,1:nCols-1]

# Transform #
mainXarr = expandWithTransformations(mainXarr)

# Standardize
standX = zScoreByColumn(mainXarr)
standY = zScoreByColumn(mainDataArr[:, nCols:nCols])
allData = hcat(standX, standY)
bCols = size(standX)[2]
nRows = size(standX)[1]
println(" \n \n \n solving Convex \n \n \n")

function solveLasso(Xtrain, Ytrain, Xpred, Ypred, gamma)
	bCols = size(Xtrain)[2]
	M = JuMP.Model(solver = GurobiSolver())
	@variables M begin
	        b[1:bCols]
	        t[1:bCols]
	        w
	end
	@objective(M,Min,0.5*w+gamma*ones(bCols)'*t)
	@constraint(M, soc, norm( [1-w;2*(Xtrain*b-Ytrain)] ) <= 1+w)
	@constraint(M,  b .<= t)
	@constraint(M, -t .<= b)

	solve(M)
	bSolved = getvalue(b)

	#In-Sample R-squared value
	errors = (Ytrain-Xtrain*bSolved)
	errorTotal = sum(errors[i]^2 for i=1:length(errors))
	isValue = sum(Ytrain[i]^2 for i=1:length(Ytrain))
	isErrorPercentage = (errorTotal / isValue) * 100
	errorsMean = Ytrain-mean(Ytrain)
	errorMeanTotal = sum(errorsMean[i]^2 for i=1:length(errorsMean))
	ISRsquared = 1-(errorTotal/errorMeanTotal)

	#OOS R-squared value
	oosErrors = Ypred - Xpred*bSolved
	oosErrorTotal = sum(oosErrors[i]^2 for i=1:length(oosErrors))
	oosValue = sum(Ypred[i]^2 for i=1:length(Ypred))
	oosErrorPercentage = (oosErrorTotal / oosValue)*100
	oosErrorsMean = Ypred - mean(Ypred)
	oosErrorsMeanTotal = sum(oosErrorsMean[i]^2 for i=1:length(oosErrorsMean))
	OOSRsquared = 1-(oosErrorTotal/oosErrorsMeanTotal)

	return isErrorPercentage, ISRsquared, oosErrorPercentage, OOSRsquared
end

testRuns = 100
nGammas = 50
trainingSize = 10

predictions = 5
ISR = zeros(testRuns, nGammas)
ISP = zeros(testRuns, nGammas)
OOSR = zeros(testRuns, nGammas)
OOSP = zeros(testRuns, nGammas)
varianceTrainArray = zeros(testRuns)
gammaArray = logspace(0, 2, nGammas)
for r = 1:testRuns
	Xtrain = allData[(1+(r-1)*trainingSize):(trainingSize+(r-1)*trainingSize), 1:bCols]
	Ytrain = allData[(1+(r-1)*trainingSize):(trainingSize+(r-1)*trainingSize), bCols+1]
	varianceTrainArray[r] = var(Ytrain)
	Xpred  = allData[(trainingSize+(r-1)*trainingSize):(trainingSize+(r-1)*trainingSize+(predictions-1)), 1:bCols]
	Ypred  = allData[(trainingSize+(r-1)*trainingSize):(trainingSize+(r-1)*trainingSize+(predictions-1)), bCols+1]
	for g = 1:nGammas
		gamma = gammaArray[g]
		ISP[r, g], ISR[r, g], OOSP[r, g], OOSR[r, g] = solveLasso(Xtrain, Ytrain, Xpred, Ypred, gamma)
	end
end

runCounter = collect(0:testRuns)
varArr = vcat("variance",varianceTrainArray)
ISR
gammaArray'
combinedArray = vcat(round.(gammaArray',3), ISR)
ISRcomb = hcat(runCounter, varArr, combinedArray)
writedlm("ISRsquaredMonthlyOther105.CSV", ISRcomb,",")

combinedArray = vcat(round.(gammaArray',3), ISP)
ISPcomb = hcat(runCounter, varArr, combinedArray)
writedlm("ISPMonthlyOther105.CSV", ISPcomb,",")

combinedArray = vcat(round.(gammaArray',3), OOSR)
OOSRcomb = hcat(runCounter, varArr, combinedArray)
writedlm("OOSRsquaredMonthlyOther105.CSV", OOSRcomb,",")

combinedArray = vcat(round.(gammaArray',3), OOSP)
OOSPcomb = hcat(runCounter, varArr, combinedArray)
writedlm("OOSPMonthlyOther105.CSV", OOSPcomb,",")


using DataFrames
using CSV
using HTTP
using Random
using LinearAlgebra
using Statistics
using Optim
using DataFramesMeta
using GLM
function create_grids()

    function xgrid(theta,xval)
        N      = length(xval)
        xub    = vcat(xval[2:N],Inf)
        xtran1  = zeros(N,N)
        xtran1c = zeros(N,N)
        lcdf   = zeros(N)
        for i=1:length(xval)
            xtran1[:,i]   = (xub[i] .>= xval).*(1 .- exp.(-theta*(xub[i] .- xval)) .- lcdf)
            lcdf        .+= xtran1[:,i]
            xtran1c[:,i] .+= lcdf
        end
        return xtran1, xtran1c
    end

    zval   = collect([.25:.01:1.25]...)
    zbin   = length(zval)
    xval   = collect([0:.125:25]...)
    xbin   = length(xval)
    tbin   = xbin*zbin
    xtran  = zeros(tbin,xbin)
    xtranc = zeros(xbin,xbin,xbin)
    for z=1:zbin
        xtran[(z-1)*xbin+1:z*xbin,:],xtranc[:,:,z] = xgrid(zval[z],xval)
    end

    return zval,zbin,xval,xbin,xtran
end
#Question 1

    url = "https://raw.githubusercontent.com/OU-PhD-Econometrics/fall-2020/master/ProblemSets/PS5-ddc/busdata.csv"
    df = CSV.read(HTTP.get(url).body,DataFrame)
    # create bus id variable
df = @transform(df, bus_id = 1:size(df,1))

#---------------------------------------------------
# reshape from wide to long (must do this twice be-
# cause DataFrames.stack() requires doing it one
# variable at a time)
#---------------------------------------------------
# first reshape the decision variable
dfy = @select(df, :bus_id,:Y1,:Y2,:Y3,:Y4,:Y5,:Y6,:Y7,:Y8,:Y9,:Y10,:Y11,:Y12,:Y13,:Y14,:Y15,:Y16,:Y17,:Y18,:Y19,:Y20,:RouteUsage,:Branded)
dfy_long = DataFrames.stack(dfy, Not([:bus_id,:RouteUsage,:Branded]))
rename!(dfy_long, :value => :Y)
dfy_long = @transform(dfy_long, time = kron(collect([1:20]...),ones(size(df,1))))
select!(dfy_long, Not(:variable))

# next reshape the odometer variable
dfx = @select(df, :bus_id,:Odo1,:Odo2,:Odo3,:Odo4,:Odo5,:Odo6,:Odo7,:Odo8,:Odo9,:Odo10,:Odo11,:Odo12,:Odo13,:Odo14,:Odo15,:Odo16,:Odo17,:Odo18,:Odo19,:Odo20)
dfx_long = DataFrames.stack(dfx, Not([:bus_id]))
rename!(dfx_long, :value => :Odometer)
dfx_long = @transform(dfx_long, time = kron(collect([1:20]...),ones(size(df,1))))
select!(dfx_long, Not(:variable))

# join reshaped df's back together
df_long = leftjoin(dfy_long, dfx_long, on = [:bus_id,:time])
sort!(df_long,[:bus_id,:time])

#question 2

theta=glm(@formula(Y ~ Odometer * Odometer^2 * RouteUsage * RouteUsage^2 * Branded * time * time^2), df_long, Binomial(), LogitLink())

#:::::::::::::::::::::::::::::::::::::::::::::::::::
# Q3:generate state transition matrices
#:::::::::::::::::::::::::::::::::::::::::::::::::::
url = "https://raw.githubusercontent.com/OU-PhD-Econometrics/fall-2020/master/ProblemSets/PS5-ddc/busdata.csv"
df = CSV.read(HTTP.get(url).body,DataFrame)
Y = Matrix(df[:,[:Y1,:Y2,:Y3,:Y4,:Y5,:Y6,:Y7,:Y8,:Y9,:Y10,:Y11,:Y12,:Y13,:Y14,:Y15,:Y16,:Y17,:Y18,:Y19,:Y20]])
X = Matrix(df[:,[:Odo1,:Odo2,:Odo3,:Odo4,:Odo5,:Odo6,:Odo7,:Odo8,:Odo9,:Odo10,:Odo11,:Odo12,:Odo13,:Odo14,:Odo15,:Odo16,:Odo17,:Odo18,:Odo19,:Odo20]])
Z = Vector(df[:,:RouteUsage])
B = Vector(df[:,:Branded])
N = size(Y,1)
T = size(Y,2)
Xstate = Matrix(df[:,[:Xst1,:Xst2,:Xst3,:Xst4,:Xst5,:Xst6,:Xst7,:Xst8,:Xst9,:Xst10,:Xst11,:Xst12,:Xst13,:Xst14,:Xst15,:Xst16,:Xst17,:Xst18,:Xst19,:Xst20]])
Zstate = Vector(df[:,:Zst])
zval,zbin,xval,xbin,xtran = create_grids()

#create data frame#
log_DF=DataFrame()
log_DF.Odometer=kron(ones(zbin),xval)
log_DF.RouteUsage =kron(ones(xbin),zval)
log_DF.Branded=zeros(size())
log_DF.time=zeros(size())
#function
function values(Xstate,Zstate,xtran,xbin,zbin,xval)
    FV=zeros(size(xtran,1),2,T+1)
   #Nested loops
    for t=2:T
        for b=0:1
FVT1[i,t] = (xtran[row1,:].-xtran[row0,:])'*FV[row0:row0+xbin-1,B[i]+1,t+1]
df_long = @transform(df_long,fv = fvt1)
theta_hat_ccp_glm = glm(@formula(Y ~ Odometer + Branded),
df_long, Binomial(), LogitLink(), offset=df_long.fv)


export cost, max_degree, depth, balance

indices(c::ContractCommand) = c.left_idxs, c.right_idxs, c.output_idxs

cost(c::AbstractCommand) = 0

"""
    cost(cmd::ContractCommand)::Number

Function to calculate the cost of the given contraction in FLOPS
"""
function cost(cmd::ContractCommand)::Number
    a, b, c = indices(cmd)
    batched_idxs = intersect(a, b, c)
    a = setdiff(a, batched_idxs)
    b = setdiff(b, batched_idxs)
    c = setdiff(c, batched_idxs)
    common = intersect(a, b)
    remaining = symdiff(a, b)
    1 << (length(common) + length(remaining) + length(batched_idxs))
end

max_degree(c::ContractCommand) = maximum(length.(indices(c)))
cost(n::ComputeNode) = cost(n.op) + sum(cost.(children(n)), init=0)
max_degree(n::ComputeNode) = maximum(max_degree.(children(n)), init=0)
max_degree(n::ComputeNode{ContractCommand}) = maximum([max_degree(n.op), max_degree.(children(n))...])
depth(n::ComputeNode) = 1 + maximum(depth.(children(n)), init=0)
Base.length(n::ComputeNode) = 1 + sum(length.(children(n)), init=0)
balance(n::ComputeNode) = ceil(log2(length(n)))/depth(n)

"""Implement show for compute node to display some useful metrics"""
function Base.show(io::IO, c::ComputeNode{T}) where T
    print(io, "ComputeNode{$(T)}: children: $(length(c)-1), ",
              "depth: $(depth(c)), balance: $(balance(c)), ",
              "max_degree: $(max_degree(c))")
end


using DelimitedFiles

input = joinpath(@__DIR__, "input")
map = readdlm(input, String)

function solve(map, slope, bounds)
    hits = 0
    pos = CartesianIndex(1, 1)
    pos += slope

    while pos[1] <= bounds[1]
        x = pos[2]
        y = pos[1]
        hits += (map[y][x] == '#' ? 1 : 0)
        pos += slope
        pos[2] > bounds[2] && (pos -= CartesianIndex(0, bounds[2]))
    end
    return hits
end

const bounds = CartesianIndex(length(map), length(map[1]))

p1 = solve(map, CartesianIndex(1, 3), bounds)

const slopes = (CartesianIndex(1, 1), CartesianIndex(1, 3), CartesianIndex(1, 5), CartesianIndex(1, 7), CartesianIndex(2, 1))

p2 = prod(solve(map, slope, bounds) for slope in slopes)

@assert(p1 == 299)
@assert(p2 == 3621285278)

println("-----------------------------------------------------------------------")
println("toboggan trajectory -- part one :: $p1")
println("toboggan trajectory -- part two :: $p2")
println("-----------------------------------------------------------------------")


abstract type AbstractPolicy end

abstract type AbstractValuePolicy <: AbstractPolicy end

struct ϵGreedy <: AbstractValuePolicy
    ϵ::Float64
end

function sample(ap::ϵGreedy, values, rng)
    if rand(rng) > ap.ϵ
        return findmax(values)[2]
    else
        return rand(rng, 1:length(values))
    end
end

function get_prob(ap::ϵGreedy, values, action)
    if action == findmax(values)[2]
        return 1 - ap.ϵ + (ap.ϵ / length(values))
    else
        return ap.ϵ / length(values)
    end
end


# connectivity functions specific to this project
using Base.Test

function TopographicWeights2D{T}(width_pre, width_post, kernel::Array{T,2}=reshape([1],1,1))
	# lists to return
	pre = Int[]
	post = Int[]
	weight = T[]

	(small_w, big_w) = sort([width_post, width_pre])

	kw=size(kernel,1)
	ks=(kw-1)/2 # num units to left or right

	big_d = [big_w big_w]
	small_d = [small_w small_w]

	# mapping from many tiles to few tiles using kernel
	for br=1:big_w
		for bc=1:big_w
			# (br, bc) = center of bigger map
			(sr, sc) = int(ceil([br, bc]/big_w * small_w)) # corresponding tile on the smaller map
			small_idx = sub2ind(small_d, sr, sc)
			# now for each element covered by kernel on the bigger tiles,
			# map back to that smaller tile
			for kr = 1:kw # [kr, kc] = location in kernel coordinates
				for kc = 1:kw
					rrange=(br-kw):(br+kw)
					crange=(bc-kw):(bc+kw)
					# location in big-layer coords
					brr = int(rrange[kr])
					bcc = int(crange[kc])
					K = kernel[kr,kc]
					if (0<brr<big_w+1) && (0<bcc<big_w+1) && K!=0 # no edge spillover, kernel is defined
						big_idx = sub2ind(big_d, brr, bcc)
						push!(pre, small_idx)
						push!(post, big_idx)
						push!(weight, K)
					end
				end
			end
		end
	end

	# big is actually pre, and small is actually post
	# - flip flop the two
	if (width_pre > width_post)
		(post, pre) = (pre, post)
	end
	return (pre, post, weight)
end

function rand_uniform(a, b)
    a + rand()*(b - a)
end

# currently unused
function all_to_all(pre_N, post_N, stype::SynapseType, stdp_enabled)
    syn = SpikingSynapse[]
    for i=1:(pre_N*post_N)
        w = 1
        delay = 4
        (s,t) = ind2sub((pre_N,post_N),i)
        s = SpikingSynapse(stype[i],s,t,float(w),float(delay))
        push!(syn,s)
    end
    SpikingConnection(syn,stdp_enabled)
end

# computes initial weights from Photoreceptor layers to LGN
function p_lgn_connectivity(pre_N, post_N, stype::SynapseType, base_weight, stdp_enabled)
	# this now acts as a pooling kernel

	(small_N, big_N) = sort([pre_N, post_N])
	kw = sqrt(big_N/small_N)
	if !isinteger(kw)
		error("kernel width ratio must be an Int64 - got $(kw) instead")
	end
	kw = int(kw)
	kernel = ones(kw, kw)
	(pre, post, w) = TopographicWeights2D(sqrt(pre_N), sqrt(post_N), kernel)
	syn = Array(SpikingSynapse,length(pre))
	for i=1:length(pre)
		delay = 1.0 # ms
		w = base_weight * rand_uniform(0.5,2)
		@test pre[i] <= pre_N
		@test post[i] <= post_N
		syn[i] = SpikingSynapse(stype,pre[i],post[i],w,delay)
	end
	SpikingConnection(syn,stdp_enabled)
end

# TODO - below generators need to be fixed!
# topographic weights from lgn layers to v1 (cortex) layers
# and from v1->v1
function forward_connectivity(pre_N, post_N, stype::SynapseType, base_weight, stdp_enabled)
	kernel = [
	    0 1 0
	    1 1 1
	    0 1 0
	];
	(pre, post, w) = TopographicWeights2D(sqrt(pre_N), sqrt(post_N), kernel)

	syn = Array(SpikingSynapse,length(pre))
	for i=1:length(pre)
		delay = 1.0
		w = base_weight * rand_uniform(0.5,2)
		@test pre[i] <= pre_N
		@test post[i] <= post_N

		syn[i] = SpikingSynapse(stype,pre[i],post[i],w,delay)
	end
	SpikingConnection(syn,stdp_enabled)
end

# lateral connections in cortex layer
function lateral_connectivity(N, w_pot::Float64, w_dep::Float64)
	# topographic lateral connectivity kernels
	# all weights should be positive, since we don't want synapses to flip-flop between e & i via STDP
	e = EXC # excitatory
	i = INH # inhibitory

	type_kernel=
	[
	    0 0 i i i 0 0
	    0 i e e e i 0
	    i e e e e e i
	    i e e 0 e e i
	    i e e e e e i
	    0 i e e e i 0
	    0 0 i i i 0 0
	]

	width = sqrt(N)
	(pre, post, stype) = TopographicWeights2D(width, width, type_kernel)

	# delays, in ms
	A = 1; B = 2; C = 3; D = 4
	delay_kernel=[
	    0 0 D D D 0 0
	    0 D C C C D 0
	    D C B A B C D
	    D C A 0 A C D
	    D C B A B C D
	    0 D C C C D 0
	    0 0 D D D 0 0
	]
	(pre, post, delay) = TopographicWeights2D(width, width, delay_kernel)

	syn = Array(SpikingSynapse,length(pre))
	for i=1:length(pre)
		w = (stype[i] == EXC) ? w_pot : w_dep
		w *= rand_uniform(0.5,2)
		syn[i] = SpikingSynapse(stype[i], pre[i], post[i], w, float(delay[i]))
	end
	SpikingConnection(syn,false)
end


# f(x,y)=x-y'

FXMYtranspose(x::Float64,y::Float64)=Fxmy(x,y)
FXMYtranspose(x::Float64,y)=Fxmy(x,y')
FXMYtranspose(x,y::Float64)=Fxmy(x,y)
FXMYtranspose(x,y)=Fxmy(x,y')
FXMYtranspose_inplace(value,auxvalue,x,y)=Fxmy_inplace(value,auxvalue,x,y')

function DXMYtranspose(derivativeIDX,f_c,faux_c,grad_c,grad_n,x::Array,y::Array)
    if derivativeIDX==1
        if size(x)==(1,1)
            axpy!(1.0,[sum(grad_c)],grad_n)
        else
            axpy!(1.0,grad_c,grad_n)
        end
    elseif derivativeIDX==2
        if size(y)==(1,1)
            axpy!(-1.0,[sum(grad_c)],grad_n)
        else
            axpy!(-1.0,grad_c',grad_n)
        end
    end
end

if PROC=="GPU"

    function FXMYtranspose_inplace(value::CudaArray,auxvalue,x::CudaArray,y::CudaArray)
        if size(x)==(1,1)
            gfill!(value,x);
            CUBLAS.geam!('N','T',1.0,value,-1.0,y,value)
        elseif size(y)==(1,1)
            gfill!(value,y)
            CUBLAS.geam!('N','N',1.0,x,-1.0,value,value)
        else
            CUBLAS.geam!('N','T',1.0,x,-1.0,y,value)
        end
    end

    function DXMYtranspose(derivativeIDX,f_c,faux_c,grad_c,grad_n,x::CudaArray,y::CudaArray)
        if derivativeIDX==1
            if size(x)==(1,1)
                axpy!(1.0,sum(grad_c),grad_n)
            else
                tmp=CudaArray(Float64,size(grad_c))
                CUBLAS.geam!('T','N',1.0,grad_c,0.0,tmp,tmp) # tmp=grad_c'
                axpy!(1.0,tmp,grad_n); free(tmp)
            end
        elseif derivativeIDX==2
            if size(y)==(1,1)
                axpy!(-1.0,sum(grad_c),grad_n)
            else
                tmp=CudaArray(Float64,size(grad_c))
                CUBLAS.geam!('T','N',1.0,grad_c,0.0,tmp,tmp) # tmp=grad_c'
                axpy!(-1.0,tmp,grad_n); free(tmp)
            end
        end
    end

end

Derivative[FXMYtranspose]=DXMYtranspose
Inplace[FXMYtranspose]=FXMYtranspose_inplace

export FXMYtranspose


# Tests for Revenue Russell DEA Models
@testset "RevenueRussellDEAModel" begin

    # Test using Book data
    X = [1; 1; 1; 1; 1; 1; 1; 1];
    Y = [7 7; 4 8; 8 4; 3 5; 3 3; 8 2; 6 4; 1.5 5];
    P = [1 1; 1 1; 1 1; 1 1; 1 1; 1 1; 1 1; 1 1];

    # Revenue Russell
    revenuerddf = dearevenuerddf(X, Y, P, :ERG)

    @test typeof(revenuerddf) == RevenueReverseDDFDEAModel

    @test nobs(revenuerddf) == 8
    @test ninputs(revenuerddf) == 1
    @test noutputs(revenuerddf) == 2
    @test ismonetary(revenuerddf) == false

    @test efficiency(revenuerddf) == efficiency(revenuerddf, :Economic)
    @test efficiency(revenuerddf, :Economic)   ≈ [0.0; 0.25; 0.25; 13/28; 4/7; 2/3; 11/35; 9/11]
    @test efficiency(revenuerddf, :Technical)  ≈ [0.0; 0.0; 0.0; 13/28; 4/7; 1/3; 11/35; 37/55]
    @test efficiency(revenuerddf, :Allocative) ≈ [0.0; 0.25; 0.25; 0.0; 0.0; 1/3; 0.0; 8/55]
    @test normfactor(revenuerddf) ≈ [14.0; 8.0; 8.0; 168/13; 14.0; 6.0; 140/11; 55/6] atol = 1e-5

    # With monetary = true
    revenurrddfmon = dearevenuerddf(X, Y, P, :ERG, monetary = true)

    @test ismonetary(revenurrddfmon) == true

    @test efficiency(revenurrddfmon) ≈ efficiency(revenuerddf) .* normfactor(revenuerddf)
    @test efficiency(revenurrddfmon, :Technical) ≈ efficiency(revenuerddf, :Technical) .* normfactor(revenuerddf)
    @test efficiency(revenurrddfmon, :Allocative) ≈ efficiency(revenuerddf, :Allocative) .* normfactor(revenuerddf)

    # With CRS (Same result in this example)
    revenurrddfcrs = dearevenuerddf(X, Y, P, :ERG, rts = :CRS)

    @test efficiency(revenurrddfcrs, :Economic)   ≈ [0.0; 0.25; 0.25; 13/28; 4/7; 2/3; 11/35; 9/11]
    @test efficiency(revenurrddfcrs, :Technical)  ≈ [0.0; 0.0; 0.0; 13/28; 4/7; 1/3; 11/35; 37/55]
    @test efficiency(revenurrddfcrs, :Allocative) ≈ [0.0; 0.25; 0.25; 0.0; 0.0; 1/3; 0.0; 8/55]

    # Print
    show(IOBuffer(), revenuerddf)

    # Test errors
    @test_throws DimensionMismatch dearevenuerddf([1; 2 ; 3], [4 ; 5], [1; 1; 1], :ERG) #  Different number of observations
    @test_throws DimensionMismatch dearevenuerddf([1; 2; 3], [4; 5; 6], [1; 2; 3; 4], :ERG) # Different number of observation in prices
    @test_throws DimensionMismatch dearevenuerddf([1 1; 2 2; 3 3 ], [4; 5; 6], [1 1 1; 2 2 2; 3 3 3], :ERG) # Different number of output prices and outputs
    @test_throws ArgumentError dearevenuerddf([1; 2; 3], [4; 5; 6], [4; 5; 6], :Error) # Invalid associated efficiency measure
    @test_throws ArgumentError dearevenuerddf([1; 2; 3], [4; 5; 6], [1; 2; 3], :ERG, rts = :Error) # Invalid returns to scale

end


I = (1,1)
G = (1,2)
αs = linspace(1/6,4,100)

N = 400
h = 1e-6
σoff = -0.2

NEV = 4
L = 20

@everywhere include("exciton-hamiltonian.jl")

######################### Just a Test #####################

@everywhere function energies(α,I,G,p0,L,N,NEV,σoff)
    BLAS.set_num_threads(1)

    σ = σoff - 2*α^2 - maximum(I)^2/2 - maximum(I)*maximum(p0)

    println("σ = $σ")
    println("p0 = $p0")
    run(`hostname`)

    @time Hm = Hamiltonian(α,I,G,p0,L,N)
    @time solution = eigs(Hm,nev=NEV,tol=1e-6,maxiter=10000,which=:LM,sigma=σ)

    return real(solution[1])
end

p0 = [((0,0,0),α) for α in αs]
px = [((h,0,0),α) for α in αs]
py = [((0,h,0),α) for α in αs]

@sync begin
    @async global E0 = pmap(x->energies(x[2],I,G,x[1],L,N,NEV,σoff),p0)
    @async global Ex = pmap(x->energies(x[2],I,G,x[1],L,N,NEV,σoff),px)
    @async global Ey = pmap(x->energies(x[2],I,G,x[1],L,N,NEV,σoff),py)
end

DEx = (Ex - E0)/h
DEy = (Ey - E0)/h

using JLD
save("outdir/derivatives-$N-$NEV.jld","αs",αs,"DEx",DEx,"DEy",DEy)


####################################################
#                                                  #
# Compute the no policy / no intervention outcome  #
#                                                  #
####################################################

"""
    nopolicy(; TT=52)
Simulate the model in Chari, Kirpalani, and Phelan (2021) for `TT` periods under no government response.
The output is a [`Dict`](https://docs.julialang.org/en/v1/base/collections/#Dictionaries) containing
the necessary variables and time-series for generating the figures.
"""
function nopolicy(; TT=52)

    # Parameters
    ISO = 0;
    TEST = 1;
    OPT = 0; #0 for fmincon, 1 for global, 2 for interp
    beta = .99; # Discount factor
    gamma = 7/18; # Exit rate
    delta = .005; # Death rate
    sigma = 2; # CES parameter on production function
    B1 = 1; # Productivity of work sector
    B2 = .1; # Productivity of home sector
    nu = 1; # Testing cost parameter
    c = (1/.01)*50*52/63000 .* YY_noep(B1, B2, sigma, 1); # Testing cost parameter
    z = 7500*52/63000 .* YY_noep(B1, B2, sigma, 1); # Treatment cost
    d = 15*3*52*(1/.96^(1/52)-1); # Death cost
    xi = .8; # Infected productivity los
    tautilde = 13/18*1/2; # Prob. of becoming symptomatic
    pi2 = .01; # Infection rate at home
    pi1 = 1.5; # Infection rate at work
    thets = .0044; #Signal probability from susceptible person
    theti = .38; #Signal probability from infected person
    ut = 1;

    # Combine into (immutable) NamedTuple
    Par = (ISO=ISO, TEST=TEST, OPT=OPT, beta=beta, gamma=gamma, delta=delta, sigma=sigma, B1=B1, B2=B2,
    nu=nu, c=c, z=z, d=d, xi=xi, tautilde=tautilde, pi2=pi2, pi1=pi1, thets=thets, theti=theti, ut=ut);

    # Empty vectors for simulation results
    I0 = 0.02;
    S0 = 1-I0/(1-Par.tautilde);
    R0  = 0;
    IT = zeros(TT);
    ST = zeros(TT);
    RT = zeros(TT);
    qT = zeros(TT);
    lamT = zeros(TT);
    muT = zeros(TT);
    #Del = zeros(TT);
    VT = zeros(TT);
    YT = zeros(TT);
    L1 = zeros(TT);
    L2 = zeros(TT);
    Rknot = zeros(TT);
    Test = zeros(TT);
    IN = zeros(TT);
    C1T = zeros(TT);
    C2T = zeros(TT);
    IT[1] = I0;
    ST[1] = S0;
    RT[1] = R0;
    VV = 0.0;
    VVC = zeros(TT);
    masstest = zeros(TT);
    mctest = zeros(TT);

    # Simulate optimal path
    for i in 1:TT
        S = ST[i];
        I = IT[i];
        R = RT[i];
        IN[i] = 1-S-R;
        qT[i] = 0;
        lamT[i] = .9901;
        muT[i] = .9901;
        Del = 0;
        LI1 = lamT[i].*(S+(1-Del).*I)+muT[i].*(1-Par.delta).*R;
        LI2 = (1-lamT[i]).*(S+(1-Del).*I)+(1-muT[i]).*(1-Par.delta).*R;
        if LI1>0
            Ipr1 = (lamT[i].^2 .* Par.pi1.*(1-Del).*I)./LI1;
            RIpr1 = (lamT[i].^2 .* Par.pi1)./LI1;
        else
            Ipr1 = 0;
            RIpr1 = 0;
        end
        if LI2>0
            Ipr2 = ((1-lamT[i]).^2 .* Par.pi2.*(1-Del).*I)./LI2;
            RIpr2 = ((1-lamT[i]).^2 .* Par.pi2)./LI2;
        else
            Ipr2 = 0;
            RIpr2 = 0;
        end
        Ipr = Ipr1 + Ipr2;
        Ipr = min(Ipr,1);
        if i<TT
            IT[i+1] = min(max((1-Del).*(1-Par.tautilde).*(1-Par.gamma).*I+S.*Ipr,0),1);
            ST[i+1] = min(max((1-Ipr).*S,0),1);
            RT[i+1] = max(min(R + Par.gamma.*IN[i],1),0);
        end

        YT[i] = YY(Par,S,I,R,qT[i],lamT[i],muT[i]);
        C1T[i],C2T[i] = CC(Par,S,I,R,qT[i],lamT[i],muT[i]);
        L1[i] = (lamT[i].*(S+Par.xi.*(1-Del).*I)+muT[i].*(1-Par.delta).*R);
        L2[i] = ((1-lamT[i]).*(S+Par.xi.*(1-Del).*I)+(1-muT[i]).*(1-Par.delta).*R);

        if I>0
            Rknot[i]=(1 ./ (Par.gamma)) .* Ipr ./ (1-S-R);
        else
            Rknot[i] = 0;
        end
        VV += Par.beta.^(i-1) * (log(YT[i]-C1T[i])-C2T[i]);
        VVC[i] = VV;
    end

    # Calculate output variables
    VV = VV + Par.beta^TT/(1-Par.beta)*VTerm(Par,ST[TT],IT[TT],RT[TT]);
    cumdeaths = cumsum(Par.gamma.*Par.delta.*max.(1 .- ST .- RT,0));
    Totaldeath =cumdeaths[TT];
    Ycumsum = YT[1]-YY_noep(Par,1);
    for i = 2:TT
        Ycumsum = Ycumsum + Par.beta.^(i-1) .* (YT[i]-YY_noep(Par,1));
    end
    Ycumsum = (1-Par.beta).*Ycumsum;
    compvar2 = 1/YY_noep(Par,1)*exp((1-Par.beta)*VV);
    compvar2 = 1-compvar2;
    #suscepss = ST[TT];
    #Iss = 1-ST[TT]-RT[TT];
    #maxI = maximum(1 .- ST .- RT);
    #R01 = Rknot[1];
    #R026 = Rknot[26];
    #R052 = Rknot[52];

    return Dict("Par" => Par, "TT" => TT, "ST" => ST, "IT" => IT, "RT" => RT,
    "YT" => YT, "C1T" => C1T, "Rknot" => Rknot, "masstest" => masstest, "mctest" => mctest)

end


# This file is a part of BAT.jl, licensed under the MIT License (MIT).

using BAT
using Test

using Distributed, Random
using ArraysOfArrays, Distributions, PDMats, StatsBase

@testset "mcmc_sample" begin
    mvec = [-0.3, 0.3]
    cmat = [1.0 1.5; 1.5 4.0]
    Σ = @inferred PDMat(cmat)
    mv_dist = MvNormal(mvec, Σ)
    density = @inferred BAT.DistributionDensity(mv_dist)
    bounds = @inferred BAT.HyperRectBounds([-5, -8], [5, 8], BAT.reflective_bounds)
    nsamples_per_chain = 50000
    nchains = 4

    # algorithmMW = @inferred MetropolisHastings() TODO: put back the @inferred
    algorithmMW = MetropolisHastings()
    @test BAT.mcmc_compatible(algorithmMW, BAT.GenericProposalDist(mv_dist), BAT.NoParamBounds(2))

    samples, stats, chains = bat_sample(
        PosteriorDensity(density, bounds), (nsamples_per_chain, nchains), algorithmMW
    )

    # ToDo: Should be able to make this exact, for MH sampler:
    @test length(samples) == nchains * nsamples_per_chain

    @test begin
        stats = @inferred bat_stats(samples)
        stats2 = @inferred bat_stats(BAT.MCMCBasicStats(samples))
        stats.mode == stats2.mode && stats.mean ≈ stats2.mean && stats.cov ≈ stats2.cov
    end

    cov_samples = cov(flatview(samples.params), FrequencyWeights(samples.weight), 2; corrected=true)
    mean_samples = mean(flatview(samples.params), FrequencyWeights(samples.weight); dims = 2)

    @test isapprox(mean_samples, mvec; rtol = 0.15)
    @test isapprox(cov_samples, cmat; rtol = 0.15)

    algorithmPW = @inferred MetropolisHastings(ARPWeighting())

    samples, stats, chains = bat_sample(
        PosteriorDensity(mv_dist, bounds), (nsamples_per_chain, nchains), algorithmPW
    )

    @test begin
        stats = @inferred bat_stats(samples)
        stats2 = @inferred bat_stats(BAT.MCMCBasicStats(samples))
        stats.mode == stats2.mode && stats.mean ≈ stats2.mean && stats.cov ≈ stats2.cov
    end

    @test samples.params[findmax(samples.logdensity)[2]] == stats.mode

    cov_samples = cov(flatview(samples.params), FrequencyWeights(samples.weight), 2; corrected=true)
    mean_samples = mean(flatview(samples.params), FrequencyWeights(samples.weight); dims = 2)

    @test isapprox(mean_samples, mvec; rtol = 0.15)
    @test isapprox(cov_samples, cmat; rtol = 0.15)
end


@generated function promote_rule(::Type{V1}, ::Type{V2}) where {V1 <: Variable, V2 <: Variable}
    if name(V1) < name(V2)
        :(Monomial{(V1(), V2())})
    else
        :(Monomial{(V2(), V1())})
    end
end

promote_rule(::Type{M}, ::Type{V}) where {V <: Variable, M <: Monomial} = promote_rule(Monomial{(V(),), (1,)}, M)
promote_rule(V::Type{<:Variable}, M::Type{<:Monomial}) = promote_rule(M, V)

function _promote_monomial(::Type{<:Monomial{V1}}, ::Type{<:Monomial{V2}}) where {V1, V2}
    varnames = union(name.(V1), name.(V2))
    sort!(varnames)
    vars = Expr(:tuple, [Expr(:call, Expr(:curly, :Variable, Expr(:quote, n))) for n in varnames]...)
    :(Monomial{$vars})
end

@generated function promote_rule(::Type{M1}, ::Type{M2}) where {M1 <: Monomial, M2 <: Monomial}
    _promote_monomial(M1, M2)
end

promote_rule(::Type{Term{T, M}}, ::Type{V}) where {V <: Variable, T, M <: Monomial} = Term{T, promote_type(V, M)}
promote_rule(V::Type{<:Variable}, T::Type{<:Term}) = promote_rule(T, V)

function promote_rule(::Type{Term{T, M2}}, ::Type{M1}) where {M1 <: Monomial, T, M2 <: Monomial}
    Term{T, promote_type(M1, M2)}
end
promote_rule(M::Type{<:Monomial}, T::Type{<:Term}) = promote_rule(T, M)

function _promote_term(::Type{Term{T1, M1}}, ::Type{Term{T2, M2}}) where {T1, M1, T2, M2}
    Term{promote_type(T1, T2), promote_type(M1, M2)}
end

promote_rule(T1::Type{<:Term}, T2::Type{<:Term}) = _promote_term(T1, T2)

function promote_rule(::Type{V}, ::Type{S}) where {S, V <: Variable}
    Term{S, Monomial{(V(),)}}
end

function promote_rule(::Type{M}, ::Type{S}) where {S, M <: Monomial}
    Term{S, M}
end

function _promote_any_term(::Type{S}, ::Type{Term{T, M}}) where {S, T, M <: Monomial}
    Term{promote_type(S, T), M}
end

promote_rule(t::Type{<:Term}, s) = _promote_any_term(s, t)


@testset "cosineknn" begin
    A = load_matrix_network("bfs_example")
    S = cosineknn(A,2)

    # A = speye(Int64,4)
    A = sparse(1.0I,4,4)
    A[4,1] = 1
    CKNN = cosineknn(A,2)
    @test ((CKNN[4,1] *2)^2 - 2) <= 1e-7

    # A = speye(Int64,4)
    A = sparse(1.0I,4,4)
    A[4,1] = 1
    CKNN = cosineknn(MatrixNetwork(A),2)
    @test ((CKNN[4,1] *2)^2 - 2) <= 1e-7
end


using LKH
using Test

@testset "LKH.jl" begin
    @testset "Symmetric TSP" begin
        M = [
            0  16   7  14
            16   0   3   5
            7   3   0  16
            14   5  16   0
        ]
        opt_tour, opt_len = solve_tsp(M)
        @test opt_len == 29

        opt_tour, opt_len = solve_tsp(M; INITIAL_TOUR_ALGORITHM="GREEDY", RUNS=5)
        @test opt_len == 29

        M = rand(1:9, 300, 300)
        M = M + M'
        for i in 1:300
            M[i, i] = 0
        end
        solve_tsp(M)
    end

    @testset "Asymmetric TSP" begin
        M = rand(1:9, 100, 100)
        solve_tsp(M)
    end

    @testset "Coordinates" begin
        n_nodes = 10
        x = rand(n_nodes) .* 10000
        y = rand(n_nodes) .* 10000
        opt_tour, opt_len = solve_tsp(x, y; dist="EUC_2D")
        opt_tour, opt_len = solve_tsp(x, y; dist="MAN_2D")
        opt_tour, opt_len = solve_tsp(x, y; dist="MAX_2D")
        opt_tour, opt_len = solve_tsp(x, y; dist="GEO")
    end

    @testset "TSP Input File" begin
        opt_tour, opt_len = solve_tsp("gr17.tsp")
        @test opt_len == 2085
    end

    @testset "CVRP Input File" begin
        opt_tour, opt_len = solve_tsp("P-n16-k8.vrp")
        @test opt_len == 450
        @show opt_tour

        opt_tour, opt_len = solve_tsp("P-n19-k2.vrp")
        @test opt_len == 212
        @show opt_tour

        opt_tour, opt_len = solve_tsp("P-n101-k4.vrp", TIME_LIMIT=1)
        @test opt_len >= 681
        @show opt_tour
    end
end


using Dates
using Distributions

#start_time = now()
#end_time = now()
@testset "Black Scholes with Banyan with size=$size" for size in [128_000_000, 256_000_000, 512_000_000]
    use_job_for_testing(sample_rate = 128) do

        # NOTE: 64M works but 128M doesn't (and also doesn't batch the I/O)
        price = BanyanArrays.fill(4.0, size)
        strike = BanyanArrays.fill(4.0, size)
        t = BanyanArrays.fill(4.0, size)
        rate = BanyanArrays.fill(4.0, size)
        vol = BanyanArrays.fill(4.0, size)

        d1 = map(
            (p, s, t, r, v) -> ((log(p / s) + (r + v^2 * 0.5) * t) / (v * sqrt(t))),
            price,
            strike,
            t,
            rate,
            vol,
        )

        d2 = map((d1, v, t) -> (d1 - (v * sqrt(t))), d1, vol, t)

        call = map(
            (d1, d2, p, s, t, r) ->
                ((cdf(Normal(), d1) * p) - (cdf(Normal(), d2) * s * exp(-r * t))),
            d1,
            d2,
            price,
            strike,
            t,
            rate,
        )

        call_sum = sum(call)
        call = nothing
        res = collect(call_sum)
        # println(res)
        @show res
        #res = collect(res)
        #end = now()
        #@test typeof(res) == Base.Vector{Float64}
        #@test all(v->v==3.999999985812889, res)
    end
end
#end_time = now()
#println(end_time - start_time)

@testset "Black Scholes without Banyan with size=$size" for size in [128_000, 128_000_000, 512_000_000]
    use_job_for_testing(sample_rate = 128) do

        price = Base.fill(4.0, size)
        strike = Base.fill(4.0, size)
        t = Base.fill(4.0, size)
        rate = Base.fill(4.0, size)
        vol = Base.fill(4.0, size)

        d1 = Base.map(
            (p, s, t, r, v) -> ((log(p / s) + (r + v^2 * 0.5) * t) / (v * sqrt(t))),
            price,
            strike,
            t,
            rate,
            vol,
        )

        d2 = Base.map((d1, v, t) -> (d1 - (v * sqrt(t))), d1, vol, t)

        call = Base.map(
            (d1, d2, p, s, t, r) ->
                ((cdf(Normal(), d1) * p) - (cdf(Normal(), d2) * s * exp(-r * t))),
            d1,
            d2,
            price,
            strike,
            t,
            rate,
        )

        call_sum = Base.sum(call)
        res = Base.collect(call_sum)
        println(res)
        # @test typeof(res) == Base.Vector{Float64}
        # @test all(v->v==3.999999985812889, res)
    end
end

@testset "Black Scholes without Banyan vectorized with size=$size" for size in [
    128_000,
    128_000_000,
    512_000_000,
]
    use_job_for_testing(sample_rate = 128) do

        price = Base.fill(4.0, size)
        strike = Base.fill(4.0, size)
        t = Base.fill(4.0, size)
        rate = Base.fill(4.0, size)
        vol = Base.fill(4.0, size)

        d1 = (log.(price ./ strike) .+ (rate .+ vol .^ 2 .* 0.5) .* t) ./ (vol .* sqrt.(t))

        d2 = d1 .- (vol .* sqrt.(t))

        call =
            (cdf.(Normal(), d1) .* price) - (cdf.(Normal(), d2) .* strike .* exp.(-rate .* t))

        call_sum = Base.sum(call)
        res = Base.collect(call_sum)
        println(res)
    end
end


"Abstract linear operator."
abstract type AbstractLinearTransform{T} end

"Operator dimensions are either integer or tuples of integers."
DimType = Union{Int,Tuple{Vararg{Int}}}

"General linear operator"
struct LinearTransform{T} <: AbstractLinearTransform{T}
    dimsRange::DimType
    dimsDomain::DimType
    prod::Function
    tprod::Function
end

"Constructor with default Float64 element type"
LinearTransform(dimsRange::DimType, dimsDomain::DimType, prod::Function, tprod::Function) = LinearTransform{Float64}(dimsRange, dimsDomain, prod, tprod)

eltype(::AbstractLinearTransform{T}) where {T} = T
isreal(op::AbstractLinearTransform) = eltype(op) <: Real


using IntervalMatrices, Test, LinearAlgebra, SparseArrays

using IntervalMatrices: _truncated_exponential_series,
                        horner,
                        scale_and_square,
                        correction_hull,
                        input_correction
include("models.jl")

include("constructors.jl")
include("arithmetic.jl")
include("setops.jl")
include("exponential.jl")
include("affine.jl")


using Distributions
using ZeroInflation
using Test
import Turing: BernoulliLogit
import Turing: BinomialLogit

@testset "ZeroInflation.jl" begin
    @testset "constructing ZeroInflated" begin
        d = zeroinflated(Normal(), 0.25)
        @test d isa ZeroInflation.ZeroInflated
        @test d.mix isa Bernoulli
        @test d.mix.p == 0.25
        @test d.original isa Normal

        d = zeroinflated(Normal(), BernoulliLogit(1.0))
        @test d.mix isa BinomialLogit
        @test d.mix.logitp == 1.0
    end

    @testset "accessing ZeroInflated" begin
        d = zeroinflated(Poisson(2.0), 0.5)
        @test params(d) == (2.0, 0.5)
        @test minimum(d) == 0
        @test maximum(d) == Inf

        d = zeroinflated(Normal(), 0.1)
        @test params(d) == (0.0, 1.0, 0.1)
        @test minimum(d) == -Inf
        @test maximum(d) == Inf

        d = zeroinflated(Uniform(-0.5, 0.5), 0.1)
        @test params(d) == (-0.5, 0.5, 0.1)
        @test minimum(d) == -0.5
        @test maximum(d) == 0.5

        d = zeroinflated(Uniform(-2.0, -1.0), 0.1)
        @test minimum(d) == -2.0
        @test maximum(d) == 0.0
    end

    @testset "rand" begin
        @test rand(zeroinflated(Normal(), 1.0)) == 0.0
        @test rand(zeroinflated(Normal(), 0.0)) != 0.0

        @test rand(zeroinflated(Normal(), 0.0)) isa AbstractFloat
        @test rand(zeroinflated(Poisson(1.0), 0.0)) isa Integer
    end

    @testset "logpdf" begin
        d = zeroinflated(Poisson(2.0), 0.2)
        samp = rand(d, 1000)
        pointwise_loglik = logpdf.(d, samp)
        @test length(pointwise_loglik) == length(samp)
        @test isapprox(sum(pointwise_loglik), logpdf(d, samp))
    end
end


@testset begin "Connectivities"
  # basic tests
  c = connect((1,2,3), Triangle)
  @test issimplex(c)
  @test paramdim(c) == 2
  @test indices(c) == (1,2,3)
  @test materialize(c, P2[(0,0),(1,0),(0,1)]) == Triangle(P2(0,0), P2(1,0), P2(0,1))

  # tuple from other collections
  c = connect(Tuple([1,2,3]), Triangle)
  @test paramdim(c) == 2
  @test indices(c) == (1,2,3)
  @test materialize(c, P2[(0,0),(1,0),(0,1)]) == Triangle(P2(0,0), P2(1,0), P2(0,1))

  # incorrect number of vertices for polytope
  @test_throws AssertionError connect((1,2,3,4), Triangle)
  @test_throws AssertionError connect((1,2,3), Quadrangle)

  # heterogeneous collections
  c = connect.([(1,2,6,5),(2,4,6),(4,3,5,6),(1,5,3)], Ngon)
  @test c[1] isa Connectivity{Quadrangle}
  @test c[2] isa Connectivity{Triangle}
  @test c[3] isa Connectivity{Quadrangle}
  @test c[4] isa Connectivity{Triangle}
  @test issimplex.(c) == [false,true,false,true]

  # ommitting polytope type means polygon or segment
  @test connect((1,2)) isa Connectivity{Segment}
  @test connect((1,2,3)) isa Connectivity{Triangle}
  @test connect((1,2,3,4)) isa Connectivity{Quadrangle}
  @test connect((1,2,3,4,5)) isa Connectivity{Pentagon}
  @test connect((1,2,3,4,5,6)) isa Connectivity{Hexagon}

  # polyhedron connectivities
  c = connect((1,2,3,4), Tetrahedron)
  @test issimplex(c)
  @test paramdim(c) == 3
  @test indices(c) == (1,2,3,4)
  points = P3[(0,0,0),(1,0,0),(0,1,0),(0,0,1)]
  @test materialize(c, points) == Tetrahedron(points...)
end


function build_gram_matrix(q::Vector{MOI.VariableIndex},
                           basis::AbstractPolynomialBasis)
    return build_gram_matrix([MOI.SingleVariable(vi) for vi in q], basis)
end
function build_gram_matrix(q::Vector,
                           basis::AbstractPolynomialBasis)
    return GramMatrix(MultivariateMoments.SymMatrix(q, length(basis)),
                      basis)
end

function add_gram_matrix(model::MOI.ModelLike, matrix_cone_type::Type,
                         basis::AbstractPolynomialBasis)
    Q, cQ = MOI.add_constrained_variables(model, matrix_cone(matrix_cone_type, length(basis)))
    q = build_gram_matrix(Q, basis)
    return q, Q, cQ
end
function add_gram_matrix(model::MOI.ModelLike, matrix_cone_type::Type,
                         bases::Vector{<:AbstractPolynomialBasis})
    qQcQ = add_gram_matrix.(model, matrix_cone_type, bases)
    return SparseGramMatrix(getindex.(qQcQ, 1)), getindex.(qQcQ, 2), getindex.(qQcQ, 3)
end

function build_moment_matrix(q::Vector,
                             basis::AbstractPolynomialBasis)
    return MomentMatrix(MultivariateMoments.SymMatrix(q, length(basis)),
                        basis)
end

struct SOSPolynomialSet{DT <: AbstractSemialgebraicSet,
                        MT <: MP.AbstractMonomial,
                        MVT <: AbstractVector{MT},
                        CT <: Certificate.AbstractCertificate} <: MOI.AbstractVectorSet
    domain::DT
    monomials::MVT
    certificate::CT
end
MOI.dimension(set::SOSPolynomialSet) = length(set.monomials)


"""
Test minimize function
"""

#push!(LOAD_PATH,"../src/")
using joptimise

function rosenbrock(x)
    f = (1 - x[1])^2 + 100*(x[2] - x[1]^2)^2
    return f
end

function rosenbrock!(g, x)
    # compute objective
    f = rosenbrock(x)
    # constraint
    g[1] = x[1]^2 + x[2]^2 - 1.0
    return f
end

# initial guess
x0 = [4.0; 4.0]
# bounds on variables
lx = [-5.0; -5.0]
ux = [5.0; 5.0]
# bounds on constriants
lg = [0.0]
ug = [0.0]
# number of constraints
ng = 1

## run minimizer with IPOPT
ip_options = Dict(
    "max_iter" => 2500,   # 1500 ~ 2500
    "tol" => 1e-6
)

xopt, fopt, info = minimize(rosenbrock!, x0, ng; lx=lx, ux=ux, lg=lg, ug=ug, solver="ipopt", options=ip_options, derivatives=ForwardAD());

println("Done with IPOPT!")
println(info)
println(xopt)

## run minimizer with SNOPT
sn_options = Dict(
    "Major feasibility tolerance" => 1.e-6,
    "Major optimality tolerance"  => 1.e-6,
    "Minor feasibility tolerance" => 1.e-6,
    "Major iterations limit" => 1000,
    "Major print level" => 1,
)

xopt, fopt, info = minimize(rosenbrock!, x0, ng; lx=lx, ux=ux, lg=lg, ug=ug, solver="snopt", options=sn_options, derivatives=ForwardAD());

println("Done with SNOPT!")
println(info)
println(xopt)


"""
    struct AngularBasis <: BasisSet

Defines a set of states with angular momentum `j`. The dimensionality if `2j+1`.

# Interface

* `m(b::AngularBasis, i)`: returns the ``m`` quantum number of the `i`th state.
"""
struct AngularBasis <: BasisSet
    j :: HalfInteger
    function AngularBasis(j)
        @assert j >= 0
        new(j)
    end
end
Base.length(b::AngularBasis) = convert(Int, 2 * b.j + 1)
m(b::AngularBasis, idx) = -b.j + (idx - 1)

"""
    struct JOperator{K} <: LinearOperator{AngularBasis}

Type to represent the angular momentum operators ``J_z``, ``J_+`` and ``J_-``, acting on
states with defined total angular momentum ([`AngularBasis`](@ref)).
"""
struct JOperator{K} <: LinearOperator{AngularBasis}
    b :: AngularBasis
end
function JOperator(b::AngularBasis, kind::Symbol)
    kind in [:z, :+, :-] || throw(ArgumentError("Bad kind: $kind"))
    JOperator{kind}(b)
end
basis(op::JOperator) = op.b
function apply(i::Integer, op::JOperator{K}, j::Integer) :: ComplexF64 where K
    angJ = convert(ComplexF64, op.b.j)
    if K == :z
        (i == j) ? m(op.b, j) : 0.0
    elseif K == :+
        (i == j + 1) ? sqrt((angJ - m(op.b, j)) * (angJ + m(op.b, j) + 1)) : 0.0
    elseif K == :-
        (i == j - 1) ? sqrt((angJ + m(op.b, j)) * (angJ - m(op.b, j) + 1)) : 0.0
    else
        error("Unreachable reached.")
    end
end

struct JOperatorSet{B <: BasisSet, T <: LinearOperator{B}, U <: LinearOperator{B}, V <: LinearOperator{B}}
    J₊ :: T
    J₋ :: U
    Jz :: V
end
function JOperatorSet(JOp, b::BasisSet)
    Jz = JOp(b, :z)
    J₊, J₋ = JOp(b, :+), JOp(b, :-)
    JOperatorSet{typeof(b),typeof(J₊),typeof(J₋),typeof(Jz)}(J₊, J₋, Jz)
end
JOperatorSet(b::BasisSet) = JOperatorSet(JOperator, b)

function JOperatorSet(f::Function, js::Vararg{JOperatorSet})
    length(js) > 0 || throw(ArgumentError("Must provide at least one J operator"))
    Angular.JOperatorSet(
        f((j.J₊ for j in js)...),
        f((j.J₋ for j in js)...),
        f((j.Jz for j in js)...)
    )
end

basis(op::JOperatorSet) = basis(op.Jz)

Jx(op::JOperatorSet) = (op.J₊ + op.J₋) / 2
Jy(op::JOperatorSet) = (op.J₊ - op.J₋) / (2im)
J2(op::JOperatorSet) = J2Operator(op)

"""
    findj(λ) -> HalfInteger

Convert and ``J^2`` operator eigenvalue into the corresponding half-integer ``j`` value.
"""
findj(λ) = HalfInteger(isqrt(1 + round(Int, real(4λ))) - 1, 2)

struct J2Operator{B <: BasisSet, JS <: JOperatorSet{B}, EOP <: LinearOperator{B}} <: LinearOperator{B}
    J :: JS
    J2 :: EOP

    function J2Operator(jops :: JOperatorSet)
        b = basis(jops)
        J2 = (jops.J₊*jops.J₋ + jops.J₋*jops.J₊)/2 + jops.Jz^2
        new{typeof(b),typeof(jops),typeof(J2)}(jops, J2)
    end
end
basis(op::J2Operator) = basis(op.J)
apply(i::Integer, op::J2Operator, j::Integer) = apply(i, op.J2, j)
function matrix(j2op::J2Operator)
    J₊ = matrix(j2op.J.J₊)
    J₋ = matrix(j2op.J.J₋)
    Jz = matrix(j2op.J.Jz)
    (J₊*J₋ .+ J₋*J₊)./2 .+ Jz^2
end


using Minibatch
using NNlib
using Base.Test

#################
# Test networks #
#################

W = rand(5, 5)
b = rand(5)
f(x) = tanh.(x)
g(x) = tanh.(x .+ b)
h(x) = tanh.(W*x .+ b)

xs = [rand(5, 3), rand(5, 4)]
x1 = VectorBatch(xs, (false, true))
x2 = MaskedBatch(xs, (false, true))

@test x1 == VectorBatch(x2)
@test x2 == MaskedBatch(x1)

for fn in (f, g, h)
    @test MaskedBatch(fn(x1)) ≈ fn(x2)
end

ys = [rand(4, 5), rand(2, 5)]
y1 = VectorBatch(ys, (true, false))
y2 = MaskedBatch(ys, (true, false))

z1 = y1 * x1; z2 = y2 * x2
@test MaskedBatch(z1) ≈ z2
@test VectorBatch(z2) ≈ z1

for axis in (1, 2)
    @test MaskedBatch(softmax(z1, axis)) ≈ softmax(z2, axis)
end

function attention(q, k, v) # q::D×N, k::D×M, v::D×M
    alpha = k'*q # ::M×N
    # TODO tri mask
    return v*softmax(alpha, 1) # ::D×N
end

d_qk = 3
d_v = 4
enc = [rand(d_v, 6), rand(d_v, 5)]
dec = [rand(d_v, 3), rand(d_v, 7)]
enc1 = VectorBatch(enc, (false, true))
enc2 = MaskedBatch(enc, (false, true))
dec1 = VectorBatch(dec, (false, true))
dec2 = MaskedBatch(dec, (false, true))

Wq = rand(d_qk, d_v)
Wk = rand(d_qk, d_v)
Wv = rand(d_v, d_v)

out1 = attention(Wq*dec1, Wk*enc1, Wv*enc1)
out2 = attention(Wq*dec2, Wk*enc2, Wv*enc2)
@test MaskedBatch(out1) ≈ out2


using MyExample
using Test

@testset "MyExample.jl" begin
    #2x + 3y
    @test my_f(2, 1) == 7
    @test my_f(2, 3) == 13
    @test my_f(1, 3) == 11
end

@testset "Derivative Tests" begin
    @test derivative_of_my_f(2, 1) == 2
end


module Wynn
    using SymPy

    # data structure to store the epsilon table
    # ϵ_ij indexed by ints i & j, stored within dict.
    struct EpsilonTable{T}
        series::T
        terms::Vector{T}
        etable::Dict{Tuple{Int,Int},T}
    end

    function EpsilonTable(terms::Vector{T}; simplified::Bool = true) where T<:Union{Float64,Sym}
        # maximum occuring i & j index
        max_ind = length(terms)

        # setting base case j = -1
        etable = Dict((i, -1) => eltype(terms)(0) for i in 0:max_ind)

        # setting base case j = 0, i = 0:max_ind-1
        merge!(etable, Dict((i-1, 0) => sum(terms[1:i]) for i in 1:max_ind))

        # setting base case i = -j-1, j even
        merge!(etable, Dict((-j-1, 2j) => eltype(terms)(0) for j in 0:max_ind))

        # setting base case for odd j
        merge!(etable, Dict((-j-1, 2j-1) => eltype(terms)(0) for j in 0:max_ind))

        # recursive calculations, j>=1, i = floor(Int,-j/2):(max_ind-j-1)
        for j in 1:2*(max_ind-1), i in floor(Int,-j/2):(max_ind-j-1)#j-1
            ϵ_ij = etable[i+1, j-2] + 1 / (etable[i+1, j-1] - etable[i, j-1])
            eltype(terms) <: Sym && simplified ? push!(etable, (i, j) =>
                simplify(ϵ_ij)) : push!(etable, (i, j) => ϵ_ij)
        end
        series = sum(terms)
        EpsilonTable(series, terms, etable)
    end

    EpsilonTable(terms::Vector{Int}; simplified::Bool = true) = EpsilonTable(convert(Vector{Float64}, terms); simplified = simplified)

    export EpsilonTable
end