import { Vector2 } from 'three';

/**
 * Edge-aware upscale + denoise for the half-resolution SAO buffer.
 *
 * A single non-separable pass (separable bilateral filters leak across diagonal silhouettes). Each
 * tap is weighted by its distance to the *plane* of the centre pixel, modelled from the screen-space
 * depth gradient (ReBLUR-style): slanted surfaces are predicted correctly, so only genuine depth
 * discontinuities are rejected -- no normal buffer and no hand-tuned raw-depth tolerance. The
 * geometry tolerance is intentionally tight so background occlusion never bleeds into the foreground
 * or vice versa.
 *
 * Taps follow a Vogel disk rotated per pixel by an R2 blue-noise angle, which both spreads the
 * gather as a dithered (non-grid) pattern and pairs with the R2-rotated sampling kernel in SAOShader.
 */

const SAOUpscaleShader = {
    defines: {
        'NUM_TAPS': 16,
        'DEPTH_PACKING': 0,
        'PERSPECTIVE_CAMERA': 1
    },
    uniforms: {

        'tAO': { value: null },
        'tDepth': { value: null },
        'tHilbert': { value: null },
        'size': { value: new Vector2(512, 512) },

        'cameraNear': { value: 1 },
        'cameraFar': { value: 100 },

        // tap-disk radius, in output (full-resolution) pixels
        'blurRadius': { value: 10.0 },

        // plane-distance tolerance as a fraction of view depth; smaller = more aggressive fg/bg split
        'depthSigma': { value: 0.02 }
    },
    vertexShader: /* glsl */`

		varying vec2 vUv;

		void main() {
			vUv = uv;
			gl_Position = projectionMatrix * modelViewMatrix * vec4( position, 1.0 );
		}`,

    fragmentShader: /* glsl */`

		#include <common>

		varying vec2 vUv;

		uniform sampler2D tAO;
		uniform sampler2D tDepth;
		uniform highp usampler2D tHilbert;
		uniform vec2 size;
		uniform float cameraNear;
		uniform float cameraFar;
		uniform float blurRadius;
		uniform float depthSigma;

		#include <packing>

		out float out_occlusion;

		float getDepth( const in vec2 uv ) {
			#if DEPTH_PACKING == 1
			return unpackRGBAToDepth( texture2D( tDepth, uv ) );
			#else
			return texture2D( tDepth, uv ).x;
			#endif
		}

		float getViewZ( const in float depth ) {
			#if PERSPECTIVE_CAMERA == 1
			return perspectiveDepthToViewZ( depth, cameraNear, cameraFar );
			#else
			return orthographicDepthToViewZ( depth, cameraNear, cameraFar );
			#endif
		}

		// Marty's R2 low-discrepancy sequence: permutation index -> 2D blue-noise-like point
		vec2 r2_marty( uint index ) {
			return fract( vec2( float( index ) ) * vec2( 0.245122333753, 0.430159709002 ) + 0.5 );
		}

		void main() {
			float centerDepth = getDepth( vUv );

			// sky / far plane -> no occlusion
			if ( centerDepth >= ( 1.0 - EPSILON ) ) {
				out_occlusion = 1.0;
				return;
			}

			float zc = getViewZ( centerDepth );

			// screen-space depth gradient defines the centre pixel's plane; predicting a tap's depth
			// from it means slanted surfaces are not rejected, only true discontinuities are
			vec2 dz = vec2( dFdx( zc ), dFdy( zc ) );

			vec2 texel = 1.0 / size;

			// plane-distance tolerance scales with view depth (footprint and depth precision grow with
			// distance); smaller depthSigma = tighter fg/bg separation
			float invSigma = 1.0 / max( EPSILON, depthSigma * abs( zc ) );

			// per-pixel rotation of the Vogel tap set -> dithered, non-separable gather. Hilbert-curve
			// permutation index through Marty's R2 sequence (matches the AO sampling dither).
			uint hilbertIndex = texelFetch( tHilbert, ivec2( gl_FragCoord.xy ) & 63, 0 ).x;
			float rot = r2_marty( hilbertIndex ).x * PI2;

			float spatialSigma = blurRadius * 0.5;
			float invSpatial2 = 1.0 / ( 2.0 * spatialSigma * spatialSigma );

			// centre tap
			float sum = texture2D( tAO, vUv ).x;
			float wsum = 1.0;

			for ( int i = 0; i < NUM_TAPS; i ++ ) {
				// Vogel-disk offset in output pixels, rotated per pixel
				float fi = float( i ) + 0.5;
				float r = sqrt( fi / float( NUM_TAPS ) ) * blurRadius;
				float theta = fi * 2.39996323 + rot;       // golden angle
				vec2 o = r * vec2( cos( theta ), sin( theta ) );

				vec2 tapUv = vUv + o * texel;

				if ( tapUv.x < 0.0 || tapUv.x > 1.0 || tapUv.y < 0.0 || tapUv.y > 1.0 ) continue;

				float zt = getViewZ( getDepth( tapUv ) );

				// distance from the tap to the centre pixel's plane
				float planeDist = abs( zt - ( zc + dot( dz, o ) ) );

				float wGeom = exp( - planeDist * planeDist * invSigma * invSigma );
				float wSpatial = exp( - dot( o, o ) * invSpatial2 );

				float w = wGeom * wSpatial;

				sum += w * texture2D( tAO, tapUv ).x;
				wsum += w;
			}

			out_occlusion = sum / wsum;
		}`

};

export { SAOUpscaleShader };
