module Interpreter
using CUDA
using StaticArrays
using ..ExpressionProcessing

export interpret

"Interprets the given expressions with the values provided.
# Arguments
 - expressions::Vector{ExpressionProcessing.PostfixType} : The expressions to execute in postfix form
 - variables::Matrix{Float32} : The variables to use. Each column is mapped to the variables x1..xn
 - parameters::Vector{Vector{Float32}} : The parameters to use. Each Vector contains the values for the parameters p1..pn. The number of parameters can be different for every expression
"
function interpret(expressions::Vector{ExpressionProcessing.PostfixType}, variables::Matrix{Float32}, parameters::Vector{Vector{Float32}})::Matrix{Float32}
	variableCols = size(variables, 2) # number of sets of variables to use for each expression
	cudaVars = CuArray(variables)
	cudaParams = create_cuda_array(parameters, NaN32) # column corresponds to data for one expression
	cudaExprs = create_cuda_array(expressions, ExpressionElement(EMPTY, 0)) # column corresponds to data for one expression
	# put into seperate cuArray, as this is static and would be inefficient to send seperatly to every kernel
	cudaStepsize = CuArray([get_max_inner_length(expressions), get_max_inner_length(parameters), size(variables, 1)]) # max num of values per expression; max nam of parameters per expression; number of variables per expression

	# each expression has nr. of variable sets (nr. of columns of the variables) results and there are n expressions
	cudaResults = CuArray{Float32}(undef, variableCols, length(expressions))

	# Start kernel for each expression to ensure that no warp is working on different expressions
	for i in eachindex(expressions)
		kernel = @cuda launch=false interpret_expression(cudaExprs, cudaVars, cudaParams, cudaResults, cudaStepsize, i)
		config = launch_configuration(kernel.fun)
		threads = min(variableCols, config.threads)
		blocks = cld(variableCols, threads)

		kernel(cudaExprs, cudaVars, cudaParams, cudaResults, cudaStepsize, i; threads, blocks)
	end

	return cudaResults
end

#TODO: Add @inbounds to all indexing after it is verified that all works https://cuda.juliagpu.org/stable/development/kernel/#Bounds-checking
const MAX_STACK_SIZE = 25 # The max number of values the expression can have. so Constant values, Variables and parameters
function interpret_expression(expressions::CuDeviceArray{ExpressionElement}, variables::CuDeviceArray{Float32}, parameters::CuDeviceArray{Float32}, results::CuDeviceArray{Float32}, stepsize::CuDeviceArray{Int}, exprIndex::Int)
	index = (blockIdx().x - 1) * blockDim().x + threadIdx().x # ctaid.x * ntid.x + tid.x
	stride = gridDim().x * blockDim().x # nctaid.x * ntid.x
	
	firstExprIndex = ((exprIndex - 1) * stepsize[1]) + 1 # Inclusive
	lastExprIndex = firstExprIndex + stepsize[1] - 1 # Inclusive
	firstParamIndex = ((exprIndex - 1) * stepsize[2]) # Exclusive
	variableCols = length(variables) / stepsize[3]

	operationStack = MVector{MAX_STACK_SIZE, Float32}(undef) # Try to get this to function with variable size too, to allow better memory usage
	operationStackTop = 0 # stores index of the last defined/valid value
	
	for varSetIndex in index:stride
		firstVariableIndex = ((varSetIndex - 1) * stepsize[3]) # Exclusive
		
		for i in firstExprIndex:lastExprIndex
			if expressions[i].Type == EMPTY
				break
			elseif expressions[i].Type == INDEX
				val = expressions[i].Value
				operationStackTop += 1

				if val > 0
					operationStack[operationStackTop] = variables[firstVariableIndex + val]
				else
					val = -val
					operationStack[operationStackTop] = parameters[firstParamIndex + val]
				end
			elseif expressions[i].Type == FLOAT32
				operationStackTop += 1
				operationStack[operationStackTop] = reinterpret(Float32, expressions[i].Value)
			elseif expressions[i].Type == OPERATOR
				type = reinterpret(Operator, expressions[i].Value)
				if type == ADD
					operationStackTop -= 1
					operationStack[operationStackTop] = operationStack[operationStackTop] + operationStack[operationStackTop + 1]
				elseif type == SUBTRACT
					operationStackTop -= 1
					operationStack[operationStackTop] = operationStack[operationStackTop] - operationStack[operationStackTop + 1]
				elseif type == MULTIPLY
					operationStackTop -= 1
					operationStack[operationStackTop] = operationStack[operationStackTop] * operationStack[operationStackTop + 1]
				elseif type == DIVIDE
					operationStackTop -= 1
					operationStack[operationStackTop] = operationStack[operationStackTop] / operationStack[operationStackTop + 1]
				elseif type == POWER
					operationStackTop -= 1
					operationStack[operationStackTop] = operationStack[operationStackTop] ^ operationStack[operationStackTop + 1]
				elseif type == ABS
					operationStack[operationStackTop] = abs(operationStack[operationStackTop])
				elseif type == LOG
					operationStack[operationStackTop] = log(operationStack[operationStackTop])
				elseif type == EXP
					operationStack[operationStackTop] = exp(operationStack[operationStackTop])
				elseif type == SQRT
					operationStack[operationStackTop] = sqrt(operationStack[operationStackTop])
				end
			else
				operationStack[operationStackTop] = NaN
				break
			end
		end
		# "(exprIndex - 1) * variableCols" -> calculates the column in which to insert the result (expression = column)
		# "+ varSetIndex" -> to get the row inside the column at which to insert the result of the variable set (variable set = row)
		resultIndex = convert(Int, (exprIndex - 1) * variableCols + varSetIndex) # Inclusive
		results[resultIndex] = operationStack[operationStackTop]
	end

	return
end


"Retrieves the number of entries for the largest inner vector"
function get_max_inner_length(vec::Vector{Vector{T}})::Int where T
	maxLength = 0
	@inbounds for i in eachindex(vec)
		if length(vec[i]) > maxLength
			maxLength = length(vec[i])
		end
	end

	return maxLength
end

"Returns a CuArray filed with the data provided. The inner vectors do not have to have the same length. All missing elements will be the value ```invalidElement```"
function create_cuda_array(data::Vector{Vector{T}}, invalidElement::T)::CuArray{T} where T
	dataCols = get_max_inner_length(data)
	dataRows = length(data)
	dataMat = convert_to_matrix(data, invalidElement)
	cudaArr = CuArray{T}(undef, dataCols, dataRows) # length(parameters) == number of expressions
	copyto!(cudaArr, dataMat)

	return cudaArr
end

"Converts a vector of vectors into a matrix. The inner vectors do not need to have the same length.

All entries that cannot be filled have ```invalidElement``` as their value
"
function convert_to_matrix(vec::Vector{Vector{T}}, invalidElement::T)::Matrix{T} where T
	vecCols = get_max_inner_length(vec)
	vecRows = length(vec)
	vecMat = fill(invalidElement, vecCols, vecRows)
	
	for i in eachindex(vec)
		vecMat[:,i] = copyto!(vecMat[:,i], vec[i])
	end

	return vecMat
end



# Kernel
function InterpretExplicit!(op::Operator, x, y)
	index = (blockIdx().x - 1) * blockDim().x + threadIdx().x
	stride = gridDim().x * blockDim().x

	if op == ADD
		# @cuprintln("Performing Addition") # Will only be displayed when the GPU is synchronized
		for i = index:stride:length(y)
			@inbounds y[i] += x[i]
		end
		return
	elseif op == SUBTRACT
		# @cuprintln("Performing Subtraction") # Will only be displayed when the GPU is synchronized
		for i = index:stride:length(y)
			@inbounds y[i] -= x[i]
		end
		return
	end
end

end