Fix snow for other aspect ratios (tablet, foldable)

uniform half intensity;

const float farthestSnowLayerWiggleSpeed = 5.8;

const float closestSnowLayerWiggleSpeed = 2.6;

const float farthestSnowLayerWiggleSpeed = 2.18;

const float closestSnowLayerWiggleSpeed = 0.9;

/**

 * Generates snow flakes.

    /* Grid. */

    // Increase the last number to make each layer more separate from the previous one.

    float depth = 0.65 + layerIndex * 0.41;

    float depth = 0.65 + layerIndex * 0.37;

    float speedAdj = 1. + layerIndex * 0.15;

    float layerR = idGenerator(layerIndex);

    snowGridSize *= depth;

    // Number of rows and columns (each one is a cell, a drop).

    float cellAspectRatio = snowGridSize.x / snowGridSize.y;

    // Aspect ratio impacts visible cells.

    snowGridSize.y /= screenAspectRatio;

    uv.y /= screenAspectRatio;

    // Skew uv.x so it goes to left or right

    uv.x += uv.y * (0.8 * layerR - 0.4);

    // scale the UV to allocate number of rows and columns.

    // Generate column id, to offset columns vertically (so snow flakes are not aligned).

    float columnId = idGenerator(floor(gridUv.x));

    // Have time affect the position of each column as well.

    gridUv.y += columnId * 2.6 + time * 0.09 * (1 - columnId);

    gridUv.y += columnId * 2.6 + time * 0.19 * (1 - columnId);

    /* Cell. */

    // Get the cell ID based on the grid position. Value from 0 to 1.

    float snowFlakePosUncorrected = (cellUv.x - horizontalWiggle);

    // Calculate snow flake.

    vec2 snowFlakeShape = vec2(1., 1.2);

    vec2 snowFlakePos = vec2(snowFlakePosUncorrected / cellAspectRatio, cellUv.y);

    snowFlakePos -= vec2(0., uv.y - 0.5) * cellId;

    vec2 snowFlakeShape = vec2(0.28, 0.26);

    vec2 snowFlakePos = vec2(snowFlakePosUncorrected, cellUv.y * cellAspectRatio);

    snowFlakePos -= vec2(

            0.,

            (uv.y - 0.5 / screenAspectRatio)  - cellUv.y / snowGridSize.y

        ) * screenAspectRatio;

    snowFlakePos *= snowFlakeShape * decreaseFactor;

    vec2 snowFlakeShapeVariation = vec2(0.055) * // max variation

        vec2((cellId * 2. - 1.), // random A based on cell ID

uniform float2 uvScaleFgd;

uniform float2 uvOffsetBgd;

uniform float2 uvScaleBgd;

uniform float2 gridSize;

uniform float time;

uniform float screenAspectRatio;

uniform float2 screenSize;

            uv,

            screenAspectRatio,

            time,

            // TODO: adjust grid size based on aspect ratio.

            /* Grid size = */ vec2(7., 1.5),

            gridSize,

            /* layer number = */ i,

            closestSnowLayerIndex,

            farthestSnowLayerIndex);

            uv,

            screenAspectRatio,

            time,

            // TODO: adjust grid size based on aspect ratio

            /* Grid size = */ vec2(7., 1.5),

            gridSize,

            /* layer number = */ i,

            closestSnowLayerIndex,

            farthestSnowLayerIndex);

        updateTextureUniforms()

        adjustCropping(surfaceSize)

        prepareColorGrading()

        updateSnowGridSize(surfaceSize)

        setIntensity(snowConfig.intensity)

        // Generate accumulated snow at the end after we updated all the uniforms.

        generateAccumulatedSnow()

    }

    override fun resize(newSurfaceSize: SizeF) = adjustCropping(newSurfaceSize)

    override fun resize(newSurfaceSize: SizeF) {

        adjustCropping(newSurfaceSize)

        updateSnowGridSize(newSurfaceSize)

    }

    override fun update(deltaMillis: Long, frameTimeNanos: Long) {

        elapsedTime += snowSpeed * deltaMillis * MILLIS_TO_SECONDS

        )

    }

    private fun updateSnowGridSize(surfaceSize: SizeF) {

        val gridSize = GraphicsUtils.computeDefaultGridSize(surfaceSize, snowConfig.pixelDensity)

        snowConfig.shader.setFloatUniform("gridSize", 7 * gridSize, 2f * gridSize)

    }

    private companion object {

        private const val MILLIS_TO_SECONDS = 1 / 1000f

    }

package com.google.android.wallpaper.weathereffects.graphics.snow

import android.content.Context

import android.content.res.AssetManager

import android.graphics.Bitmap

import android.graphics.RuntimeShader

import androidx.annotation.FloatRange

    val foreground: Bitmap,

    /** A bitmap containing the background of the image. */

    val background: Bitmap,

    /** Pixel density of the display. Used for dithering. */

    val pixelDensity: Float,

    /** The amount of the snow flakes. This contributes to the color grading as well. */

    @FloatRange(from = 0.0, to = 1.0) val intensity: Float,

    /** The intensity of the color grading. 0: no color grading, 1: color grading in full effect. */

     * @param context the application context.

     * @param foreground a bitmap containing the foreground of the image.

     * @param background a bitmap containing the background of the image.

     * @param pixelDensity pixel density of the display.

     * @param intensity initial intensity that affects the amount of snow flakes and color grading.

     *   Expected range is [0, 1]. You can always change the intensity dynamically. Defaults to 1.

     */

    constructor(

        context: Context,

        assets: AssetManager,

        foreground: Bitmap,

        background: Bitmap,

        pixelDensity: Float,

        intensity: Float = DEFAULT_INTENSITY,

    ) : this(

        shader = GraphicsUtils.loadShader(context.assets, SHADER_PATH),

        accumulatedSnowShader =

            GraphicsUtils.loadShader(context.assets, ACCUMULATED_SNOW_SHADER_PATH),

        colorGradingShader = GraphicsUtils.loadShader(context.assets, COLOR_GRADING_SHADER_PATH),

        noiseTexture = GraphicsUtils.loadTexture(context.assets, NOISE_TEXTURE_PATH)

        shader = GraphicsUtils.loadShader(assets, SHADER_PATH),

        accumulatedSnowShader = GraphicsUtils.loadShader(assets, ACCUMULATED_SNOW_SHADER_PATH),

        colorGradingShader = GraphicsUtils.loadShader(assets, COLOR_GRADING_SHADER_PATH),

        noiseTexture = GraphicsUtils.loadTexture(assets, NOISE_TEXTURE_PATH)

            ?: throw RuntimeException("Noise texture is missing."),

        lut = GraphicsUtils.loadTexture(context.assets, LOOKUP_TABLE_TEXTURE_PATH),

        lut = GraphicsUtils.loadTexture(assets, LOOKUP_TABLE_TEXTURE_PATH),

        foreground,

        background,

        pixelDensity,

        intensity,

        COLOR_GRADING_INTENSITY,

        MAX_SNOW_THICKNESS

/** Contains functions for rendering. */

object GraphicsUtils {

    /* Default width dp is calculated as default_display_width / default_display_density. */

    private const val DEFAULT_WIDTH_DP = 1080 / 2.625f

    /**

     * Loads a shader from an asset file.

        allocationOut.copyTo(blurredImage)

        return blurredImage

    }

    /**

     * @return the [Float] representing the aspect ratio of width / height, -1 if either width or

     *   height is equal to or less than 0.

        } else width / height

    }

    /**

     * Compute the weather effect default grid size. This takes into consideration the different

     * display densities and aspect ratio so the effect looks good on displays with different sizes.

     * @param surfaceSize the size of the surface where the wallpaper is being rendered.

     * @param density the current display density.

     * @return a [Float] representing the default size.

     */

    fun computeDefaultGridSize(surfaceSize: SizeF, density: Float): Float {

        val displayWidthDp = surfaceSize.width / density

        val adjustedScale = when {

            // "COMPACT"

            displayWidthDp < 600 -> 1f

            // "MEDIUM"

            displayWidthDp >= 600 && displayWidthDp < 840 -> 0.9f

            // "EXPANDED"

            else -> 0.8f

        }

        return adjustedScale * displayWidthDp / DEFAULT_WIDTH_DP

    }

    private fun resolveShaderIncludes(assetManager: AssetManager, string: String): String {

        val match = Regex("[ \\t]*#include +\"([\\w\\d./]+)\"")

        return string.replace(match) { m ->