import { Fn, If, PI, clamp, cos, cross, dot, equirectUV, float, log, max, mix, normalize, pow, reflect, sin, sqrt, struct, vec3 } from 'three/tsl';

/**
 * Specular / microfacet BRDF helpers: VNDF sampling, GTR distribution, Smith geometry,
 * Fresnel, reflection importance sampling, parallax-corrected ray-length terms, and
 * equirectangular environment sampling / MIS helpers.
 * Pure TSL functions of their inputs (no scene/camera state).
 */

/**
 * Sentinel ray length the SSR pass writes for environment misses (no screen-space hit), set far above
 * any real hit distance so a single magnitude test separates misses from hits and survives `.max( 0 )`.
 *
 * @type {number}
 */
export const ENV_RAY_LENGTH = 1e4;

/**
 * Classification threshold for {@link ENV_RAY_LENGTH}: above this is an env miss, below a real hit.
 * An order of magnitude under the sentinel, robust to fp16 storage and bilinear blending at borders.
 *
 * @type {number}
 */
export const ENV_RAY_LENGTH_THRESHOLD = 1e3;

// Bounded-VNDF sampler (Eto & Tokuyoshi 2023; spherical-cap form, Dupuy & Benyoub 2023)
const SampleGGXVNDF = Fn( ( [ V, ax, ay, r1, r2 ] ) => {

	// Warp the view direction to the hemisphere ("standard") configuration.
	const wiStd = normalize( vec3( ax.mul( V.x ), ay.mul( V.y ), V.z ) ).toVar();

	// Isotropic bound on the spherical cap (Eto & Tokuyoshi eq. 5). alpha ∈ [0,1] here,
	// so the sign term in `s` is always +1 and is dropped.
	const a = ax.min( ay ).toVar();
	const s = float( 1.0 ).add( V.xy.length() ).toVar();
	const a2 = a.mul( a ).toVar();
	const s2 = s.mul( s ).toVar();
	const k = a2.oneMinus().mul( s2 ).div( s2.add( a2.mul( V.z ).mul( V.z ) ) ).toVar();

	// Tighten the cap with the bound (upper hemisphere; N·V ≥ 0 in our usage).
	const b = wiStd.z.mul( k ).toVar();

	// Sample the (bounded) spherical cap.
	const phi = float( 6.283185307179586 ).mul( r1 ).toVar(); // 2*pi
	const z = r2.oneMinus().mul( float( 1.0 ).add( b ) ).sub( b ).toVar();
	const sinTheta = sqrt( max( float( 0.0 ), float( 1.0 ).sub( z.mul( z ) ) ) ).toVar();
	const c = vec3( sinTheta.mul( cos( phi ) ), sinTheta.mul( sin( phi ) ), z ).toVar();

	// Microfacet normal in the standard config, then warp back to the ellipsoid (unstretch).
	const wmStd = c.add( wiStd ).toVar();
	const Ne = normalize( vec3(
		ax.mul( wmStd.x ),
		ay.mul( wmStd.y ),
		max( float( 0.0 ), wmStd.z )
	) ).toVar();

	return Ne;

}, {
	name: 'SampleGGXVNDF',
	type: 'vec3',
	inputs: [
		{ name: 'V', type: 'vec3' },
		{ name: 'ax', type: 'float' },
		{ name: 'ay', type: 'float' },
		{ name: 'r1', type: 'float' },
		{ name: 'r2', type: 'float' },
	]
} );

// Generalized Trowbridge-Reitz (GTR). For GGX set k=2.
// D_GTR(roughness, NoH, k) where roughness = α (not α²).
export const D_GTR = Fn( ( [ roughness, NoH, k ] ) => {

	// see: Filament - Normal distribution function (specular D) - 4.4.1
	const a2 = roughness.mul( roughness ).toVar(); // α²
	const NoH2 = NoH.mul( NoH ).toVar();
	const base = NoH2.mul( a2.sub( float( 1.0 ) ) ).add( float( 1.0 ) ).toVar();
	// a² / (π * base^k)
	return a2.div( PI.mul( pow( base, k ) ) ).toVar(); // float

} );

// Smith G1 (Heitz): expects alpha (not squared); it squares internally
export const SmithG = Fn( ( [ NDotX, alpha ] ) => {

	// see: Filament - Geometric shadowing (specular G) - 4.4.2

	const a2 = alpha.mul( alpha ).toVar(); // α²
	const NDotX2 = NDotX.mul( NDotX ).toVar(); // (N·X)²
	return float( 2.0 ).mul( NDotX ).div(
		NDotX.add( sqrt(
			a2.add( a2.oneMinus().mul( NDotX2 ) )
		) )
	);

} );

// Geometry term G = G1(N·L, α_G) * G1(N·V, α_G)  (α_G is NOT squared here)
export const GeometryTerm = Fn( ( [ NoL, NoV, alphaG ] ) => {

	const G1v = SmithG( NoV, alphaG ).toVar();
	const G1l = SmithG( NoL, alphaG ).toVar();
	return G1v.mul( G1l ).toVar();

} );

// Bounded VNDF direction PDF (reflection mapping), matching SampleGGXVNDF above.
// p(L) = D_GTR(roughness, NoH, 2) / ( 2 * (k * N·V + t) )   (Eto & Tokuyoshi eq. 8)
// with the isotropic cap bound k and t = ‖(α·V.xy, V.z)‖. Here 'roughness' is α, not α².
const GGXVNDFPdf = Fn( ( [ NoH, NoV, roughness ] ) => {

	const D = D_GTR( roughness, NoH, float( 2.0 ) ).toVar();
	const a2 = roughness.mul( roughness ).toVar();
	const sinV2 = max( float( 0.0 ), float( 1.0 ).sub( NoV.mul( NoV ) ) ).toVar(); // ‖V.xy‖²
	const s = float( 1.0 ).add( sqrt( sinV2 ) ).toVar();
	const s2 = s.mul( s ).toVar();
	const k = float( 1.0 ).sub( a2 ).mul( s2 ).div( s2.add( a2.mul( NoV ).mul( NoV ) ) ).toVar();
	const t = sqrt( a2.mul( sinV2 ).add( NoV.mul( NoV ) ) ).toVar();
	return D.div( max( float( 1e-6 ), float( 2.0 ).mul( k.mul( NoV ).add( t ) ) ) ).toVar();

} );

/**
 * Fresnel reflectance for the Schlick approximation.
 */
export const F_Schlick = Fn( ( [ f0, theta ] ) => {

	const oneMinus = float( 1.0 ).sub( theta ).toVar();
	const oneMinus2 = oneMinus.mul( oneMinus ).toVar();
	const oneMinus5 = oneMinus2.mul( oneMinus2 ).mul( oneMinus ).toVar();
	return f0.add( vec3( 1.0 ).sub( f0 ).mul( oneMinus5 ) ).toVar(); // vec3

} );

/**
 * Specular dominant factor for parallax-corrected ray length.
 * From REBLUR: A Hierarchical Recurrent Denoiser (NRD).
 */
export const getSpecularDominantFactor = Fn( ( [ NoV, roughness ] ) => {

	const a = float( 0.298475 ).mul(
		log( float( 39.4115 ).sub( float( 39.0029 ).mul( roughness ) ) )
	);
	const f = float( 1.0 ).sub( NoV ).pow( 10.8649 )
		.mul( float( 1.0 ).sub( a ) )
		.add( a );
	return clamp( f );

} ).setLayout( {
	name: 'getSpecularDominantFactor',
	type: 'float',
	inputs: [
		{ name: 'NoV', type: 'float' },
		{ name: 'roughness', type: 'float' }
	]
} );

/**
 * Everything a single GGX reflection sample produces. `reflectDir` and `sampleWeight`
 * drive the SSR ray-march and compositing; `pdf`, `NdotV`, `alpha` and `f0` are the GGX
 * terms the env-miss MIS fallback needs so the caller never re-derives microfacet math.
 */
const ggxReflectionStruct = struct( {
	reflectDir: 'vec3', // view-space reflected ray direction
	sampleWeight: 'vec3', // chromatic weight (incl. Fresnel tint) to multiply gathered radiance by
	pdf: 'float', // VNDF direction pdf (for MIS against the env CDF)
	NdotV: 'float',
	alpha: 'float', // GGX alpha (roughness²), clamped
	f0: 'vec3' // Fresnel reflectance at normal incidence
} );

/**
 * Importance-samples the GGX/VNDF specular lobe for one pixel and returns the reflected
 * ray direction plus the Monte-Carlo weight to apply to the gathered radiance, along with
 * the GGX terms the SSR env-miss MIS fallback needs.
 *
 * @param {Node<vec3>} N - View-space shading normal (normalized).
 * @param {Node<vec3>} V - View-space surface→camera direction (normalized).
 * @param {Node<float>} roughness - Perceptual roughness in `[0,1]`.
 * @param {Node<float>} metalness - Metalness in `[0,1]`.
 * @param {Node<vec3>} albedo - Surface base color; tints the metal Fresnel reflectance (`f0`).
 * @param {Node<vec4>} Xi - Per-pixel random numbers; only `.xy` are used.
 * @return {ggxReflectionStruct}
 */
export const ggxReflectionSample = Fn( ( [ N, V, roughness, metalness, albedo, Xi ] ) => {

	// GGX alpha (use r^2, clamp to avoid degenerate)
	const a = roughness.mul( roughness ).max( 0.001 ).toVar();
	const ax = a.toVar();
	const ay = a.toVar();

	// TBN from view-space normal
	const up = vec3( 0, 0, 1 );
	let T = cross( up, N ).toVar();
	T = T.normalize().toVar();
	If( T.length().lessThan( 1e-3 ), () => {

		T.assign( cross( vec3( 0, 1, 0 ), N ).normalize() );

	} );
	const B = cross( N, T ).normalize().toVar();

	// transform V to LOCAL frame (N = +Z)
	const Vlocal = vec3( dot( T, V ), dot( B, V ), dot( N, V ) ).toVar();

	// VNDF sample **in local frame**
	const Hlocal = SampleGGXVNDF( Vlocal, ax, ay, Xi.x, Xi.y ).toVar();
	If( Hlocal.z.lessThan( 0 ), () => {

		Hlocal.assign( Hlocal.negate() );

	} );

	// transform H back to VIEW space
	const h = normalize( T.mul( Hlocal.x ).add( B.mul( Hlocal.y ) ).add( N.mul( Hlocal.z ) ) ).toVar();

	// reflect with V (surface->camera) and H
	const viewReflectDir = reflect( V.negate(), h ).normalize().toVar();

	// BRDF/PDF evaluation for the sampled direction
	// V: surface -> camera, L: reflected direction, N: normal, H: half-vector
	const L = viewReflectDir.toVar();
	const H = normalize( V.add( L ) ).toVar(); // ~h; recomputed for robustness

	const NdotV = max( float( 0.0 ), dot( N, V ) ).toVar();
	const NdotL = max( float( 0.0 ), dot( N, L ) ).toVar();
	const NdotH = max( float( 0.0 ), dot( N, H ) ).toVar();
	const VdotH = max( float( 0.0 ), dot( V, H ) ).toVar();

	const f0 = mix( vec3( 0.04 ), albedo, metalness ).toVar();
	// Chromatic Fresnel reflectance: for metals f0 = albedo, so the reflection is tinted and desaturates
	// toward white at grazing angles. Kept as vec3 so colored metals reflect with the correct chroma.
	const fresnelWeight = F_Schlick( f0, VdotH ).toVar(); // vec3

	// Bounded-VNDF direction pdf — still needed for the env-miss MIS path.
	const pdf = GGXVNDFPdf( NdotH, NdotV, ax ).toVar();

	// Numerically stable importance weight: brdf·NdotL/pdf ≡ fresnel·G2·(k·NdotV + t)/(2·NdotV), which
	// cancels the GGX D analytically. Evaluating D explicitly is catastrophic at low roughness
	// (D → 3e5 at α = 0.001 wrecks f32 precision); the cancelled form stays stable down to a mirror.
	// (k·NdotV + t) is the bounded-cap normalization; k shrinks the cap to drop below-horizon samples.
	const a2 = ax.mul( ax ).toVar();
	const sinV2 = NdotV.mul( NdotV ).oneMinus().max( 0.0 ).toVar(); // ‖V.xy‖²
	const sB = float( 1.0 ).add( sqrt( sinV2 ) ).toVar();
	const s2B = sB.mul( sB ).toVar();
	const kB = a2.oneMinus().mul( s2B ).div( s2B.add( a2.mul( NdotV ).mul( NdotV ) ) ).toVar();
	const tB = sqrt( a2.mul( sinV2 ).add( NdotV.mul( NdotV ) ) ).toVar();
	const glossyWeight = fresnelWeight
		.mul( GeometryTerm( NdotL, NdotV, ax ) )
		.mul( kB.mul( NdotV ).add( tB ) )
		.div( float( 2.0 ).mul( NdotV ).max( 1e-4 ) ).toVar();

	return ggxReflectionStruct( viewReflectDir, glossyWeight, pdf, NdotV, ax, f0 );

} );

// Equirectangular environment sampling

/**
 * Equirectangular direction / UV / PDF helpers and MIS weighting shared by environment sampling code.
 * Env-miss MIS integration lives in {@link ImportanceSampledEnvironment}.
 *
 * Equirectangular parameterization helpers used with CDF importance sampling are adapted from
 * [three-gpu-pathtracer](https://github.com/gkjohnson/three-gpu-pathtracer).
 *
 * @see {@link https://github.com/gkjohnson/three-gpu-pathtracer}
 */

// uv -> direction (equirectangular)
export const equirectUvToDir = Fn( ( [ uvIn ] ) => {

	const phi = uvIn.x.mul( Math.PI * 2 ).sub( Math.PI );
	const lat = uvIn.y.sub( 0.5 ).mul( Math.PI );
	const cosLat = cos( lat );
	return normalize( vec3(
		cosLat.mul( cos( phi ) ),
		sin( lat ),
		cosLat.mul( sin( phi ) )
	) );

} ).setLayout( {
	name: 'equirectUvToDir',
	type: 'vec3',
	inputs: [ { name: 'uv', type: 'vec2' } ]
} );

// Solid-angle PDF of a direction under equirectangular parameterization.
export const equirectDirPdf = Fn( ( [ direction ] ) => {

	const uvDir = equirectUV( direction );
	const sinTheta = sin( uvDir.y.mul( Math.PI ) );
	return sinTheta.abs().lessThan( float( 1e-6 ) ).select(
		float( 0 ),
		float( 1 ).div( float( 2 * Math.PI * Math.PI ).mul( sinTheta ) )
	);

} ).setLayout( {
	name: 'equirectDirPdf',
	type: 'float',
	inputs: [ { name: 'direction', type: 'vec3' } ]
} );

/**
 * MIS power heuristic with β = 2: `pdfA² / (pdfA² + pdfB²)`.
 * Weights the contribution of the strategy that produced `pdfA` against the other strategy.
 *
 * @see Eric Veach, *Optimally Combining Sampling Techniques for Monte Carlo Rendering*
 * @tsl
 */
export const misPowerHeuristic = Fn( ( [ pdfA, pdfB ] ) => {

	const pdfASq = pdfA.mul( pdfA );
	const pdfBSq = pdfB.mul( pdfB );
	return pdfASq.div( pdfASq.add( pdfBSq ) );

} ).setLayout( {
	name: 'misPowerHeuristic',
	type: 'float',
	inputs: [
		{ name: 'pdfA', type: 'float' },
		{ name: 'pdfB', type: 'float' }
	]
} );