#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)

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;
};

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;
    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 depth01) {
    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;

    // Bake depth is linear in ortho: 0=near surface, 1=far surface.
    // 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);
}

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(float4 vertexOS, inout float3 worldPos) {
    float3 center = mul(unity_ObjectToWorld, float4(0,0,0,1)).xyz;

    float3 scale = float3(
        length(unity_ObjectToWorld._m00_m10_m20),
        length(unity_ObjectToWorld._m01_m11_m21),
        length(unity_ObjectToWorld._m02_m12_m22));

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

    // Billboard facing the camera direction
    float3 viewWS = normalize(camPos - center);
    float3 right, up;
    BillboardBasis(viewWS, right, up);
    worldPos = center + vertexOS.x * right * scale.x + vertexOS.y * up * scale.y;
}

// 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);

    float3 scale = float3(
        length(unity_ObjectToWorld._m00_m10_m20),
        length(unity_ObjectToWorld._m01_m11_m21),
        length(unity_ObjectToWorld._m02_m12_m22));
    float radiusWS = _Impostors_Sphere_Radius * max(scale.x, max(scale.y, scale.z));

    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;
    }

    ImpostorSample s0 = SampleImpostorCell(cell0, uv0, gridRes);
    ImpostorSample s1 = SampleImpostorCell(cell1, uv1, gridRes);
    ImpostorSample s2 = SampleImpostorCell(cell2, uv2, 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) s0 = SampleImpostorCell(cell0, uvBase0, gridRes);
        if (s1.albedo.a < _Impostors_Cutoff && baseAlpha1 > _Impostors_Cutoff) s1 = SampleImpostorCell(cell1, uvBase1, gridRes);
        if (s2.albedo.a < _Impostors_Cutoff && baseAlpha2 > _Impostors_Cutoff) s2 = SampleImpostorCell(cell2, uvBase2, 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;

    return result;
}

#endif  // _IMPOSTORS

#endif  // __IMPOSTOR_INC