Mesh Gradient

import {
  Shader,
  ShaderComponentProps,
  ShaderMotionParams,
  ShaderPreset,
  vec4,
} from "@/components/Shader";
import {
  colorBandingFix,
  colorPropsAreEqual,
  declarePI,
  declareRotate,
  defaultObjectSizing,
  getShaderColorFromString,
  ShaderFitOptions,
  ShaderSizingParams,
  ShaderSizingUniforms,
  sizingVariablesDeclaration,
} from "@/lib/shader-utils";
import { cn } from "@/lib/utils";
import { memo } from "react";

export const meshGradientMeta = {
  maxColorCount: 10,
} as const;

/**
 * A composition of N color spots (one per color) with 2 types of
 * distortions applied to the coordinate space
 *
 * Uniforms:
 * - u_colors (vec4[]), u_colorsCount (float used as integer)
 * - u_distortion: warp distortion
 * - u_swirl: vortex distortion
 *
 */

// language=GLSL
export const meshGradientFragmentShader: string = `#version 300 es
precision mediump float;

uniform float u_time;

uniform vec4 u_colors[${meshGradientMeta.maxColorCount}];
uniform float u_colorsCount;

uniform float u_distortion;
uniform float u_swirl;

${sizingVariablesDeclaration}

out vec4 fragColor;

${declarePI}
${declareRotate}

vec2 getPosition(int i, float t) {
  float a = float(i) * .37;
  float b = .6 + mod(float(i), 3.) * .3;
  float c = .8 + mod(float(i + 1), 4.) * 0.25;

  float x = sin(t * b + a);
  float y = cos(t * c + a * 1.5);

  return .5 + .5 * vec2(x, y);
}

void main() {
  vec2 shape_uv = v_objectUV;

  shape_uv += .5;

  float t = .5 * u_time;

  float radius = smoothstep(0., 1., length(shape_uv - .5));
  float center = 1. - radius;
  for (float i = 1.; i <= 2.; i++) {
    shape_uv.x += u_distortion * center / i * sin(t + i * .4 * smoothstep(.0, 1., shape_uv.y)) * cos(.2 * t + i * 2.4 * smoothstep(.0, 1., shape_uv.y));
    shape_uv.y += u_distortion * center / i * cos(t + i * 2. * smoothstep(.0, 1., shape_uv.x));
  }

  vec2 uvRotated = shape_uv;
  uvRotated -= vec2(.5);
  float angle = 3. * u_swirl * radius;
  uvRotated = rotate(uvRotated, -angle);
  uvRotated += vec2(.5);

  vec3 color = vec3(0.);
  float opacity = 0.;
  float totalWeight = 0.;

  for (int i = 0; i < ${meshGradientMeta.maxColorCount}; i++) {
    if (i >= int(u_colorsCount)) break;

    vec2 pos = getPosition(i, t);
    vec3 colorFraction = u_colors[i].rgb * u_colors[i].a;
    float opacityFraction = u_colors[i].a;

    float dist = length(uvRotated - pos);

    dist = pow(dist, 3.5);
    float weight = 1. / (dist + 1e-3);
    color += colorFraction * weight;
    opacity += opacityFraction * weight;
    totalWeight += weight;
  }

  color /= totalWeight;
  opacity /= totalWeight;

  ${colorBandingFix}

  fragColor = vec4(color, opacity);
}
`;

export interface MeshGradientUniforms extends ShaderSizingUniforms {
  u_colors: vec4[];
  u_colorsCount: number;
  u_distortion: number;
  u_swirl: number;
}

export interface MeshGradientParams
  extends ShaderSizingParams,
    ShaderMotionParams {
  colors?: string[];
  distortion?: number;
  swirl?: number;
}

export interface MeshGradientProps
  extends ShaderComponentProps,
    MeshGradientParams {}

type MeshGradientPreset = ShaderPreset<MeshGradientParams>;

export const defaultPreset: MeshGradientPreset = {
  name: "Default",
  params: {
    ...defaultObjectSizing,
    speed: 1,
    frame: 0,
    colors: ["#e0eaff", "#241d9a", "#f75092", "#9f50d3"],
    distortion: 0.8,
    swirl: 0.1,
  },
};

export const purplePreset: MeshGradientPreset = {
  name: "Purple",
  params: {
    ...defaultObjectSizing,
    speed: 0.6,
    frame: 0,
    colors: ["#aaa7d7", "#3c2b8e"],
    distortion: 1,
    swirl: 1,
  },
};

export const beachPreset: MeshGradientPreset = {
  name: "Beach",
  params: {
    ...defaultObjectSizing,
    speed: 0.1,
    frame: 0,
    colors: ["#bcecf6", "#00aaff", "#00f7ff", "#ffd447"],
    distortion: 0.8,
    swirl: 0.35,
  },
};

export const inkPreset: MeshGradientPreset = {
  name: "Ink",
  params: {
    ...defaultObjectSizing,
    speed: 1,
    frame: 0,
    colors: ["#ffffff", "#000000"],
    distortion: 1,
    swirl: 0.2,
    rotation: 90,
  },
};

export const meshGradientPresets: MeshGradientPreset[] = [
  defaultPreset,
  inkPreset,
  purplePreset,
  beachPreset,
];

export const MeshGradient: React.FC<MeshGradientProps> = memo(
  function MeshGradientImpl({
    // Own props
    speed = defaultPreset.params.speed,
    frame = defaultPreset.params.frame,
    colors = defaultPreset.params.colors,
    distortion = defaultPreset.params.distortion,
    swirl = defaultPreset.params.swirl,

    // Sizing props
    fit = defaultPreset.params.fit,
    rotation = defaultPreset.params.rotation,
    scale = defaultPreset.params.scale,
    originX = defaultPreset.params.originX,
    originY = defaultPreset.params.originY,
    offsetX = defaultPreset.params.offsetX,
    offsetY = defaultPreset.params.offsetY,
    worldWidth = defaultPreset.params.worldWidth,
    worldHeight = defaultPreset.params.worldHeight,
    ...props
  }: MeshGradientProps) {
    const uniforms = {
      // Own uniforms
      u_colors: colors.map(getShaderColorFromString),
      u_colorsCount: colors.length,
      u_distortion: distortion,
      u_swirl: swirl,

      // Sizing uniforms
      u_fit: ShaderFitOptions[fit],
      u_rotation: rotation,
      u_scale: scale,
      u_offsetX: offsetX,
      u_offsetY: offsetY,
      u_originX: originX,
      u_originY: originY,
      u_worldWidth: worldWidth,
      u_worldHeight: worldHeight,
    } satisfies MeshGradientUniforms;

    return (
      <Shader
        {...props}
        className={cn("w-full h-full", props.className)}
        speed={speed}
        frame={frame}
        fragmentShader={meshGradientFragmentShader}
        uniforms={uniforms}
      />
    );
  },
  colorPropsAreEqual,
);