#pragma kernel CSMain

RWTexture2D<float4> Destination;

float4x4 _CameraToWorld;
float4x4 _CameraInverseProjection;
float3 _LightDirection;

float power;
float darkness;
float blackAndWhite;
float3 colourAMix;
float3 colourBMix;

static const float epsilon = 0.001f;
static const float maxDst = 200;
static const int maxStepCount = 250;

struct Ray {
    float3 origin;
    float3 direction;
};

Ray CreateRay(float3 origin, float3 direction) {
    Ray ray;
    ray.origin = origin;
    ray.direction = direction;
    return ray;
}

Ray CreateCameraRay(float2 uv) {
    float3 origin = mul(_CameraToWorld, float4(0, 0, 0, 1)).xyz;
    float3 direction = mul(_CameraInverseProjection, float4(uv, 0, 1)).xyz;
    direction = mul(_CameraToWorld, float4(direction, 0)).xyz;
    direction = normalize(direction);
    return CreateRay(origin, direction);
}

float potential(in float3 pos)
{
        float3 z = pos;
        for (int i = 1; i < 15; i++)
        {
            return log(length(z)) / pow(power, float(i));
        }
    return 0.0;
}

float2 SceneInfo(float3 p) {
    float pot = potential(p);
    if (pot == 0.0) return 0.0;
    gradient = (float3(potential(p + xDir * EPS), potential(p + yDir * EPS), potential(p + zDir * EPS)) - pot) / EPS;
    return (0.5 / exp(pot))*sinh(pot) / length(gradient);
}



float3 EstimateNormal(float3 p) {
    float x = SceneInfo(float3(p.x + epsilon, p.y, p.z)).y - SceneInfo(float3(p.x - epsilon, p.y, p.z)).y;
    float y = SceneInfo(float3(p.x, p.y + epsilon, p.z)).y - SceneInfo(float3(p.x, p.y - epsilon, p.z)).y;
    float z = SceneInfo(float3(p.x, p.y, p.z + epsilon)).y - SceneInfo(float3(p.x, p.y, p.z - epsilon)).y;
    return normalize(float3(x, y, z));
}

[numthreads(8, 8, 1)]
void CSMain(uint3 id : SV_DispatchThreadID)
{
    uint width, height;
    Destination.GetDimensions(width, height);

    float2 uv = id.xy / float2(width, height);

    // Background gradient
    float4 result = lerp(float4(51, 3, 20, 1), float4(16, 6, 28, 1), uv.y) / 255.0;

    // Raymarching:
    Ray ray = CreateCameraRay(uv * 2 - 1);
    float rayDst = 0;
    int marchSteps = 0;

    while (rayDst < maxDst && marchSteps < maxStepCount) {
        marchSteps++;
        float2 sceneInfo = SceneInfo(ray.origin);
        float dst = sceneInfo.y;

        // Ray has hit a surface
        if (dst <= epsilon) {
            float escapeIterations = sceneInfo.x;
            float3 normal = EstimateNormal(ray.origin - ray.direction*epsilon * 2);

            float colourA = saturate(dot(normal*.5 + .5, -_LightDirection));
            float colourB = saturate(escapeIterations / 16.0);
            float3 colourMix = saturate(colourA * colourAMix + colourB * colourBMix);

            result = float4(colourMix.xyz, 1);
            break;
        }
        ray.origin += ray.direction * dst;
        rayDst += dst;
    }

    float rim = marchSteps / darkness;
    Destination[id.xy] = lerp(result, 1, blackAndWhite) * rim;
}