# ============================
# shaders.py — GLSL Shader Sources
# ============================

# -- Sphere Vertex Shader --
SPHERE_VS = r"""#version 330 core
layout(location=0) in vec2 aQuad;
uniform mat4 uView; uniform mat4 uProj;
uniform samplerBuffer uAtomPosRad; uniform samplerBuffer uAtomColor;
uniform samplerBuffer uReplicaOffsets;
uniform usamplerBuffer uVisiblePairs; // .r = atomIdx, .g = repIdx
uniform usamplerBuffer uSelectionMaskFull; // .r = 1 if selected, 0 otherwise
uniform bool uSelectionOnly;
uniform int uVisibleCount;
uniform bool uIsOrtho;
uniform float uAtomScale;
uniform float uProjScale;
uniform float uViewportHeight;
uniform float uTime;
uniform bool uPrettyRender;
uniform float uPulseSpeed;
uniform float uPulseAmp;
out VS_OUT { 
    vec3 centerVS; 
    float radiusVS; 
    vec4 color; 
    vec3 posVS; 
    float pixelRadius; 
    flat int atomIdx;
} v;
void main(){
    int inst = gl_InstanceID;
    uvec2 pair = texelFetch(uVisiblePairs, inst).rg;
    int atomIdx = int(pair.x);
    int repIdx = int(pair.y);
    
    vec4 posRad = texelFetch(uAtomPosRad, atomIdx);
    vec4 color = texelFetch(uAtomColor, atomIdx);
    vec3 offset = texelFetch(uReplicaOffsets, repIdx).xyz;

    if (uSelectionOnly) {
        uint selected = texelFetch(uSelectionMaskFull, atomIdx).r;
        if (selected == 0u) {
            gl_Position = vec4(2.0, 2.0, 2.0, 1.0); // Discard by position
            return;
        }
    }

    vec3 centerWorld = posRad.xyz + offset;
    float radius = posRad.w * uAtomScale;
    if (uPrettyRender && uPulseAmp > 0.0) {
        radius *= (1.0 + uPulseAmp * sin(uTime * uPulseSpeed + float(atomIdx) * 0.5));
    }
    vec4 centerView4 = uView * vec4(centerWorld, 1.0);
    vec3 centerView = centerView4.xyz;
    
    float pixel_radius;
    if (uIsOrtho) {
        pixel_radius = radius * uProjScale; 
    } else {
        pixel_radius = radius * uProjScale / abs(centerView.z);
    }
    
    vec3 posVS = centerView + vec3(aQuad * radius, 0.0);
    v.centerVS = centerView; v.radiusVS = radius; v.color = color; v.posVS = posVS; 
    v.pixelRadius = pixel_radius;
    v.atomIdx = atomIdx;
    gl_Position = uProj * vec4(posVS, 1.0);
}
"""

# -- Sphere Fragment Shader --
SPHERE_FS = r"""#version 330 core
in VS_OUT { 
    vec3 centerVS; 
    float radiusVS; 
    vec4 color; 
    vec3 posVS; 
    float pixelRadius; 
    flat int atomIdx;
} v;
uniform mat4 uProj; 
uniform mat4 uViewInv;
uniform mat4 uShadowMatrix;
uniform bool uIsOrtho;
uniform usamplerBuffer uSelectedAtoms; 
uniform usamplerBuffer uSelectionMaskFull;
uniform int uSelectedCount;
out vec4 FragColor;

const int MAX_LIGHTS = 4;
uniform vec3 uLightDir[MAX_LIGHTS];
uniform vec3 uLightCol[MAX_LIGHTS];
uniform int uNumLights;
uniform samplerBuffer uAtomPosRad; 
uniform int uAtomCount;
uniform samplerBuffer uAtomMaterial; 
uniform sampler2DShadow uShadowMap;
uniform float uExposure;
uniform float uHeadlight;
uniform bool uShowShadows;
uniform bool uShowGroundShadow;
uniform float uShadowIntensity;
uniform float uShadowSoftness;
uniform float uShadowBias;
uniform float uSkipPx;
uniform float uPointPx;
uniform bool uPrettyRender;
uniform float uReflectionIntensity;
uniform float uReflectionGloss;
uniform float uFogDensity;
uniform samplerBuffer uAtomMaterial2;
uniform float uSelectionGlow;
uniform float uTime;

uniform bool uSelectionOnly;

void main(){
    if (uSelectionOnly) { FragColor = vec4(1.0); return; }
    if (v.pixelRadius < uSkipPx) discard;

    vec2 coord = v.posVS.xy - v.centerVS.xy;
    float distSq = dot(coord, coord);
    float rSq = v.radiusVS * v.radiusVS;

    if (v.pixelRadius < uPointPx) {
        if (distSq > rSq) discard; 
        gl_FragDepth = gl_FragCoord.z;
        FragColor = vec4(v.color.rgb, v.color.a);
        return;
    }

    if (distSq > rSq) discard;

    float zw = sqrt(rSq - distSq);
    vec3 N = normalize(vec3(coord, zw)); 
    vec3 surfVS = v.centerVS + vec3(coord, zw);
    vec3 V = normalize(-surfVS);
    
    vec4 clipPos = uProj * vec4(surfVS, 1.0);
    gl_FragDepth = (clipPos.z / clipPos.w) * 0.5 + 0.5;

    // --- Global Variable Initialization ---
    float occlusion = 1.0;
    float shadow = 0.0;
    float visibility = 1.0;
    // --------------------------------------

    // 0. Contact Shadows (Cheap Inter-atomic depth)
    if (uPrettyRender && uShowShadows) {
        float horizon = max(dot(N, vec3(0, 1, 0)), 0.0);
        // Stronger local contact occlusion
        occlusion = mix(1.0, 0.4 + 0.6 * horizon, uShadowIntensity);
        float contact = 1.0 - pow(dot(N, vec3(0,0,1)), 3.0);
        occlusion *= (1.0 - 0.4 * contact * uShadowIntensity);
    }

    // 1. Shadow Map Lookup (Optimized PCF with early exit)
    if (uPrettyRender && uShowShadows) {
        vec3 worldPos = (uViewInv * vec4(surfVS, 1.0)).xyz;
        vec4 sc = uShadowMatrix * vec4(worldPos, 1.0);
        // Slope-scaled bias
        float ndotl = max(dot(N, uLightDir[0]), 0.0);
        float bias = max(uShadowBias * 0.1, uShadowBias * (1.0 - ndotl));
        sc.z -= bias * sc.w;
        
        // 4-corner test for early exit
        float c0 = textureProj(uShadowMap, sc + vec4(vec2(-1.5, -1.5) * uShadowSoftness, 0, 0));
        float c1 = textureProj(uShadowMap, sc + vec4(vec2(1.5, -1.5) * uShadowSoftness, 0, 0));
        float c2 = textureProj(uShadowMap, sc + vec4(vec2(-1.5, 1.5) * uShadowSoftness, 0, 0));
        float c3 = textureProj(uShadowMap, sc + vec4(vec2(1.5, 1.5) * uShadowSoftness, 0, 0));
        
        if ((c0 + c1 + c2 + c3) > 3.99) {
            shadow = 0.0; 
        } else if ((c0 + c1 + c2 + c3) < 0.01) {
            shadow = 1.0;
        } else {
            float v_sum = c0 + c1 + c2 + c3;
            for(float x = -1.5; x <= 1.5; x += 1.0){
                for(float y = -1.5; y <= 1.5; y += 1.0){
                    if ((abs(x) > 1.4 && abs(y) > 1.4)) continue;
                    v_sum += textureProj(uShadowMap, sc + vec4(vec2(x, y) * uShadowSoftness, 0, 0));
                }
            }
            shadow = 1.0 - (v_sum / 16.0);
        }
    }
    
    if (uPrettyRender && uShowShadows) {
        visibility = clamp(mix(1.0, 1.0 - shadow, uShadowIntensity * 2.5), 0.0, 1.0);
    }

    // Per-Atom Materials (Expanded to 8 Parameters)
    vec4 m1 = texelFetch(uAtomMaterial, v.atomIdx);
    vec4 m2 = texelFetch(uAtomMaterial2, v.atomIdx);
    
    float ambient = uPrettyRender ? m1.r : 0.4;
    float specular = uPrettyRender ? m1.g : 0.6;
    float shininess = uPrettyRender ? m1.b : 48.0;
    float rim = uPrettyRender ? m1.a : 0.1;
    
    float m_sheen = uPrettyRender ? m2.r : 0.0;
    float m_clearcoat = uPrettyRender ? m2.g : 0.0;
    float m_reflection = uPrettyRender ? m2.b : 0.3;
    float m_glossiness = uPrettyRender ? m2.a : 0.5;

    vec3 totalDiff = vec3(0.0);
    vec3 totalSpec = vec3(0.0);
    
    // 1. Headlight
    {
        vec3 L = vec3(0, 0, 1);
        float d = max(dot(N, L), 0.0);
        totalDiff += d * uHeadlight * vec3(1.0);
        float s = pow(max(dot(N, normalize(L + V)), 0.0), shininess);
        totalSpec += s * specular * uHeadlight * vec3(1.0);
    }

    // 2. Directional Lights
    for(int i = 0; i < uNumLights; i++){
        vec3 L = normalize(uLightDir[i]); 
        vec3 H = normalize(L + V);
        float NdotL = dot(N, L);
        float d = NdotL * 0.5 + 0.5;
        d = d * d; 
        
        float l_vis = (i == 0) ? visibility : 1.0;
        totalDiff += l_vis * d * uLightCol[i];
        float s = pow(max(dot(N, H), 0.0), shininess);
        totalSpec += l_vis * s * specular * uLightCol[i];
    }
    
    // 3. Cheap Environment Reflection (Fake MatCap)
    vec3 reflectionCol = vec3(0.0);
    float rimEffect = pow(1.0 - max(dot(N, V), 0.0), 3.0);
    
    if (uPrettyRender && m_reflection > 0.01) {
        vec3 R = reflect(-V, N);
        float sky = smoothstep(-0.4, 0.4, R.y);
        float sun = pow(max(dot(R, normalize(vec3(0.5, 1.0, 0.5))), 0.0), 4.0 + m_glossiness * 60.0);
        vec3 env = mix(vec3(0.05, 0.05, 0.1), vec3(0.8, 0.9, 1.0), sky) + sun;
        reflectionCol = env * m_reflection;
    }

    vec3 finalCol;
    if (uPrettyRender) {
        // Shadow-clamped lighting (apply visibility to BOTH diffuse and ambient for deeper shadows)
        vec3 lightFactor = vec3(ambient) + totalDiff * visibility;
        finalCol = v.color.rgb * lightFactor * occlusion + totalSpec * visibility * occlusion;
        finalCol += rim * rimEffect * vec3(1.0) * occlusion * visibility;
        
        // Darken shadows even further based on intensity
        finalCol *= (1.0 - shadow * uShadowIntensity * 1.5); 
        
        // Sheen (Soft velvet glow)
        if (m_sheen > 0.0) {
            float sheen_val = pow(1.0 - max(dot(N, V), 0.0), 4.0);
            finalCol += m_sheen * sheen_val * v.color.rgb;
        }

        // Clearcoat (Secondary sharp highlight)
        if (m_clearcoat > 0.0) {
            vec3 H_cc = normalize(normalize(vec3(1.0, 1.0, 1.0)) + V);
            float specular_cc = pow(max(dot(N, H_cc), 0.0), 256.0);
            finalCol += m_clearcoat * specular_cc * vec3(1.0);
        }

        finalCol += reflectionCol * (0.3 + 0.7 * rimEffect);
        finalCol *= uExposure;

        // Exponential Fog (Subtler)
        float z = abs(surfVS.z);
        float fog = exp(-z * uFogDensity * 2.0);
        vec3 fogColor = vec3(0.01, 0.01, 0.02);
        finalCol = mix(fogColor, finalCol, fog);
        
        float fade = smoothstep(150.0, 600.0, z); // Broader fade
        finalCol = mix(finalCol, fogColor, fade);
    } else {
        // Balanced lighting for standard view
        float d = max(dot(N, normalize(vec3(0.5, 0.5, 1.0))), 0.0);
        finalCol = v.color.rgb * (0.3 + 0.7 * d);
    }

    // Selection Glow (O(1) bitmask-based lookup)
    bool is_selected = uint(texelFetch(uSelectionMaskFull, v.atomIdx).r) > 0u;
    if (is_selected) {
        // High-contrast vibrant cyan-yellow mix for visibility
        finalCol = mix(finalCol, vec3(0.0, 1.0, 1.0), 0.4); 
        finalCol += vec3(0.2, 0.2, 0.0); // Slight golden boost
        finalCol *= 1.4;
    }

    FragColor = vec4(finalCol, v.color.a);
}
"""

SHADOW_VS = r"""#version 330 core
layout(location=0) in vec2 aQuad;
uniform mat4 uView; uniform mat4 uProj;
uniform samplerBuffer uAtomPosRad;
uniform samplerBuffer uReplicaOffsets;
uniform usamplerBuffer uVisiblePairs;
uniform int uVisibleCount;
uniform float uAtomScale;
out VS_OUT { vec3 centerVS; float radiusVS; vec3 posVS; } v;
void main(){
    int inst = gl_InstanceID;
    uvec2 pair = texelFetch(uVisiblePairs, inst).rg;
    int atomIdx = int(pair.x);
    int repIdx = int(pair.y);
    vec4 posRad = texelFetch(uAtomPosRad, atomIdx);
    vec3 offset = texelFetch(uReplicaOffsets, repIdx).xyz;
    vec3 centerWorld = posRad.xyz + offset;
    float radius = posRad.w * uAtomScale;
    vec4 centerView4 = uView * vec4(centerWorld, 1.0);
    v.centerVS = centerView4.xyz;
    v.radiusVS = radius;
    v.posVS = v.centerVS + vec3(aQuad * radius, 0.0);
    gl_Position = uProj * vec4(v.posVS, 1.0);
}
"""

SHADOW_FS = r"""#version 330 core
in VS_OUT { 
    vec3 centerVS; 
    float radiusVS; 
    vec3 posVS; 
} v;
uniform mat4 uProj; 
void main() {
    vec2 coord = v.posVS.xy - v.centerVS.xy;
    float distSq = dot(coord, coord);
    float rSq = v.radiusVS * v.radiusVS;
    if (distSq > rSq) discard;
    float zw = sqrt(rSq - distSq);
    vec3 surfVS = v.centerVS + vec3(coord, zw);
    vec4 clipPos = uProj * vec4(surfVS, 1.0);
    gl_FragDepth = (clipPos.z / clipPos.w) * 0.5 + 0.5;
}
"""

# -- Ground Shaders --
# -- Line Shaders --
LINE_VS = r"""#version 330 core
layout(location=0) in vec3 aPos;
uniform mat4 uView; uniform mat4 uProj; uniform samplerBuffer uReplicaOffsets;
void main(){
    int inst = gl_InstanceID;
    vec3 offset = texelFetch(uReplicaOffsets, inst).xyz;
    gl_Position = uProj * uView * vec4(aPos + offset, 1.0);
}
"""
LINE_FS = r"""#version 330 core
uniform vec4 uRGBA;
out vec4 FragColor;
void main(){
    FragColor = uRGBA;
}
"""

SHADOW_CYLINDER_VS = r"""#version 330 core
layout(location=0) in vec3 aPos;     
uniform mat4 uView; uniform mat4 uProj;
uniform samplerBuffer uReplicaOffsets;
uniform samplerBuffer uBondData; 
uniform usamplerBuffer uVisiblePairs; // .r = bondIdx, .g = repIdx
uniform float uBondScale; 
void main(){
    int inst = gl_InstanceID;
    uvec2 pair = texelFetch(uVisiblePairs, inst).rg;
    int bondIdx = int(pair.x);
    int repIdx = int(pair.y);
    
    vec3 offset = texelFetch(uReplicaOffsets, repIdx).xyz;
    vec3 iStart = texelFetch(uBondData, 4 * bondIdx + 0).xyz;
    vec3 iEnd   = texelFetch(uBondData, 4 * bondIdx + 1).xyz;
    vec3 dir = iEnd - iStart;
    float len = length(dir);
    if(len < 0.0001) { gl_Position = vec4(0,0,0,1); return; }
    vec3 yAxis = normalize(dir);
    vec3 up = (abs(yAxis.y) > 0.999) ? vec3(1, 0, 0) : vec3(0, 1, 0);
    vec3 axis = cross(up, yAxis);
    float c = dot(up, yAxis);
    mat3 rot;
    if (abs(c + 1.0) < 0.0001) rot = mat3(-1, 0, 0,  0, -1, 0,  0, 0, 1);
    else if (abs(c - 1.0) < 0.0001) rot = mat3(1,0,0, 0,1,0, 0,0,1);     
    else {
        mat3 K = mat3(0, axis.z, -axis.y,  -axis.z, 0, axis.x,  axis.y, -axis.x, 0);
        rot = mat3(1,0,0, 0,1,0, 0,0,1) + K + K*K * (1.0/(1.0+c));
    }
    vec3 localPos = aPos;
    localPos.x *= uBondScale; localPos.z *= uBondScale; localPos.y *= len; 
    gl_Position = uProj * uView * vec4(iStart + offset + rot * localPos, 1.0);
}
"""
SHADOW_CYLINDER_FS = r"""#version 330 core
void main() {}
"""

# -- Cylinder Shader (Instanced) --
CYLINDER_VS = r"""#version 330 core
layout(location=0) in vec3 aPos;     
layout(location=1) in vec3 aNormal;
uniform mat4 uView; uniform mat4 uProj;
uniform samplerBuffer uReplicaOffsets;
uniform samplerBuffer uBondData; 
uniform usamplerBuffer uVisiblePairs; // .r = bondIdx, .g = repIdx
uniform float uBondScale;   
uniform usamplerBuffer uSelectionMaskFull; // Reuse atom mask for bonds (Start Atom Index)
uniform bool uSelectionOnly;
out vec3 vNormal; out vec4 vColor; out vec3 vPos;
flat out int vAtomIdx;
void main(){
    int inst = gl_InstanceID;
    uvec2 pair = texelFetch(uVisiblePairs, inst).rg;
    int bondIdx = int(pair.x);
    int repIdx = int(pair.y);
    
    vec3 offset = texelFetch(uReplicaOffsets, repIdx).xyz;
    vec3 iStart = texelFetch(uBondData, 4 * bondIdx + 0).xyz;
    vec3 iEnd   = texelFetch(uBondData, 4 * bondIdx + 1).xyz;
    vec3 iColor = texelFetch(uBondData, 4 * bondIdx + 2).xyz;
    int atomIdx = int(texelFetch(uBondData, 4 * bondIdx + 3).x);
    
    if (uSelectionOnly) {
        uint selected = texelFetch(uSelectionMaskFull, atomIdx).r;
        if (selected == 0u) {
            gl_Position = vec4(2.0, 2.0, 2.0, 1.0);
            return;
        }
    }

    vec3 p0 = iStart + offset;
    vec3 p1 = iEnd + offset;
    vec3 dir = p1 - p0;
    float len = length(dir);
    if(len < 0.0001) { gl_Position = vec4(0,0,0,1); return; }
    vec3 yAxis = normalize(dir);
    vec3 up = (abs(yAxis.y) > 0.999) ? vec3(1, 0, 0) : vec3(0, 1, 0);
    vec3 axis = cross(up, yAxis);
    float c = dot(up, yAxis);
    mat3 rot;
    if (abs(c + 1.0) < 0.0001) {
        rot = mat3(-1, 0, 0,  0, -1, 0,  0, 0, 1);
    } else if (abs(c - 1.0) < 0.0001) {
        rot = mat3(1,0,0, 0,1,0, 0,0,1);     
    } else {
        mat3 K = mat3(0, axis.z, -axis.y,  -axis.z, 0, axis.x,  axis.y, -axis.x, 0);
        rot = mat3(1,0,0, 0,1,0, 0,0,1) + K + K*K * (1.0/(1.0+c));
    }
    vec3 localPos = aPos;
    localPos.x *= uBondScale;
    localPos.z *= uBondScale;
    localPos.y *= len; 
    vec3 worldPos = p0 + rot * localPos;
    vNormal = rot * aNormal; 
    vColor = texelFetch(uBondData, 4 * bondIdx + 2); // Fetch RGBA
    vPos = worldPos;
    vAtomIdx = atomIdx;
    gl_Position = uProj * uView * vec4(worldPos, 1.0);
}
"""

CYLINDER_FS = r"""#version 330 core
in vec3 vNormal; in vec4 vColor; in vec3 vPos;
uniform mat4 uView;
uniform mat4 uViewInv;
uniform mat4 uShadowMatrix;
uniform sampler2DShadow uShadowMap;
uniform bool uPrettyRender;
uniform bool uShowShadows;
uniform float uShadowIntensity;
uniform float uShadowSoftness;
uniform float uShadowBias;
uniform float uReflectionIntensity;
uniform float uFogDensity;
uniform vec3 uLightDirVS; 
uniform samplerBuffer uAtomMaterial;
uniform samplerBuffer uAtomMaterial2;
uniform usamplerBuffer uSelectedAtoms; 
uniform usamplerBuffer uSelectionMaskFull;
uniform int uSelectedCount;
out vec4 FragColor;
flat in int vAtomIdx;
uniform bool uSelectionOnly;

void main(){
    if (uSelectionOnly) { FragColor = vec4(1.0); return; }
    vec3 normal = normalize(mat3(uView) * vNormal); 
    vec3 L = normalize(uLightDirVS);
    vec3 V = normalize(-(mat3(uView) * vPos));
    vec3 H = normalize(L + V);
    
    // --- Global Variable Initialization ---
    float shadow = 0.0;
    float visibility = 1.0;
    // --------------------------------------
    
    // 1. Shadow Map Lookup (16-sample PCF)
    if (uPrettyRender && uShowShadows) {
        vec4 sc = uShadowMatrix * vec4(vPos, 1.0);
        // Simplified bias for calibration
        sc.z -= uShadowBias * sc.w;
        float c0 = textureProj(uShadowMap, sc + vec4(vec2(-1.5, -1.5) * uShadowSoftness, 0, 0));
        float c1 = textureProj(uShadowMap, sc + vec4(vec2(1.5, -1.5) * uShadowSoftness, 0, 0));
        float c2 = textureProj(uShadowMap, sc + vec4(vec2(-1.5, 1.5) * uShadowSoftness, 0, 0));
        float c3 = textureProj(uShadowMap, sc + vec4(vec2(1.5, 1.5) * uShadowSoftness, 0, 0));
        
        if ((c0 + c1 + c2 + c3) > 3.99) {
            shadow = 0.0;
        } else if ((c0 + c1 + c2 + c3) < 0.01) {
            shadow = 1.0;
        } else {
            float v_sum = c0 + c1 + c2 + c3;
            for(float x = -1.5; x <= 1.5; x += 1.0) {
                for(float y = -1.5; y <= 1.5; y += 1.0) {
                    if (abs(x) > 1.4 && abs(y) > 1.4) continue;
                    v_sum += textureProj(uShadowMap, sc + vec4(vec2(x,y) * uShadowSoftness, 0, 0));
                }
            }
            shadow = 1.0 - (v_sum / 16.0);
        }
    }
    
    if (uPrettyRender && uShowShadows) {
        visibility = mix(1.0, 1.0 - shadow, uShadowIntensity * 1.5);
    }
    
    float diff = max(dot(normal, L), 0.2) * visibility;
    vec3 baseColor = vColor.rgb * (0.3 + 0.7 * diff);
    
    if (uPrettyRender) {
        // Fetch 8-parameter material set from TBOs
        vec4 m1 = texelFetch(uAtomMaterial, vAtomIdx);
        vec4 m2 = texelFetch(uAtomMaterial2, vAtomIdx);
        
        float ambient   = m1.r;
        float specular  = m1.g;
        float shininess = m1.b;
        float rim       = m1.a;
        
        float m_sheen      = m2.r;
        float m_clearcoat  = m2.g;
        float m_reflection = m2.b;
        float m_glossiness = m2.a;

        // Stronger local contact occlusion for Bonds
        float horizon = max(normal.y, 0.0);
        float occlusion = mix(1.0, 0.4 + 0.6 * horizon, uShadowIntensity);
        float contact = 1.0 - pow(dot(normal, vec3(0,0,1)), 3.0);
        occlusion *= (1.0 - 0.4 * contact * uShadowIntensity);

        float diffuse = max(dot(normal, L), 0.0);
        float spec = pow(max(dot(normal, H), 0.0), shininess);
        
        float lightFactor = ambient + diffuse * visibility;
        baseColor = vColor.rgb * lightFactor * occlusion + vec3(spec * specular * visibility * occlusion);
        baseColor *= (1.0 - shadow * uShadowIntensity * 1.5);
        
        // Sheen
        if (m_sheen > 0.0) {
            float sheen_val = pow(1.0 - max(dot(normal, V), 0.0), 4.0);
            baseColor += m_sheen * sheen_val * vColor.rgb;
        }
        
        // Clearcoat
        if (m_clearcoat > 0.0) {
            vec3 H_cc = normalize(normalize(vec3(1.0, 1.0, 1.0)) + V);
            float specular_cc = pow(max(dot(normal, H_cc), 0.0), 256.0);
            baseColor += m_clearcoat * specular_cc * vec3(1.0);
        }

        if (m_reflection > 0.0) {
            vec3 R = reflect(-V, normal);
            vec3 worldR = normalize(vec3(uViewInv * vec4(R, 0.0)));
            float sky = smoothstep(-0.4, 0.4, worldR.y);
            float sun = pow(max(dot(R, normalize(vec3(0.5, 1.0, 0.5))), 0.0), 4.0 + m_glossiness * 60.0);
            vec3 env = mix(vec3(0.05, 0.05, 0.1), vec3(0.8, 0.9, 1.0), sky) + sun;
            baseColor += env * m_reflection * (0.3 + 0.7 * pow(1.0 - max(dot(normal, V), 0.0), 3.0));
        }
        
        // Exponential Fog
        float z = (uView * vec4(vPos, 1.0)).z;
        float fog = exp(-abs(z) * uFogDensity * 2.0);
        vec3 fogColor = vec3(0.01, 0.01, 0.02);
        baseColor = mix(fogColor, baseColor, fog);
        
        float fade = smoothstep(150.0, 600.0, abs(z));
        baseColor = mix(baseColor, fogColor, fade);
    }

    // Selection Glow (O(1) bitmask-based lookup)
    bool is_selected = uint(texelFetch(uSelectionMaskFull, vAtomIdx).r) > 0u;
    if (is_selected) {
        // High-contrast vibrant cyan-yellow mix for visibility
        baseColor = mix(baseColor, vec3(0.0, 1.0, 1.0), 0.4); 
        baseColor += vec3(0.2, 0.2, 0.0); // Slight golden boost
        baseColor *= 1.4;
    }

    FragColor = vec4(baseColor, vColor.a);
}
"""

# -- Post-Processing Shaders --

FULLSCREEN_VS = r"""#version 330 core
layout(location=0) in vec2 aPos;
out vec2 vUV;
void main(){
    vUV = aPos * 0.5 + 0.5;
    gl_Position = vec4(aPos, 0.0, 1.0);
}
"""

# 1. SSAO Shader (Fast 8-sample
SSAO_FS = r"""#version 330 core
in vec2 vUV;
uniform sampler2D uDepthTex;
uniform mat4 uProj;
uniform mat4 uProjInv;
uniform float uNear;
uniform float uFar;
uniform float uRadius;
uniform float uStrength;
uniform float uBias;
out vec4 FragColor;

uniform int uIsOrtho;

vec3 getViewPos(vec2 uv) {
    float z = texture(uDepthTex, uv).r;
    vec4 clip = vec4(uv * 2.0 - 1.0, z * 2.0 - 1.0, 1.0);
    vec4 view = uProjInv * clip;
    if (uIsOrtho == 0) {
        return view.xyz / max(view.w, 1e-6);
    } else {
        return view.xyz; // w is 1.0 in ortho
    }
}

float rand(vec2 co){
    return fract(sin(dot(co, vec2(12.9898, 78.233))) * 43758.5453);
}

void main() {
    float zRaw = texture(uDepthTex, vUV).r;
    if (zRaw >= 0.999999) {
        FragColor = vec4(1.0);
        return;
    }

    vec3 fragPos = getViewPos(vUV);

    vec2 texel = 1.0 / vec2(textureSize(uDepthTex, 0));
    float rotation = rand(vUV) * 6.2831853;
    mat2 R = mat2(cos(rotation), -sin(rotation),
                  sin(rotation),  cos(rotation));

    const int SAMPLES = 16;
    float occ = 0.0;

    // Derived world-space radius for distance check
    float worldRadius;
    if (uIsOrtho == 0) {
        worldRadius = uRadius * abs(fragPos.z) * 2.0;
    } else {
        worldRadius = uRadius * (2.0 / uProj[0][0]);
    }

    for (int i = 0; i < SAMPLES; ++i) {
        float t = (float(i) + 0.5) / float(SAMPLES);
        float ang = float(i) * 2.39996323;
        vec2 dir2 = R * vec2(cos(ang), sin(ang));

        vec2 sampleUV = vUV + dir2 * (uRadius * t);
        if (sampleUV.x < 0.0 || sampleUV.x > 1.0 || sampleUV.y < 0.0 || sampleUV.y > 1.0)
            continue;

        vec3 samplePos = getViewPos(sampleUV);
        float diff = samplePos.z - fragPos.z;
        float dist = distance(samplePos, fragPos);

        // Simple robust occlusion: Is the sample closer than frag?
        if (diff > uBias && dist < worldRadius * 2.0) {
            float rangeCheck = smoothstep(0.0, 1.0, worldRadius / dist);
            occ += 1.0 * rangeCheck;
        }
    }

    float ao = 1.0 - uStrength * (occ / float(SAMPLES));
    FragColor = vec4(vec3(clamp(ao, 0.0, 1.0)), 1.0);
}
"""

# 2. Bloom Extract & Blur
BLOOM_EXTRACT_FS = r"""#version 330 core
in vec2 vUV;
uniform sampler2D uColorTex;
uniform float uThreshold;
out vec4 FragColor;
void main(){
    vec3 col = texture(uColorTex, vUV).rgb;
    float brightness = dot(col, vec3(0.2126, 0.7152, 0.0722));
    if (brightness > uThreshold) FragColor = vec4(col, 1.0);
    else FragColor = vec4(0, 0, 0, 1);
}
"""

GAUSSIAN_BLUR_FS = r"""#version 330 core
in vec2 vUV;
uniform sampler2D uImage;
uniform bool uHorizontal;
uniform float uSize;
out vec4 FragColor;
void main(){
    float weight[5] = float[](0.227027, 0.1945946, 0.1216216, 0.054054, 0.016216);
    vec2 tex_offset = 1.0 / textureSize(uImage, 0);
    vec3 result = texture(uImage, vUV).rgb * weight[0];
    if(uHorizontal){
        for(int i=1; i<5; ++i){
            result += texture(uImage, vUV + vec2(tex_offset.x * i, 0.0)).rgb * weight[i];
            result += texture(uImage, vUV - vec2(tex_offset.x * i, 0.0)).rgb * weight[i];
        }
    } else {
        for(int i=1; i<5; ++i){
            result += texture(uImage, vUV + vec2(0.0, tex_offset.y * i)).rgb * weight[i];
            result += texture(uImage, vUV - vec2(0.0, tex_offset.y * i)).rgb * weight[i];
        }
    }
    FragColor = vec4(result, 1.0);
}
"""
GAUSSIAN_BILATERAL_FS = r"""#version 330 core
in vec2 vUV;
uniform sampler2D uImage;
uniform sampler2D uDepthTex;
uniform bool uHorizontal;
uniform float uNear;
uniform float uFar;
out vec4 FragColor;

float getLinearDepth(vec2 uv) {
    float z = texture(uDepthTex, uv).r;
    return (2.0 * uNear * uFar) / (uFar + uNear - (z * 2.0 - 1.0) * (uFar - uNear));
}

void main(){
    float weight[5] = float[](0.227027, 0.1945946, 0.1216216, 0.054054, 0.016216);
    vec2 tex_offset = 1.0 / textureSize(uImage, 0);
    float centerAO = texture(uImage, vUV).r;
    float centerDepth = getLinearDepth(vUV);
    
    float result = centerAO * weight[0];
    float normalization = weight[0];
    
    for(int i=1; i<5; ++i){
        vec2 off = uHorizontal ? vec2(tex_offset.x * i, 0.0) : vec2(0.0, tex_offset.y * i);
        
        // Positive side
        float depthP = getLinearDepth(vUV + off);
        float weightP = weight[i] * max(0.0, 1.0 - abs(centerDepth - depthP) * 2.0);
        result += texture(uImage, vUV + off).r * weightP;
        normalization += weightP;
        
        // Negative side
        float depthN = getLinearDepth(vUV - off);
        float weightN = weight[i] * max(0.0, 1.0 - abs(centerDepth - depthN) * 2.0);
        result += texture(uImage, vUV - off).r * weightN;
        normalization += weightN;
    }
    FragColor = vec4(vec3(result / normalization), 1.0);
}
"""

# 3. Final Resolve (Combine + FXAA + Gamma)
RESOLVE_FS = r"""#version 330 core
in vec2 vUV;
uniform sampler2D uSceneTex;
uniform sampler2D uSSAO;
uniform sampler2D uBloom;
uniform float uBloomIntensity;
uniform bool uEnableSSAO;
uniform bool uEnableBloom;
uniform bool uEnableFXAA;
uniform bool uDebugAO;
uniform sampler2D uDepthTex;
uniform sampler2D uSelectionMask;
uniform int uDebugMode; // 0: None, 1: AO, 2: Raw Depth, 3: Linear Depth
uniform float uGamma;
uniform float uNear;
uniform float uFar;
out vec4 FragColor;

float getLinearDepth(float z) {
    return (2.0 * uNear * uFar) / (uFar + uNear - (z * 2.0 - 1.0) * (uFar - uNear));
}

void main(){
    if (uDebugMode == 1) { // AO Only
        float ao = texture(uSSAO, vUV).r;
        FragColor = vec4(vec3(pow(ao, 2.0)), 1.0);
        return;
    }
    if (uDebugMode == 2) { // Raw Depth (Non-linear)
        float z = texture(uDepthTex, vUV).r;
        FragColor = vec4(vec3(z), 1.0);
        return;
    }
    if (uDebugMode == 3) { // Linear Depth (Normalized)
        float z = texture(uDepthTex, vUV).r;
        float d = getLinearDepth(z);
        float d_norm = clamp((d - uNear) / max(uFar - uNear, 1e-6), 0.0, 1.0);
        FragColor = vec4(vec3(d_norm), 1.0);
        return;
    }
    if (uDebugMode == 4) { // Amplified depth contrast
        float z = texture(uDepthTex, vUV).r;
        float d = getLinearDepth(z);
        float d_norm = clamp((d - uNear) / max(uFar - uNear, 1e-6), 0.0, 1.0);
        float vis = pow(1.0 - d_norm, 6.0);
        FragColor = vec4(vec3(vis), 1.0);
        return;
    }
    
    vec3 color = texture(uSceneTex, vUV).rgb;
    
    if (uEnableSSAO) {
        float ao = texture(uSSAO, vUV).r;
        color *= pow(ao, 4.0); // Extreme AO contrast for 'punchy' look
    }
    
    if (uEnableBloom) {
        vec3 bloom = texture(uBloom, vUV).rgb;
        color += bloom * uBloomIntensity;
    }
    
    if (uEnableFXAA) {
        vec2 off = 1.0 / textureSize(uSceneTex, 0);
        vec3 n = texture(uSceneTex, vUV + vec2(0, off.y)).rgb;
        vec3 s = texture(uSceneTex, vUV - vec2(0, off.y)).rgb;
        vec3 e = texture(uSceneTex, vUV + vec2(off.x, 0)).rgb;
        vec3 w = texture(uSceneTex, vUV - vec2(off.x, 0)).rgb;
        color = (color + n + s + e + w) * 0.2;
    }

    // color = color / (color + vec3(1.0)); // Removed Reinhard compression
    
    // ACES Filmic Tonemapping
    float a = 2.51;
    float b = 0.03;
    float c = 2.43;
    float d = 0.59;
    float e = 0.14;
    vec3 toneMapped = clamp((color*(a*color+b))/(color*(c*color+d)+e), 0.0, 1.0);
    
    // Vignette
    vec2 dist = (vUV - 0.5) * 2.0;
    float vignette = clamp(1.0 - dot(dist, dist) * 0.3, 0.0, 1.0);
    toneMapped *= vignette;

    // Dithering
    float noise = fract(sin(dot(vUV, vec2(12.9898, 78.233))) * 43758.5453) * 0.005;
    toneMapped += noise;

    // Gamma
    vec3 finalColor = pow(toneMapped, vec3(1.0/uGamma));
    
    // 4. Selection Silhouette (Outline)
    // Sobel-ish filter on selection mask
    vec2 tSz = 1.0 / textureSize(uSelectionMask, 0);
    float mask = texture(uSelectionMask, vUV).r;
    float edge = 0.0;
    if (mask == 0.0) {
        float m1 = texture(uSelectionMask, vUV + vec2(tSz.x, 0)).r;
        float m2 = texture(uSelectionMask, vUV - vec2(tSz.x, 0)).r;
        float m3 = texture(uSelectionMask, vUV + vec2(0, tSz.y)).r;
        float m4 = texture(uSelectionMask, vUV - vec2(0, tSz.y)).r;
        if (m1 > 0.0 || m2 > 0.0 || m3 > 0.0 || m4 > 0.0) edge = 1.0;
    }
    
    if (edge > 0.5) {
        finalColor = mix(finalColor, vec3(0.0), 1.0); // Black outline
    }

    FragColor = vec4(finalColor, 1.0);
}
"""

# -- Picking Shaders --
PICK_VS = r"""#version 330 core
layout(location=0) in vec2 aQuad;
uniform mat4 uView; uniform mat4 uProj;
uniform samplerBuffer uAtomPosRad;
uniform samplerBuffer uReplicaOffsets;
uniform int uAtomCount;
out vec2 vCoord;
out vec4 vColor;
void main(){
    int inst = gl_InstanceID;
    int atomIdx = inst % uAtomCount;
    int repIdx = inst / uAtomCount;
    vec4 posRad = texelFetch(uAtomPosRad, atomIdx);
    vec3 offset = texelFetch(uReplicaOffsets, repIdx).xyz;
    vec3 center = posRad.xyz + offset;
    float r = posRad.w;
    vec4 centerView = uView * vec4(center, 1.0);
    vec3 posVS = centerView.xyz + vec3(aQuad * r, 0.0);
    gl_Position = uProj * vec4(posVS, 1.0);
    vCoord = aQuad; // From -1 to 1
    int id = atomIdx;
    float R = float(id & 0xFF) / 255.0;
    float G = float((id >> 8) & 0xFF) / 255.0;
    float B = float((id >> 16) & 0xFF) / 255.0;
    vColor = vec4(R, G, B, 1.0);
}
"""
PICK_FS = r"""#version 330 core
in vec4 vColor;
in vec2 vCoord;
out vec4 FragColor;
void main() {
    if (dot(vCoord, vCoord) > 1.0) discard;
    FragColor = vColor;
}
"""

# -- Volume Render Shaders (Raymarching) --
VOLUME_VS = r"""#version 330 core
layout(location=0) in vec3 aPos;
uniform mat4 uView; uniform mat4 uProj;
uniform vec3 uOrigin; uniform vec3 uSize;
out vec3 vPos;
void main(){
    vec3 worldPos = uOrigin + aPos * uSize;
    vPos = worldPos;
    gl_Position = uProj * uView * vec4(worldPos, 1.0);
}
"""

VOLUME_FS = r"""#version 330 core
in vec3 vPos;
uniform vec3 uCamPos;
uniform vec3 uOrigin;
uniform vec3 uSize;
uniform float uIsoLevel;
uniform float uOpacity;
uniform float uStepSize;
uniform int uMode; // 0=Isosurface, 1=MIP
uniform vec3 uBaseColor;
uniform sampler3D uVolume;
out vec4 FragColor;

void main(){
    vec3 dir = normalize(vPos - uCamPos);
    
    // Ray-box intersection for exit point
    vec3 invDir = 1.0 / (dir + 1e-9);
    vec3 t0 = (uOrigin - uCamPos) * invDir;
    vec3 t1 = (uOrigin + uSize - uCamPos) * invDir;
    vec3 tmax = max(t0, t1);
    float tExit = min(min(tmax.x, tmax.y), tmax.z);
    
    float tStart = length(vPos - uCamPos);
    float tEnd = tExit;
    
    vec4 accum = vec4(0.0);
    float maxVal = 0.0;
    float rayLen = tEnd - tStart;
    if (rayLen <= 0.0) discard;

    int maxSteps = int(min(256.0, rayLen / uStepSize));
    for(int i=0; i < maxSteps; i++){
        vec3 p = vPos + dir * (float(i) * uStepSize);
        vec3 texCoord = (p - uOrigin) / uSize;
        
        float val = texture(uVolume, texCoord).r;
        
        if (uMode == 0) { // Isosurface
            if(val >= uIsoLevel){
                accum = vec4(uBaseColor, uOpacity);
                break;
            }
        } else { // MIP
            maxVal = max(maxVal, val);
        }
    }
    
    if (uMode == 1) { // MIP Finalize
        if (maxVal < 0.01) discard;
        accum = vec4(uBaseColor * maxVal, maxVal * uOpacity);
    }
    
    if(accum.a < 0.01) discard;
    FragColor = accum;
}
"""
