#ifndef __IMPOSTOR_INC
#define __IMPOSTOR_INC

#include "UnityCG.cginc"
#include "globals.cginc"
#include "vertex_deformation.hlsl"

// Utility functions for hemispherical octahedral mapping
float2 HemiOctEncode(float3 N) {
    N.y = max(N.y, 1e-4);
    float3 p = hemi_octahedron_to_plane(normalize(N), 0, float3(1,0,0), float3(0,1,0), 1);
    return p.xz;
}

float3 HemiOctDecode(float2 uv) {
    return normalize(plane_to_hemi_octahedron(float3(uv.x, 0, uv.y), 0, float3(1,0,0), float3(0,1,0), 1));
}

void FrameBasis(float3 frameDir, out float3 planeX, out float3 planeY, out float3 planeN) {
    planeN = normalize(frameDir);
    float3 up = abs(planeN.y) > 0.999 ? float3(0,0,1) : float3(0,1,0);
    planeX = normalize(cross(planeN, up));
    planeY = normalize(cross(planeX, planeN));
}

void BillboardBasis(float3 fwd, out float3 right, out float3 up) {
    right = abs(fwd.y) > 0.999 ? float3(-1,0,0) : normalize(cross(float3(0,1,0), fwd));
    up = cross(fwd, right);
}

float2 GridFromDir(float3 viewDir, float gridRes) {
    float2 uv = HemiOctEncode(viewDir) * 0.5 + 0.5;
    return clamp(uv * (gridRes - 1), 0, gridRes - 1);
}

float3 DirFromCell(float2 cell, float gridRes) {
    float2 uv = cell / max(1.0, gridRes - 1) * 2.0 - 1.0;
    return HemiOctDecode(uv);
}

// Branchless barycentric weights for 3 points
float3 BarycentricWeights3(float2 gridFrac, bool isBottomRight) {
    if (isBottomRight) {
        // Bottom-Right Triangle: (0,0), (1,0), (1,1)
        return float3(1.0 - gridFrac.x, gridFrac.x - gridFrac.y, gridFrac.y);
    } else {
        // Top-Left Triangle: (0,0), (0,1), (1,1)
        return float3(1.0 - gridFrac.y, gridFrac.y - gridFrac.x, gridFrac.x);
    }
}

// Compute UV on a virtual plane facing frameDir
float2 VirtualPlaneUV(float3 frameDir, float3 pivotToCam, float3 vertexToCam) {
    float3 planeX, planeY, planeN;
    FrameBasis(frameDir, planeX, planeY, planeN);

    float projPivot = dot(planeN, pivotToCam);
    float projVertex = dot(planeN, vertexToCam);
    projVertex = (abs(projVertex) < 1e-4) ? (projVertex < 0 ? -1e-4 : 1e-4) : projVertex;
    float ratio = projPivot / projVertex;

    float3 offset = vertexToCam * ratio - pivotToCam;

    float2 uv = float2(dot(planeX, offset), dot(planeY, offset));
    return uv * -1.0 + 0.5;
}

#if defined(_IMPOSTORS)

// Decode linear 0-1 depth (0=near clip, 1=far clip) to world space distance from camera.
float DecodeImpostorDepth(float linearDepth) {
    return lerp(_Impostors_Near_Clip, _Impostors_Far_Clip, linearDepth);
}

float2 ClampUvInCell(float2 uv) {
    uv = saturate(uv);
    float2 halfTexelInCell = 0.5 * _Impostors_Atlas_TexelSize.xy * (float)_Impostors_Grid_Resolution;
    return clamp(uv, halfTexelInCell, 1.0 - halfTexelInCell);
}

struct ImpostorSample {
    float4 albedo;
    float4 normal;
    float4 metallicGloss;
};

struct ImpostorResult {
    float4 albedo;
    float3 normal;
    float metallic;
    float smoothness;
    float3 objPos;
    // TODO rm
    float debug;
};

float SampleImpostorDepthCell(float2 cell, float2 uvInCell, float gridRes) {
    uvInCell = ClampUvInCell(uvInCell);
    float invGridRes = rcp(gridRes);
    float2 atlasUv = (cell + uvInCell) * invGridRes;
    float2 gradX = ddx(uvInCell) * invGridRes;
    float2 gradY = ddy(uvInCell) * invGridRes;
    // 0 = near clip plane, 1 = far clip plane
    return _Impostors_Depth_Atlas.SampleGrad(bilinear_clamp_s, atlasUv, gradX, gradY).r;
}

float SampleImpostorAlphaCell(float2 cell, float2 uvInCell, float gridRes) {
    uvInCell = ClampUvInCell(uvInCell);
    float invGridRes = rcp(gridRes);
    float2 atlasUv = (cell + uvInCell) * invGridRes;
    float2 gradX = ddx(uvInCell) * invGridRes;
    float2 gradY = ddy(uvInCell) * invGridRes;
    return _Impostors_Atlas.SampleGrad(bilinear_clamp_s, atlasUv, gradX, gradY).a;
}

ImpostorSample SampleImpostorCell(float2 cell, float2 uvInCell, float gridRes) {
    uvInCell = ClampUvInCell(uvInCell);
    float invGridRes = rcp(gridRes);
    float2 atlasUv = (cell + uvInCell) * invGridRes;
    float2 gradX = ddx(uvInCell) * invGridRes;
    float2 gradY = ddy(uvInCell) * invGridRes;

    ImpostorSample s;
    s.albedo = _Impostors_Atlas.SampleGrad(bilinear_clamp_s, atlasUv, gradX, gradY);
    s.normal = _Impostors_Normal_Atlas.SampleGrad(bilinear_clamp_s, atlasUv, gradX, gradY);
    s.metallicGloss = _Impostors_Metallic_Gloss_Atlas.SampleGrad(bilinear_clamp_s, atlasUv, gradX, gradY).xzzy;
    return s;
}

float2 ImpostorParallaxOffsetForFrame(float3 frameDir, float3 pivotToCamOS, float2 uvBase, float encodedDepth) {
    float3 planeX, planeY, planeN;
    FrameBasis(frameDir, planeX, planeY, planeN);

    float2 camXY = float2(dot(pivotToCamOS, planeX), dot(pivotToCamOS, planeY));
    float camZ = dot(pivotToCamOS, planeN);
    camZ = (abs(camZ) < 1e-4) ? (camZ < 0 ? -1e-4 : 1e-4) : camZ;

    // Convert world-space depth to object space to match scaled impostor.
    float objScale = max(2.0 * _Impostors_Sphere_Radius, 1e-4);
    float objRadius = _Impostors_Sphere_Radius / objScale;
    float objNear = _Impostors_Near_Clip / objScale;
    float worldSpaceDepth = DecodeImpostorDepth(encodedDepth);
    float objectSpaceDepth = worldSpaceDepth / objScale;
    float depth01 = (objectSpaceDepth - objNear) / (2.0 * objRadius);
    // Convert to signed "height" where nearer pixels shift more (matches typical bump-offset convention).
    float zSurface = (depth01 - 0.5);
    float2 planeCoord = 0.5 - uvBase;

    return (planeCoord - camXY) * (zSurface / camZ);
}

// Reconstruct object-space offset from impostor center.
// frameDir: object-space direction from center to baking camera (same space as UV computation)
// UV encodes X,Y on the virtual plane (0.5 = center), depth is linear 0-1 (near to far clip).
float3 ReconstructObjectOffsetFromFrame(float3 frameDir, float2 uv, float encodedDepth) {
    float3 planeX, planeY, planeN;
    FrameBasis(frameDir, planeX, planeY, planeN);

    float objScale = max(2.0f * _Impostors_Sphere_Radius, 1e-4);
    float objRadius = _Impostors_Sphere_Radius / objScale;
    float objNear = _Impostors_Near_Clip / objScale;
    float objFar = _Impostors_Far_Clip / objScale;

    float2 offsetXY = (0.5f - uv) * 2.0f * objRadius;
    float objectSpaceDepth = lerp(objNear, objFar, encodedDepth);
    float offsetZ = (objRadius + objNear) - objectSpaceDepth;

    return offsetXY.x * planeX + offsetXY.y * planeY + offsetZ * planeN;
}

ImpostorSample BlendImpostorSamples(ImpostorSample s0, ImpostorSample s1, ImpostorSample s2, float3 bw) {
    ImpostorSample result;
    float3 alpha = float3(s0.albedo.a, s1.albedo.a, s2.albedo.a);
    float alphaOut = dot(alpha, bw);
    float3 premul =
        s0.albedo.rgb * (alpha.x * bw.x) +
        s1.albedo.rgb * (alpha.y * bw.y) +
        s2.albedo.rgb * (alpha.z * bw.z);
    float3 rgbOut = premul / max(alphaOut, 1e-4);
    result.albedo = float4(rgbOut, alphaOut);

    // Weight normal/metallicGloss by alpha to avoid blending with transparent (zero) pixels
    float3 alphaBw = alpha * bw;
    alphaBw /= max(alphaBw.x + alphaBw.y + alphaBw.z, 0.001);

    result.normal = s0.normal * alphaBw.x + s1.normal * alphaBw.y + s2.normal * alphaBw.z;
    result.metallicGloss = s0.metallicGloss * alphaBw.x + s1.metallicGloss * alphaBw.y + s2.metallicGloss * alphaBw.z;
    return result;
}

// Billboard vertex transformation for impostors
void impostor_vert(inout float3 vertexOS) {
    float3 center = mul(unity_ObjectToWorld, float4(0,0,0,1)).xyz;

#ifdef SHADOW_CASTER_PASS
    float3 camPos = _Impostors_Main_Camera_Pos;
#else
    float3 camPos = _WorldSpaceCameraPos;
#endif

    // Billboard facing the camera direction (world space, then convert back to object space).
    float3 viewWS = normalize(camPos - center);
    float3 right, up;
    BillboardBasis(viewWS, right, up);
    float radiusScale = _Impostors_Sphere_Radius * 2.0;
    float3 worldPos = center + vertexOS.x * right * radiusScale + vertexOS.y * up * radiusScale;
    vertexOS = mul(unity_WorldToObject, float4(worldPos, 1)).xyz;
}

// Sample impostor atlas with view-dependent blending
ImpostorResult impostor_frag(float3 worldPos) {
    ImpostorResult result = (ImpostorResult)0;

    // Calculate center in fragment shader to avoid extra interpolator
    float3 center = mul(unity_ObjectToWorld, float4(0,0,0,1)).xyz;

    // Sphere culling first
#ifdef SHADOW_CASTER_PASS
    float3 camPos = _Impostors_Main_Camera_Pos;
#else
    float3 camPos = _WorldSpaceCameraPos;
#endif
    float3 viewDir = normalize(worldPos - camPos);

    float radiusWS = _Impostors_Sphere_Radius;

    float3 originToRo = camPos - center;
    float b = dot(originToRo, viewDir);
    float c = dot(originToRo, originToRo) - radiusWS * radiusWS;
    clip(b * b - c);

    // For lattice lookup, use the camera-to-impostor-center direction (matches billboard orientation).
    float3x3 worldToObject = (float3x3)unity_WorldToObject;
    float3 viewOS = normalize(mul(worldToObject, normalize(camPos - center)));

    // Get continuous grid position and find the 3 nearest frames
    float gridRes = (float)_Impostors_Grid_Resolution;
    float2 grid = GridFromDir(viewOS, gridRes);
    float2 gridFloor = floor(grid);
    float2 gridFrac = frac(grid);
    
    // Determine triangle and weights
    bool isBottomRight = gridFrac.x > gridFrac.y;
    float3 bw = BarycentricWeights3(gridFrac, isBottomRight);

    // Frame cells
    float2 cell0 = clamp(gridFloor, 0, gridRes - 1);
    float2 cell1 = clamp(gridFloor + (isBottomRight ? float2(1,0) : float2(0,1)), 0, gridRes - 1);
    float2 cell2 = clamp(gridFloor + float2(1,1), 0, gridRes - 1);

    // Compute shared vectors for virtual plane UV calculation
    float3 pivotToCamOS = mul(worldToObject, camPos - center);
    float3 vertexPosOS = mul(worldToObject, worldPos - center);
    float3 vertexToCamOS = pivotToCamOS - vertexPosOS;

    float3 frameDir0 = DirFromCell(cell0, gridRes);
    float3 frameDir1 = DirFromCell(cell1, gridRes);
    float3 frameDir2 = DirFromCell(cell2, gridRes);

    float2 uvBase0 = ClampUvInCell(VirtualPlaneUV(frameDir0, pivotToCamOS, vertexToCamOS));
    float2 uvBase1 = ClampUvInCell(VirtualPlaneUV(frameDir1, pivotToCamOS, vertexToCamOS));
    float2 uvBase2 = ClampUvInCell(VirtualPlaneUV(frameDir2, pivotToCamOS, vertexToCamOS));

    float baseAlpha0 = SampleImpostorAlphaCell(cell0, uvBase0, gridRes);
    float baseAlpha1 = SampleImpostorAlphaCell(cell1, uvBase1, gridRes);
    float baseAlpha2 = SampleImpostorAlphaCell(cell2, uvBase2, gridRes);
    float baseAlphaBlended = baseAlpha0 * bw.x + baseAlpha1 * bw.y + baseAlpha2 * bw.z;
    float parallaxStrength = _Impostors_Parallax * smoothstep(_Impostors_Cutoff, 1.0, baseAlphaBlended);

    float depth0 = SampleImpostorDepthCell(cell0, uvBase0, gridRes);
    float depth1 = SampleImpostorDepthCell(cell1, uvBase1, gridRes);
    float depth2 = SampleImpostorDepthCell(cell2, uvBase2, gridRes);
    float depthBlended = depth0 * bw.x + depth1 * bw.y + depth2 * bw.z;

    if (_Impostors_Debug_Depth > 0.5) {
        result.albedo = float4(depthBlended.xxx, 1);
        return result;
    }

    float2 uv0 = uvBase0;
    float2 uv1 = uvBase1;
    float2 uv2 = uvBase2;
    if (parallaxStrength > 0.001) {
        uv0 = uvBase0 + ImpostorParallaxOffsetForFrame(frameDir0, pivotToCamOS, uvBase0, depthBlended) * parallaxStrength;
        uv1 = uvBase1 + ImpostorParallaxOffsetForFrame(frameDir1, pivotToCamOS, uvBase1, depthBlended) * parallaxStrength;
        uv2 = uvBase2 + ImpostorParallaxOffsetForFrame(frameDir2, pivotToCamOS, uvBase2, depthBlended) * parallaxStrength;
    }

    float2 uv0Final = uv0;
    float2 uv1Final = uv1;
    float2 uv2Final = uv2;

    ImpostorSample s0 = SampleImpostorCell(cell0, uv0Final, gridRes);
    ImpostorSample s1 = SampleImpostorCell(cell1, uv1Final, gridRes);
    ImpostorSample s2 = SampleImpostorCell(cell2, uv2Final, gridRes);

    // Parallax can push UVs into transparent pixels at silhouettes; fall back to unshifted UVs when that happens.
    if (parallaxStrength > 0.001) {
        if (s0.albedo.a < _Impostors_Cutoff && baseAlpha0 > _Impostors_Cutoff) {
            uv0Final = uvBase0;
            s0 = SampleImpostorCell(cell0, uv0Final, gridRes);
        }
        if (s1.albedo.a < _Impostors_Cutoff && baseAlpha1 > _Impostors_Cutoff) {
            uv1Final = uvBase1;
            s1 = SampleImpostorCell(cell1, uv1Final, gridRes);
        }
        if (s2.albedo.a < _Impostors_Cutoff && baseAlpha2 > _Impostors_Cutoff) {
            uv2Final = uvBase2;
            s2 = SampleImpostorCell(cell2, uv2Final, gridRes);
        }
    }

    // Blend 3 samples
    ImpostorSample blended = BlendImpostorSamples(s0, s1, s2, bw);

    if (_Impostors_Debug_Mode > 0.5) {
        // Debug view of weights
        result.albedo = float4(bw, 1);
        return result;
    }

    clip(blended.albedo.a - _Impostors_Cutoff);

    result.albedo = blended.albedo;
    result.normal = normalize(blended.normal.xyz * 2.0 - 1.0);
    result.metallic = blended.metallicGloss.r;
    result.smoothness = blended.metallicGloss.a;

    // Reconstruct object-space position from each frame's UV and depth, then blend.
    float depth0Pos = SampleImpostorDepthCell(cell0, uv0Final, gridRes);
    float depth1Pos = SampleImpostorDepthCell(cell1, uv1Final, gridRes);
    float depth2Pos = SampleImpostorDepthCell(cell2, uv2Final, gridRes);

    // Use object-space frame directions (same space as UV computation)
    float3 offset0 = ReconstructObjectOffsetFromFrame(frameDir0, uv0Final, depth0Pos);
    float3 offset1 = ReconstructObjectOffsetFromFrame(frameDir1, uv1Final, depth1Pos);
    float3 offset2 = ReconstructObjectOffsetFromFrame(frameDir2, uv2Final, depth2Pos);
    float3 objOffset = offset0 * bw.x + offset1 * bw.y + offset2 * bw.z;

    // Return object-space position (offset from origin)
    result.objPos = objOffset;

    return result;
}

#endif  // _IMPOSTORS

#endif  // __IMPOSTOR_INC