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"
#mainData = loadConcrete(path)
#fileName = path*"/Results/Concrete/Run1/Concrete_Run"
mainData = loadHousingData(path)
#fileName = path*"/Results/HousingData/HousingData"
#mainData = loadCPUData(path)
path = "/Users/SkipperAfRosenborg/Google Drive/DTU/10. Semester/Thesis/GitHubCode/Results"
fileName = path*"/HousingData/Housing"

#Reset HPC path
#path = "/zhome/9f/d/88706/SpecialeCode/Thesis/ML"
#cd(path)

mainDataArr = Array(mainData)
halfRows = Int64(floor(size(mainDataArr)[1]/2))
train, test = splitDataIn2(mainDataArr, halfRows, size(mainData)[1])
halfRows = Int64(floor(size(test)[1]/2))
test, valid = splitDataIn2(mainDataArr, halfRows, size(test)[1])

colNames = names(mainData)
trainingData = Array(train)
testData = Array(test)
valiData = Array(valid)
nRows = size(trainingData)[1]
nCols = size(trainingData)[2]

#=
mainData, testData = loadElevatorData(path)
dataSize = size(mainData)
colNames = names(mainData)
trainingData = Array(mainData)
testData = Array(testData)
nRows = size(trainingData)[1]
nCols = size(trainingData)[2]
=#

### STAGE 1 ###
println("STAGE 1 INITIATED")
#Define solve
m = JuMP.Model(solver = GurobiSolver(OutputFlag =0 ))

#Add binary variables variables
@variable(m, 0 <= z[1:nCols] <= 1, Bin )

#Calculate highly correlated matix
HC = cor(trainingData)

#Define objective function
@objective(m, Max, sum(z[i] for i=1:length(z)))

#Define max correlation and constraints
rho = 0.8
for i=1:length(z)
	for j=1:length(z)
		if i != j
			if HC[i,j] >= rho
				@constraint(m,z[i]+z[j] <= 1)
			end
		end
	end
end

#Get solution status
status = solve(m);

#Get objective value
println("Objective value kMax: ", getobjectivevalue(m))
kmax = getobjectivevalue(m)

#Get solution value
#zSolved = getvalue(z)
println("STAGE 1 DONE")

### STAGE 2 ###
println("STAGE 2 INITIATED")
println("Standardizing data")

#Split data by X and y
y = trainingData[:,nCols]
X = trainingData[:,1:nCols-1]
yTest = testData[:,nCols]
XTest = testData[:,1:nCols-1]
yVali = valiData[:,nCols]
XVali = valiData[:,1:nCols-1]

# TRANSFORMATIONS ###
X = expandWithTransformations(X)
XTest = expandWithTransformations(XTest)
XVali = expandWithTransformations(XVali)

# Standardize
standX = zScoreByColumn(X)
standY = zscore(y)
standXTest = zScoreByColumn(XTest)
standYTest = zscore(yTest)
standXVali = zScoreByColumn(XVali)
standYVali = zscore(yVali)

# Output for Lasso regression in R
writedlm(fileName*"XTrain.CSV",standX,",")
writedlm(fileName*"YTrain.CSV",standY,",")
writedlm(fileName*"XTest.CSV",standXTest,",")
writedlm(fileName*"YTest.CSV",standYTest,",")
writedlm(fileName*"XVali.CSV",standXVali,",")
writedlm(fileName*"YVali.CSV",standYVali,",")

#Initialise values for check later
bCols = size(X)[2]
kValue = []
RsquaredValue = []
bSolved = []
bestBeta = []
warmStartBeta = []
warmstartZ = []

startIter = 1

bigM = 100

stage2Model = JuMP.Model(solver = GurobiSolver(TimeLimit = 30))
SSTO = sum((standY[i]-mean(standY))^2 for i=1:length(standY))
amountOfGammas = 2
#Spaced between 0 and half the SSTO since this would then get SSTO*absSumOfBeta which would force everything to 0
gammaArray = log10.(logspace(0, log10.(SSTO), amountOfGammas))

bbdata = NodeData[]

function buildStage2(standX, standY, kmax)
	HC = cor(X)
	HCPairCounter = 0
	#println("Building model")
	#Define parameters and model
	stage2Model = JuMP.Model(solver = GurobiSolver(TimeLimit = 20, OutputFlag = 0));
	gamma = 10

	#Define variables
	@variable(stage2Model, b[1:bCols]) #Beta values

	#Define binary variable (5b)
	@variable(stage2Model, 0 <= z[1:bCols] <= 1, Bin )

	@variable(stage2Model, v[1:bCols]) # auxiliary variables for abs

	@variable(stage2Model, T) #First objective term
	@variable(stage2Model, G) #Second objective term

	#Define objective function (5a)
	@objective(stage2Model, Min, T+G)

	#println("Trying to implement new constraint")
	xSquareExpr = @expression(stage2Model, 0*b[1]^2)
	for l = 1:bCols
		coef = 0
		for j = 1:nRows
		   coef += standX[j,l]^(2)
		end
		append!(xSquareExpr, @expression(stage2Model, coef*b[l]^2))
	end
	#println("Implemented x^2")

	ySquaredExpr = @expression(stage2Model, 0*b[1])
	for j = 1:nRows
		append!(ySquaredExpr,@expression(stage2Model, standY[j,1]^2))
	end
	#println("Implemented y^2")

	simpleBetaExpr = @expression(stage2Model, 0*b[1])
	for l = 1:bCols
		coef = 0
		for j = 1:nRows
			coef += -1*2*standX[j, l]*standY[j]
		end
		append!(simpleBetaExpr, @expression(stage2Model, coef*b[l]))
	end
	#println("Implemented simpleBetaExpr")

	crossBetaExpr = @expression(stage2Model, 0*b[1]*b[2])
	iter = 1
	for l = 1:bCols
		for k = (l + 1):bCols
			coef = 0
			for j = 1:nRows
				coef += 2*standX[j,l]*standX[j,k]
			end
			append!(crossBetaExpr, @expression(stage2Model, coef*b[l,1]*b[k,1]))
		end
		#println("Finished loop $l/$bCols")
	end
	#println("Implemented crossBetaExpr")
	totalExpr = @expression(stage2Model, crossBetaExpr+simpleBetaExpr+xSquareExpr+ySquaredExpr)
	@constraint(stage2Model, quadConst, totalExpr <= T)
	#println("Successfully added quadratic constraints")

	#Define constraints (5c)
	@constraint(stage2Model, conBigMN, -1*b .<= bigM*z) #from 0 to bCols -1
	@constraint(stage2Model, conBigMP,  1*b .<= bigM*z) #from bCols to 2bCols-1

	#Second objective term
	@constraint(stage2Model, 1*b .<= v) #from 2bCols to 3bCols-1
	@constraint(stage2Model, -1*b .<= v) #from 3bCols to 4bCols-1
	oneNorm = sum(v[i] for i=1:bCols)

	#gamma[g]*oneNorm <= G ---> -G <= -gamma[g]*oneNorm --> G >= gamma[g]*oneNorm
	g=5
	@constraint(stage2Model, gammaConstr, 1*oneNorm <= G) #4bCols

	#Define kmax constraint (5d)
	@constraint(stage2Model, kMaxConstr, sum(z[j] for j=1:bCols) <= kmax) #4bCols+1

	#Constraint 5f - can only select one of a pair of highly correlated features
	rho = 0.8
	for k=1:bCols
		for j=1:bCols
			if k != j
				if HC[k,j] >= rho
					HCPairCounter += 1
					@constraint(stage2Model,z[k]+z[j] <= 1) #from 4bCols+1 to (4bCols+1+HCPairCounter)
				end
			end
		end
	end

	#Constraint (5g) - only one transformation allowed (x, x^2, log(x) or sqrt(x))
	for j=1:(nCols-1)
		@constraint(stage2Model, z[j]+z[j+(nCols-1)]+z[j+2*(nCols-1)]+z[j+3*(nCols-1)] <= 1) #from (4bCols+1+HCPairCounter) to (4bCols+1+HCPairCounter+nCols)
	end

	#Adding empty cuts to possibly contain cuts. Can generate a maximum of 6 cuts for
	#each stage 3 and default value is 3 iterations between stage 2 and stage 3,
	#so 18 empty constraints are added + 1 for fixing the solution
	nEmptyCuts = 19
	zOnes = ones(bCols)
	for c = 1:nEmptyCuts
		@constraint(stage2Model, sum(zOnes[i]*z[i] for i=1:bCols) <= bCols) #(4bCols+1+HCPairCounter+nCols) to (4bCols+1+HCPairCounter+nCols+18)
	end

	#addinfocallback(stage2Model, infocallback, when = :Intermediate)

	JuMP.build(stage2Model)

	return internalmodel(stage2Model), HCPairCounter
end

function buildLasso(standX, standY)
	#Define parameters and model
	lassoModel = JuMP.Model(solver = GurobiSolver(TimeLimit = 30, OutputFlag = 0));
	gamma = 10

	#Define variables
	@variable(lassoModel, b[1:bCols]) #Beta values

	@variable(lassoModel, v[1:bCols]) # auxiliary variables for abs

	@variable(lassoModel, T) #First objective term
	@variable(lassoModel, G) #Second objective term

	#Define objective function (5a)
	@objective(lassoModel, Min, T+G)

	#println("Trying to implement new constraint")
	xSquareExpr = @expression(lassoModel, 0*b[1]^2)
	for l = 1:bCols
		coef = 0
		for j = 1:nRows
		   coef += standX[j,l]^(2)
		end
		append!(xSquareExpr, @expression(lassoModel, coef*b[l]^2))
	end
	#println("Implemented x^2")

	ySquaredExpr = @expression(lassoModel, 0*b[1])
	for j = 1:nRows
		append!(ySquaredExpr,@expression(lassoModel, standY[j,1]^2))
	end
	#println("Implemented y^2")

	simpleBetaExpr = @expression(lassoModel, 0*b[1])
	for l = 1:bCols
		coef = 0
		for j = 1:nRows
			coef += -1*2*standX[j, l]*standY[j]
		end
		append!(simpleBetaExpr, @expression(lassoModel, coef*b[l]))
	end
	#println("Implemented simpleBetaExpr")

	crossBetaExpr = @expression(lassoModel, 0*b[1]*b[2])
	iter = 1
	for l = 1:bCols
		for k = (l + 1):bCols
			coef = 0
			for j = 1:nRows
				coef += 2*standX[j,l]*standX[j,k]
			end
			append!(crossBetaExpr, @expression(lassoModel, coef*b[l,1]*b[k,1]))
		end
		#println("Finished loop $l/$bCols")
	end
	#println("Implemented crossBetaExpr")
	totalExpr = @expression(lassoModel, crossBetaExpr+simpleBetaExpr+xSquareExpr+ySquaredExpr)
	@constraint(lassoModel, quadConst, totalExpr <= T)
	#println("Successfully added quadratic constraints")

	#Second objective term
	@constraint(lassoModel, 1*b .<= v) #from 2bCols to 3bCols-1
	@constraint(lassoModel, -1*b .<= v) #from 3bCols to 4bCols-1
	oneNorm = sum(v[i] for i=1:bCols)

	#gamma[g]*oneNorm <= G ---> -G <= -gamma[g]*oneNorm --> G >= gamma[g]*oneNorm
	@constraint(lassoModel, gammaConstr, 1*oneNorm <= G) #4bCols

	JuMP.build(lassoModel)

	return internalmodel(lassoModel)
end

function solveLasso(model)
	bestNBeta = zeros(bCols+1,bCols+4)
	for i=1:bCols+1
		bestNBeta[i,1] = i-1
	end

	println("Solving Lasso for all gamma")
	gammaArray = logspace(-1, 3, 100)
	gamma = 0
	tol = 1e-6

	open(fileName*"LassoLog.csv", "w") do f
		allColNames = expandedColNamesToString(colNames)
		write(f, "RsquaredValue, gamma, nNZ, $allColNames\n")
		for j in 1:length(gammaArray)
			gamma = gammaArray[j]
			changeGammaLasso(model, gamma)

			#Gurobi.writeproblem(model, "testproblem.lp")

			#solve problem
			status = Gurobi.optimize!(model)

			#Get solution
			sol = Gurobi.getsolution(model)

			#Get parameters
			bSolved = sol[1:bCols]

			#Shrink values to 0 if within tolerance "tol"
			for i=1:length(bSolved)
				if bSolved[i] < tol
					if bSolved[i] > -tol
						bSolved[i] = 0
					end
				end
			end

			#Out of sample test
			Rsquared = getRSquared(standXVali,standYVali,bSolved)

			#Count number of non zero elements
			numNZ = countnz(bSolved)

			if Rsquared > bestNBeta[numNZ+1,2]
				bestNBeta[numNZ+1,2] = getRSquared(standXVali,standYVali,bSolved)
				bestNBeta[numNZ+1,3] = getRSquared(standXTest,standYTest,bSolved)
				bestNBeta[numNZ+1,4] = gamma
				bestNBeta[numNZ+1,5:bCols+4] = bSolved
			end

			write(f, "$Rsquared, $gamma, $numNZ, $bSolved\n")
			#println("Rsquared =$Rsquared \t gamma =$j \t n NZ=$numNZ")
		end
	end
	f = open(fileName*"LassoBestK.csv", "w")
	write(f, "k,Rsquared Vali, Rsquared Test,gamma,"*expandedColNamesToString(colNames)*"\n")
	writecsv(f,bestNBeta)
	close(f)
	println("Solve all Lasso problems")
	return bestNBeta
end

lassoM = buildLasso(standX, standY)
bestNLASSO = solveLasso(lassoM)

stage2Model, HCPairCounter = buildStage2(standX,standY, kmax)

# best3Beta[:,1] = Rsquared
# best3Beta[:,2] = kmax
# best3Beta[:,3] = gamma
# best3Beta[:,4:bCols+3] = beta

function solveForAllK(model, kmax)
	best3Beta = zeros(3,bCols+3)
	solArr = zeros(kmax*length(gammaArray),3)
	println("Solving for all k and gamma")
	for i in 1:1:kmax
		#println("Calculating warmstart solution")

		#=
		#Calculate warmstart
		warmStartError = 1e6
		for j in 1:5
			warmStartBetaTemp = gradDecent(standX, standY, 30000, 1e-3, i, HC, bSolved)
			tempError = norm(standY- standX*warmStartBetaTemp)^2
			println("Iteration $j of 5: Warmstart error is:", tempError)
			if tempError < warmStartError
				warmStartError = copy(tempError)
				warmStartBeta = copy(warmStartBetaTemp)
				println("Iteration $j of 5: Changed warmstartBeta")
			end
		end

		println("Warmstart error is:", warmStartError)

		#Set warmstart values
		startVar = zeros(length(warmStartBeta)*3+2)
		warmstartZ = zeros(warmStartBeta)
		for j in 1:bCols
			startVar[j] = warmStartBeta[j]
			startVar[j+2*bCols] = abs(warmStartBeta[j])
			if isequal(warmStartBeta[j],0)
				startVar[j+bCols] = 0
			else
				startVar[j+bCols] = 1
				warmstartZ[j] = 1
			end
		end
		startVar[bCols*3+1] = 10000
		startVar[bCols*3+2] = 10000

		Gurobi.setwarmstart!(m2,startVar)
		=#

		#Set kMax rhs constraint
		curUB = Gurobi.getconstrUB(model) #Get current UBbounds
		curUB[bCols*4+2] = i #Change upperbound in current bound vector
		Gurobi.setconstrUB!(model, curUB) #Push bound vector to model
		Gurobi.updatemodel!(model)

		#Set new Big M
		newBigM = 3#tau*norm(warmStartBeta, Inf)
		changeBigM(model,newBigM)
		for j in 1:length(gammaArray)
			changeGamma(model, gammaArray[j])

			#println("Starting to solve stage 2 model with kMax = $i and gamma = $j")

			#Gurobi.writeproblem(model, "testproblem.lp")

			#Solve Stage 2 model
			status = Gurobi.optimize!(model)
			#println("Objective value: ", Gurobi.getobjval(model))

			sol = Gurobi.getsolution(model)

			#printSolutionToCSV("lars.csv", bbdata)
			#Get solution and calculate R^2
			bSolved = sol[1:bCols]
			zSolved = sol[1+bCols:2*bCols]

			#Out of sample test
			Rsquared = getRSquared(standXVali,standYVali,bSolved)

			if Rsquared > best3Beta[1,1] #largest than 3rd largest Rsquared
				if Rsquared > best3Beta[2,1] #largest than 2nd largest Rsquared
					if Rsquared > best3Beta[3,1] #largest than largest Rsquared
						best3Beta[1,:] = best3Beta[2,:]
						best3Beta[2,:] = best3Beta[3,:]
						best3Beta[3,1] = Rsquared
						best3Beta[3,2] = i
						best3Beta[3,3] = j #potentially store actual gamma value
						best3Beta[3,4:bCols+3] = bSolved
					else
						best3Beta[1,:] = best3Beta[2,:]
						best3Beta[2,1] = Rsquared
						best3Beta[2,2] = i
						best3Beta[2,3] = j #potentially store actual gamma value
						best3Beta[2,4:bCols+3] = bSolved
					end
				else
					best3Beta[1,1] = Rsquared
					best3Beta[1,2] = i
					best3Beta[1,3] = j #potentially store actual gamma value
					best3Beta[1,4:bCols+3] = bSolved
				end
			end

			if any(Rsquared .> RsquaredValue) || isempty(RsquaredValue)
				bestBeta = bSolved
			end
			solArr[Int64((i-1)*length(gammaArray)+j),1] = i
			solArr[Int64((i-1)*length(gammaArray)+j),2] = j
			solArr[Int64((i-1)*length(gammaArray)+j),3] = Rsquared
			push!(kValue, (i-1)*length(gammaArray)+j)
			push!(RsquaredValue, Rsquared)
			#println("Rsquared =$Rsquared\t kMax =$i \t gamma =$j")
		end
		#printNonZeroValues(bSolved)
	end

	return best3Beta, solArr, model
end

function solveAndLogForAllK(model, kmax)
	best3Beta = zeros(3,bCols+3)
	solArr = zeros(kmax*length(gammaArray),3)
	println("Solving for all k and gamma")
	open(fileName*"AALRLog.csv", "w") do f
		allColNames = expandedColNamesToString(colNames)
		write(f, "RsquaredValue,gamma,k, $allColNames\n")
		for i in 1:1:kmax
			#println("Calculating warmstart solution")

			#=
			#Calculate warmstart
			warmStartError = 1e6
			for j in 1:5
				warmStartBetaTemp = gradDecent(standX, standY, 30000, 1e-3, i, HC, bSolved)
				tempError = norm(standY- standX*warmStartBetaTemp)^2
				println("Iteration $j of 5: Warmstart error is:", tempError)
				if tempError < warmStartError
					warmStartError = copy(tempError)
					warmStartBeta = copy(warmStartBetaTemp)
					println("Iteration $j of 5: Changed warmstartBeta")
				end
			end

			println("Warmstart error is:", warmStartError)

			#Set warmstart values
			startVar = zeros(length(warmStartBeta)*3+2)
			warmstartZ = zeros(warmStartBeta)
			for j in 1:bCols
				startVar[j] = warmStartBeta[j]
				startVar[j+2*bCols] = abs(warmStartBeta[j])
				if isequal(warmStartBeta[j],0)
					startVar[j+bCols] = 0
				else
					startVar[j+bCols] = 1
					warmstartZ[j] = 1
				end
			end
			startVar[bCols*3+1] = 10000
			startVar[bCols*3+2] = 10000

			Gurobi.setwarmstart!(m2,startVar)
			=#

			#Set kMax rhs constraint
			curUB = Gurobi.getconstrUB(model) #Get current UBbounds
			curUB[bCols*4+2] = i #Change upperbound in current bound vector
			Gurobi.setconstrUB!(model, curUB) #Push bound vector to model
			Gurobi.updatemodel!(model)

			#Set new Big M
			newBigM = 10#tau*norm(warmStartBeta, Inf)
			changeBigM(model,newBigM)
			for j in 1:length(gammaArray)
				gamma=gammaArray[j]
				changeGamma(model, gammaArray[j])

				#println("Starting to solve stage 2 model with kMax = $i and gamma = $j")

				#Gurobi.writeproblem(model, "testproblem.lp")

				#Solve Stage 2 model
				status = Gurobi.optimize!(model)
				#println("Objective value: ", Gurobi.getobjval(model))

				sol = Gurobi.getsolution(model)

				#printSolutionToCSV("lars.csv", bbdata)
				#Get solution and calculate R^2
				bSolved = sol[1:bCols]
				zSolved = sol[1+bCols:2*bCols]

				#Out of sample test
				Rsquared = getRSquared(standXVali,standYVali,bSolved)

				if Rsquared > best3Beta[1,1] #largest than 3rd largest Rsquared
					if Rsquared > best3Beta[2,1] #largest than 2nd largest Rsquared
						if Rsquared > best3Beta[3,1] #largest than largest Rsquared
							best3Beta[1,:] = best3Beta[2,:]
							best3Beta[2,:] = best3Beta[3,:]
							best3Beta[3,1] = Rsquared
							best3Beta[3,2] = i
							best3Beta[3,3] = j #potentially store actual gamma value
							best3Beta[3,4:bCols+3] = bSolved
						else
							best3Beta[1,:] = best3Beta[2,:]
							best3Beta[2,1] = Rsquared
							best3Beta[2,2] = i
							best3Beta[2,3] = j #potentially store actual gamma value
							best3Beta[2,4:bCols+3] = bSolved
						end
					else
						best3Beta[1,1] = Rsquared
						best3Beta[1,2] = i
						best3Beta[1,3] = j #potentially store actual gamma value
						best3Beta[1,4:bCols+3] = bSolved
					end
				end

				if any(Rsquared .> RsquaredValue) || isempty(RsquaredValue)
					bestBeta = bSolved
				end
				solArr[Int64((i-1)*length(gammaArray)+j),1] = i
				solArr[Int64((i-1)*length(gammaArray)+j),2] = j
				solArr[Int64((i-1)*length(gammaArray)+j),3] = Rsquared
				push!(kValue, (i-1)*length(gammaArray)+j)
				push!(RsquaredValue, Rsquared)
				println("Rsquared =$Rsquared\t kMax =$i \t gamma =$j")
				write(f, "$Rsquared,$gamma,$i,$bSolved\n")
			end
			#printNonZeroValues(bSolved)
		end
	end
	return best3Beta, solArr, model
end

#="""
For each of the three beta sets produced, we will test for significance
and condition number of model to see if cuts are necessary
High or low condition number doesn't mean that one correlation matrix is "better"
than the other. All it means is that variables are more correlated or less.
Whether it's good or not depends on the application.
"""=#

function stageThree(best3Beta, X, Y, allCuts)
	#=Condition Number
	  A high condition number indicates a multicollinearity problem. A condition
	  number greater than 15 is usually taken as evidence of
	  multicollinearity and a condition number greater than 30 is
	  usually an instance of severe multicollinearity
	=#

	bCols = size(X)[2]
	nRows = size(X)[1]
	cuts = Matrix(0, bCols+1)
	rowsPerSample = nRows #All of rows in training data to generate beta estimates, but selected with replacement
	totalSamples = 25 #25 different times we will get a beta estimate
	nBoot = 10000

	#For loop start
	for i = 1:size(best3Beta)[1]
		if signifBoolean[i] == 1 #if the previous solution was already completely significant, skip the bootstrapping
			continue
		end
		bestBeta = best3Beta[i,4:bCols+3]
		bestK = best3Beta[i,2]
		bestGamma = best3Beta[i,3]

		xColumns = []
		bSample = Matrix(totalSamples, bCols)

		bZeros = zeros(bCols)
		for j = 1:bCols
			if !isequal(bestBeta[j],0)
				push!(xColumns, j)
			end
		end

		selectedX = X[:,xColumns]
		condNumber = cond(selectedX)
		if condNumber >= 15
			bZeros[xColumns] = 1
			subsetSize = size(xColumns)[1]
			newCut = [bZeros' subsetSize]
			cuts = [cuts; newCut]
			println("A cut based on Condition number = $condNumber has been created from Beta$i")
		end

		# test significance
		bestZ = zeros(bestBeta)
		for l=1:size(bestBeta)[1]
			if bestBeta[l] != 0
				bestZ[l] = 1
			end
		end

		bZeros = zeros(bCols)
		createBetaDistribution(bSample, X, Y, bestK, totalSamples, rowsPerSample,  bestGamma, allCuts, bestZ) #standX, standY, k, sampleSize, rowsPerSample

		confArray99 = createConfidenceIntervalArray(bSample, nBoot, 0.99)
		confArray95 = createConfidenceIntervalArray(bSample, nBoot, 0.95)
		confArray90 = createConfidenceIntervalArray(bSample, nBoot, 0.90)

		significanceResult = testSignificance(confArray99, confArray95, confArray90, bestBeta)
		significanceResultNONSI = [] # = significanceResult[xColumns]
		subsetSize = size(xColumns)[1]
		for n = 1:size(significanceResult)[1]
			for s = 1:subsetSize
				if significanceResult[n] == 0 && xColumns[s] == n
					push!(significanceResultNONSI,xColumns[s])
					println("Parameter $n is selected, but NOT significant")
				elseif significanceResult[n] > 0 && xColumns[s] == n
					println("Parameter $n is significant with ", significanceResult[n])
				end
			end
		end

		if !isempty(significanceResultNONSI)
			bZeros[significanceResultNONSI] = 1
			subsetSize = size(significanceResultNONSI)[1]
			newCut = [bZeros' subsetSize]
			cuts = [cuts; newCut]
			println("A cut based on parameters being non-significant in Beta$i has been created")
		end
		if isempty(significanceResultNONSI)
			signifBoolean[i] = 1
		end
	end
	return cuts
end

bSample = []
allCuts = []
signifBoolean = zeros(3)

stage2Model, HCPairCounter = buildStage2(standX,standY, 2)

#Gurobi.writeproblem(stage2Model, "testproblem1.lp")

best3Beta, solArr, stage2Model = solveAndLogForAllK(stage2Model, 2)

cuts = stageThree(best3Beta, standX, standY, allCuts)

#cuts = stageThree(best3Beta[3,4:bCols+3], best3Beta[3,2], best3Beta[3,3],
# 				  best3Beta[2,4:bCols+3], best3Beta[2,2], best3Beta[2,3],
#				  best3Beta[1,4:bCols+3], best3Beta[1,2], best3Beta[1,3],
#				  standX, standY, allCuts)
#writedlm("betaOut.CSV", cuts,",")

cutCounter = 0
runCount = 1
cutMatrix = []

while !isempty(cuts)
	#Function to add cuts to problem
	#addCuts(stage2Model, cuts)
	preCutCounter = copy(cutCounter)
	allCuts = vcat(allCuts,cuts)
	cutMatrix = cuts
	cutRows = size(cutMatrix)[1]
	cutCols = size(cutMatrix)[2]
	cutCounter += cutRows
	addCuts(stage2Model, cutMatrix, preCutCounter)

	#Resolve problem
	Gurobi.updatemodel!(stage2Model)
	#Gurobi.writeproblem(stage2Model, "testproblem2.lp")
	#println(stage2Model)
	best3Beta, solArr = solveAndLogForAllK(stage2Model, 2)

	#Stage 3
	cuts = stageThree(best3Beta, standX, standY, allCuts)
	println("Finished iteration $runCount")
	runCount += 1
	if runCount == 4
		break
	end
end

best3Beta[1,1] = getRSquared(standXTest, standYTest, best3Beta[1,4:end])
best3Beta[2,1] = getRSquared(standXTest, standYTest, best3Beta[2,4:end])
best3Beta[3,1] = getRSquared(standXTest, standYTest, best3Beta[3,4:end])

best3Beta = cat(2,signifBoolean,best3Beta)

f = open(fileName*"AALRBestK.csv", "w")
write(f, "Significant,Rsquared,k,gamma,"*expandedColNamesToString(colNames)*"\n")
writecsv(f,best3Beta)
close(f)

bSolvedLinReg = inv(standX'*standX)*standX'*standY
standLinReg = getRSquared(standXTest,standYTest,bSolvedLinReg)
f = open(fileName*"standardLinReg.csv", "w")
write(f, "R² \n")
writecsv(f,standLinReg)
close(f)

#println(getRMSE(standXTest, standYTest, best3Beta[1,4:111]))
#println(getRMSE(standXTest, standYTest, best3Beta[2,4:111]))
#println(getRMSE(standXTest, standYTest, best3Beta[3,4:111]))
#Gurobi.writeproblem(stage2Model, "testproblem3.lp")
println("Complete")


module Permutations
using LinearAlgebra

import Base: length, show, inv, reverse, ==, getindex, *, ^, sign, hash, getindex,
                Matrix, Array, AbstractMatrix, AbstractArray, Array, adjoint

#                SparseMatrixCSC, AbstractSparseMatrix, AbstractSparseArray, sparse

import Combinatorics: nthperm
import Random: randperm

export Permutation, RandomPermutation
export length, getindex, array, two_row
export inv, cycles, cycle_string
export order, matrix, fixed_points
export longest_increasing, longest_decreasing, reverse, sign
export hash, dict
export CoxeterDecomposition  # no reason to expose CoxeterGenerator

# Defines the Permutation class. Permutations are bijections of 1:n.

abstract type AbstractPermutation end

"""
* `Permutation(list)` creates a new `Permutation`. Here `list` must be a rearrangement of `1:n`.
* `Permutation(n)` creates the identity `Permutation` of `1:n`.
* `Permutation(n,k)` creates the `k`'th `Permutation` of `1:n`.
"""
struct Permutation <: AbstractPermutation
    data::Vector{Int}
    function Permutation(dat::Vector{Int})
        n = length(dat)
        if sort(dat) != collect(1:n)
            error("Improper array: must be a permutation of 1:n")
        end
        new(dat)
    end
end

Permutation(v::AbstractVector) = Permutation(Vector{Int}(v))

# create the k'th permutation of 1:n
function Permutation(n,k)
    return Permutation(nthperm(collect(1:n),k))
end

# create the identity permutation
Permutation(n::Int) = Permutation(collect(1:n))

"""
`RandomPermutation(n)` creates a random permutation of `1:n`,
each with probability `1/factorial(n)`.
"""
RandomPermutation(n::Int) = Permutation(randperm(n))

# Returns the number of elements in the Permtuation

"""
`length(p)` is the number of elements in the `Permutation` `p`.
"""
function length(p::Permutation)
    return length(p.data)
end

# Check for equality of permutations
==(p::Permutation, q::Permutation) = p.data==q.data

# Apply the Permutation to an element: p[k] or p(k)
getindex(p::Permutation, k::Int) = p.data[k]
(p::Permutation)(k::Int) = p.data[k]

# Create a two-row representation of this permutation
"""
`two_row(p)` creates a two-row representation of the `Permutation` `p`
in which the first row is `1:n` and the second row are the
values `p(1), p(2), ..., p(n)`.
"""
function two_row(p::Permutation)
    n = length(p)
    return [ (1:n)'; p.data']
end

# Composition of two permutations
*(p::Permutation) = p
function *(p::Permutation, q::Permutation)
    n = length(p)
    if n != length(q)
        error("Cannot compose Permutations of different lengths")
    end
    @inbounds dat = [ p.data[q.data[k]] for k in 1:n ]
    return Permutation(dat)
end

# Inverse of a permutation
"""
`inv(p)` gives the inverse of `Permutation` `p`. This may
also be computed with `p'`.
"""
function inv(p::Permutation)
    n = length(p)
    data = zeros(Int,n)
    for k=1:n
        @inbounds j = p.data[k]
        @inbounds data[j] = k
    end
    return Permutation(data)
end

adjoint(p::Permutation) = inv(p)

# Find the cycles in a permutation
"""
`cycles(p)` returns a list of the cycles in `Permutation` `p`.
"""
function cycles(p::Permutation)
    n = length(p)
    result = Array{Int,1}[]
    todo = trues(n)
    while any(todo)
        k = findall(todo)[1]
        todo[k] = false
        cycle = [k]
        j = p[k]
        while j != k
            append!(cycle,[j])
            todo[j] = false
            j = p[j]
        end
        push!(result, cycle)
    end
    return result
end

"""
`cycle_string(p)` creates a nice, prinatble string representation
from the cycle structure of the permutation `p`.
"""
function cycle_string(p::Permutation)
    if length(p)==0
        return "()"
    end
    str = ""
    cc = cycles(p)
    for c in cc
        str *= array2string(c)
    end
    return str
end

# helper function for cycle_string (not exposed). Converts an integer
# array such as [1,3,5,2] into the string "(1,3,5,2)".
function array2string(a::Array{Int,1})
    n = length(a)
    res = "("
    for k=1:n
        res *= string(a[k])
        if k<n
            res *= ","
        else
            res *= ")"
        end
    end
    return res
end

"""
`sign(p)` is `+1` is `p` is an even `Permtuation` and `-1` if p is odd.
"""
function sign(p::Permutation)
    cc = cycles(p)
    szs = [ length(c)+1 for c in cc ]
    par = sum(szs)%2
    return (par>0) ? -1 : 1
end

# When we print a permutation, it appears in disjoint cycle format.
function show(io::IO, p::Permutation)
    print(io, cycle_string(p))
end

# Find the smallest positive n such that p^n is the identity
"""
`order(p)` is the smallest positive integer `n`
such that `p^n` is the identity `Permutation`.
"""
function order(p::Permutation)
    result = 1
    clist = cycles(p)
    for c in clist
        result = lcm(result, length(c))
    end
    return result
end

# Extend p^n so negative exponents work too
function ^(p::Permutation, n::Int)
    if n==0
        return Permutation(length(p))
    end
    if n<0
        return inv(p)^(-n)
    end
    if n==1
        return p
    end

    m::Int = round(Int,floor(n/2)) # int(floor(n/2))  # m = floor(n/2)
    q::Permutation = p^m

    if n%2 == 0   # if even
        return q*q
    end

    return p*q*q
end

# Represent as a permtuation matrix.
"""
`Matrix(p)` returns the permutation matrix for the `Permutation` `p`.
"""
function Matrix{T}(p::Permutation) where T
    n = length(p)
    A = Matrix{T}(I,n,n)     # A = eye(T,n)   #  int(eye(n))
    return A[:,p.data]
end

Matrix(p::Permutation) = Matrix{Int}(p)
Array(p::Permutation) = Matrix(p)
AbstractMatrix(p::Permutation) = Matrix(p)
AbstractArray(p::Permutation) = Matrix(p)

Array{T}(p::Permutation) where T = Matrix{T}(p)
AbstractMatrix{T}(p::Permutation) where T = Matrix{T}(p)
AbstractArray{T}(p::Permutation) where T = Matrix{T}(p)

#
# """
# `SparseMatrixCSC(p)` returns the permutation matrix for the `Permutation` `p`.
# """
# function SparseMatrixCSC{T}(p::Permutation) where T
#     n = length(p)
#     A = speye(T,n)   #  int(eye(n))
#     return A[p.data,:]
# end
#
# SparseMatrixCSC(p::Permutation) = SparseMatrixCSC{Int}(p)
# AbstractSparseMatrix(p::Permutation) = SparseMatrixCSC{Int}(p)
# AbstractSparseArray(p::Permutation) = SparseMatrixCSC{Int}(p)
# AbstractSparseMatrix{T}(p::Permutation) where T = SparseMatrixCSC{T}(p)
# AbstractSparseArray{T}(p::Permutation) where T = SparseMatrixCSC{T}(p)
#
# sparse(p::Permutation) = SparseMatrixCSC(p)

# find the fixed points of a Permutation
"""
`fixed_points(p)` returns the list of values `k` for which `p(k)==k`.
"""
fixed_points(p::Permutation) = findall([ p[k]==k for k in 1:length(p)])

# Find a longest monotone subsequence of p in a given direction. This
# is not exposed, but it is used by longest_increasing and
# longest_increasing.
function longest_monotone(p::Permutation, order=<)
    n = length(p)
    if n==0
        return Int[]
    end
    scores = ones(Int,n)
    scores[n] = 1

    for k=n-1:-1:1
        for i=k+1:n
            if order(p[k],p[i]) && scores[k] <= scores[i]
                scores[k] = scores[i]+1
            end
        end
    end

    seq = Int[]
    mx = maximum(scores)
    a = findfirst([s == mx for s in scores])
    push!(seq,a)

    while scores[a] > 1
        val = scores[a]-1
        idx = findfirst([scores[k]==val && order(p[a],p[k]) for k=a+1:n ])
        a += idx
        push!(seq,a)
    end

    return [p[i] for i in seq]

end

"""
Let `p` be a `Permutation`. Thinking of this as a list of values,
`longest_increasing(p)` returns a longest subsequence of values
that are in ascending order
"""
longest_increasing(p::Permutation) = longest_monotone(p,<)

"""
Let `p` be a `Permutation`. Thinking of this as a list of values,
`longest_decreasing(p)` returns a longest subsequence of values
that are in descending order
"""
longest_decreasing(p::Permutation) = longest_monotone(p,>)

"""
Let `p` be a `Permutation`. Thinking of `p` as a list of values,
`reverse(p)` creates a `Permutation` with those values in reverse
sequence.
"""
function reverse(p::Permutation)
    d = reverse(p.data)
    return Permutation(d)
end

# hash function so Permutations can be keys in dictionaries, etc.
hash(p::Permutation, h::UInt64) = hash(p.data,h)

"""
`dict(p)` converts a permutation into a dictionary. If `d` is the
result then `d[k]` equals `p(k)` for all `k` in the domain of the
permutation.
"""
function dict(p::Permutation)
    n = length(p)
    d = Dict{Int,Int}()
    for i=1:n
        d[i] = p(i)
    end
    return d
end

#####
# Decomposing into Coxeter generators
#####

struct CoxeterGenerator <: AbstractPermutation
    n::Int
    i::Int
    function CoxeterGenerator(n::Int, i::Int)
        1 ≤ i ≤ n-1 || throw(ArgumentError("$i must be between 1 and $n-1"))
        new(n, i)
    end
end

length(sᵢ::CoxeterGenerator) = sᵢ.n
Permutation(P::CoxeterGenerator) = Permutation([1:P.i-1; P.i+1; P.i; P.i+2:P.n])

function show(io::IO, p::CoxeterGenerator)
    print(io, "length $(p.n) permutation: s_$(p.i)")
end

# _coxeter_reduce reduces a product of simple transpositions using the relationships
# https://en.wikipedia.org/wiki/Symmetric_group#Generators_and_relations
# combined with a sorting for uniqueness
function _coxeter_reduce!(terms::Vector{Int})
    for i = 1:length(terms)-1
        # sᵢ^2 = I
        if terms[i] == terms[i+1]
            return _coxeter_reduce!(deleteat!(terms, i:i+1))
        end
        # sort using s_is_j = s_js_i
        if terms[i+1] ≠ terms[i]-1 && terms[i+1] ≠ terms[i]+1 && terms[i] > terms[i+1]
            terms[i], terms[i+1] = terms[i+1], terms[i]
            return _coxeter_reduce!(terms)
        end
    end
    # (s_is_{i+1})^3 = I
    for i = 1:length(terms)-5
        if terms[i]+1 == terms[i+1] == terms[i+2]+1 == terms[i+3] == terms[i+4]+1 == terms[i+5]
            return _coxeter_reduce!(deleteat!(terms, i:i+5))
        end
    end

    terms
end

"""
`CoxeterDecomposition(p)` expresses the `Permutation` `p` as a
composition of transpositions of the form `(k,k+1)`.
"""
struct CoxeterDecomposition <: AbstractPermutation
    n::Int
    terms::Vector{Int}
    function CoxeterDecomposition(n::Int, terms::Vector{Int})
        for t in terms
            1 ≤ t ≤ n-1 || throw(ArgumentError("$t must be between 1 and $n-1"))
        end
        new(n, _coxeter_reduce!(terms))
    end
end

CoxeterDecomposition(n::Int, terms::AbstractVector) = CoxeterDecomposition(n, Vector{Int}(terms))
CoxeterDecomposition(sᵢ::CoxeterGenerator) = CoxeterDecomposition(sᵢ.n, [sᵢ.i])

length(P::CoxeterDecomposition) = P.n

function Permutation(P::CoxeterDecomposition)
    if isempty(P.terms)
        Permutation(1:P.n)
    else
        *(Permutation.(CoxeterGenerator.(P.n,P.terms))...)
    end
end
CoxeterDecomposition(P::Permutation) = CoxeterDecomposition!(Permutation(copy(P.data)))
CoxeterDecomposition!(P::Permutation) = CoxeterDecomposition(length(P), _coxeterdecomposition!(P))
function _coxeterdecomposition!(P::Permutation)
    n = length(P)
    data = P.data
    ret = Int[]
    while !issorted(data)
        for k=1:n-1
            if data[k] > data[k+1]
                data[k], data[k+1] =  data[k+1], data[k]
                push!(ret, k)
            end
        end
    end
    reverse!(ret)
end

==(A::CoxeterDecomposition, B::CoxeterDecomposition) = A.terms == B.terms

function *(A::CoxeterGenerator, B::CoxeterGenerator)
    length(A) == length(B) || throw(ArgumentError("Permutations must have same length to multiply"))
    CoxeterDecomposition(length(A), [A.i,B.i])
end
function *(A::CoxeterGenerator, B::CoxeterDecomposition)
    length(A) == length(B) || throw(ArgumentError("Permutations must have same length to multiply"))
    CoxeterDecomposition(length(A), [A.i; B.terms])
end
function *(A::CoxeterDecomposition, B::CoxeterGenerator)
    length(A) == length(B) || throw(ArgumentError("Permutations must have same length to multiply"))
    CoxeterDecomposition(length(A), [A.terms; B.i])
end
function *(A::CoxeterDecomposition, B::CoxeterDecomposition)
    length(A) == length(B) || throw(ArgumentError("Permutations must have same length to multiply"))
    CoxeterDecomposition(length(A), [A.terms; B.terms])
end

inv(A::CoxeterDecomposition) = CoxeterDecomposition(A.n, reverse(A.terms))
adjoint(A::CoxeterDecomposition) = inv(A)

function show(io::IO, p::CoxeterDecomposition)
    print(io, "Permutation of 1:$(length(p)): ")
    if length(p.terms)==0
        print(io,"()")
    end
    for i in p.terms
        # print(io, "s_$(i)")
        print(io,"($i,$(i+1))")
    end
end

# deprecations
@deprecate array(p::Permutation) p.data
@deprecate matrix(p::Permutation, sparse::Bool = false) sparse ? sparse(p) : Matrix(p)

include("sqrt.jl")

end # end of module Permutations


# using DataFrames, CSV

# include("../src/LinearRegression.jl")

@testset "Cook's Distance" begin
        st_df = DataFrame(
                Y=[6.4, 7.4, 10.4, 15.1, 12.3 , 11.4],
                XA=[1.5, 6.5, 11.5, 19.9, 17.0, 15.5],
                XB=[1.8, 7.8, 11.8, 20.5, 17.3, 15.8],
                XC=[3., 13., 23., 39.8, 34., 31.],
                # values from SAS proc reg
                CooksD_base=[1.4068501943, 0.176809102, 0.0026655177, 1.0704009915, 0.0875726457, 0.1331183932],
                CooksD_multi=[1.7122291956, 18.983407026, 0.000118078, 0.8470797843, 0.0715921999, 0.1105843157],
                )

        t_lm_base = regress(@formula(Y ~ 1+ XA), st_df)
        results = predict_in_sample(t_lm_base, st_df, req_stats=["all"])
        @test isapprox(st_df.CooksD_base, results.cooksd)

        t_lm_multi = regress(@formula(Y ~ 1+ XA + XB), st_df)
        results = predict_in_sample(t_lm_multi, st_df, req_stats=["all"])
        @test isapprox(st_df.CooksD_multi, results.cooksd)
end


#============= SHE related functions ============#

export RHS_SHE, Jacobian_SHE, RHS_SHE1D, Jacobian_SHE1D

function RHS_SHE(S::State)
    @unpack ϵ, ν = S.gp.p
    u = S.u
    ϵ * u .- u .^ 3 .- (ν * ∇²(S)) .- ∇⁴(S)
end

# Vectorized (flattened) version of the RHS
function RHS_SHE(V::Array, p::Dict, N::Int, Δx::Real; dim::Int=2)
    @unpack ϵ, ν = p
    ϵ * V .- V .^3 .- (ν * ∇²(V, N, Δx; dim=dim)) .- ∇⁴(V, N, Δx; dim=dim)
end

# Jacobian will go here
function Jacobian_SHE(V::Array, p::Dict, N::Int, Δx::Real; dim::Int=2)
    @unpack ϵ, ν = p
    spdiagm(0 => ϵ .- 3 .* V .^ 2 ) .- ν * ∇²(N, Δx; dim=dim) .- ∇⁴(N, Δx; dim=dim)
end

## One dimensional case ##

RHS_SHE1D(V::Array, p::Dict, N::Int, Δx::Real) = RHS_SHE(V, p, N, Δx; dim=1)
Jacobian_SHE1D(V::Array, p::Dict, N::Int, Δx::Real) = Jacobian_SHE(V, p, N, Δx; dim=1)


immutable NoncentralF <: ContinuousUnivariateDistribution
    ndf::Float64
    ddf::Float64
    ncp::Float64
    function NoncentralF(n::Real, d::Real, nc::Real)
	n > zero(n) && d > zero(d) && nc >= zero(nc) ||
	    error("ndf and ddf must be > 0 and ncp >= 0")
	new(float64(n), float64(d), float64(nc))
    end
end

@_jl_dist_3p NoncentralF nf

@continuous_distr_support NoncentralF 0.0 Inf

mean(d::NoncentralF) = d.ddf > 2.0 ? d.ddf / (d.ddf - 2.0) * (d.ndf + d.ncp) / d.ndf : NaN

var(d::NoncentralF) = d.ddf > 4.0 ? 2.0 * d.ddf^2 *
		       ((d.ndf+d.ncp)^2 + (d.ddf - 2.0)*(d.ndf + 2.0*d.ncp)) /
		       (d.ndf * (d.ddf - 2.0)^2 * (d.ddf - 4.0)) : NaN

function rand(d::NoncentralF)
    rn = rand(NoncentralChisq(d.ndf,d.ncp)) / d.ndf
    rd = rand(Chisq(d.ddf)) / d.ddf
    rn / rd
end


"""
    Normal(μ,σ)

The *Normal distribution* with mean `μ` and standard deviation `σ≥0` has probability density function

```math
f(x; \\mu, \\sigma) = \\frac{1}{\\sqrt{2 \\pi \\sigma^2}}
\\exp \\left( - \\frac{(x - \\mu)^2}{2 \\sigma^2} \\right)
```

Note that if `σ == 0`, then the distribution is a point mass concentrated at `μ`.
Though not technically a continuous distribution, it is allowed so as to account for cases where `σ` may have underflowed,
and the functions are defined by taking the pointwise limit as ``σ → 0``.

```julia
Normal()          # standard Normal distribution with zero mean and unit variance
Normal(μ)         # Normal distribution with mean μ and unit variance
Normal(μ, σ)      # Normal distribution with mean μ and variance σ^2

params(d)         # Get the parameters, i.e. (μ, σ)
mean(d)           # Get the mean, i.e. μ
std(d)            # Get the standard deviation, i.e. σ
```

External links

* [Normal distribution on Wikipedia](http://en.wikipedia.org/wiki/Normal_distribution)

"""
struct Normal{T<:Real} <: ContinuousUnivariateDistribution
    μ::T
    σ::T
    Normal{T}(µ::T, σ::T) where {T<:Real} = new{T}(µ, σ)
end

function Normal(μ::T, σ::T; check_args=true) where {T <: Real}
    check_args && @check_args(Normal, σ >= zero(σ))
    return Normal{T}(μ, σ)
end

#### Outer constructors
Normal(μ::Real, σ::Real) = Normal(promote(μ, σ)...)
Normal(μ::Integer, σ::Integer) = Normal(float(μ), float(σ))
Normal(μ::T) where {T <: Real} = Normal(μ, one(T))
Normal() = Normal(0.0, 1.0, check_args=false)

const Gaussian = Normal

# #### Conversions
convert(::Type{Normal{T}}, μ::S, σ::S) where {T <: Real, S <: Real} = Normal(T(μ), T(σ))
convert(::Type{Normal{T}}, d::Normal{S}) where {T <: Real, S <: Real} = Normal(T(d.μ), T(d.σ), check_args=false)

@distr_support Normal -Inf Inf

#### Parameters

params(d::Normal) = (d.μ, d.σ)
@inline partype(d::Normal{T}) where {T<:Real} = T

location(d::Normal) = d.μ
scale(d::Normal) = d.σ

Base.eltype(::Type{Normal{T}}) where {T} = T

#### Statistics

mean(d::Normal) = d.μ
median(d::Normal) = d.μ
mode(d::Normal) = d.μ

var(d::Normal) = abs2(d.σ)
std(d::Normal) = d.σ
skewness(d::Normal{T}) where {T<:Real} = zero(T)
kurtosis(d::Normal{T}) where {T<:Real} = zero(T)

entropy(d::Normal) = (log2π + 1)/2 + log(d.σ)

function kldivergence(p::Normal, q::Normal)
    μp = mean(p)
    σ²p = var(p)
    μq = mean(q)
    σ²q = var(q)
    σ²p_over_σ²q = σ²p / σ²q
    return (abs2(μp - μq) / σ²q - logmxp1(σ²p_over_σ²q)) / 2
end

#### Evaluation

# Helpers
"""
    xval(d::Normal, z::Real)

Computes the x-value based on a Normal distribution and a z-value.
"""
function xval(d::Normal, z::Real)
    if isinf(z) && iszero(d.σ)
        d.μ + one(d.σ) * z
    else
        d.μ + d.σ * z
    end
end
"""
    zval(d::Normal, x::Real)

Computes the z-value based on a Normal distribution and a x-value.
"""
zval(d::Normal, x::Real) = (x - d.μ) / d.σ

gradlogpdf(d::Normal, x::Real) = -zval(d, x) / d.σ

# logpdf
_normlogpdf(z::Real) = -(abs2(z) + log2π)/2

function logpdf(d::Normal, x::Real)
    μ, σ = d.μ, d.σ
    if iszero(d.σ)
        if x == μ
            z = zval(Normal(μ, one(σ)), x)
        else
            z = zval(d, x)
            σ = one(σ)
        end
    else
        z = zval(Normal(μ, σ), x)
    end
    return _normlogpdf(z) - log(σ)
end

# pdf
_normpdf(z::Real) = exp(-abs2(z)/2) * invsqrt2π

function pdf(d::Normal, x::Real)
    μ, σ = d.μ, d.σ
    if iszero(σ)
        if x == μ
            z = zval(Normal(μ, one(σ)), x)
        else
            z = zval(d, x)
            σ = one(σ)
        end
    else
        z = zval(Normal(μ, σ), x)
    end
    return _normpdf(z) / σ
end

# logcdf
function _normlogcdf(z::Real)
    if z < -one(z)
        return log(erfcx(-z * invsqrt2)/2) - abs2(z)/2
    else
        return log1p(-erfc(z * invsqrt2)/2)
    end
end

function logcdf(d::Normal, x::Real)
    if iszero(d.σ) && x == d.μ
        z = zval(Normal(zero(d.μ), d.σ), one(x))
    else
        z = zval(d, x)
    end
    return _normlogcdf(z)
end

# logccdf
function _normlogccdf(z::Real)
    if z > one(z)
        return log(erfcx(z * invsqrt2)/2) - abs2(z)/2
    else
        return log1p(-erfc(-z * invsqrt2)/2)
    end
end

function logccdf(d::Normal, x::Real)
    if iszero(d.σ) && x == d.μ
        z = zval(Normal(zero(d.μ), d.σ), one(x))
    else
        z = zval(d, x)
    end
    return _normlogccdf(z)
end

# cdf
_normcdf(z::Real) = erfc(-z * invsqrt2)/2

function cdf(d::Normal, x::Real)
    if iszero(d.σ) && x == d.μ
        z = zval(Normal(zero(d.μ), d.σ), one(x))
    else
        z = zval(d, x)
    end
    return _normcdf(z)
end

# ccdf
_normccdf(z::Real) = erfc(z * invsqrt2)/2

function ccdf(d::Normal, x::Real)
    if iszero(d.σ) && x == d.μ
        z = zval(Normal(zero(d.μ), d.σ), one(x))
    else
        z = zval(d, x)
    end
    return _normccdf(z)
end

# quantile
function quantile(d::Normal, p::Real)
    # Promote to ensure that we don't compute erfcinv in low precision and then promote
    _p, _μ, _σ = map(float, promote(p, d.μ, d.σ))
    q = xval(d, -erfcinv(2*_p) * sqrt2)
    if isnan(_p)
        return oftype(q, _p)
    elseif iszero(_σ)
        # Quantile not uniquely defined at p=0 and p=1 when σ=0
        if iszero(_p)
            return oftype(q, -Inf)
        elseif isone(_p)
            return oftype(q, Inf)
        else
            return oftype(q, _μ)
        end
    end
    return q
end

# cquantile
function cquantile(d::Normal, p::Real)
    # Promote to ensure that we don't compute erfcinv in low precision and then promote
    _p, _μ, _σ = map(float, promote(p, d.μ, d.σ))
    q = xval(d, erfcinv(2*_p) * sqrt2)
    if isnan(_p)
        return oftype(q, _p)
    elseif iszero(_σ)
        # Quantile not uniquely defined at p=0 and p=1 when σ=0
        if iszero(_p)
            return oftype(q, Inf)
        elseif isone(_p)
            return oftype(q, -Inf)
        else
            return oftype(q, _μ)
        end
    end
    return q
end

mgf(d::Normal, t::Real) = exp(t * d.μ + d.σ^2 / 2 * t^2)
cf(d::Normal, t::Real) = exp(im * t * d.μ - d.σ^2 / 2 * t^2)

#### Affine transformations

Base.:+(d::Normal, c::Real) = Normal(d.μ + c, d.σ)
Base.:*(c::Real, d::Normal) = Normal(c * d.μ, abs(c) * d.σ)

#### Sampling

rand(rng::AbstractRNG, d::Normal{T}) where {T} = d.μ + d.σ * randn(rng, float(T))

#### Fitting

struct NormalStats <: SufficientStats
    s::Float64    # (weighted) sum of x
    m::Float64    # (weighted) mean of x
    s2::Float64   # (weighted) sum of (x - μ)^2
    tw::Float64    # total sample weight
end

function suffstats(::Type{<:Normal}, x::AbstractArray{T}) where T<:Real
    n = length(x)

    # compute s
    s = x[1]
    for i = 2:n
        @inbounds s += x[i]
    end
    m = s / n

    # compute s2
    s2 = abs2(x[1] - m)
    for i = 2:n
        @inbounds s2 += abs2(x[i] - m)
    end

    NormalStats(s, m, s2, n)
end

function suffstats(::Type{<:Normal}, x::AbstractArray{T}, w::AbstractArray{Float64}) where T<:Real
    n = length(x)

    # compute s
    tw = w[1]
    s = w[1] * x[1]
    for i = 2:n
        @inbounds wi = w[i]
        @inbounds s += wi * x[i]
        tw += wi
    end
    m = s / tw

    # compute s2
    s2 = w[1] * abs2(x[1] - m)
    for i = 2:n
        @inbounds s2 += w[i] * abs2(x[i] - m)
    end

    NormalStats(s, m, s2, tw)
end

# Cases where μ or σ is known

struct NormalKnownMu <: IncompleteDistribution
    μ::Float64
end

struct NormalKnownMuStats <: SufficientStats
    μ::Float64      # known mean
    s2::Float64     # (weighted) sum of (x - μ)^2
    tw::Float64     # total sample weight
end

function suffstats(g::NormalKnownMu, x::AbstractArray{T}) where T<:Real
    μ = g.μ
    s2 = abs2(x[1] - μ)
    for i = 2:length(x)
        @inbounds s2 += abs2(x[i] - μ)
    end
    NormalKnownMuStats(g.μ, s2, length(x))
end

function suffstats(g::NormalKnownMu, x::AbstractArray{T}, w::AbstractArray{Float64}) where T<:Real
    μ = g.μ
    s2 = abs2(x[1] - μ) * w[1]
    tw = w[1]
    for i = 2:length(x)
        @inbounds wi = w[i]
        @inbounds s2 += abs2(x[i] - μ) * wi
        tw += wi
    end
    NormalKnownMuStats(g.μ, s2, tw)
end

struct NormalKnownSigma <: IncompleteDistribution
    σ::Float64

    function NormalKnownSigma(σ::Float64)
        σ > 0 || throw(ArgumentError("σ must be a positive value."))
        new(σ)
    end
end

struct NormalKnownSigmaStats <: SufficientStats
    σ::Float64      # known std.dev
    sx::Float64      # (weighted) sum of x
    tw::Float64     # total sample weight
end

function suffstats(g::NormalKnownSigma, x::AbstractArray{T}) where T<:Real
    NormalKnownSigmaStats(g.σ, sum(x), Float64(length(x)))
end

function suffstats(g::NormalKnownSigma, x::AbstractArray{T}, w::AbstractArray{T}) where T<:Real
    NormalKnownSigmaStats(g.σ, dot(x, w), sum(w))
end

# fit_mle based on sufficient statistics

fit_mle(::Type{<:Normal}, ss::NormalStats) = Normal(ss.m, sqrt(ss.s2 / ss.tw))
fit_mle(g::NormalKnownMu, ss::NormalKnownMuStats) = Normal(g.μ, sqrt(ss.s2 / ss.tw))
fit_mle(g::NormalKnownSigma, ss::NormalKnownSigmaStats) = Normal(ss.sx / ss.tw, g.σ)

# generic fit_mle methods

function fit_mle(::Type{<:Normal}, x::AbstractArray{T}; mu::Float64=NaN, sigma::Float64=NaN) where T<:Real
    if isnan(mu)
        if isnan(sigma)
            fit_mle(Normal, suffstats(Normal, x))
        else
            g = NormalKnownSigma(sigma)
            fit_mle(g, suffstats(g, x))
        end
    else
        if isnan(sigma)
            g = NormalKnownMu(mu)
            fit_mle(g, suffstats(g, x))
        else
            Normal(mu, sigma)
        end
    end
end

function fit_mle(::Type{<:Normal}, x::AbstractArray{T}, w::AbstractArray{Float64}; mu::Float64=NaN, sigma::Float64=NaN) where T<:Real
    if isnan(mu)
        if isnan(sigma)
            fit_mle(Normal, suffstats(Normal, x, w))
        else
            g = NormalKnownSigma(sigma)
            fit_mle(g, suffstats(g, x, w))
        end
    else
        if isnan(sigma)
            g = NormalKnownMu(mu)
            fit_mle(g, suffstats(g, x, w))
        else
            Normal(mu, sigma)
        end
    end
end


# ------------ #
# Original API #
# ------------ #

# lindef.m -- DONE
function lindef(breaks::AbstractVector, evennum::Int=0)
    p = LinParams(breaks, evennum)
    return (
        length(p.breaks),
        minimum(p.breaks),
        maximum(p.breaks),
        Any[p.breaks, evennum],
        p
    )
end

# linnode.m -- DONE
linnode(breaks, evennum) = nodes(LinParams(breaks, evennum))

# lindop.m
function lindop(breaks, evennum, order)
    D, params = derivative_op(LinParams(breaks, evennum), order)
    n, a, b = length(params.breaks), params.breaks[1], params.breaks[end]
    D, n, a, b, params
end

# linbase.m -- DONE
function linbase(breaks, evennum=0, x=breaks, order=0)
    B = evalbase(LinParams(breaks, evennum), x, order)
    B, x
end


module Day9
using DataStructures
export parse_input, part_1

function parse_input(input::String)::Tuple{Int, Int}
    m = match(r"(\d+) players; last marble is worth (\d+) points", input)
    return tuple(map(x -> parse(Int, x), m.captures)...)
end

function rotate!(deque::Deque, times::Int)
    if times < 0
        for i in 1:abs(times)
            pushfirst!(deque, pop!(deque))
        end
    else
        for i in 1:times
            push!(deque, popfirst!(deque))
        end
    end
end

function part_1(players, max_points)::Int
    marbles = Deque{Int}()
    push!(marbles, 0)
    points = zeros(Int, players)

    for value in 1:max_points
        if value % 23 == 0
            player = mod1(value, players)
            points[player] += value

            rotate!(marbles, -7)
            points[player] += pop!(marbles)
            rotate!(marbles, 1)
        else
            rotate!(marbles, 1)
            push!(marbles, value)
        end
    end

    return maximum(points)
end

end


struct CFMap{T,N}
    result::T
    cf_type::Symbol
    img_size::Tuple{Vararg{Int,N}}
    zero_index::Tuple{Vararg{Int,N}}
end

function CFMap(img, cf_type)
    img_size = size(img)

    if cf_type == :full
        zero_index = img_size
        result_size = @. 2 * img_size - 1

    elseif cf_type == :central_symmetry
        zero_index = (img_size[1:end - 1]..., 1)
        result_size = (2 .* img_size[1:end - 1] .- 1 ..., img_size[end])

    elseif cf_type == :periodic_point_point
        zero_index = map(s -> 1, img_size)
        result_size = img_size

    else
        error("$cf_type is not supported")
    end

    result = similar(img, Float32, result_size)

    CFMap(result, cf_type, img_size, zero_index)
end

function dir_ixs(cfmap, dir)
    n = maximum(abs, dir)

    way = Float32.(collect(0:n) / (n == 0 ? 1 : n))

    if cfmap.cf_type == :central_symmetry && dir[end] < 0
        dir = .- dir
    end

    ixs = map((z, d) -> way .* d, cfmap.zero_index, dir)

    if cfmap.cf_type == :periodic_point_point
        for (ix, s) in zip(ixs, cfmap.img_size)
            out_ix = ix .< 0
            ix[out_ix] .+= s
        end
    end
    ixs
end

function dir_ixs(cfmap, dir::Symbol)
    tdir = to_tuple_dir(cfmap.img_size, dir)
    dir_ixs(cfmap, tdir)
end

function shiftmap(cfmap::CFMap)
    res = cfmap.result
    ns = size(res)
    ms = ns .÷ 2
    raw_ixs = map(m -> -m:m, ms)
    ixs = map((rix, n) -> mod.(rix, n) .+ 1, raw_ixs, ns)
    res[ixs...], raw_ixs
end

function itp_map(cfmap)
    if cfmap.cf_type == :periodic_point_point
        shiftres, xs = shiftmap(cfmap)
        itp = interpolate(xs, shiftres, Gridded(Linear()))
        extrapolate(itp, Periodic())
    else
        xs = map((ix, z) -> ix .- z, axes(cfmap.result), cfmap.zero_index)
        itp = interpolate(xs, cfmap.result, Gridded(Linear()))
    end
end

"""
    dir_from_map(cfmap, dir)

Return CF on one direcion using interpolations.

`CF.(r...)`, where `r = Map.dir_ixs(cfmap, dir)`.
"""
function dir_from_map(cfmap, dir)
    way = dir_ixs(cfmap, dir)
    itp = itp_map(cfmap)
    itp.(way...)
end

function surface_dirs(cfmap, dim, ispos)
    ixs = map(s -> -(s - 1):(s - 1), cfmap.img_size)
    map(ixs, 1:length(ixs)) do ix, d
        if d == dim
            if ispos
                ix[end]
            else
                ix[1]
            end
        else
            if d < dim
                ix[:]
            else
                ix[2:end-1]
            end
        end
    end
end

function all_dirs(cfmap)
    dirs = Vector{Tuple{Vararg{Int, ndims(cfmap.result)}}}()
    for i in 1:ndims(cfmap.result)
        for ispos in [true, false]

            s_dirs = surface_dirs(cfmap, i, ispos)
            for dir in Iterators.product(s_dirs...)
                push!(dirs, dir)
            end
        end
    end
    dirs
end

"""
    mean_dir(cfmap::CFMap)

Return averaged correlation function for r = 0, 1, 2, ...
"""
function mean_dir(cfmap::CFMap)
    data = cfmap.result

    m = cfmap.img_size |> norm |> floor |> Int
    m += ndims(data) > 1

    cfdir = zeros(Float32, m)
    weights = zeros(Float32, m)

    for ix in CartesianIndices(data)
        tix = Tuple(ix)
        r = norm(tix .- cfmap.zero_index)

        q = floor(Int, r) + 1
        α = q - r

        cfdir[q] += data[ix] * α
        cfdir[q + 1] += data[ix] * (1 - α)
        weights[q] += α
        weights[q + 1] += (1 - α)
    end

    cfdir ./= weights
end

"""
    restore_full_map(cfmap::CFMap)

Return centered full correlation function map.
"""
function restore_full_map(cfmap::CFMap{T,N}) where T where N
    res = cfmap.result

    if cfmap.cf_type == :full
        res

    elseif cfmap.cf_type == :central_symmetry
        n = size(res, N)
        reversed = reverse(res)
        [selectdim(reversed, N, 1:n - 1) res]

    elseif cfmap.cf_type == :periodic_point_point
        shiftmap(cfmap)[1]

    else
        error()
    end
end

function to_tuple_dir(img_size::Tuple{Int}, dir)
    if dir == :x
        (img_size[1] - 1,)
    else
        error(":$dir is not supported for 1D images")
    end
end

function to_tuple_dir(img_size::Tuple{Int,Int}, dir)
    if dir == :x
        (img_size[1] - 1, 0)
    elseif dir == :y
        (0, img_size[2] - 1)
    elseif dir == :xy
        q = minimum(img_size) - 1
        (q, q)
    elseif dir == :yx
        q = minimum(img_size) - 1
        (-q, q)
    else
        error(":$dir is not supported for 2D images")
    end
end

function to_tuple_dir(img_size::Tuple{Int,Int,Int}, dir)
    if dir == :x
        (img_size[1] - 1, 0, 0)
    elseif dir == :y
        (0, img_size[2] - 1, 0)
    elseif dir == :z
        (0, 0, img_size[3] - 1)
    else
        error(":$dir is not supported for 3D images")
    end
end

function cut_cfmap(cfmap, mask)
    if cfmap.cf_type == :central_symmetry && mask[end]
        error("negative direction in last dimension is not supported for $(cfmap.cf_type)")
    elseif cfmap.cf_type == :periodic_point_point && any(mask)
        error("negative direction is not supported for $(cfmap.cf_type)")
    end

    ixs = map(cfmap.zero_index, cfmap.img_size, mask) do z, s, m
        if m
            z:-1:1
        else
            z:z + s - 1
        end
    end

    view(cfmap.result, ixs...)
end


using HTTP
using CSV
using DataFrames
using GLMakie
using LinearAlgebra
GLMakie.inline!(false)
url = "https://raw.githubusercontent.com/plotly/datasets/master/vortex.csv"
df = CSV.File(HTTP.get(url).body)|> DataFrame;

uvw = Vec3f.(df.u, df.v, df.w)
len = norm.(uvw)

arrows(df.x, df.y, df.z, df.u, df.v, df.w, linewidth=0, color=len, arrowsize = Vec3f(0.5, 0.5, 1)) |> display

# surface parameterization  (u, v)----> (f(u,v), g(u,v), h(u,v)):
f = (u, v) -> cos(v)*(6 - (5/4 +sin(3*u))*sin(u-3*v))
g = (u, v) -> sin(v)*(6 - (5/4 +sin(3*u))*sin(u-3*v))
h = (u, v) -> -cos(u-3*v) * (5/4 +sin(3*u));

u = range(0, stop=2π, length=150)
v = range(0, stop=2π, length=150)

tr = surface(f.(u,v'),
             g.(u,v'),
             h.(u,v'),
             ambient=Vec3f(0.5),
             diffuse=Vec3f(1),
             specular=0.5,
             colormap=:balance) |> display


"Z-normalize the input array"
function znorm(x::AbstractVector)
    x = x .- mean(x)
    x ./= std(x, mean=0, corrected=false)
end

function znorm!(x::AbstractVector)
    x .= x .- mean(x)
    x ./= std(x, mean=0, corrected=false)
end

"""
    meanstd(x::AbstractArray)

Return mean and std of `x`. Is faster but less accurate than built-in functions.
"""
function meanstd(x::AbstractArray{T}) where T
    s = ss = zero(float(T))
    n = length(x)
    @avx for i in eachindex(x)
        s  += x[i]
        ss += x[i]^2
    end
    m   = s/n
    sig = √(ss/n - m^2)
    m, sig
end

import Base: @propagate_inbounds, @boundscheck, getindex
import LinearAlgebra: normalize

abstract type AbstractNormalizer{T,N} <: AbstractArray{T,N} end

advance!(x) = 0 # Catches everything that's not a normalizer

setup_normalizer(n::Type{Nothing}, q, y) = q, y

normalize(::Type{Nothing}, q) = q

(z::AbstractNormalizer)(q) = normalize(z, q)
(::Type{N})(q) where N <: AbstractNormalizer = normalize(N, q)

abstract type AbstractZNormalizer{T,N} <: AbstractNormalizer{T,N} end

"""
    ZNormalizer{T} <: AbstractNormalizer{T, 1}

Utility object that normalizes an input array on the fly.
Works like a vector, index into window with `z[i]`, index into normalized window with `z[!,i]`.

# Arguments:
- `x::Vector{T}`: The vecetor to operate on
- `n::Int`: The lenght of a window
- `μ::T`: mean over window
- `σ::T`: std over window
- `s::T`: sum over window
- `ss::T`: sum^2 over window
- `i::Int`: the current index
- `buffer::Vector{T}`: stores calculated normalized values, do not access this, it might not be complete
- `bufi::Int`: index of last normalized value
"""
mutable struct ZNormalizer{T,N} <: AbstractZNormalizer{T,N}
    x::Array{T,N}
    n::Int
    μ::T
    σ::T
    s::T
    ss::T
    i::Int
    buffer::Array{T,N}
    bufi::Int
end

function ZNormalizer(x::AbstractVector{T}, n) where T
    @assert length(x) >= n
    s = ss = zero(T)
    @avx for i in 1:n
        s += x[i]
        ss += x[i]^2
    end
    μ = s/n
    σ = sqrt(max(0, ss/n - μ^2))
    buffer = similar(x, n)
    ZNormalizer(x, n, μ, σ, s, ss, 0, buffer, 0)
end

function ZNormalizer(x::AbstractMatrix{T}, n) where T
    @assert length(x) >= n
    m = size(x,1)*n
    s = ss = zero(T)
    for i in 1:n
        s += sum(x[!,i])
        ss += sum(abs2, x[!,i])
    end
    μ = s/m
    σ = sqrt(max(0, ss/m - μ^2))
    buffer = similar(x, size(x,1), n)
    ZNormalizer(x, n, μ, σ, s, ss, 0, buffer, 0)
end

function normalize(::Type{ZNormalizer}, q::AbstractArray)
    q = q .- mean(q)
    q ./= std(q, corrected=false, mean=0)
end

setup_normalizer(z::Type{ZNormalizer}, q::AbstractVecOrMat, y::AbstractVecOrMat) = normalize(z, q), ZNormalizer(y, lastlength(q))

@propagate_inbounds function advance!(z::ZNormalizer{T,1}) where T
    if z.i == 0
        return z.i = 1
    end
    @boundscheck if z.i + z.n > length(z.x)
        return z.i += 1
    end
    z.bufi = 0

    # Remove old point
    x = z.x[z.i]
    z.s -= x
    z.ss -= x^2

    # Add new point
    x = z.x[z.i+z.n]
    z.s += x
    z.ss += x^2
    z.μ = z.s/z.n
    z.σ = sqrt(max(z.ss/z.n - z.μ^2, 0))
    z.i += 1
end

@propagate_inbounds function advance!(z::ZNormalizer{T,2}) where T
    if z.i == 0
        return z.i = 1
    end
    @boundscheck if z.i + z.n > lastlength(z.x)
        return z.i += 1
    end
    z.bufi = 0

    # Remove old point
    x = z.x[!, z.i]
    z.s -= sum(x)
    z.ss -= sum(abs2, x)

    # Add new point
    x = z.x[!, z.i+z.n]
    z.s += sum(x)
    z.ss += sum(abs2, x)
    m = length(z)
    z.μ = z.s/m
    z.σ = sqrt(max(z.ss/m - z.μ^2, 0))
    z.i += 1
end

@inline @propagate_inbounds function getindex(z::ZNormalizer{<:Any, 1}, i)
    @boundscheck 1 <= i <= z.n || throw(BoundsError(z,i))
    xi = i+z.i-1
    @boundscheck xi <= length(z.x) || throw(BoundsError(z,i))
    z.x[xi]
end

@inline @propagate_inbounds function getindex(z::ZNormalizer{<:Any, 1}, ::typeof(!), i::Int, inorder = i == z.bufi + 1)
    y = (z[i]-z.μ) / z.σ
    if inorder
        z.bufi = i
        z.buffer[i] = y
    end
    y
end

@inline @propagate_inbounds function getindex(z::ZNormalizer{<:Any, 1}, ::typeof(!), i::AbstractRange)
    @boundscheck (i[1] == z.i && length(i) == z.n) || throw(ArgumentError("ZNormalizers can only be indexed by ranges corresponding to their current state. Got range $i but state was $(z.i) corresponding to range $(z.i):$(z.i+z.n-1)"))
    z
end

@inline @propagate_inbounds function getindex(z::ZNormalizer{T,2}, ::typeof(!), i::Int, inorder = i == z.bufi + 1) where T
    j = inorder ? i : z.n
    xj = z.i + i - 1
    μ,σ = z.μ, z.σ + eps(T)
    @avx for k = 1:size(z.x, 1)
        z.buffer[k, j] = (z.x[k, xj] - μ) / σ
    end
    if inorder
        z.bufi = i
    end
    z.buffer[!, j]
end

Statistics.mean(z::AbstractZNormalizer) = z.μ
Statistics.std(z::AbstractZNormalizer) = z.σ

SlidingDistancesBase.lastlength(z::AbstractZNormalizer) = z.n
Base.length(z::AbstractZNormalizer) = length(z.buffer)
Base.size(z::ZNormalizer, args...) = size(z.buffer, args...)
actuallastlength(x::AbstractZNormalizer) = lastlength(x.x)

# Multi dim ==================================================================================
mutable struct DiagonalZNormalizer{T} <: AbstractZNormalizer{T,2}
    x::Array{T,2}
    n::Int
    μ::Array{T,1}
    σ::Array{T,1}
    s::Array{T,1}
    ss::Array{T,1}
    i::Int
    buffer::Array{T,2}
    bufi::Int
end

function DiagonalZNormalizer(x::AbstractArray{T,2}, n) where T
    @assert length(x) >= n
    s  = zeros(T, size(x,1))
    ss = zeros(T, size(x,1))
    @inbounds @simd for i in 1:n
        s  .+= x[!, i]
        ss .+= x[!, i].^2
    end
    μ = s./n
    σ = sqrt.(max.(ss./n .- μ.^2, 0))
    buffer = similar(x, size(x,1), n)
    DiagonalZNormalizer(x, n, μ, σ, s, ss, 0, buffer, 0)
end

function normalize(::Type{DiagonalZNormalizer}, q::AbstractMatrix)
    q = q .- mean(q, dims=2) # TODO: this will cause a ton of allocations
    q ./= (std(q, dims=2, corrected=false) .+ eps(eltype(q)))
end

setup_normalizer(z::Type{DiagonalZNormalizer}, q, y) = normalize(z, q), DiagonalZNormalizer(y, lastlength(q))

@propagate_inbounds function advance!(z::DiagonalZNormalizer{T}) where T

    if z.i == 0
        return z.i = 1
    end
    @boundscheck if z.i + z.n > length(z.x)
        return z.i += 1
    end
    z.bufi = 0

    # Remove old point
    x = z.x[!, z.i]
    @avx z.s .-= x
    @avx z.ss .-= x.^2

    # Add new point
    x = z.x[!, z.i+z.n]
    @avx z.s .+= x
    @avx z.ss .+= x.^2
    @avx z.μ .= z.s./z.n
    @avx z.σ .= sqrt.(z.ss./z.n .- z.μ.^2)
    z.i += 1
end

@inline @propagate_inbounds function getindex(z::AbstractNormalizer{<:Any, 2}, i::Union{Number, AbstractRange}, j)
    @boundscheck 1 <= j <= z.n || throw(BoundsError(z,j))
    @boundscheck 1 <= 1 <= size(z.x, 1) || throw(BoundsError(z,i))
    xj = j+z.i-1
    @boundscheck xj <= lastlength(z.x) || throw(BoundsError(z,j))
    z.x[i, xj]
end

@inline @propagate_inbounds function getindex(z::DiagonalZNormalizer{T}, ::typeof(!), i::Int, inorder = i == z.bufi + 1) where T
    j = inorder ? i : z.n
    xj = z.i + i - 1
    @avx for k = 1:size(z.x, 1)
        z.buffer[k, j] = (z.x[k, xj] - z.μ[k]) / (z.σ[k] + eps(T))
    end
    if inorder
        z.bufi = i
    end
    z.buffer[!, j]
end

@inline @propagate_inbounds function getindex(z::AbstractNormalizer{<:Any, 2}, ::typeof(!), i::AbstractRange)
    @boundscheck (i[1] == z.i && length(i) == z.n) || throw(ArgumentError("Normalizers can only be indexed by ranges corresponding to their current state. Got range $i but state was $(z.i) corresponding to range $(z.i):$(z.i+z.n-1)"))
    z
end

Base.Matrix(z::AbstractNormalizer{<:Any, 2}) = z.x[:,z.i:z.i+z.n-1]

Base.length(z::DiagonalZNormalizer) = size(z.x,1) * z.n
Base.size(z::AbstractNormalizer{<:Any, 2}) = (size(z.x,1), z.n)

# Norm normalizer ==================================================================================

abstract type AbstractNormNormalizer{T,N} <: AbstractNormalizer{T,N} end

mutable struct NormNormalizer{T,N} <: AbstractNormNormalizer{T,N}
    x::Array{T,N}
    n::Int
    σ::T
    ss::T
    i::Int
    buffer::Array{T,N}
    bufi::Int
end

mutable struct SqNormNormalizer{T,N} <: AbstractNormNormalizer{T,N}
    x::Array{T,N}
    n::Int
    σ::T
    ss::T
    i::Int
    buffer::Array{T,N}
    bufi::Int
end

function NormNormalizer(x::AbstractArray{T,N}, n) where {T,N}
    @assert length(x) >= n
    ss = zero(T)
    @inbounds for i in 1:n
        ss += sum(abs2, x[!, i])
    end
    σ = sqrt(ss)
    if N == 1
        buffer = similar(x, n)
    else
        buffer = similar(x, size(x,1), n)
    end
    NormNormalizer(x, n, σ, ss, 0, buffer, 0)
end

function SqNormNormalizer(x::AbstractArray{T,N}, n) where {T,N}
    @assert length(x) >= n
    ss = zero(T)
    @inbounds for i in 1:n
        ss += sum(abs2, x[!, i])
    end
    σ = ss
    if N == 1
        buffer = similar(x, n)
    else
        buffer = similar(x, size(x,1), n)
    end
    SqNormNormalizer(x, n, σ, ss, 0, buffer, 0)
end

function normalize(::Type{N}, q::AbstractArray) where N <: AbstractNormNormalizer
    power = N == NormNormalizer ? 1 : 2
    q ./ (norm(q)^power .+ eps(eltype(q)))
end

setup_normalizer(z::Type{N}, q, y) where N <: AbstractNormNormalizer = normalize(N, q), N(y, lastlength(q))

@propagate_inbounds function advance!(z::AbstractNormNormalizer{T}) where T

    if z.i == 0
        return z.i = 1
    end
    @boundscheck if z.i + z.n > length(z.x)
        return z.i += 1
    end
    z.bufi = 0

    # Remove old point
    x = z.x[!, z.i]
    z.ss -= sum(abs2, x)

    # Add new point
    x = z.x[!, z.i+z.n]
    z.ss += sum(abs2, x)
    z.σ = z isa SqNormNormalizer ? z.ss : sqrt(z.ss)
    z.i += 1
end

@inline @propagate_inbounds function getindex(z::AbstractNormNormalizer{T}, ::typeof(!), i::Int, inorder = i == z.bufi + 1) where T
    j = inorder ? i : z.n
    xj = z.i + i - 1
    σ = z.σ + eps(T)
    @avx for k = 1:size(z.x, 1)
        z.buffer[k, j] = z.x[k, xj] / σ
    end
    if inorder
        z.bufi = i
    end
    z.buffer[!, j]
end

LinearAlgebra.norm(z::NormNormalizer) = z.σ
LinearAlgebra.norm(z::SqNormNormalizer) = sqrt(z.σ)

Base.length(z::AbstractNormNormalizer{<:Any,2}) = size(z.x,1) * z.n
Base.length(z::AbstractNormNormalizer{<:Any,1}) = z.n
SlidingDistancesBase.lastlength(z::AbstractNormNormalizer) = z.n
actuallastlength(x::AbstractNormNormalizer) = lastlength(x.x)

# ============================================================================================

for T in [ZNormalizer, DiagonalZNormalizer, NormNormalizer, SqNormNormalizer]
    @eval @inline @propagate_inbounds function normalize(::Type{$T}, z::$T)
        if z.bufi == z.n
            return z.buffer
        end
        for i = z.bufi+1:z.n
            z[!, i, true] # This populates the buffer
        end
        return z.buffer
    end
end


module TestSlices
using Test

using Catlab.Theories, Catlab.CategoricalAlgebra, Catlab.Graphs

"""
A product in a slice category is equivalent to a pullback. Here we work with
discrete graphs:

PB     B=2
        ↓ id
A=2 ⟶ C=3
   swap

The pullback of A and B has two elements:
- one element mapping to (a₂,b₁) - which both map to c₁
- one element mapping to (a₁,b₂) - which both map to c₂

Given slice morphisms into A and B (say, a single element set which maps to c₁,
which forces its map to A to point to a₂ and its map to B to point to b₁), the
universal property gives a map into the limit object, so it picks the element
that maps to (a₂,b₁).
"""
# create limit
G2,G1,G3 = Graph.([2,1,3])
A, B = Slice.([ACSetTransformation(G2, G3;V=x) for x in [[1,2],[2,1]]])
@test force(id(A)) == force(SliceHom(A,A,id(G2)))
slice_dia = FreeDiagram(DiscreteDiagram([A,B], SliceHom))
slice_lim = limit(slice_dia)

# test universal property
mA, mB = [ACSetTransformation(G1,G2;V=x) for x in [[1],[2]]]
map_to_base = Slice(ACSetTransformation(G1,G3;V=[1]))
toA = SliceHom(map_to_base, A, mA)
toB = SliceHom(map_to_base, B, mB)
u = universal(slice_lim, Span(toA, toB))
@test dom(u) == map_to_base
@test codom(u) == apex(slice_lim)
@test force(compose(u, legs(slice_lim)[1]).f) == force(toA.f)

# Pushout complement in Petri, computed as slice in Grph
two = @acset Graph begin V=2; E=2; src=[1,2]; tgt=[2,1] end
a,b,c = Graph(1), path_graph(Graph, 2), path_graph(Graph, 3)
A = Slice(ACSetTransformation(a, two, V=[2])) # a transition
B = Slice(ACSetTransformation(b, two, V=[2,1], E=[2]))
C = Slice(ACSetTransformation(c, two, V=[1,2,1], E=[1,2]))
ff = ACSetTransformation(a,b,V=[1])
gg = ACSetTransformation(b,c,V=[2,3], E=[2])
f = SliceHom(A,B,ff)
g = SliceHom(B,C,gg)
pc1, pc2 = pushout_complement(ComposablePair(ff,gg))
p1, p2 = pushout_complement(f,g)

# Colimit: a pushout of slices
#------------------------------
# do the pushout associated with the pushout complement above
slice_dia = FreeDiagram{Slice,SliceHom}(Multispan(A, [f, p1]))
clim = colimit(slice_dia)
@test is_isomorphic(dom(apex(clim)), c)
pushout(([f, p1]))

# DPO of slices: do the pushout complement and add an extra state
R = Slice(ACSetTransformation(Graph(2), two, V=[2, 1]))
r = SliceHom(A,R,ACSetTransformation(a, Graph(2), V=[1]))
res = rewrite_match(f, r, g)
@test is_isomorphic(dom(res), apex(coproduct(b, Graph(1))))

end #module


"""
- cdfE calculated empirical cumulative distribution function F(y<=yi)
- To get the CCDF of Y (complementary cumulative distribution function):
	- ccdf(Y) = 1- cdf(Y);
Example:
    [cdfY,y] = cdfE(Y,nbins); % nbins: number of bins (i.e. number of y points)
    [cdfY,y] = cdfE(Y,y); % y is the mid-points of the edges of bins. % This is expected to be the fastest.
                                               % Noted that the 'EdgeScale' has no effect on output results since y is specified.
    [cdfY,y] = cdfE(Y);
"""
function cdfE(Y; nedges::Int=100)
	edges(l,r,pt) = collect(range(l,r,length = pt));
	y = edges(minimum(Y), maximum(Y), nedges);
	cdfY = cdfE_core(y, Y, nedges);
	(y, cdfY);
end

function cdfE(Y, y::Vector; edgescale=:linear)
	nedges = length(y);
	cdfY = cdfE_core(y, Y, nedges)
end

function cdfE_core(y, Y, nedges)
	Y = [Y...];
	cdfY = fill(NaN,nedges)
	invlenY = 1/length(Y);
	for i = 1:nedges
		cdfY[i] = sum(Y .<= y[i])*invlenY;
	end
	cdfY
end


using Base.Test

include("MNL.jl")

### Generate data

# make it deterministic
srand(10)

m, n, k = 100, 100, 3
Theta = randn(m, k) * randn(k, n)
Theta = Theta .- (Theta * ones(n,1))

nobs = int(10*m*n)
K = 2 #int(.05*n) # size of sets
sets = Array(Array{Int,1}, nobs) # set presented
rows = Array(Int, nobs)      # customer who arrived
cols = Array(Int, nobs)      # product they picked
for iobs = 1:nobs
	it = sample(1:m)
	St = sample(1:n, K, replace=false, ordered=true)
	wv = WeightVec(Float64[exp(-Theta[it, j]) for j in St])
	jt = sample(wv)
	rows[iobs], cols[iobs], sets[iobs] = it, jt, St
end
data = MNLdata(nobs, m, n, sets, rows, cols)

### Test invariances

ThetaTest = randn(m, n)
u,s,v = svd(ThetaTest)
ThetaTest = u[:,1:k]*diagm(s[1:k])*v[:,1:k]'
UTest = factor(ThetaTest, k=k)

@test_approx_eq negloglik(ThetaTest, data) negloglik(ThetaTest + rand(m,1)*ones(1,n), data)
@test_approx_eq grad_negloglik(ThetaTest + rand(m,1)*ones(1,n), data) grad_negloglik(ThetaTest + rand(m,1)*ones(1,n), data)

@test_approx_eq negloglik(ThetaTest, data) negloglik(UTest, data)

@test_approx_eq nucnorm(ThetaTest) nucnorm(UTest)

# gradients satisfy some relationship too...
# @test_approx_eq grad_negloglik(ThetaTest, data) grad_negloglik(UTest, data)
# @test_approx_eq grad_nucnorm(ThetaTest) grad_nucnorm(UTest)


using StaticCompiler
using StaticTools
using LoopVectorization

@inline function mul!(C::MallocArray, A::MallocArray, B::MallocArray)
    @turbo for n ∈ indices((C,B), 2), m ∈ indices((C,A), 1)
        Cmn = zero(eltype(C))
        for k ∈ indices((A,B), (2,1))
            Cmn += A[m,k] * B[k,n]
        end
        C[m,n] = Cmn
    end
    return C
end

function loopvec_matrix(argc::Int, argv::Ptr{Ptr{UInt8}})
    argc == 3 || return printf(stderrp(), c"Incorrect number of command-line arguments\n")
    rows = parse(Int64, argv, 2)            # First command-line argument
    cols = parse(Int64, argv, 3)            # Second command-line argument

    # LHS
    A = MallocArray{Float64}(undef, rows, cols)
    @turbo for i ∈ axes(A, 1)
        for j ∈ axes(A, 2)
           A[i,j] = i*j
        end
    end

    # RHS
    B = MallocArray{Float64}(undef, cols, rows)
    @turbo for i ∈ axes(B, 1)
        for j ∈ axes(B, 2)
           B[i,j] = i*j
        end
    end

    # # Matrix multiplication
    C = MallocArray{Float64}(undef, cols, cols)
    mul!(C, B, A)

    # Print to stdout
    printf(C)
    # Also print to file
    fp = fopen(c"table.tsv",c"w")
    printf(fp, C)
    fclose(fp)
    # Clean up matrices
    free(A)
    free(B)
    free(C)
end

# Attempt to compile
path = compile_executable(loopvec_matrix, (Int64, Ptr{Ptr{UInt8}}), "./")


using ITensors

# Define indices
a = Index(2, "a")
b = Index(2, "b")
c = Index(2, "c")

# Note that Index tags don't have
# to match variable names
i = Index(3, "Red,Link")

# Define 3 order 2 tensors (matrices)
Z = ITensor(a, b)
X = ITensor(b, c)
Y = ITensor(b, c)

# Set some elements
Z[a(1), b(1)] = 1.0
Z[a(2), b(2)] = -1.0

X[b(1), c(2)] = 1.0
X[b(2), c(1)] = 1.0

Y[b(1), c(1)] = 1.0
Y[b(2), c(2)] = 1.0

# Operations with tensors
R = Z * X

S = Y + X

T = Y - X

# Print results
println("Z =\n", Z, "\n")
println("X =\n", X, "\n")
println("Y =\n", Y, "\n")
println("R = Z * X =\n", R, "\n")
println("S = Y + X =\n", S, "\n")
println("S = Y - X =\n", T, "\n")

# Check that adding incompatible tensors cause an error
try
  U = Z + X
catch
  println("Cannot add Z and X")
end

# Compare calculations to Julia arrays
jZ = [
  1.0 0.0
  0.0 -1.0
]
jX = [
  0.0 1.0
  1.0 0.0
]
jY = [
  1.0 0.0
  0.0 1.0
]
@assert Array(R, a, c) == jZ * jX
@assert Array(S, b, c) == jY + jX
@assert Array(T, b, c) == jY - jX


mean(d::LogitNormal; kwargs...) = estimateMean(d,kwargs...) # moved to NormalTransforms

function mode(d::LogitNormal)
    (μ, σ) = params(d)
    # if mu<0 then maximum is left of median, if mu>0 right of median
    # if mu=0 the maximum is either at mu or there are two maxima of the same height
    # here, report the left mode
    interval = μ <= 0.0 ? (0.0,logistic(μ)) : (logistic(μ),1.0)
    resOpt = optimize(x -> -pdf(d, x), interval[1], interval[2])
    Optim.minimizer(resOpt)
end

function var(d::LogitNormal; mean=missing, kwargs...)
    # estimateVariance does not pass kwargs to mean, need to do it here
    m = ismissing(mean) ? estimateMean(d,kwargs...) : mean
    estimateVariance(d; kwargs..., mean=m)
end

### Helpers

# opjective function minimized in mode(d)
# function objectiveModeLogitNormal(mode::Real, mu::Real, sd2::Real)
#     mode <= 0.0 || mode >= 1.0 && return prevfloat(T(Inf))
#     diff =  mu - sd2*(1 - 2*mode) - logit(mode)
#     diff^2
# end


export Linear

const WeightType{T} = AbstractMatrix{T}
const BiasType{T} = Union{Compat.Nothing, AbstractVector{T}}

"""
    Linear{T, WeightType, BiasType, InputType} <: AbstractBlock

Linear layer.

``z^{l+1} = W^l * z^{l} + b^l``

expression of gradients:

``\\delta^{l} = (W^l)^T\\delta^{l+1}``
``\\delta^{l}_w = \\delta^{l+1}(z^l)^T``
``\\delta^{l}_b = \\delta^{l+1}``
"""
mutable struct Linear{
                        T, # Numeric Type
                        WType<:WeightType{T},
                        BType<:BiasType{T},
                        IType<:AbstractVecOrMat{T}
                    } <: AbstractBlock

    weight::WType
    bias::BType

    grad_weight::WType
    grad_bias::BType

    input::IType

    function Linear{T, W, B, I}(
        weight::W,
        bias::B,
        grad_weight::W,
        grad_bias::B
        ) where {T, W, B, I}
        new{T, W, B, I}(weight, bias, grad_weight, grad_bias)
    end
end

include("Utils.jl")
include("Constructor.jl")
include("Forward.jl")
include("Backward.jl")

import Base: show

function show(io::IO, op::Linear)
    dname = string(device(op.weight))
    print(io, "Linear{$(dname)} ($(size(op, 1)) -> $(size(op, 2)))")
end


using Base.Test
using Distributions
# using Gadfly
using Klara

# Target distribution shape, unscaled
g(x) = exp(-dot(x, x)/2)
log_g(x) = -dot(x, x)/2

# Candidate (starting) density, use Normal(0, 2)
g0(x) = 1/(sqrt(2*pi)*2)*exp(-dot(x, x)/8)
log_g0(x) = -dot(x, x)/8-log(2 * sqrt(2*pi))

# Compute M0 such that candidate dominates target
h0(x) = -g(x)/g0(x)

#res0 = optimize(h0, ones(L))
#M0 = -res0.f_minimum

M0 = 10.0
log_M0 = log(M0)

# Candidate (starting) Distribution, use Normal(0, 2)
g1(x) = exp(sum(logpdf(Normal(0, 2), x)))
log_g1(x) = sum(logpdf(Normal(0, 2), x))

# Compute M1 such that candidate dominates target
h1(x) = -g(x)/sum(g1(x))
#res1 = optimize(h1, ones(L))
#M1 = -res1.f_minimum
M1 = 10.0
log_M1 = log(M1)

model = likelihood_model(BasicContUnvParameter(:p, logtarget=log_g), false)

sampler = ARS(log_g0, proposalscale=log_M0)

mcrange = BasicMCRange(nsteps=10000, burnin=1000)

v0 = Dict(:p=>1.)

job = BasicMCJob(model, sampler, mcrange, v0)

run(job)

chain1 = output(job)

sampler = ARS(log_g1, proposalscale=log_M1)

mcrange = BasicMCRange(nsteps=10000, burnin=1000)

v0 = Dict(:p=>1.)

job = BasicMCJob(model, sampler, mcrange, v0)

run(job)

chain2= output(job)

println()
println("M: ", [M0 M1])
println()

println("Acceptance rate: ", acceptance(chain1, diagnostics=false))

println("Acceptance rate: ", acceptance(chain2, diagnostics=false))

theta = hcat(chain1.value, chain2.value);
for i in sample(1:length(theta), 200)
  @test isapprox(h0(theta[i])/M0, h1(theta[i])/M1)
end

# plot(x=chain1.value, Geom.histogram(bincount=20))

# plot(x=chain2.value, Geom.histogram(bincount=20))


struct DSASignature{T} <: AbstractSignature where T<:Integer
    hash::T
    r::T
    s::T
    pubkey::T
end

==(x::DSASignature,y::DSASignature) = x.hash==y.hash && x.r==y.r && x.s==y.s && x.pubkey==y.pubkey

import Base.mod
mod(G::CyclicGroup) = mod(value(G),order(G))
mod(n::Integer,G::CyclicGroup) = mod(n,order(G))

# https://github.com/zydeon/DSA/blob/master/DSA.py
function modinv_euclidean(z::Integer,a::Integer)
    if 0<z<a
	i = a
	j = z
	y1 = BigInt(1)
	y2 = BigInt(0)
	while j > 0
	    q = div(i,j)
	    r = i-j*q
	    y = y2 - y1*q
	    i, j = j, r
	    y2, y1 = y1, y
	    if i == 1
		return mod(y2,a)
            end
        end
    else
        error("Inverse Error")
    end
end

modinv_fermat(k::Integer,a::Integer) = powermod(k,a-2,a)

modinv(k::Integer,G::CyclicGroup) = modinv_euclidean(k,order(G))

function DSASignature(hash,signer::AbstractSigner)

    h = BigInt(hash)
    x = BigInt(signer.privkey)
    G = signer.G

    k = rand(1:order(G)) ### chooses a number from 0 to q
    r = mod(G^k)
    kinv = modinv(k,G)
    s = mod(kinv*(h + x*r),G)

    # This always works for modinv_fermat
    # s < order(G)/2 || s = order(G) - s

    if s==0
        return DSASignature(hash,signer,G)
    else
        return DSASignature(hash,r,s,signer.pubkey)
    end
end

function verify(sr::DSASignature,G::CyclicGroup)
    r,s = sr.r,sr.s
    h = BigInt(sr.hash)
    Y = typeof(G)(sr.pubkey,G)

    0<r<order(G) || return false
    0<s<order(G) || return false

    sinv = modinv(s,G)
    w = mod(sinv,G)
    u1 = mod(h*w,G)
    u2 = mod(r*w,G)
    v = mod(G^u1 * Y^u2)

    if v==r
        return true
    else
        return false
    end
end

import Base.Dict
function Dict(signature::DSASignature)
    config = Dict()
    config["hash"] = string(signature.hash,base=16)
    config["pub"] = string(signature.pubkey,base=16)
    config["r"] = string(signature.r,base=16)
    config["s"] = string(signature.s,base=16)
    return config
end

function DSASignature{BigInt}(dict::Dict)
    hash = parse(BigInt,dict["hash"],base=16)
    r = parse(BigInt,dict["r"],base=16)
    s = parse(BigInt,dict["s"],base=16)
    pub = parse(BigInt,dict["pub"],base=16)
    DSASignature(hash,r,s,pub)
end


function solve_plug_everything(input)
	arr = sort(input)
	ones = arr[1] == 1 ? 1 : 0
	twos = arr[1] == 2 ? 1 : 0
	threes = arr[1] == 3 ? 2 : 1
	for i in 2:length(arr)
		if arr[i] - arr[i - 1] == 1
			ones += 1
		elseif arr[i] - arr[i - 1] == 2
			twos += 1
		else
			threes += 1
		end
	end
	println("twos: " * string(twos)) # <- Note there are no twos!
	ones * threes
end

function solve_2(input)
	arr = sort(input)
	contiguous_ones = []
	ones = arr[1] == 1 ? 0 : -1
	for i in 2:length(arr)
		if arr[i] - arr[i - 1] == 1
			ones += 1
		else
			push!(contiguous_ones, ones)
			ones = -1
		end
	end
	push!(contiguous_ones, ones)
	reduce(*, calc_configs.(filter(x -> x > 0, contiguous_ones)), init = 1)
end

function calc_configs(n)::Int64
	2^n - (n - 2)*(n - 1)/2
end

@assert solve_plug_everything(parse.(Int64, readlines("input/test/test_day10.txt"))) == 220 ["test case 1 failed"]
@assert solve_2(parse.(Int64, readlines("input/test/test_day10.txt"))) == 19208 ["test case 2 failed"]

input = parse.(Int64, readlines("input/day10.txt"))
println("First star: " * string(solve_plug_everything(input)))
println("Second star: " * string(solve_2(input)))


using Test

include("trinary.jl")

@testset "trinary 1 is decimal 1" begin
    @test trinary_to_decimal("1") == 1
end

@testset "trinary 2 is decimal 2" begin
    @test trinary_to_decimal("2") == 2
end

@testset "trinary 10 is decimal 3" begin
    @test trinary_to_decimal("10") == 3
end

@testset "trinary 11 is decimal 4" begin
    @test trinary_to_decimal("11") == 4
end

@testset "trinary 100 is decimal 9" begin
    @test trinary_to_decimal("100") == 9
end

@testset "trinary 112 is decimal 14" begin
    @test trinary_to_decimal("112") == 14
end

@testset "trinary 222 is decimal 26" begin
    @test trinary_to_decimal("222") == 26
end

@testset "trinary 1122000120 is decimal 32091" begin
    @test trinary_to_decimal("1122000120") == 32091
end

@testset "invalid trinary digits returns 0" begin
    @test trinary_to_decimal("1234") == 0
end

@testset "invalid word as input returns 0" begin
    @test trinary_to_decimal("carrot") == 0
end

@testset "invalid numbers with letters as input returns 0" begin
    @test trinary_to_decimal("0a1b2c") == 0
end


include("../src/dependencies.jl")

function plot_data_callbacks(time::Float64, state::NetworkState, event::TimedEvent)
    return nothing
end

function proportion_test(params::NetworkParameters; max_time::Float64, lambda_range::Vector{Float64})
    print("Begin proportion test [")
    mean_system_job_totals = zeros(Float64, length(lambda_range))
    proportion_orbiting_jobs = zeros(Float64, length(lambda_range))

    for i in 1:length(lambda_range)
        print("-")
        params.λ = lambda_range[i]
        Random.seed!(1)
        sim_states = simulate(NetworkState(0, zeros(Int8, params.L), 0), TimedEvent(ArrivalEvent(),0.0, 0, 0.0), params, max_time = max_time, callback = plot_data_callbacks)

        # Test Criteria
        proportion_data = sim_states.orbiting_jobs/sim_states.number_in_system
        !isnan(proportion_data) ? proportion_orbiting_jobs[i] = proportion_data : proportion_orbiting_jobs[i] = 0.0
        @assert proportion_orbiting_jobs[i] >= 0.0 && proportion_orbiting_jobs[i] <= 1.0

    end
    print("]")
    println(proportion_orbiting_jobs)
    println("proportion_test - PASSED!")
end

"""
Call simulate function below for lambdas in lambda_range
"""

test_params = NetworkParameters(  L=3,
                                gamma_shape = 3.0,
                                λ = NaN,
                                η = 4.0,
                                μ_vector = ones(3),
                                P = [0 1.0 0;
                                    0 0 1.0;
                                    0.5 0 0],
                                Q = [0 0.5 0;
                                     0 0 0.5;
                                     0.5 0 0],
                                p_e = AnalyticWeights([1.0, 0, 0]),
                                K = fill(5,3))

lambdas = [i for i in 0.2:0.2:1.2]

proportion_test(test_params, max_time = 1000.0, lambda_range = lambdas)


# # Advection-diffusion of tracer by cellular flow
#
#md # This example can be viewed as a Jupyter notebook via [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/literated/cellularflow.ipynb).
#
# An example demonstrating the advection-diffusion of a tracer by a cellular flow.
#
# ## Install dependencies
#
# First let's make sure we have all required packages installed.

# ```julia
# using Pkg
# pkg"add PassiveTracerFlows, Plots, Printf"
# ```

# ## Let's begin
# Let's load `PassiveTracerFlows.jl` and some other needed packages.
#
using PassiveTracerFlows, Plots, Printf

# ## Choosing a device: CPU or GPU

dev = CPU()     # Device (CPU/GPU)
nothing # hide

# ## Numerical parameters and time-stepping parameters

      n = 128            # 2D resolution = n²
stepper = "RK4"          # timestepper
     dt = 0.02           # timestep
 nsteps = 800            # total number of time-steps
 nsubs  = 25             # number of time-steps for intermediate logging/plotting (nsteps must be multiple of nsubs)
nothing # hide

# ## Numerical parameters and time-stepping parameters

L = 2π        # domain size
κ = 0.002     # diffusivity
nothing # hide

# ## Set up cellular flow
# We create a two-dimensional grid to construct the cellular flow. Our cellular flow is derived
# from a streamfunction ``ψ(x, y) = ψ₀ \cos(x) \cos(y)`` as ``(u, v) = (-∂_y ψ, ∂_x ψ)``.
grid = TwoDGrid(n, L)

ψ₀ = 0.2
mx, my = 1, 1

ψ = [ψ₀ * cos(mx * grid.x[i]) * cos(my * grid.y[j]) for i in 1:grid.nx, j in 1:grid.ny]

uvel(x, y) =  ψ₀ * my * cos(mx * x) * sin(my * y)
vvel(x, y) = -ψ₀ * mx * sin(mx * x) * cos(my * y)
nothing # hide

# ## Problem setup
# We initialize a `Problem` by providing a set of keyword arguments.
prob = TracerAdvectionDiffusion.Problem(dev; nx=n, Lx=L, κ=κ, steadyflow=true, u=uvel, v=vvel,
                                          dt=dt, stepper=stepper)
nothing # hide

# and define some shortcuts
sol, clock, vars, params, grid = prob.sol, prob.clock, prob.vars, prob.params, prob.grid
x, y = grid.x, grid.y
nothing # hide

# ## Setting initial conditions

# Our initial condition for the tracer ``c`` is a gaussian centered at ``(x, y) = (L_x/5, 0)``.

gaussian(x, y, σ) = exp(-(x^2 + y^2) / (2σ^2))

amplitude, spread = 0.5, 0.15
c₀ = [amplitude * gaussian(x[i] - 0.2 * grid.Lx, y[j], spread) for i=1:grid.nx, j=1:grid.ny]

TracerAdvectionDiffusion.set_c!(prob, c₀)
nothing # hide

# Let's plot the initial tracer concentration and streamlines. Note that when plotting, we decorate
# the variable to be plotted with `Array()` to make sure it is brought back on the CPU when
# `vars` live on the GPU.

function plot_output(prob)
  c = prob.vars.c

  p = heatmap(x, y, Array(vars.c'),
         aspectratio = 1,
              c = :balance,
         legend = :false,
           clim = (-0.2, 0.2),
          xlims = (-grid.Lx/2, grid.Lx/2),
          ylims = (-grid.Ly/2, grid.Ly/2),
         xticks = -3:3,
         yticks = -3:3,
         xlabel = "x",
         ylabel = "y",
          title = "initial tracer concentration (shading) + streamlines",
     framestyle = :box)

  contour!(p, x, y, Array(ψ'),
     levels=0.0125:0.025:0.2,
     lw=2, c=:black, ls=:solid, alpha=0.7)

  contour!(p, x, y, Array(ψ'),
     levels=-0.1875:0.025:-0.0125,
     lw=2, c=:black, ls=:dash, alpha=0.7)

  return p
end
nothing # hide

# ## Time-stepping the `Problem` forward

# We time-step the `Problem` forward in time.

startwalltime = time()

p = plot_output(prob)

anim = @animate for j = 0:round(Int, nsteps/nsubs)

 if j % (200 / nsubs) == 0
    log = @sprintf("step: %04d, t: %d, walltime: %.2f min",
                   clock.step, clock.t, (time()-startwalltime)/60)

    println(log)
  end

  p[1][1][:z] = Array(vars.c)
  p[1][:title] = "concentration, t=" * @sprintf("%.2f", clock.t)

  stepforward!(prob, nsubs)
  TracerAdvectionDiffusion.updatevars!(prob)
end

mp4(anim, "cellularflow.mp4", fps=12)


cumtrapz(y, x, x0=0) = cumsum(vcat(0, diff(x).*(y[1:end-1] .+ y[2:end])/2)) .+ x0
invcumtrapz(y, x, x0=0) = cumsum(vcat(0, 2*diff(x)./(y[1:end-1] .+ y[2:end]))) .+ x0

@inline function distance2(line::Tuple{<:AbstractVector,<:AbstractVector}, x)
    v = line[2] - line[1]
    λ = clamp(v⋅(x - line[1])/(v⋅v), 0, 1)
    p = (1 - λ)*line[1] + λ*line[2]
    (p - x)⋅(p - x)
end


using Pkg
using JuMP;
using Plots;
using Distributions
using Random
using StatsBase

function create_serie()
    serie = zeros(120).+7
    d = Normal(0.0, 0.3)
    for t in 1:120
        serie[t] += cos(2*pi*t/12 + 6) + rand(d, 1)[1]
    end
    return serie
end

function create_dummy()
    A = zeros(120, 12)

    A[:,1] .= 1

    for i in 1:120
        col = mod(i,12)
        if col != 0
            A[i,col+1] = 1
        end
    end
    return A
end

function create_trig()
    A = zeros(120, 12)

    A[:,1] .= 1

    for t in 1:120, i = 2:12
        use_cos = (mod(i,2) == 0)

        j = (use_cos ? i/2 : i/2 - 0.5)
        if use_cos
            A[t,i] = cos(pi/6*j*t)
        else
            A[t,i] = sin(pi/6*j*t)
        end
    end
    return A
end

function solve_MQO(A,b)
    x = A'A \ A'b
    return x
end

b = create_serie()
Ad = create_dummy()
At = create_trig()
xd = solve_MQO(Ad,b)
xt = solve_MQO(At,b)

yd = Ad*xd
yt = At*xt

minimum(isapprox.(yd,yt))

p = plot(autocov(yd - b))

savefig(p,"Modelos Estruturais\\Lista 1\\imagens\\FAC.png")


using Combinatorics

const data = parse.(Int, readlines())

function solve_1()
    n = 150
    dp = fill(0, (n + 1,))
    dp[1] = 1

    for d in data
        for i in n+1:-1:1
            if i+d <= n+1
                dp[i+d] += dp[i]
            end
        end
    end
    dp[end]
end

function solve_2()
    n = 150
    m = length(data)
    dp = zeros(Int, n+1, m+1)

    dp[1, 1] = 1

    for d in data
        for i in n+1:-1:1
            for j in 1:m
                if i+d <= n+1
                    dp[i+d, j+1] += dp[i, j]
                end
            end
        end
    end
    filter(>(0), dp[end, :])[1]
end

solve_1() |> println

solve_2() |> println


# ------------------------------------------ #
#                  Constants                 #
# ------------------------------------------ #
Pr = 13.5
k = 0.57  # thermal conductivity of water, W / (K*m)
nu = mu/rhow # kin. viscosity

# ------------------------------------------ #
#       Nusselt number parameterisations     #
# ------------------------------------------ #

# Where A, a and b are defined: Dittus-Boelter correlation A * Pr^a * Re^b
coefs = (
    standard = (A=0.023, a=2/5, b=0.8),
    sommers = (A=0.0025, a=0.4, b=0.95),
    lunardini = (A=0.0078, a=1/3, b=0.927),
    vincent = (A=0.332, a=1/3, b=0.74),
    ogier = (A=1.78, a=1/3, b=0.58),
    gnielinski = ()
)

# ------------------------------------------ #
#                  Formulas                  #
# ------------------------------------------ #

r(S) = sqrt(S/pi)
Re(Q,S) = Q/S * 2 * r(S) / nu
Nu_Dittus_Boelter(Re, p) = p.A * Pr^p.a .* Re.^p.b
Nu_Gnielinski(f, Re) = f./8 .* (Re .- 1000) .* Pr ./
                       (1 .+ 12.7 .* (f./8).^0.5 .* (Pr^(2/3) .- 1))
# select between the two depending on parameter:
Nu_fn(f, Re, p::Tuple{}) = Nu_Gnielinski(f, Re)
Nu_fn(f, Re, p) = Nu_Dittus_Boelter(Re, p)

melt_energy_eq(Q, dphi_dz, dp_dz) = -Q * ( dphi_dz + rhow*cw*ct*dp_dz ) # Energy

# from Sommers and RAJARAM 2020, Eq. 29
z_eq_fn(Q, Nu) = rhow * cw * Q ./ (pi * k * Nu) # e-folding length, m
tau_eq_fn(Q, dphi_dz, dp_dz, Nu) = melt_energy_eq(Q, dphi_dz, dp_dz) / (k * pi * Nu ) # equilibrium offset-temperature: T_eq - T_melt, °C

# same paper, modification of Eq. 21
tau_w_diff(z_eq, dtauw_dz, tau_eq) = tau_eq - z_eq * dtauw_dz # water offset-temperature, T_w - T_melt, °C, calculated from differential equation
tau_w_sol(z, tau_eq, t0, z_eq) = tau_eq + (t0-tau_eq)*exp(-z/z_eq) # water offset-temperature, T_w - T_melt, °C, calculated from solution of differential equation as a function of z

# ------------------------------------------ #
#          Initialise dictionaries           #
# ------------------------------------------ #

z_eq = Dict()
tau_eq = Dict()
tau_w = Dict()
tau_diff = Dict() # Tw - Teq = tau_w - tau_eq
T_surf = Dict()
tauw_measured = Dict()
dtauw_dz = Dict()
Nus = Dict()

# ------------------------------------------ #
#    Calculate heat transfer parameters      #
# ------------------------------------------ #

days = ["0908", "2108"]

for date in days

    # Quantities derived from field experiments
    Qs = mid_309_265[date][:Q]
    Ss = mid_309_265[date][:S]
    f = mid_309_265[date][:f]
    dphi_dzs = mid_309_265[date][:dphi_dz]
    dp_dzs = mid_309_265[date][:dp] ./ mid_309_265[date][:dz]
    T_ct = ct * mid_309_265[date][:p_mean]

    # Data using output of free-gradient model: dtauw_dz = dTw_dz - dTmelt_dz, the first term comes from the MCMC runs
    local dtauw_dz = bootstrap(Particles(model_runs[date][:dTdz_MCMC]), n_partcl) .-  model_runs[date][:dTdz_ct]  # first term has to be converted from array to particles

    # Temperature measurements
    i = indexin(DateTime.(mid_309_265[date][:t_inj]), DateTime.(ctd309[date][:t]))  # extract the indices where there was a tracer experiment
    tauw_measured[date] = Particles(n_partcl, Normal(                             # compute the mean value over the date since it does not vary much
            mean(
            [mean(ctd309[date][:temp_PT][i]),
             mean(ctd265[date][:temp_PT][i])]), e_T)
            ) .- mean(T_ct, dims=1)

    # mean depth of test-section
    d = mean([ctd309[date][:depth], ctd265[date][:depth]], dims=1)

    # compute Reynolds number
    Res = Re.(Qs, Ss)

    # e-folding length
    z_eq[date] = Dict()
    for n = 1:length(coefs)
        c = coefs[n]
        z_eq[date][keys(coefs)[n]] = z_eq_fn.(Qs, Nu_fn(f, Res, c))
    end

    # equilibrium offset-temperature and Nusselt number
    tau_eq[date] = Dict()
    Nus[date] = Dict()
    for n = 1:length(coefs)
        c = coefs[n]
        tau_eq[date][keys(coefs)[n]] = tau_eq_fn.(Qs, dphi_dzs, dp_dzs, Nu_fn(f, Res, c))
        Nus[date][keys(coefs)[n]] = Nu_fn(f, Res, c)
    end

    # water offset-temperature, its difference to tau_eq and the estimated surface temperature
    tau_w[date] = Dict() # Tw - Tmelt
    tau_diff[date] = Dict()
    T_surf[date] = Dict()
    for k = keys(coefs)
        c = coefs[k]
        tau_w[date][k] = tau_w_diff.(z_eq[date][k], dtauw_dz, tau_eq[date][k])
        tau_diff[date][k] = tau_w[date][k] .- tau_eq[date][k]
        T_surf[date][k] = tau_w_sol.(-d, tau_eq[date][k], tau_w[date][k], z_eq[date][k])
    end
end


struct Householder <: SRSolver end

function impliedVolatilitySRHalley(
    isCall::Bool,
    price::T,
    f::T,
    strike::T,
    tte::T,
    df::T,
    ftolrel::T,
    maxEval::Int,
    solver::Householder,
)::T where {T}
    c, ex = normalizePrice(isCall, price, f, strike, df)
    if c >= 1 / ex || c <= 0
        throw(DomainError(c, string("Price out of range, must be < ", 1 / ex, " and > 0")))
    end
    x = log(ex)
    guess = T(impliedVolatilitySRGuessUndiscountedCall(c, ex,x))
    b = guess
    xtolrel = 32 * eps(T)
    xtol = T(0)
    ftolrel = max(ftolrel, eps(T))
    ftol = ftolrel * max(1, c)
    logc = log(c)
    fb, fbOverfpb, fp2bOverfpb, fp3bOverfpb = objectiveHouseholderLog(x, ex, b, logc)
    if abs(fb) < ftol
        return b / sqrt(tte)
    end
    for iteration = 1:maxEval
        x0 = b
        lf = (fbOverfpb) * (fp2bOverfpb)
        hn = -fbOverfpb
        num = (1 + fp2bOverfpb * hn / 2)
        denom = (1 + fp2bOverfpb * hn + fp3bOverfpb / 6 * hn * hn)
        x1 = x0 + hn * num / denom
        a = x0
        fb, fbOverfpb, fp2bOverfpb, fp3bOverfpb = objectiveHouseholderLog(x, ex, x1, logc)
        b = x1
        xtol_ = xtol + max(1, abs(b)) * xtolrel
        if abs(b - a) <= xtol_ || abs(fb) < ftol
            break
        end
    end
    return b / sqrt(tte)
end

function objectiveHouseholderLog(x::T, ex::T, v::T, logc::T) where {T}
    v = abs(v)
    h = x / v
    t = v / 2
    h2 = h^2
    t2 = t^2
    sqrt2 = sqrt(T(2))
    Np = erfcx(-(h + t) / sqrt2)
    Nm = erfcx(-(h - t) / sqrt2)
    eh2t2 = exp(-(h2 + t2) / 2)
    norm = 1 / (2 * sqrt(ex)) * eh2t2
    cEstimate = norm * (Np - Nm)
    logcEstimate = log(cEstimate)
    twopi = 2 * T(pi)
    logvega = (2 /sqrt(twopi)) / (Np - Nm)  #u'/u
    volgaOverVega = (h + t) * (h - t) / v
    logvolgaOverVega = volgaOverVega - logvega  #u'' u - u'2 /u^2
    h2mt2 = h2 - t2
    c3OverVega = (-3 * h2 - t2 + h2mt2 * h2mt2) / (v^2)   #
    logc3overVega = c3OverVega - 3 * logvega * volgaOverVega + 2 * logvega * logvega
    return logcEstimate - logc, (logcEstimate - logc) / logvega, logvolgaOverVega, logc3overVega
end

function objectiveHouseholder(x::T, ex::T, v::T, c::T) where {T}
    v = abs(v)
    h = x / v
    t = v / 2
    h2 = h^2
    t2 = t^2
    sqrt2 = sqrt(T(2))
    Np = erfcx(-(h + t) / sqrt2)
    Nm = erfcx(-(h - t) / sqrt2)
    eh2t2 = exp(-(h2 + t2) / 2)
    norm = 1 / (2 * sqrt(ex)) * eh2t2
    cEstimate = norm * (Np - Nm)
    eht = exp(-h * t)
    twopi = 2 * T(pi)
    vega = eh2t2 * eht / sqrt(twopi)    #normpdf(h + t)  #this is exp(-0.5* (h^2+t^2 +2*h*t)) = eh2t2*exp(-h*t)
    cEstimateOverVega = (Np - Nm) / (2 * eht * sqrt(ex / (twopi))) - c / vega
    volgaOverVega = (h + t) * (h - t) / v
    h2mt2 = h2 - t2

    c3OverVega = (-3 * h2 - t2 + h2mt2 * h2mt2) / (v^2)
    return cEstimate - c, cEstimateOverVega, volgaOverVega, c3OverVega
end