#include "lib/shared/point.hlsl"
#include "lib/shared/quat-functions.hlsl"
#include "lib/shared/hash-functions.hlsl"

static const float3 Corners[] =
    {
        float3(-1, -1, 0),
        float3(1, -1, 0),
        float3(1, 1, 0),
        float3(1, 1, 0),
        float3(-1, 1, 0),
        float3(-1, -1, 0),
};

cbuffer Transforms : register(b0)
{
    float4x4 CameraToClipSpace;
    float4x4 ClipSpaceToCamera;
    float4x4 WorldToCamera;
    float4x4 CameraToWorld;
    float4x4 WorldToClipSpace;
    float4x4 ClipSpaceToWorld;
    float4x4 ObjectToWorld;
    float4x4 WorldToObject;
    float4x4 ObjectToCamera;
    float4x4 ObjectToClipSpace;
};

cbuffer Params : register(b1)
{
    float Scale;
    float2 Stretch;
    float __padding;

    float3 Offset;
    float OrientationMode;

    float Rotate;
    float3 RotationAxis;

    float Randomize;
    float RandomPhase;
    float RandomRotate;
    float __padding0;

    float3 RandomPosition;
    float RandomScale;

    float4 Color;

    float ColorVariationMode;
    float ScaleDistribution;
    float SpreadLength;
    float SpreadPhase;

    float SpreadPingPong;
    float SpreadRepeat;
    float2 AtlasSize;

    float TextureAtlasMode;
    float FxTextureMode;
    float AlphaCutOff;
    float IsFxTextureConnected;

    // float __padding1; // order adjusted for 4x4 padding

    float4 FxTextureAmount;

    float UseRotationAsRgba;

    float UseWFoScale;
    float UseStretch;
    float TooCloseFadeOut;
};

cbuffer FogParams : register(b2)
{
    float4 FogColor;
    float FogDistance;
    float FogBias;
}

struct psInput
{
    float4 position : SV_POSITION;
    float4 color : COLOR;
    float2 texCoord : TEXCOORD;
    float fog : FOG;
};

sampler texSampler : register(s0);

StructuredBuffer<Point> Points : t0;
Texture2D<float4> SpriteTexture : register(t1);
Texture2D<float4> FxTexture : register(t2);
Texture2D<float4> ColorOverW : register(t3);
Texture2D<float4> SizeOverW : register(t4);

inline float fmod(float x, float y)
{
    return (x - y * floor(x / y));
}

inline float GetUFromMode(float mode, int id, float f, float4 scatter, float w, float fog)
{
    switch ((int)(mode + 0.5))
    {

    case 0:
        return scatter.w;

    case 1:
        return hash11u(id);

    case 2:
        float f1 = (f + SpreadPhase) / SpreadLength;
        f1 = SpreadRepeat > 0.5 ? fmod(f1, 1) : f1;
        return SpreadPingPong > 0.5 ? (1 - abs(f1 * 2 - 1)) : f1;

    case 3:
        float w1 = (w + SpreadPhase) / SpreadLength;
        w1 = SpreadRepeat > 0.5 ? fmod(w1, 1) : w1;
        return SpreadPingPong > 0.5 ? (1 - abs(w1 * 2 - 1)) : w1;

    default:
        return fog;
    }
}

psInput vsMain(uint id
               : SV_VertexID)
{
    uint particleCount, stride;
    Points.GetDimensions(particleCount, stride);

    uint width, height, mips;
    SpriteTexture.GetDimensions(0, width, height, mips);
    float2 textureAspect = width > height ? float2((float)width / height, 1)
                                          : float2(1, (float)height / width);

    psInput output;

    uint quadIndex = id % 6;
    uint pointId = id / 6;
    float3 cornerFactors = Corners[quadIndex];

    Point p = Points[pointId]; 

    float4 pRotation = normalize(p.Rotation); 
    float f = pointId / (float)particleCount;

    float phase = RandomPhase + 133.1123 * f;
    int phaseId = (int)phase; 

    float4 normalizedScatter = lerp(hash41u(pointId * 12341 + phaseId),
                                    hash41u(pointId * 12341 + phaseId + 1),
                                    smoothstep(0, 1, phase - phaseId));

    float3 scatterForScale = normalizedScatter.xyx * 2 - 1;

    output.fog = 0;

    int2 altasSize = (int2)AtlasSize;

    float textureUx = GetUFromMode(TextureAtlasMode, pointId, (f * altasSize.x) % 1, normalizedScatter, p.W, output.fog) % altasSize.x;
    float textureUy = GetUFromMode(TextureAtlasMode, pointId, f, normalizedScatter.wxyz, p.W, output.fog); 
    
    int textureCelX =  textureUx * altasSize.x;
    int textureCelY =  textureUy * altasSize.y;
    output.texCoord = (cornerFactors.xy * float2(-1, 1) * 0.5 + 0.5);
    output.texCoord /= altasSize;
    output.texCoord += float2(textureCelX, textureCelY) / altasSize;

    float4 posInObject = float4(p.Position, 1);

    float3 randomOffset = qRotateVec3((normalizedScatter.xyz - 0.5) * 2 * RandomPosition * Randomize, pRotation);
    posInObject.xyz += randomOffset;

    if (OrientationMode <= 1.5)
    {
            posInObject.xyz += qRotateVec3(float3(0,0,Offset.z), pRotation);
    }

    float4 quadPosInCamera = mul(posInObject, ObjectToCamera);

    // Shrink too close particles
    float4 posInCamera = mul(posInObject, ObjectToCamera);
    float tooCloseFactor = saturate(-posInCamera.z / 0.1 - 1);
    
    output.fog = pow(saturate(-posInCamera.z / FogDistance), FogBias);

    float4 colorFromPoint = (UseRotationAsRgba > 0.5) ? pRotation : 1;

    float colorFxU = GetUFromMode(ColorVariationMode, pointId, f, normalizedScatter, p.W, output.fog);
    if (TooCloseFadeOut > 0.5){
        p.Color.a *= tooCloseFactor; // fade out the points too close to the camera
    }
    
    output.color = p.Color * Color * ColorOverW.SampleLevel(texSampler, float2(colorFxU, 0), 0) * colorFromPoint;

    float adjustedRotate = Rotate;
    float adjustedScale = Scale;
    float adjustedRandomize = Randomize;

    if (IsFxTextureConnected)
    {
        float4 centerPos = mul(float4(quadPosInCamera.xyz, 1), CameraToClipSpace);
        centerPos.xyz /= centerPos.w;

        float4 fxColor = FxTexture.SampleLevel(texSampler, (centerPos.xy * float2(1, -1) + 1) / 2, 0);

        if(FxTextureMode < 0.5) 
        {
            output.color *= fxColor;
        }
        else {
            adjustedRotate += FxTextureAmount.r * fxColor.r * fxColor.a * 360;
            adjustedScale += FxTextureAmount.g * fxColor.g * fxColor.a;
            adjustedRandomize += FxTextureAmount.b * fxColor.b * fxColor.a;
        }
    }

    // Scale and stretch
    float scaleFxU = GetUFromMode(ScaleDistribution, pointId, f, normalizedScatter, p.W, output.fog);
    float scaleFromCurve = SizeOverW.SampleLevel(texSampler, float2(scaleFxU, 0), 0).r;
    float hideUndefinedPoints = isnan(p.W) ? 0 : (UseWFoScale > 0.5 ? p.W : 1 ) ;
    //float computedScale = adjustedScale * (RandomScale * scatterForScale.y *adjustedRandomize + 1) * tooCloseFactor * scaleFromCurve * hideUndefinedPoints; 
    float computedScale = adjustedScale * (RandomScale * scatterForScale.y *adjustedRandomize + 1) * scaleFromCurve * hideUndefinedPoints;
    output.position = 0;

    if (OrientationMode <= 1.5)
    {
        float2 corner = float2((cornerFactors.xy + Offset.xy) * 0.010 * Stretch * (UseStretch > 0.5 ? p.Stretch.xy : 1) * textureAspect) * float2(-1, -1);

        float4 rot = qFromAngleAxis((adjustedRotate + RandomRotate * scatterForScale.x * adjustedRandomize) * 3.141578 / 180, RotationAxis);

        if ((int)OrientationMode == 1)
        {
            float3 yRotated = qRotateVec3(float3(0, 1, 0), pRotation);
            float4 yRotatedInCam = mul(float4(yRotated, 1), ObjectToCamera);
            float a = atan2(yRotatedInCam.x, yRotatedInCam.y);
            float4 xx = qFromAngleAxis(-a, float3(0, 0, 1));
            rot = qMul(xx, rot);
        }

        corner = qRotateVec3(float3(corner, 0), rot).xy;
        output.position = mul(quadPosInCamera + float4(corner * computedScale, 0, 0), CameraToClipSpace);
    }
    else
    {
        float3 axis = ( cornerFactors + Offset) * 0.010 * float3(Stretch * textureAspect,1);
        float4 rotation = qMul(normalize(pRotation), qFromAngleAxis((adjustedRotate + 180 + RandomRotate * scatterForScale.x * Randomize) / 180 * PI, RotationAxis));
        axis = qRotateVec3(axis, rotation) * computedScale;
        output.position = mul(posInObject + float4(axis, 0), ObjectToClipSpace);
    }
    return output;
}

float4 psMain(psInput input) : SV_TARGET
{
    float4 imgColor = SpriteTexture.Sample(texSampler, input.texCoord);
    float4 color = input.color * imgColor;
    
    if (color.a < AlphaCutOff)
        discard;

    color.rgb = lerp(color.rgb, FogColor.rgb, input.fog * FogColor.a);
    return clamp(color, 0, float4(100, 100, 100, 1));
}
