module Interpreter
using CUDA
using StaticArrays
using ..ExpressionProcessing
using ..Utils

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{Expr}, variables::Matrix{Float32}, parameters::Vector{Vector{Float32}})::Matrix{Float32}
	
	exprs = Vector{ExpressionProcessing.PostfixType}(undef, length(expressions))
	@inbounds for i in eachindex(expressions)
		exprs[i] = ExpressionProcessing.expr_to_postfix(expressions[i])
	end
	
	variableCols = size(variables, 2) # number of variable sets to use for each expression
	cudaVars = CuArray(variables)
	cudaParams = Utils.create_cuda_array(parameters, NaN32) # column corresponds to data for one expression
	cudaExprs = Utils.create_cuda_array(exprs, 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([Utils.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(exprs))

	# Start kernel for each expression to ensure that no warp is working on different expressions
	@inbounds for i in eachindex(exprs)
		kernel = @cuda launch=false fastmath=true interpret_expression(cudaExprs, cudaVars, cudaParams, cudaResults, cudaStepsize, i)
		# config = launch_configuration(kernel.fun)
		threads = min(variableCols, 128)
		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 depth of the stack to store the values and intermediate results
function interpret_expression(expressions::CuDeviceArray{ExpressionElement}, variables::CuDeviceArray{Float32}, parameters::CuDeviceArray{Float32}, results::CuDeviceArray{Float32}, stepsize::CuDeviceArray{Int}, exprIndex::Int)
	varSetIndex = (blockIdx().x - 1) * blockDim().x + threadIdx().x # ctaid.x * ntid.x + tid.x (1-based)
	@inbounds variableCols = length(variables) / stepsize[2]

	if varSetIndex > variableCols
		return
	end

	# firstExprIndex = ((exprIndex - 1) * stepsize[1]) + 1 # Inclusive
	# lastExprIndex = firstExprIndex + stepsize[1] - 1 # Inclusive
	@inbounds firstParamIndex = ((exprIndex - 1) * stepsize[1]) # Exclusive

	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
	
	@inbounds firstVariableIndex = ((varSetIndex-1) * stepsize[2]) # Exclusive
	
	@inbounds for expr in expressions
		if expr.Type == EMPTY
			break
		elseif expr.Type == INDEX
			val = expr.Value
			operationStackTop += 1

			if val > 0
				operationStack[operationStackTop] = variables[firstVariableIndex + val]
			else
				val = abs(val)
				operationStack[operationStackTop] = parameters[firstParamIndex + val]
			end
		elseif expr.Type == FLOAT32
			operationStackTop += 1
			operationStack[operationStackTop] = reinterpret(Float32, expr.Value)
		elseif expr.Type == OPERATOR
			type = reinterpret(Operator, expr.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] = NaN32
			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
	@inbounds results[resultIndex] = operationStack[operationStackTop]

	return
end

end