struct Settings {
  moveRate: f32,
  turnRate: f32,
  sensorAngle: f32,
  sensorOffset: f32,
  turnWhenLost: f32,
  individualTrailWeight: f32,
  agentCount: f32,
  introProgress: f32,
  color1ToColor1: f32,
  color1ToColor2: f32,
  color1ToColor3: f32,
  color2ToColor1: f32,
  color2ToColor2: f32,
  color2ToColor3: f32,
  color3ToColor1: f32,
  color3ToColor2: f32,
  color3ToColor3: f32,
};

@group(1) @binding(0) var<uniform> settings: Settings;
@group(1) @binding(2) var trailMapIn: texture_2d<f32>;
@group(1) @binding(3) var trailMapOut: texture_storage_2d<rgba16float, write>;
@group(1) @binding(4) var sourceMap: texture_2d<f32>;

@compute @workgroup_size(64)
fn main(
  @builtin(global_invocation_id) global_id: vec3<u32>,
  @builtin(num_workgroups) workgroup_count: vec3<u32>
) {
  let id = get_id(global_id, workgroup_count);

  if id >= u32(settings.agentCount) {
    return;
  }

  var agent = agents[id];
  if agent.colorIndex < 0.0 {
    return;
  }

  let hasIntroTarget =
    settings.introProgress < 0.999 &&
    agent.targetPosition.x >= 0.0 &&
    agent.targetPosition.y >= 0.0;

  if hasIntroTarget && settings.introProgress < agent.introDelay {
    return;
  }

  let random = textureSampleLevel(
    noise,
    noiseSampler,
    fract(vec2(f32(id) * 0.7548777, state.time * 0.00017 + f32(id) * 0.5698403)),
    0
  );

  let forwardSensor = sensor_position(agent.position, agent.angle, settings.sensorOffset, 0);
  let leftSensor = sensor_position(agent.position, agent.angle, settings.sensorOffset, settings.sensorAngle);
  let rightSensor = sensor_position(agent.position, agent.angle, settings.sensorOffset, -settings.sensorAngle);

  let trailForward = textureLoad(trailMapIn, forwardSensor, 0);
  let trailLeft = textureLoad(trailMapIn, leftSensor, 0);
  let trailRight = textureLoad(trailMapIn, rightSensor, 0);
  let sourceForwardSample = textureLoad(sourceMap, forwardSensor, 0);
  let sourceLeftSample = textureLoad(sourceMap, leftSensor, 0);
  let sourceRightSample = textureLoad(sourceMap, rightSensor, 0);

  let channelMask = get_channel_mask(agent.colorIndex);
  let reactionMask = get_reaction_mask(agent.colorIndex);

  let trailForwardWeight = dot(trailForward.rgb, reactionMask);
  let trailLeftWeight = dot(trailLeft.rgb, reactionMask);
  let trailRightWeight = dot(trailRight.rgb, reactionMask);

  let sourceForwardWeight = dot(sourceForwardSample.rgb, reactionMask);
  let sourceLeftWeight = dot(sourceLeftSample.rgb, reactionMask);
  let sourceRightWeight = dot(sourceRightSample.rgb, reactionMask);

  let weightForward = trailForwardWeight + sourceForwardWeight * 24.0;
  let weightLeft = trailLeftWeight + sourceLeftWeight * 24.0;
  let weightRight = trailRightWeight + sourceRightWeight * 24.0;

  var rotation = (random.r - 0.5) * settings.turnWhenLost;
  if weightForward >= weightLeft && weightForward >= weightRight {
    rotation = rotation * 0.25;
  } else {
    rotation += sign(weightLeft - weightRight) * settings.turnRate;
  }

  let sourceAtAgent = textureLoad(sourceMap, vec2<i32>(agent.position), 0);
  let positiveReactionMask = max(reactionMask, vec3<f32>(0.0));
  let sourceAtAgentStrength = clamp(dot(sourceAtAgent.rgb, positiveReactionMask), 0.0, 1.0);
  var moveRate = settings.moveRate * mix(1.0, 0.08, sourceAtAgentStrength);
  var introTargetOffset = vec2<f32>(0.0, 0.0);
  var introTargetDistance = 0.0;

  if hasIntroTarget {
    introTargetOffset = agent.targetPosition - agent.position;
    introTargetDistance = length(introTargetOffset);
    let targetAngle = atan2(introTargetOffset.y, introTargetOffset.x);
    let nearTitle = 1.0 - smoothstep(4.0, max(28.0, settings.sensorOffset * 0.75), introTargetDistance);
    let desiredAngle = mix(targetAngle, agent.targetAngle, nearTitle * 0.2);
    let introTurn = angle_delta(agent.angle, desiredAngle);

    rotation = clamp(introTurn, -settings.turnRate * 3.4, settings.turnRate * 3.4)
      + (random.g - 0.5) * settings.turnWhenLost * 0.18;
    moveRate = min(settings.moveRate * mix(2.65, 0.01, nearTitle), introTargetDistance);
  }

  var step = vec2(cos(agent.angle), sin(agent.angle)) * moveRate;
  if hasIntroTarget {
    step = vec2<f32>(0.0, 0.0);
    if introTargetDistance > 0.5 {
      step = introTargetOffset / introTargetDistance * moveRate;
    }
  }

  let maxPosition = state.size - vec2<f32>(1.0, 1.0);
  let nextPosition = clamp(agent.position + step, vec2<f32>(0, 0), maxPosition);
  if nextPosition.x == 0 || nextPosition.x == maxPosition.x || nextPosition.y == 0 || nextPosition.y == maxPosition.y {
    rotation = 3.14159265359 + random.a - 0.5;
  }

  let sourceBelow = textureLoad(sourceMap, vec2<i32>(nextPosition), 0);
  let sourceBelowStrength = clamp(dot(sourceBelow.rgb, positiveReactionMask), 0.0, 1.0);
  let trailWeight = settings.individualTrailWeight * (1.0 + sourceBelowStrength * 16.0);
  var trailBelow = textureLoad(trailMapIn, vec2<i32>(nextPosition), 0);
  trailBelow = vec4<f32>(
    trailBelow.rgb + channelMask * trailWeight,
    max(trailBelow.a, 0.0)
  );

  agent.angle += rotation;
  agent.position = nextPosition;

  textureStore(trailMapOut, vec2<i32>(nextPosition), trailBelow);
  agents[id] = agent;
}

fn sensor_position(agentPosition: vec2<f32>, agentAngle: f32, sensorOffset: f32, sensorOffsetAngle: f32) -> vec2<i32> {
  let sensorAngle = agentAngle + sensorOffsetAngle;
  return vec2<i32>(clamp(
    agentPosition + vec2(cos(sensorAngle), sin(sensorAngle)) * sensorOffset,
    vec2<f32>(0, 0),
    state.size - vec2<f32>(1, 1)
  ));
}

fn get_channel_mask(colorIndex: f32) -> vec3<f32> {
  if colorIndex < 0.5 {
    return vec3<f32>(1, 0, 0);
  }
  if colorIndex < 1.5 {
    return vec3<f32>(0, 1, 0);
  }
  return vec3<f32>(0, 0, 1);
}

fn get_reaction_mask(colorIndex: f32) -> vec3<f32> {
  if colorIndex < 0.5 {
    return vec3<f32>(
      settings.color1ToColor1,
      settings.color1ToColor2,
      settings.color1ToColor3
    );
  }
  if colorIndex < 1.5 {
    return vec3<f32>(
      settings.color2ToColor1,
      settings.color2ToColor2,
      settings.color2ToColor3
    );
  }
  return vec3<f32>(
    settings.color3ToColor1,
    settings.color3ToColor2,
    settings.color3ToColor3
  );
}

fn angle_delta(sourceAngle: f32, targetAngle: f32) -> f32 {
  return atan2(sin(targetAngle - sourceAngle), cos(targetAngle - sourceAngle));
}