#ifndef __IMPOSTOR_INC
#define __IMPOSTOR_INC

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

struct ImpostorResult {
    float4 albedo;
    float3 normal;
    float metallic;
    float smoothness;
    float3 objPos;
};

#if defined(_IMPOSTORS)
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 = cross(planeX, planeN);  // Already normalized since planeX and planeN are orthonormal
}

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), 0.0, gridRes - 1.0);
}

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

float3 BarycentricWeights3(float2 gridFrac, bool isBottomRight) {
    float3 br = float3(1.0 - gridFrac.x, gridFrac.x - gridFrac.y, gridFrac.y);
    float3 tl = float3(1.0 - gridFrac.y, gridFrac.y - gridFrac.x, gridFrac.x);
    return isBottomRight ? br : tl;
}

float2 VirtualPlaneUV(float3 planeX, float3 planeY, float3 planeN, float3 pivotToCam, float3 vertexToCam) {
    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;

    return 0.5 - float2(dot(planeX, offset), dot(planeY, offset));
}

float2 ClampUvInCell(float2 uv, float2 halfTexelInCell) {
    return clamp(saturate(uv), halfTexelInCell, 1.0 - halfTexelInCell);
}

struct ImpostorSample {
    float4 albedo;
    float4 normal;
    float2 metallicGloss;
    float depth;
};

ImpostorSample SampleImpostorCell(float2 cell, float2 uvInCell, float invGridRes) {
    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);
    float4 mgd = _Impostors_Metallic_Gloss_Depth_Atlas.SampleGrad(bilinear_clamp_s, atlasUv, gradX, gradY);
    s.metallicGloss = mgd.rg;
    s.depth = mgd.b;
    return s;
}

float2 ImpostorParallaxOffset(float3 planeX, float3 planeY, float3 planeN, float3 pivotToCamOS, float encodedDepth,
                              float impostorNear, float impostorFar, float impostorRadius) {
    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;

    float worldSpaceDepth = lerp(impostorNear, impostorFar, encodedDepth);
    float depth01 = (worldSpaceDepth - impostorNear) / (2.0 * impostorRadius);
    float height = 0.5 - depth01;

    return (camXY / camZ) * height;
}

float3 ReconstructObjectOffset(float3 planeX, float3 planeY, float3 planeN, float2 uv, float encodedDepth,
                               float impostorNear, float impostorFar, float impostorRadius) {
    float2 offsetXY = (0.5 - uv) * (2.0 * impostorRadius);
    float worldSpaceDepth = lerp(impostorNear, impostorFar, encodedDepth);
    float offsetZ = (impostorRadius + impostorNear) - worldSpaceDepth;

    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 alphaBw = alpha * bw * rcp(max(alphaOut, 1e-4));

    result.albedo = float4(s0.albedo.rgb * alphaBw.x + s1.albedo.rgb * alphaBw.y + s2.albedo.rgb * alphaBw.z, alphaOut);
    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;
}
#endif  // _IMPOSTORS

void impostor_vert(inout float3 vertexOS) {
#if defined(_IMPOSTORS)
    float3 center = mul(unity_ObjectToWorld, float4(0,0,0,1)).xyz;

#if defined(_INSTANCE_TEXTURE_OFFSET) && defined(UNITY_INSTANCING_ENABLED)
    float3 instanceCenter;
    float4 instanceRotation;
    float3 instanceScale;
    if (get_instance_transform(instanceCenter, instanceRotation, instanceScale)) {
        center = mul(unity_ObjectToWorld, float4(instanceCenter, 1.0)).xyz;
    }
#endif  // _INSTANCE_TEXTURE_OFFSET && UNITY_INSTANCING_ENABLED

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

    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 + vertexOS.y * up) * radiusScale;
    vertexOS = mul(unity_WorldToObject, float4(worldPos, 1.0)).xyz;
#endif  // _IMPOSTORS
}

ImpostorResult impostor_frag(float3 worldPos) {
    ImpostorResult result = (ImpostorResult)0;
#if defined(_IMPOSTORS)

    float3 center = mul(unity_ObjectToWorld, float4(0,0,0,1)).xyz;
    float4 instanceRotation = float4(0, 0, 0, 1);
    float instanceScale = 1.0;
    bool hasInstanceTransform = false;

#if defined(_INSTANCE_TEXTURE_OFFSET) && defined(UNITY_INSTANCING_ENABLED)
    float3 instanceCenter;
    float3 instanceScale3;
    if (get_instance_transform(instanceCenter, instanceRotation, instanceScale3)) {
        center = mul(unity_ObjectToWorld, float4(instanceCenter, 1.0)).xyz;
        instanceScale = max(instanceScale3.x, max(instanceScale3.y, instanceScale3.z));
        hasInstanceTransform = true;
    }
#endif  // _INSTANCE_TEXTURE_OFFSET && UNITY_INSTANCING_ENABLED

#ifdef SHADOW_CASTER_PASS
    float3 camPos = _Impostors_Main_Camera_Pos;
#else
    float3 camPos = _WorldSpaceCameraPos;
#endif
    float3 camToCenter = camPos - center;
    float3 viewDir = normalize(worldPos - camPos);

    float impostorRadius = _Impostors_Sphere_Radius * instanceScale;
    float impostorNear = _Impostors_Near_Clip * instanceScale;
    float impostorFar = _Impostors_Far_Clip * instanceScale;

    float b = dot(camToCenter, viewDir);
    float c = dot(camToCenter, camToCenter) - impostorRadius * impostorRadius;
    clip(b * b - c);

    float3x3 worldToObject = (float3x3)unity_WorldToObject;
    float3 viewOS = mul(worldToObject, normalize(camToCenter));

    float3 pivotToCamOS = mul(worldToObject, camToCenter);
    float3 vertexPosOS = mul(worldToObject, worldPos - center);

    if (hasInstanceTransform) {
        float4 instanceRotationInv = qconj(instanceRotation);
        viewOS = qrotate(instanceRotationInv, viewOS);
        pivotToCamOS = qrotate(instanceRotationInv, pivotToCamOS);
        vertexPosOS = qrotate(instanceRotationInv, vertexPosOS);
    }
    float3 vertexToCamOS = pivotToCamOS - vertexPosOS;

    float gridRes = (float)_Impostors_Grid_Resolution;
    float invGridRes = rcp(gridRes);
    float2 halfTexelInCell = 0.5 * _Impostors_Atlas_TexelSize.xy * gridRes;

    float2 grid = GridFromDir(viewOS, gridRes);
    float2 gridFloor = floor(grid);
    float2 gridFrac = frac(grid);

    bool isBottomRight = gridFrac.x > gridFrac.y;
    float3 bw = BarycentricWeights3(gridFrac, isBottomRight);

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

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

    float3 planeX0, planeY0, planeN0;
    float3 planeX1, planeY1, planeN1;
    float3 planeX2, planeY2, planeN2;
    FrameBasis(frameDir0, planeX0, planeY0, planeN0);
    FrameBasis(frameDir1, planeX1, planeY1, planeN1);
    FrameBasis(frameDir2, planeX2, planeY2, planeN2);

    float2 uvBase0 = ClampUvInCell(VirtualPlaneUV(planeX0, planeY0, planeN0, pivotToCamOS, vertexToCamOS), halfTexelInCell);
    float2 uvBase1 = ClampUvInCell(VirtualPlaneUV(planeX1, planeY1, planeN1, pivotToCamOS, vertexToCamOS), halfTexelInCell);
    float2 uvBase2 = ClampUvInCell(VirtualPlaneUV(planeX2, planeY2, planeN2, pivotToCamOS, vertexToCamOS), halfTexelInCell);

    ImpostorSample s0 = SampleImpostorCell(cell0, uvBase0, invGridRes);
    ImpostorSample s1 = SampleImpostorCell(cell1, uvBase1, invGridRes);
    ImpostorSample s2 = SampleImpostorCell(cell2, uvBase2, invGridRes);

    float baseAlpha0 = s0.albedo.a, depth0 = s0.depth;
    float baseAlpha1 = s1.albedo.a, depth1 = s1.depth;
    float baseAlpha2 = s2.albedo.a, depth2 = s2.depth;

    float baseAlphaBlended = dot(float3(baseAlpha0, baseAlpha1, baseAlpha2), bw);
    float parallaxStrength = _Impostors_Parallax * smoothstep(_Impostors_Cutoff, 1.0, baseAlphaBlended);
    float depthBlended = dot(float3(depth0, depth1, depth2), bw);

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

    float2 parallaxOffset0 = 0, parallaxOffset1 = 0, parallaxOffset2 = 0;
    bool needsParallaxResample = false;

    [branch]
    if (parallaxStrength > 0.001) {
        parallaxOffset0 = ImpostorParallaxOffset(planeX0, planeY0, planeN0, pivotToCamOS, depthBlended,
                                                 impostorNear, impostorFar, impostorRadius) * parallaxStrength;
        parallaxOffset1 = ImpostorParallaxOffset(planeX1, planeY1, planeN1, pivotToCamOS, depthBlended,
                                                 impostorNear, impostorFar, impostorRadius) * parallaxStrength;
        parallaxOffset2 = ImpostorParallaxOffset(planeX2, planeY2, planeN2, pivotToCamOS, depthBlended,
                                                 impostorNear, impostorFar, impostorRadius) * parallaxStrength;

        float maxOffsetSq = max(max(dot(parallaxOffset0, parallaxOffset0), dot(parallaxOffset1, parallaxOffset1)), dot(parallaxOffset2, parallaxOffset2));
        needsParallaxResample = maxOffsetSq > 0.00005;
    }

    float2 finalUv0 = uvBase0, finalUv1 = uvBase1, finalUv2 = uvBase2;

    [branch]
    if (needsParallaxResample) {
        float2 uv0 = ClampUvInCell(uvBase0 + parallaxOffset0, halfTexelInCell);
        float2 uv1 = ClampUvInCell(uvBase1 + parallaxOffset1, halfTexelInCell);
        float2 uv2 = ClampUvInCell(uvBase2 + parallaxOffset2, halfTexelInCell);

        ImpostorSample ps0 = SampleImpostorCell(cell0, uv0, invGridRes);
        ImpostorSample ps1 = SampleImpostorCell(cell1, uv1, invGridRes);
        ImpostorSample ps2 = SampleImpostorCell(cell2, uv2, invGridRes);

        if (ps0.albedo.a >= _Impostors_Cutoff || baseAlpha0 <= _Impostors_Cutoff) {
            s0 = ps0;
            finalUv0 = uv0;
        }
        if (ps1.albedo.a >= _Impostors_Cutoff || baseAlpha1 <= _Impostors_Cutoff) {
            s1 = ps1;
            finalUv1 = uv1;
        }
        if (ps2.albedo.a >= _Impostors_Cutoff || baseAlpha2 <= _Impostors_Cutoff) {
            s2 = ps2;
            finalUv2 = uv2;
        }
    }

    ImpostorSample blended = BlendImpostorSamples(s0, s1, s2, bw);

    [branch]
    if (_Impostors_Debug_Mode > 0.5) {
        result.albedo = float4(bw, 1);
        return result;
    }

    clip(blended.albedo.a - _Impostors_Cutoff);

    result.albedo = blended.albedo;
    float3 normalOS = blended.normal.xyz * 2.0 - 1.0;
    if (hasInstanceTransform) {
        normalOS = qrotate(instanceRotation, normalOS);
    }
    result.normal = normalize(mul((float3x3)unity_ObjectToWorld, normalOS));
    result.metallic = blended.metallicGloss.r;
    result.smoothness = blended.metallicGloss.g;

#if defined(_IMPOSTORS_DEPTH)
    float3 offset0 = ReconstructObjectOffset(planeX0, planeY0, planeN0, finalUv0, depth0,
                                             impostorNear, impostorFar, impostorRadius);
    float3 offset1 = ReconstructObjectOffset(planeX1, planeY1, planeN1, finalUv1, depth1,
                                             impostorNear, impostorFar, impostorRadius);
    float3 offset2 = ReconstructObjectOffset(planeX2, planeY2, planeN2, finalUv2, depth2,
                                             impostorNear, impostorFar, impostorRadius);
    result.objPos = offset0 * bw.x + offset1 * bw.y + offset2 * bw.z;
#endif

    return result;
#endif  // _IMPOSTORS
}

#endif  // __IMPOSTOR_INC