#version 440 core
#define NR_POINT_LIGHTS 2
#define NR_DIR_LIGHTS 1
#define NR_SPOT_LIGHTS 1

//based on tutorial https://learnopengl.com/PBR/Lighting

layout (location = 0) out vec4 FragColor;
layout (location = 1) out vec4 BrightColor;


struct DirectionalLight {
  vec3 direction;
  vec3 color;
  float intensity;

  float ambientLight;
  float specularLight;
};

struct PointLight {
  vec3 position;
  vec3 color;
  float intensity;
  float a;          // Constant attenuation
  float b;          // Linear attenuation

  float ambientLight;
  float specularLight;
};

struct SpotLight {
  vec3 position;
  vec3 direction;
  float intensity;
  vec3 color;
  float cutOff;
  float outerCutOff;
};


in vec2 TexCoords;
in vec3 WorldPos;
in vec3 Normal;

// material parameters
uniform sampler2D albedoMap;
uniform sampler2D normalMap;
uniform sampler2D metallicMap;
uniform sampler2D roughnessMap;

uniform DirectionalLight dirLights[NR_DIR_LIGHTS];
uniform PointLight pointLights[NR_POINT_LIGHTS];
uniform SpotLight spotLights[NR_SPOT_LIGHTS];
uniform SpotLight flashLight;


uniform vec3 cameraPos;

uniform bool nightVision;

const float PI = 3.14159265359;




// ----------------------------------------------------------------------------
vec3 getNormalFromMap()
{
    vec3 tangentNormal = texture(normalMap, TexCoords).xyz * 2.0 - 1.0;

    vec3 Q1  = dFdx(WorldPos);
    vec3 Q2  = dFdy(WorldPos);
    vec2 st1 = dFdx(TexCoords);
    vec2 st2 = dFdy(TexCoords);

    vec3 N   = normalize(Normal);
    vec3 T  = normalize(Q1*st2.t - Q2*st1.t);
    vec3 B  = -normalize(cross(N, T));
    mat3 TBN = mat3(T, B, N);

    return normalize(TBN * tangentNormal);
}

// ----------------------------------------------------------------------------
float DistributionGGX(vec3 N, vec3 H, float roughness)
{
    float a = roughness*roughness;
    float a2 = a*a;
    float NdotH = max(dot(N, H), 0.0);
    float NdotH2 = NdotH*NdotH;

    float nom   = a2;
    float denom = (NdotH2 * (a2 - 1.0) + 1.0);
    denom = PI * denom * denom;

    return nom / denom;
}
// ----------------------------------------------------------------------------
float GeometrySchlickGGX(float NdotV, float roughness)
{
    float r = (roughness + 1.0);
    float k = (r*r) / 8.0;

    float nom   = NdotV;
    float denom = NdotV * (1.0 - k) + k;

    return nom / denom;
}
// ----------------------------------------------------------------------------
float GeometrySmith(vec3 N, vec3 V, vec3 L, float roughness)
{
    float NdotV = max(dot(N, V), 0.0);
    float NdotL = max(dot(N, L), 0.0);
    float ggx2 = GeometrySchlickGGX(NdotV, roughness);
    float ggx1 = GeometrySchlickGGX(NdotL, roughness);

    return ggx1 * ggx2;
}
// ----------------------------------------------------------------------------
vec3 fresnelSchlick(float cosTheta, vec3 F0)
{
    return F0 + (1.0 - F0) * pow(clamp(1.0 - cosTheta, 0.0, 1.0), 5.0);
}
// ----------------------------------------------------------------------------

vec3 CalculatePointLight(PointLight light, vec3 N, vec3 V, vec3 WorldPos, vec3 albedo, float metallic, float roughness)
{
  // Calculate reflectance at normal incidence (F0)
  vec3 F0 = vec3(0.04); 
  F0 = mix(F0, albedo, metallic);

  // reflectance equation
  vec3 Lo = vec3(0.0);

  // calculate per-light radiance
  vec3 L = normalize(light.position - WorldPos);
  vec3 H = normalize(V + L);
  float distance = length(light.position - WorldPos);
  float attenuation = 1.0 / (light.a * distance * distance + light.b * distance + 1.0);
  vec3 radiance = light.color * attenuation;

  // Cook-Torrance BRDF
  float NDF = DistributionGGX(N, H, roughness);   
  float G   = GeometrySmith(N, V, L, roughness);      
  vec3 F    = fresnelSchlick(clamp(dot(H, V), 0.0, 1.0), F0);

  vec3 numerator    = NDF * G * F; 
  float denominator = 4.0 * max(dot(N, V), 0.0) * max(dot(N, L), 0.0) + 0.0001; // + 0.0001 to prevent divide by zero
  vec3 specular = numerator / denominator;

  // kS is equal to Fresnel
    vec3 kS = F;
    // for energy conservation, the diffuse and specular light can't
    // be above 1.0 (unless the surface emits light); to preserve this
    // relationship the diffuse component (kD) should equal 1.0 - kS.
    vec3 kD = vec3(1.0) - kS;
    // multiply kD by the inverse metalness such that only non-metals 
    // have diffuse lighting, or a linear blend if partly metal (pure metals
    // have no diffuse light).
    kD *= 1.0 - metallic;	  

    // scale light by NdotL
    float NdotL = max(dot(N, L), 0.0);        

    // add to outgoing radiance Lo
    Lo += (kD * albedo / PI + specular) * radiance * NdotL;

    return Lo;
}

vec3 CalculateDirectionalLight(DirectionalLight light, vec3 N, vec3 V, vec3 albedo, float metallic, float roughness)
{
    // Calculate reflectance at normal incidence (F0)
    vec3 F0 = vec3(0.04); 
    F0 = mix(F0, albedo, metallic);

    // reflectance equation
    vec3 Lo = vec3(0.0);

    // calculate per-light radiance
    vec3 L = normalize(-light.direction);
    vec3 H = normalize(V + L);
    vec3 radiance = light.color * light.intensity;

    // Cook-Torrance BRDF
    float NDF = DistributionGGX(N, H, roughness);   
    float G   = GeometrySmith(N, V, L, roughness);      
    vec3 F    = fresnelSchlick(clamp(dot(H, V), 0.0, 1.0), F0);

    vec3 numerator    = NDF * G * F; 
    float denominator = 4.0 * max(dot(N, V), 0.0) * max(dot(N, L), 0.0) + 0.0001; // + 0.0001 to prevent divide by zero
    vec3 specular = numerator / denominator;

    // kS is equal to Fresnel
    vec3 kS = F;
    // for energy conservation, the diffuse and specular light can't
    // be above 1.0 (unless the surface emits light); to preserve this
    // relationship the diffuse component (kD) should equal 1.0 - kS.
    vec3 kD = vec3(1.0) - kS;
    // multiply kD by the inverse metalness such that only non-metals 
    // have diffuse lighting, or a linear blend if partly metal (pure metals
    // have no diffuse light).
    kD *= 1.0 - metallic;	  

    // scale light by NdotL
    float NdotL = max(dot(N, L), 0.0);        

    // add to outgoing radiance Lo
    Lo += (kD * albedo / PI + specular) * radiance * NdotL;

    return Lo;
}

vec3 CalculateSpotLight(SpotLight light, vec3 N, vec3 V, vec3 WorldPos, vec3 albedo, float metallic, float roughness)
{
    // Calculate reflectance at normal incidence (F0)
    vec3 F0 = vec3(0.04); 
    F0 = mix(F0, albedo, metallic);

    // reflectance equation
    vec3 Lo = vec3(0.0);

    // calculate per-light radiance
    vec3 L = normalize(light.position - WorldPos);
    vec3 H = normalize(V + L);
    float distance = length(light.position - WorldPos);
    float attenuation = 1.0 / (distance * distance); // Simple attenuation for now
    vec3 radiance = light.color * attenuation;

    // spotlight (soft edges)
    float theta = dot(L, normalize(-light.direction));
    float epsilon = light.cutOff - light.outerCutOff;
    float intensity = clamp((theta - light.outerCutOff) / epsilon, 0.0, 1.0);

    // combine radiance with intensity
    radiance *= intensity;

    // Cook-Torrance BRDF
    float NDF = DistributionGGX(N, H, roughness);   
    float G   = GeometrySmith(N, V, L, roughness);      
    vec3 F    = fresnelSchlick(clamp(dot(H, V), 0.0, 1.0), F0);

    vec3 numerator    = NDF * G * F; 
    float denominator = 4.0 * max(dot(N, V), 0.0) * max(dot(N, L), 0.0) + 0.0001; // + 0.0001 to prevent divide by zero
    vec3 specular = numerator / denominator;

    // kS is equal to Fresnel
    vec3 kS = F;
    // for energy conservation, the diffuse and specular light can't
    // be above 1.0 (unless the surface emits light); to preserve this
    // relationship the diffuse component (kD) should equal 1.0 - kS.
    vec3 kD = vec3(1.0) - kS;
    // multiply kD by the inverse metalness such that only non-metals 
    // have diffuse lighting, or a linear blend if partly metal (pure metals
    // have no diffuse light).
    kD *= 1.0 - metallic;	  

    // scale light by NdotL
    float NdotL = max(dot(N, L), 0.0);        

    // add to outgoing radiance Lo
    Lo += (kD * albedo / PI + specular) * radiance * NdotL;

    return Lo;
}





void main()
{		

    vec3 albedo     = pow(texture(albedoMap, TexCoords).rgb, vec3(2.2));
    float metallic  = texture(metallicMap, TexCoords).r;
    float roughness = texture(roughnessMap, TexCoords).r;

    vec3 N = getNormalFromMap();
    vec3 V = normalize(cameraPos - WorldPos);

    vec3 Lo = vec3(0.0);

    // Directional lights
    for (int i = 0; i < NR_DIR_LIGHTS; i++) {
        Lo += CalculateDirectionalLight(dirLights[i], N, V, albedo, metallic, roughness);
    }

    // Point lights
    for (int i = 0; i < NR_POINT_LIGHTS; i++)
    {
        Lo += CalculatePointLight(pointLights[i], N, V, WorldPos, albedo, metallic, roughness);
    }

    // Spot lights (including flashlight)
    Lo += CalculateSpotLight(flashLight, N, V, WorldPos, albedo, metallic, roughness);

    for (int i = 0; i < NR_SPOT_LIGHTS; i++) {
        Lo += CalculateSpotLight(spotLights[i], N, V, WorldPos, albedo, metallic, roughness);
    }
    
    
    
    
    // ambient lighting (note that the next IBL tutorial will replace 
    // this ambient lighting with environment lighting).
    vec3 ambient = vec3(0.03) * albedo;

    vec3 color = ambient + Lo;

    if (nightVision == true)
    {
        // fog parameters
        vec3 fogColor = vec3(0.f, 0.f, 0.f);
        float fogStart = 25.0f;
        float fogEnd = 30.0f;
        // Fog calculation
        float distance = length(cameraPos - WorldPos);
        float fogFactor = clamp((fogEnd - distance) / (fogEnd - fogStart), 0.0, 1.0);
        color = mix(fogColor, color, fogFactor);
    }
    else
    {
        // fog parameters
        vec3 fogColor = vec3(0.f, 0.f, 0.f);
        float fogStart = 10.0f;
        float fogEnd = 30.0f;
        // Fog calculation
        float distance = length(cameraPos - WorldPos);
        float fogFactor = clamp((fogEnd - distance) / (fogEnd - fogStart), 0.0, 1.0);
        color = mix(fogColor, color, fogFactor);
    }

    FragColor = vec4(color, 1.0);
    BrightColor = vec4(0.0, 0.0, 0.0, 1.0);


}