#version 450
#extension GL_ARB_separate_shader_objects : enable

layout(location = 0) out vec4 outColour;
layout(location = 1) out vec4 pickOutColour;
layout(location = 0) in vec3 fragColour;
layout(location = 1) in vec2 fragTexCoord;
layout(location = 2) in flat uint fragTexIndex;
layout(location = 3) in vec3 fragExtraColour;
layout(location = 4) in vec3 fragNormal;
layout(location = 5) in flat uint fragTileIndex;
layout(location = 6) in vec3 fragCoord;
layout(set = 1, binding = 1) uniform sampler2D texSampler;

layout(binding = 0) uniform UniformBufferObject {
    mat4 normalMat;
    mat4 view;
    mat4 proj;
    uint mode;
} ubo;

// for more information look at lightinfo.h
struct LightInfo {
    uint type;
    float intensity;
    float coneWidth;
    uint isActive;
    vec3 position;
    vec3 direction;
    vec3 colour;
};

struct TileRoomInfo {
    uint roomId;
    float temperature;
    LightInfo lightInfo;
    uint lightIndices[10];
};

layout(std430, set = 0, binding = 1) buffer WorldInfo
{
    uint xSize;
    uint ySize;
    TileRoomInfo tileRoomInfos[ ];
};

layout(std430, set = 0, binding = 2) buffer LightInfos
{
    LightInfo lightInfos[ ];
};

float calculateColourLuminance(vec3 colour) {
    return colour.r * 0.21 + colour.g * 0.71 + colour.b * 0.07;
}

void main() {
    bool drawSolid = (ubo.mode >> 2) == 2 || (ubo.mode >> 2) == 0;

    // pure colour-only output used for picking walls
    float selectionAlpha = (drawSolid) ? 1.0 : 0;
    pickOutColour = vec4(fragColour, selectionAlpha);

    // regular render output
    float alpha = (drawSolid) ? 1 : 0.5;
    vec3 textureColour = texture(texSampler, vec2(fragTexCoord[0] + (128 * fragTexIndex), fragTexCoord[1])).rgb;
    if (fragExtraColour[0] != 0.0 || fragExtraColour[1] != 0.0 || fragExtraColour[2] != 0.0) {
        textureColour = (textureColour + fragExtraColour) / 2;
        outColour = vec4(textureColour, alpha);
        return;
    }
    outColour = vec4(textureColour, alpha);

    // set the environment light as a default for the final light info
    LightInfo ambientLight = tileRoomInfos[fragTileIndex].lightInfo;
    float intensity = ambientLight.intensity;
    // factor in the ambient light direction as intensity
    ambientLight.intensity = clamp(max(dot(fragNormal, ambientLight.direction), intensity * 0.80), 0.0, intensity);

    LightInfo diffuseLight = { 0, 0.0, 0.0, 1, { 0.0, 0.0, 0.0 }, { 0.0, 0.0, 0.0 }, { 0.0, 0.0, 0.0 }};
    float lightIntensities[10];
    float maxIntensity = 0.0;
    float maxAbsoluteIntensity = 0.0;
    // first pass, determine the number of light sources as well as the maximum value
    uint numElements = 0;
    for (uint i = 0; i < 10; i++) {
        if (tileRoomInfos[fragTileIndex].lightIndices[i] != 0) {
            numElements++;
            LightInfo dynamicLightInfo = lightInfos[tileRoomInfos[fragTileIndex].lightIndices[i]];

            // do not do anything if the light source is not currently active
            if (dynamicLightInfo.isActive == 0) {
                lightIntensities[i] = 0;
                continue;
            }

            vec3 lightPos = dynamicLightInfo.position;
            float lightIntensity = dynamicLightInfo.intensity;
            vec3 direction = normalize(lightPos - fragCoord);
            // have a fall-off rate of 20% *after* the first tile
            float distanceOffset = max(0, distance(lightPos, fragCoord) - 1) / 5;
            // keep track of the absolute maximum intensity
            float absoluteIntensity = max(lightIntensity - distanceOffset, 0);
            // calculate the intensity given the angle surfae and light direction
            lightIntensities[i] = max(dot(fragNormal, direction) - distanceOffset, 0.0) * lightIntensity;

            // if this a spot light (type 2) make sure to take the cone width into consideration
            if (
            dynamicLightInfo.type == 2
            && ((abs(acos(dot(direction, dynamicLightInfo.direction))) / 3.14) * 180) > dynamicLightInfo.coneWidth
            ) {
                lightIntensities[i] = 0;
                lightIntensity = 0;
                absoluteIntensity = 0;
            }

            diffuseLight.intensity += lightIntensities[i];
            if (maxIntensity < lightIntensities[i]) {
                maxIntensity = lightIntensities[i];
            }
            if (maxAbsoluteIntensity < absoluteIntensity) {
                maxAbsoluteIntensity = absoluteIntensity;
            }
        }
    }

    // second pass: calculate the final colour
    vec3 tempColour = { 0, 0, 0 };
    float allIntensity = 0;
    for (uint i = 0; i < numElements && maxIntensity > 0; i++) {
        float relativeIntensity = lightIntensities[i] / maxIntensity;
        allIntensity += relativeIntensity;
        tempColour = lightInfos[tileRoomInfos[fragTileIndex].lightIndices[i]].colour;
        diffuseLight.colour.r += tempColour.r * relativeIntensity;
        diffuseLight.colour.g += tempColour.g * relativeIntensity;
        diffuseLight.colour.b += tempColour.b * relativeIntensity;
    }

    // check for the edge case of no light source at all
    if (numElements > 0 && allIntensity > 0) {
        diffuseLight.colour /= allIntensity;
    }

    // make sure not to create overly bright light sources
    diffuseLight.intensity = clamp(diffuseLight.intensity, 0, maxAbsoluteIntensity);

    // determine the factor in which the light source colour and ambient light is used
    float sourceFactor = clamp(diffuseLight.intensity / ambientLight.intensity, 0, 1);

    float luminanceDiffuse = calculateColourLuminance(diffuseLight.colour);
    float luminanceAmbient = calculateColourLuminance(ambientLight.colour);
    diffuseLight.colour = (diffuseLight.colour * sourceFactor) + (ambientLight.colour * (1 - sourceFactor));
    float newLuminance =  calculateColourLuminance(diffuseLight.colour);
    if (newLuminance < max(newLuminance, luminanceAmbient)) {
        diffuseLight.colour *= 1 + (max(newLuminance, luminanceAmbient) - newLuminance);
    }
    // make sure not to calculate a light intensity that is lower than either of its components
    float targetIntensity = max(maxAbsoluteIntensity, ambientLight.intensity);
    diffuseLight.intensity = max(targetIntensity, (diffuseLight.intensity * sourceFactor) + (ambientLight.intensity * (1 - sourceFactor)));

    // get the weighted ambient colour - if we get too close to a neutral "weight"
    // make sure to gradually tone it down, as things otherwise look too pale
    float sum = dot(diffuseLight.colour, vec3(1, 1, 1));
    float factor  = (sum > 2.5) ? 0.2 * ((3 - sum) / 0.5) : 0.2;
    outColour.rgb = outColour.rgb * (1 - factor) + diffuseLight.colour * factor;
    outColour.rgb *= diffuseLight.intensity;
}
