#version 450

// This variable is set by the pipeline
layout (constant_id = 0) const uint BUFFER_ELEMENTS = 32;

layout(binding = 0) buffer Pos {
   uint values[ ];
};

layout (local_size_x = 4, local_size_y = 1, local_size_z = 1) in;

float learningRate = 0.01;

uint fibonacci(uint n) {
	if(n <= 1){
		return n;
	}
	uint curr = 1;
	uint prev = 1;
	for(uint i = 2; i < n; ++i) {
		uint temp = curr;
		curr += prev;
		prev = temp;
	}
	return curr;
}

float sigmoid(float z) {
    return 1.0 / (1.0 + exp(-z));
}

float forwardProp(float x, float w, float b) {
    float z = dot(w, x) + b;
    float yHat = sigmoid(z);
    return yHat;
}

void backwardProp(float x, float yHat, float y, float m, out float w, out float b) {
    float dZ = yHat - y;
    float dW = (1/m) * x * dZ;
    float dB = (1/m) * dZ;
    w -= (learningRate * dW);
    b -= (learningRate * dB);
}

float calculateLoss(float yHat, float y) {
    return -(y * log(yHat) + (1.0 - y) * log(1.0 - yHat));
}

shared uint sharedTotal[1];

void main() 
{
	uint index = gl_GlobalInvocationID.x;
	if (index >= BUFFER_ELEMENTS) 
		return;	

//    float m = int(BUFFER_ELEMENTS);
//    float w = 0.001;
//    float b = 0.0;
//
//    float x = values[index];
//
//    float yHat = forwardProp(x, w, b);
//    float loss = calculateLoss(yHat, x);
//
//    backwardProp(x, yHat, x, m, w, b);
//
//	values[index] = fibonacci(values[index]);

    sharedTotal[0] = 0;

    barrier();
    memoryBarrierShared();

    atomicAdd(sharedTotal[0], values[index]);

    barrier();
    memoryBarrierShared();

    values[index] = int(sharedTotal[0]);
}