/*
 * Copyright (C) 2023 The Android Open Source Project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

uniform shader foreground;
uniform shader background;
uniform shader accumulatedSnow;
uniform shader noise;
uniform float2 gridSize;
uniform float time;
uniform float screenAspectRatio;
uniform float2 screenSize;
uniform mat3 transformMatrixBitmap;
uniform mat3 transformMatrixWeather;

#include "shaders/constants.agsl"
#include "shaders/utils.agsl"
#include "shaders/snow.agsl"

// Snow tint.
const vec4 snowColor = vec4(1., 1., 1., 0.95);
// Background tint
const vec4 bgdTint = vec4(0.8, 0.8, 0.8, 0.07);

// Indices of the different snow layers.
const float farthestSnowLayerIndex = 6;
const float midSnowLayerIndex = 2;
const float closestSnowLayerIndex = 0;

vec4 main(float2 fragCoord) {
    /**
     * The effect is consisted of 2 image textures (foreground and background) + 10 layers of
     * snow + 1 layer of snow accumulation. Below describes the rendering order (back to front):
     * 1. Background
     * 2. Background snow layers (from farthest layer to mid layer)
     * 3. Foreground
     * 4. Snow accumulation layer (on subject)
     * 5. Foreground snow layers (from mid layer to closest layer)
     */

    // Apply transform matrix to fragCoord
    float2 adjustedUv = transformPoint(transformMatrixBitmap, fragCoord);
    // Calculate uv for snow based on transformed coordinates
    float2 uv = transformPoint(transformMatrixWeather, fragCoord) / screenSize;
    float2 uvAdjusted = vec2(uv.x, uv.y / screenAspectRatio);

    vec4 colorForeground = foreground.eval(adjustedUv);
    vec4 colorBackground = background.eval(adjustedUv);

     // Adjusts contrast and brightness.
    float noiseT = triangleNoise(fragCoord.xy + vec2(12.31, 1024.1241));
    colorBackground.rgb =
        imageRangeConversion(colorBackground.rgb, 0.88, 0.02, noiseT * 0.025, intensity);
    colorForeground.rgb =
        imageRangeConversion(colorForeground.rgb, 0.88, 0.02, noiseT * 0.025, intensity);

    // 1. Draw background.
    vec4 color = colorBackground;

    // Add slight tint to the background.
    color.rgb = normalBlendNotPremultiplied(color.rgb, bgdTint.rgb, bgdTint.a);

    // 2. Generate snow layers behind the subject.
    if (colorForeground.a == 0) {
        for (float i = farthestSnowLayerIndex; i > midSnowLayerIndex; i--) {
            Snow snow = generateSnow(
                uv,
                screenAspectRatio,
                time,
                gridSize,
                /* layer number = */ i,
                closestSnowLayerIndex,
                farthestSnowLayerIndex);

            color.rgb =
                normalBlendNotPremultiplied(color.rgb, snowColor.rgb, snowColor.a * snow.flakeMask);
        }
    }

    // 3. Add the foreground layer. Any effect from here will be in front of the subject.
    color.rgb = normalBlend(color.rgb, colorForeground.rgb, colorForeground.a);

    // 4. Add accumulated snow layer.
    // Load noise texture to give "fluffy-ness" to the snow. Displace the sampling of the noise.
    vec3 cloudsNoise = noise.eval(uvAdjusted * 7000 + vec2(fragCoord.y, -fragCoord.x)).rgb;
    // Add dither to give texture to the snow and ruffle the edges.
    float dither = abs(triangleNoise(fragCoord * 0.01));

    // Get the accumulated snow buffer. r contains its mask, g contains some random noise.
    vec2 accSnow = accumulatedSnow.eval(adjustedUv).rg;
    // Sharpen the mask of the accumulated snow, but not in excess.
    float accSnowMask = smoothstep( (1.-intensity), 1.0, /* mask= */accSnow.r);
    if (accSnowMask > 0) {
        // Makes the edges of the snow layer accumulation rougher.
        accSnowMask = map(accSnowMask, 1. - cloudsNoise.b - 0.3 * dither, 1., 0., 1.);
        // Load snow texture and dither. Make it have gray-ish values.
        float accSnowTexture = smoothstep(0.2, 0.7, /* noise= */ accSnow.g) * 0.7;
        accSnowTexture = map(accSnowTexture, dither - 1, 1, 0, 1);
        // Adjust snow texture coverage/shape.
        accSnowTexture = map(accSnowTexture, 0.67, 0.8, 0, 1);
        accSnowMask = map(accSnowMask, 0., 1., 0., 1.- 0.6 * accSnowTexture - 0.35 * dither);

        color.rgb = normalBlendNotPremultiplied(color.rgb, snowColor.rgb, snowColor.a * accSnowMask);
    }

    // 5. Generate snow in front of the subject.
    for (float i = midSnowLayerIndex; i >= closestSnowLayerIndex; i--) {
        Snow snow = generateSnow(
            uv,
            screenAspectRatio,
            time,
            gridSize,
            /* layer number = */ i,
            closestSnowLayerIndex,
            farthestSnowLayerIndex);

        color.rgb =
            normalBlendNotPremultiplied(color.rgb, snowColor.rgb, snowColor.a * snow.flakeMask);
    }

    return color;
}