using CSV, Tables, Glob


const TEST_DATA_DIR = joinpath(@__DIR__, "..", "data")


function csv_to_matrix(file_path, types=Float32)
    return CSV.read(file_path, Tables.matrix, header=0, delim=',', types=types)
end


function load_test_data()
    inputs_path = joinpath(TEST_DATA_DIR, "inputs.csv")
    inputs = vec(csv_to_matrix(inputs_path))
    weight_matrices, bias_vectors = Matrix{Float32}[], Vector{Float32}[]
    for file_path in glob("weights*.csv", TEST_DATA_DIR)
        push!(weight_matrices, csv_to_matrix(file_path))
    end
    for file_path in glob("biases*.csv", TEST_DATA_DIR)
        push!(bias_vectors, vec(csv_to_matrix(file_path)))
    end
    return inputs, [zip(weight_matrices, bias_vectors)...]
end


function sigmoid(x)
    return 1 ./ (1 .+ exp.(-x))
end


function layer_func(input_vector, weight_matrix, bias_vector)
    return sigmoid(weight_matrix * input_vector + bias_vector)
end


function feed_forward(x, layers...)
    for (w, b) in layers
        x = layer_func(x, w, b)
    end
    return x
end


function main()
    n = parse(Int32, ARGS[1])
    x, weights_biases = load_test_data()
    
    # # To get info about memory allocations etc.:
    # @time feed_forward(x, weights_biases...)

    # Call once to compile, then again to time
    @elapsed feed_forward(x, weights_biases...)
    t = @elapsed begin
        for _ in 1:n
            feed_forward(x, weights_biases...)
        end
    end
    println(round(t, digits=5))
end


if abspath(PROGRAM_FILE) == @__FILE__
    main()
end