import { Matrix4, Vector2 } from 'three';

/**
 * Point-sampled screen-space volumetric obscurance.
 *
 * A hemisphere of radius `kernelRadius` is sampled in the volume above the surface (oriented by the
 * surface normal, lifted slightly off it to avoid self-occlusion). Each sample is counted as
 * occluded when it lies behind the depth buffer *and* the surface it tapped is itself within
 * `kernelRadius` (so distant geometry seen through the depth buffer does not occlude). The
 * occluded fraction is a physically plausible obscurance estimate, so `intensity = 1.0` produces a
 * correct-looking result with no scene-dependent multiplier; higher values are an artistic
 * darkening control.
 *
 * Sample placement uses a fixed R3 low-discrepancy kernel rotated per pixel by an R2 blue-noise
 * dither (interleaved sampling). The blue-noise rotation pushes the sampling error into high
 * frequencies that the edge-aware upscale downstream removes cleanly, leaving little residual.
 */

const SAOShader = {
    defines: {
        'NUM_SAMPLES': 7,
        'NORMAL_TEXTURE': 0,
        'DEPTH_PACKING': 1,
        'PERSPECTIVE_CAMERA': 1
    },
    uniforms: {

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

        'cameraNear': { value: 1 },
        'cameraFar': { value: 100 },
        'cameraProjectionMatrix': { value: new Matrix4() },
        'cameraInverseProjectionMatrix': { value: new Matrix4() },

        // artistic darkening control; 1.0 == physically plausible
        'intensity': { value: 1.0 },

        // radius of the sampling hemisphere, in world/view-space units; also the maximum distance at
        // which a tapped surface counts as an occluder
        'kernelRadius': { value: 0.5 }
    },

    //language=GLSL
    vertexShader:`

		varying vec2 vUv;

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

    //language=GLSL
    fragmentShader:        `

		#include <common>

		varying vec2 vUv;

		uniform sampler2D tDepth;

		#if NORMAL_TEXTURE == 1
		uniform sampler2D tNormal;
		#endif

		uniform highp usampler2D tHilbert;

		uniform float cameraNear;
		uniform float cameraFar;
		uniform mat4 cameraProjectionMatrix;
		uniform mat4 cameraInverseProjectionMatrix;

		uniform float intensity;
		uniform float kernelRadius;
		uniform vec2 size;

		// RGBA depth

		#include <packing>

		out float out_occlusion;

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

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

		vec3 getViewPosition( const in vec2 screenPosition, const in float depth, const in float viewZ ) {
			float clipW = cameraProjectionMatrix[2][3] * viewZ + cameraProjectionMatrix[3][3];
			vec4 clipPosition = vec4( ( vec3( screenPosition, depth ) - 0.5 ) * 2.0, 1.0 );
			clipPosition *= clipW; // unprojection.

			return ( cameraInverseProjectionMatrix * clipPosition ).xyz;
		}

		vec3 getViewNormal( const in vec3 viewPosition, const in vec2 screenPosition ) {
			#if NORMAL_TEXTURE == 1
			return unpackRGBToNormal( texture2D( tNormal, screenPosition ).xyz );
			#else
			// Accurate normal reconstruction (Drobot): for each axis pick the nearer-in-depth
			// neighbour, so the tangent basis never straddles a silhouette. dFdx-based normals read
			// garbage across depth discontinuities, which mis-orients the sampling hemisphere and
			// makes flat, camera-facing surfaces self-occlude along their edges.
			vec2 texel = 1.0 / size;

			float dc = getDepth( screenPosition );
			float dl = getDepth( screenPosition - vec2( texel.x, 0.0 ) );
			float dr = getDepth( screenPosition + vec2( texel.x, 0.0 ) );
			float dd = getDepth( screenPosition - vec2( 0.0, texel.y ) );
			float du = getDepth( screenPosition + vec2( 0.0, texel.y ) );

			vec3 dpdx;
			if ( abs( dl - dc ) < abs( dr - dc ) ) {
				vec3 left = getViewPosition( screenPosition - vec2( texel.x, 0.0 ), dl, getViewZ( dl ) );
				dpdx = viewPosition - left;
			} else {
				vec3 right = getViewPosition( screenPosition + vec2( texel.x, 0.0 ), dr, getViewZ( dr ) );
				dpdx = right - viewPosition;
			}

			vec3 dpdy;
			if ( abs( dd - dc ) < abs( du - dc ) ) {
				vec3 down = getViewPosition( screenPosition - vec2( 0.0, texel.y ), dd, getViewZ( dd ) );
				dpdy = viewPosition - down;
			} else {
				vec3 up = getViewPosition( screenPosition + vec2( 0.0, texel.y ), du, getViewZ( du ) );
				dpdy = up - viewPosition;
			}

			return normalize( cross( dpdx, dpdy ) );
			#endif
		}

		// Marty's R2 low-discrepancy sequence: maps a permutation index to a 2D blue-noise-like point.
		// (GLSL ES has no fma; the expanded multiply-add is equivalent here.)
		vec2 r2_marty( uint index ) {
			return fract( vec2( float( index ) ) * vec2( 0.245122333753, 0.430159709002 ) + 0.5 );
		}

		// Concentric disk mapping (Shirley & Chiu): low-distortion uniform disk sample from [0,1)^2.
		vec2 sampleUniformDiskConcentric( const in vec2 u ) {
			vec2 offset = 2.0 * u - 1.0;

			if ( offset.x == 0.0 && offset.y == 0.0 ) return vec2( 0.0 );

			float r;
			float theta;

			if ( abs( offset.x ) > abs( offset.y ) ) {
				r = offset.x;
				theta = ( PI * 0.25 ) * ( offset.y / offset.x );
			} else {
				r = offset.y;
				theta = ( PI * 0.5 ) - ( PI * 0.25 ) * ( offset.x / offset.y );
			}

			return r * vec2( cos( theta ), sin( theta ) );
		}

		// Cosine-weighted hemisphere direction in tangent space (z is up), via Malley's method.
		vec3 sampleCosineWeightedHemisphere( const in vec2 u ) {
			vec2 d = sampleUniformDiskConcentric( u );
			float z = sqrt( max( 0.0, 1.0 - dot( d, d ) ) );
			return vec3( d.x, d.y, z );
		}

		// Branchless orthonormal basis from a unit normal (Duff et al. 2017). Columns are T, B, N, so
		// basis * v maps a tangent-space vector (z up) into the normal's frame.
		mat3 buildOrthonormalBasis( const in vec3 n ) {
			vec3 T;
			vec3 B;

			if ( n.z < 0.0 ) {
				float a = 1.0 / ( 1.0 - n.z );
				float b = n.x * n.y * a;
				T = vec3( 1.0 - n.x * n.x * a, -b, n.x );
				B = vec3( b, n.y * n.y * a - 1.0, -n.y );
			} else {
				float a = 1.0 / ( 1.0 + n.z );
				float b = -n.x * n.y * a;
				T = vec3( 1.0 - n.x * n.x * a, b, -n.x );
				B = vec3( b, 1.0 - n.y * n.y * a, -n.y );
			}

			return mat3( T, B, n );
		}

		// Roberts' R3 low-discrepancy additive recurrence: ( 1/g, 1/g^2, 1/g^3 ), g^4 = g + 1
		const vec3 R3 = vec3( 0.8191725134, 0.6710436067, 0.5497004779 );

		// normal-bias as a fraction of kernelRadius: lifts the sample origin off the surface to avoid
		// self-occlusion. Larger removes acne but lets near-surface contact occlusion leak.
		const float SURFACE_BIAS = 0.05;

		// push each sample this many pixels further along its screen-space direction so it never
		// reads the centre pixel (self-occlusion); folded into the offset instead of a hard reject
		const float SELF_OCCLUSION_MARGIN_PIXELS = 2.0;

		// radial distribution bias: normalized radius -> pow( radius, power ) pulls samples toward
		// the origin, where occlusion contributes most (small crevices over large open surfaces)
		const float sampleDistributionPower = 2.0;

		// distant samples contribute less (less reliable / more likely to be noise): fraction of the
		// radius over which an occluder's contribution fades to zero at the edge. Must be in [0, 1]. (GTAO)
		const float GTAO_EFFECT_FALLOFF_RANGE = 0.615;

		float getOccludedFraction( const in vec3 centerViewPosition, const in vec2 centerUv ) {
			vec3 centerViewNormal = getViewNormal( centerViewPosition, centerUv );

			// sample origin lifted slightly off the surface along the normal (see SURFACE_BIAS)
			vec3 sampleOrigin = centerViewPosition + centerViewNormal * ( kernelRadius * SURFACE_BIAS );

			// tangent->view basis, built once and reused for every sample
			mat3 tbn = buildOrthonormalBasis( centerViewNormal );

			// per-pixel 2D blue noise: a Hilbert-curve permutation index (spatially well distributed)
			// fed through Marty's R2 sequence. .x rotates the kernel around the normal, .y offsets the
			// sampling radius -- decorrelating neighbours so the edge-aware upscale resolves them.
			uint hilbertIndex = texelFetch( tHilbert, ivec2( gl_FragCoord.xy ) & 63, 0 ).x;
			vec2 noise = r2_marty( hilbertIndex );

			float ca = cos( noise.x * PI2 );
			float sa = sin( noise.x * PI2 );
			vec3 rT = tbn[0] * ca + tbn[1] * sa;
			vec3 rB = tbn[1] * ca - tbn[0] * sa;
			tbn[0] = rT;
			tbn[1] = rB;

			// distance falloff (GTAO): full weight until falloffFrom, then linear to zero at the radius
			float falloffRange = GTAO_EFFECT_FALLOFF_RANGE * kernelRadius;
			float falloffFrom  = kernelRadius - falloffRange;
			float falloffMul   = -1.0 / falloffRange;
			float falloffAdd   = falloffFrom / falloffRange + 1.0;

			float occlusion = 0.0;
			float weightSum = 0.0;

			for ( int i = 0; i < NUM_SAMPLES; i ++ ) {
				// fixed R3 low-discrepancy point; per-pixel variation comes from the rotated frame
				// and the radius offset below
				vec3 q = fract( float( i ) * R3 );

				// cosine-weighted hemisphere direction, oriented to the (rotated) surface frame:
				// concentrates samples near the normal where occlusion matters most
				vec3 dir = normalize( tbn * sampleCosineWeightedHemisphere( q.xy ) );

				// normalized radius with a per-pixel offset, biased toward the origin (see
				// sampleDistributionPower) so the near occluders that dominate AO are favoured
				float radius = pow( fract( q.z + noise.y ), sampleDistributionPower );

				// falloff weight from the probe's own distance -- defined for every sample, so the
				// weighted normaliser is well posed; distant probes are less reliable (GTAO)
				float weight = saturate( radius * kernelRadius * falloffMul + falloffAdd );

				vec3 samplePositionVS = sampleOrigin + dir * ( radius * kernelRadius );

				// project the sample into screen space
				vec4 clip = cameraProjectionMatrix * vec4( samplePositionVS, 1.0 );
				vec2 sampleUv = ( clip.xy / clip.w ) * 0.5 + 0.5;

				// screen-space offset from the centre, in pixels
				vec2 offsetPixels = ( sampleUv - centerUv ) * size;

				// push forward along the screen direction by the self-occlusion margin (extends the
				// screen-space radius), then snap to the nearest pixel centre so the depth fetch is
				// exact and neighbouring pixels stay coherent (GTAO)
				float offsetLength = length( offsetPixels );
				offsetPixels += offsetPixels / max( offsetLength, EPSILON ) * SELF_OCCLUSION_MARGIN_PIXELS;
				offsetPixels = round( offsetPixels );

				sampleUv = centerUv + offsetPixels / size;

				// samples outside the screen have no depth to test against
				if ( sampleUv.x < 0.0 || sampleUv.x > 1.0 || sampleUv.y < 0.0 || sampleUv.y > 1.0 ) {
					continue;
				}

				// a valid (on-screen) sample contributes its weight to the normaliser, occluded or not
				weightSum += weight;

				float sampleDepth = getDepth( sampleUv );

				// skip the sky / far plane
				if ( sampleDepth >= ( 1.0 - EPSILON ) ) {
					continue;
				}

				float sceneViewZ = getViewZ( sampleDepth );

				// ...and the tapped surface must actually be within the sampling radius, otherwise we
				// have hit distant geometry through the depth buffer that does not occlude this point
				vec3 occluderVS = getViewPosition( sampleUv, sampleDepth, sceneViewZ );
				vec3 to_hit = occluderVS - centerViewPosition;
				
				float distance_to_hit = length( to_hit );
				
				if ( distance_to_hit > kernelRadius ) {
					continue;
				}
				
				vec3 direction_to_hit = normalize(occluderVS - sampleOrigin);
				
				if( dot( direction_to_hit, centerViewNormal) < 0.06 ){ // ~sqrt( 1/255 )
				    // unable to occlude
				    continue;
				}
				

				// occluded: add this sample's weight to the numerator
				occlusion += weight;
			}

			// weighted occlusion fraction; off-screen samples were excluded so screen borders are unbiased
			return weightSum > 0.0 ? occlusion / weightSum : 0.0;
		}

		void main() {
			// snap to the centre of the full-res depth texel we sample. The half-res pass otherwise
			// lands on texel boundaries, where NEAREST biases consistently to one side and the
			// reconstructed position disagrees with the sampled depth -> asymmetric edge self-occlusion.
			vec2 centerUv = ( floor( vUv * size ) + 0.5 ) / size;

			float centerDepth = getDepth( centerUv );
			if( centerDepth >= ( 1.0 - EPSILON ) ) {
				// background / far plane is fully visible. Write white rather than discarding, so the
				// half-res buffer never leaves the clear value for the upscale to read back as occlusion
				// (that is the foreground-edge halo over background).
				out_occlusion = 1.0;
				return;
			}

			float centerViewZ = getViewZ( centerDepth );
			vec3 viewPosition = getViewPosition( centerUv, centerDepth, centerViewZ );

			float occlusion = getOccludedFraction( viewPosition, centerUv );

			// intensity is an artistic power: 1.0 leaves the physical estimate untouched, higher
			// darkens occluded areas while leaving fully-lit areas at 1.0
			float visibility = pow( 1.0 - occlusion, intensity );

			// never fully occlude: if the receiver can see a sample at all it is not truly black, and
			// the R8 target cannot represent anything below 1/255 anyway
			out_occlusion = max( 1.0 / 255.0, visibility );
		}`

};

export { SAOShader };
