#extension GL_ARB_bindless_texture : require
#extension GL_KHR_shader_subgroup_ballot : enable
#extension GL_KHR_shader_subgroup_vote : enable

#include "camera.glsl"
#include "draw_cmds_data.glsl"

// Keep in sync with cpu
#define MAX_POINT_LIGHTS 8
#define PI 3.1415926535897932384626433832795
#define CSM_SPLITS 4

layout(location = 16, bindless_sampler) uniform sampler2DArray shadow_maps;
layout(location = 17, bindless_sampler) uniform samplerCubeArray cube_shadow_maps;
layout(location = 22, bindless_sampler) uniform sampler2D brdfLut;


// Input, output blocks

VERTEX_EXPORT VertexData {
  vec3 vPos;
  vec2 uv;
  mat3 vTBN;
  vec3 wPos;
  vec3 vUp;
  vec3 wNormal;
  vec3 vNormal;
} VertexOut;

VERTEX_EXPORT flat uint DrawID;

float random(vec4 seed4) {
  float dot_product = dot(seed4, vec4(12.9898,78.233,45.164,94.673));
  return fract(sin(dot_product) * 43758.5453);
}

#if VERTEX_SHADER

layout(location = 0) in vec3 aPos;
layout(location = 1) in vec3 aNormal;
layout(location = 2) in vec2 aUV;
layout(location = 3) in vec3 aTangent;

void main() {
    DrawID = gl_BaseInstance + gl_InstanceID;
    mat4 model = draw_data[DrawID].transform;
    mat4 viewModel = view * model;
    vec4 vPos = viewModel * vec4(aPos.xyz, 1.0);
    gl_Position = projection * vPos;

    VertexOut.vUp = normalize(mat3(viewModel) * vec3(0.0, 1.0, 0.0));

    VertexOut.vPos = vPos.xyz / vPos.w;
    VertexOut.uv = aUV;
    vec3 aBitangent = cross(aTangent, aNormal);
    vec3 T = normalize(vec3(viewModel * vec4(aTangent, 0.0)));
    vec3 B = normalize(vec3(viewModel * vec4(aBitangent, 0.0)));
    vec3 N = normalize(vec3(viewModel * vec4(aNormal, 0.0)));
    VertexOut.vTBN = mat3(T, B, N);
    VertexOut.wPos = (model * vec4(aPos.xyz, 1.0)).xyz;
    VertexOut.wNormal = normalize(model * vec4(aNormal, 0.0)).xyz;
    VertexOut.vNormal = normalize(mat3(viewModel) * aNormal).xyz;
}
#endif // VERTEX_SHADER

#if FRAGMENT_SHADER

#include "material.glsl"

// Types
struct Light {
  vec4 vPos;
  vec4 color;

  mat4 view_mat;

  // for directional lights contains view projection matrices for each split
  // TODO: compress this somehow
  mat4[4] view_proj_mats;

  // x, y = near, far for point lights
  // z = shadow map index
  // w = csm split count for directional lights
  vec4 params;
  vec4 csm_split_points; // TODO: Maybe increase to 8, though it's probably too many
};

layout(std430, binding = 1) readonly buffer Lights {
  uint lights_count;
  Light lights[];
};

out vec4 FragColor;

int getShadowMapIndex(int lightIdx) {
  return int(lights[lightIdx].params.z);
}

int getCSMSplitCount(int lightIdx) {
  return clamp(int(lights[lightIdx].params.w), 0, 4);
}

int getCSMSplit(int lightIdx, float depth) {
  // int totalSplits = getCSMSplitCount(lightIdx);

  for (int i = 0; i < CSM_SPLITS; i++) {
    if (depth >= lights[lightIdx].csm_split_points[i]) {
      return i;
    }
  }

  return CSM_SPLITS - 1;
}

vec3 schlickFresnel(int matIdx, float NDotV) {
  vec3 f0 = mix(vec3(0.04), getAlbedo(matIdx).rgb, getMetallic(matIdx));

  return f0 + (1.0 - f0) * pow(1.0 - NDotV, 5.0);
}

vec3 schlickFresnelRoughness(int matIdx, float NDotV) {
  vec3 f0 = mix(vec3(0.04), getAlbedo(matIdx).rgb, getMetallic(matIdx));

  return f0 + (max(vec3(1.0 - getRoughness(matIdx)), f0) - f0) * pow(1.0 - NDotV, 5.0);
}

const float eps = 0.0001;

float geomSmith(int matIdx, float DotVal) {
  float k = ((getRoughness(matIdx) + 1.0) * (getRoughness(matIdx) + 1.0)) / 8.0;
  float denom = DotVal * (1 - k) + k;
  return DotVal / denom;
}

float ggxDistribution(int matIdx, float NDotH) {
  float a = getRoughness(matIdx) * getRoughness(matIdx);
  float alpha2 = a * a;
  float NDotH2 = NDotH * NDotH;
  float nom = alpha2;
  float denom = (NDotH2 * (alpha2 - 1.0) + 1.0);
  denom = PI * denom * denom;
  return nom / denom;
}

float lightAttenuation(bool point, float dist, float radius) {
  float d = max(dist - radius, 0) * (point ? 1.0 : 0.0);

  float denom = d/radius + 1;
  float att = 1 / (denom * denom);
  // TODO: cutoff
  att = max(att, 0);

  return att;
}

vec2 poissonDisk[4] = vec2[](
  vec2( -0.94201624, -0.39906216 ),
  vec2( 0.94558609, -0.76890725 ),
  vec2( -0.094184101, -0.92938870 ),
  vec2( 0.34495938, 0.29387760 )
);

const vec3 csm_split_colors[4] = vec3[](
  vec3(0f, 1f, 0.17f),
  vec3(1f, 0.2f, 0.6f),
  vec3(0.17f, 0.78f, 1f),
  vec3(0.91f, 0.93f, 0.64f)
);

float map(float value, float min1, float max1, float min2, float max2) {
  return min2 + (value - min1) * (max2 - min2) / (max1 - min1);
}
float linestep(float min, float max, float v) {
  return clamp((v - min) / (max - min), 0, 1);
}
float ReduceLightBleeding(float p_max, float amount) {
  return linestep(amount, 1, p_max);
}
float g_min_variance = 0.00001;
float ChebyshevUpperBound(vec2 moments, float t) {
  // // One-tailed inequality valid if t > Moments.x
  float p = t <= moments.x ? 1.0 : 0.0;
  // // Compute variance.
  float variance = moments.y - (moments.x * moments.x);
  // // Compute probabilistic upper bound.
  variance = max(variance, g_min_variance);
  float d = t - moments.x;
  float p_max = variance / (variance + d * d);
  return max(p, ReduceLightBleeding(p_max, 0.6));
}

float weight(float t, float log2radius, float gamma) {
    return exp(-gamma*pow(log2radius-t,2.0f));
}

vec2 sampleDirectShadowBlurred(vec3 uv, float radius, float gamma) {
    vec2 pix = vec2(0);
    float norm = 0;
    // weighted integration over mipmap levels
    for(float i = 0; i < 10; i += 0.5) {
        float k = weight(i, log2(radius), gamma);
        pix += k*textureLod(shadow_maps, uv, i).xy;
        norm += k;
    }
    // nomalize, and a bit of brigtness hacking 
    return pix/norm;
}

vec3 microfacetModel(int matIdx, int light_idx, vec3 P, vec3 N) {
  // Light light = lights[light_idx];

  vec3 diffuseBrdf = vec3(0); // metallic
  if (getMetallic(matIdx) < 1.0) {
    diffuseBrdf = getAlbedo(matIdx).rgb / PI;
  }

  // 0 - means directional, 1 - means point light
  bool point = subgroupAll(lights[light_idx].vPos.w == 1);
  vec3 lightI = lights[light_idx].color.rgb;
  vec3 L = mix(-lights[light_idx].vPos.xyz, lights[light_idx].vPos.xyz - P, lights[light_idx].vPos.w);
  vec3 V = normalize(-P);
  vec3 H = normalize(V + L);

  float NDotH = max(dot(N, H), 0.0);
  float NDotL = max(dot(N, L), 0.0);
  float NDotV = max(dot(N, V), 0.0);

  float normal_offset_scale = clamp(1 - NDotL, 0, 1);
  normal_offset_scale *= 10; // constant

  int shadow_map_idx = subgroupBroadcastFirst(getShadowMapIndex(light_idx));

  float constant_bias = 0.003;
  float shadow_mult = 1;
  vec4 shadow_offset = vec4(VertexOut.wNormal * normal_offset_scale, 0);
  if (point) {
    float dist = max(length(L), eps);
    L /= dist;
    float lightRadius = max(lights[light_idx].color.a, eps);
    float att = lightAttenuation(
      point,
      dist,
      lightRadius
    );
    lightI *= att;

    vec2 shadow_map_texel_size = 1.0 / vec2(textureSize(cube_shadow_maps, 0));
    shadow_offset *= shadow_map_texel_size.x;
    vec3 shadow_dir = (lights[light_idx].view_mat * vec4(VertexOut.wPos, 1.0)).xyz;
    float world_depth = length(shadow_dir.xyz);
    shadow_dir = normalize((lights[light_idx].view_mat * (vec4(VertexOut.wPos, 1.0) + shadow_offset)).xyz);
    float mapped_depth = map(world_depth, lights[light_idx].params.x, lights[light_idx].params.y, 0, 1);

    vec4 texcoord;
    texcoord.xyz = shadow_dir;
    texcoord.w = float(light_idx);

    float sum = 0;

    // sum += texture(cube_shadow_maps, vec4(normalize(texcoord.xyz), texcoord.w), mapped_depth - constant_bias);
    // for (float y = -1.5; y <= 1.5; y += 1) {
    //   for (float x = -1.5; x <= 1.5; x += 1) {
    //     // sum += texture(cube_shadow_maps, vec4(normalize(texcoord.xyz + vec3(x, y, z) * shadow_map_texel_size.x), texcoord.w), mapped_depth - constant_bias);
    //   }
    // }

    shadow_mult = sum / 1.0;
  } else {
    int csm_split_idx = subgroupBroadcastFirst(getCSMSplit(light_idx, P.z));
    // Visualize CSM splits
    // mat.albedo = vec4(mix(mat.albedo.rgb, csm_split_colors[csm_split_idx], 0.8), mat.albedo.a);
    // Directional shadow
    vec2 shadow_map_texel_size = 1.0 / vec2(textureSize(shadow_maps, 0));
    shadow_offset *= shadow_map_texel_size.x;
    vec4 shadow_pos = lights[light_idx].view_proj_mats[csm_split_idx] * vec4(VertexOut.wPos, 1.0);
    // shadow_pos.xy = (lights[light_idx].view_proj_mats[csm_split_idx] * (vec4(VertexOut.wPos, 1.0) + shadow_offset)).xy;
    shadow_pos /= shadow_pos.w;
    shadow_pos.xy = shadow_pos.xy * 0.5 + 0.5; // [-1, 1] to [0, 1]
    // shadow_pos.z = min(shadow_pos.z, 1);
    // shadow_pos.z -= constant_bias;

    vec4 texcoord;
    texcoord.xyw = shadow_pos.xyz; // sampler2DArrayShadow strange texcoord mapping
    texcoord.z = shadow_map_idx + csm_split_idx;


    // vec4 xs = textureGather(shadow_maps, texcoord.xyz, 0);
    // vec4 ys = textureGather(shadow_maps, texcoord.xyz, 1);

    // vec2 moments = vec2(
    //   xs.x + xs.y + xs.z + xs.w,
    //   ys.x + ys.y + ys.z + ys.w
    // );

    // for (int i = 0; i < 4; i++) {
    //   moments += texture(shadow_maps, vec3(texcoord.xy + poissonDisk[i] * shadow_map_texel_size, texcoord.z)).xy;
    // }
    // moments /= 4.0;

    float sum = 0;

    // moments = sampleDirectShadowBlurred(texcoord.xyz, shadow_pos.z * 10, 0.5);

    for (float y = -1.5; y <= 1.5; y += 1) {
      for (float x = -1.5; x <= 1.5; x += 1) {
        vec2 moments = texture(shadow_maps, vec3(texcoord.xy + vec2(x, y) * shadow_map_texel_size, texcoord.z)).xy;
        sum += ChebyshevUpperBound(moments, shadow_pos.z);
      }
    }

    shadow_mult = sum / 16.0;;
  }
  shadow_mult = clamp(shadow_mult, 0.0, 1.0);

  vec3 F = schlickFresnelRoughness(matIdx, NDotV);
  vec3 specBrdf = F * (ggxDistribution(matIdx, NDotH) * geomSmith(matIdx, NDotV) * geomSmith(matIdx, NDotL));
  specBrdf /= 4.0 * NDotL * NDotV + 0.0001;

  vec3 kS = F;
  vec3 kD = vec3(1.0) - kS;
  kD *= 1.0 - getMetallic(matIdx);


  return (kD * diffuseBrdf + specBrdf) * lightI * (NDotL * shadow_mult);
}

vec3 ibl(int matIdx, vec3 N, vec3 V) {
  float NDotV = max(dot(N, V), 0.0);
  vec3 F = schlickFresnelRoughness(matIdx, NDotV);
  vec3 kS = F;
  vec3 kD = 1.0 - kS;

  vec3 R = reflect(-V, N);
  float ambient_diff = dot(N, VertexOut.vUp) * 0.5 + 0.5;
  float ambient_spec = dot(R, VertexOut.vUp) * 0.5 + 0.5;

  // TODO: don't hardcode this
  vec3 irradiance = vec3(1.0, 0.9764705882352941, 0.9921568627450981) * 79 * ambient_diff * 0.02;
  vec3 diffuse = irradiance * getAlbedo(matIdx).rgb;

  vec3 reflectedColor = vec3(0.9, 0.9064705882352941, 0.9921568627450981) * 79 * ambient_spec * 0.02;
  vec2 envBRDF = textureLod(brdfLut, vec2(NDotV, getRoughness(matIdx)), 0).rg;
  vec3 specular = reflectedColor * (F * envBRDF.x + envBRDF.y);
  return kD * diffuse + specular;
}

void main() {
  int matIdx = draw_data[DrawID].materialIdx;
#if OVERRIDE_COLOR
  FragColor = getAlbedo(matIdx);
#else
  if (getAlbedo(matIdx).a < 0.5) {
    FragColor = vec4(0);
    return;
  }
  sampler2D normal_map = materials[matIdx].normal_map;
  vec2 normal_map_uv_scale = materials[matIdx].normal_map_uv_scale;
  
  vec3 N = textureSize(normal_map, 0) == ivec2(0) ? vec3(0.5) : vec3(texture(normal_map, VertexOut.uv * normal_map_uv_scale).xy, 0);
  N = N * 2.0 - 1.0;
  N.z = sqrt(clamp(1 - N.x * N.x - N.y * N.y, 0, 1));
  N = normalize(N);
  N = normalize(VertexOut.vTBN * N);
  // vec3 N = VertexOut.vNormal;

  vec3 finalColor = vec3(0);

  int n_lights = clamp(int(lights_count), 0, MAX_POINT_LIGHTS);
  for (int i = 0; i < MAX_POINT_LIGHTS; i++) {
    if (i > lights_count) break;
    finalColor += microfacetModel(matIdx, i, VertexOut.vPos, N);
  }

  vec3 V = normalize(-VertexOut.vPos);
  // ambient
  finalColor += ibl(matIdx, N, V);
  finalColor += getEmission(matIdx);

  FragColor = vec4(finalColor, getAlbedo(matIdx).a);
#endif
}


#endif // FRAGMNET_SHADER