Untitled

import bpy
from mathutils import Vector as vec
from math import exp, cos, sin, pi, atan2, fmod
from random import uniform

wall = 9.0
speed = 0.25
repulsion_strength=0.25
wall_repulsion_strength=0.5
alignment_strength=0.1
repulsion_zone_width=1.5
alignment_zone_width=1.0
keyframe_interval=10
frames=2000
print("hello")

Fish = bpy.data.collections['Fish'].objects
velocities = [vec([0,0,0])] * len(Fish)

for i,f in enumerate(Fish):
    ang = uniform(0, 2*pi)
    v = speed*vec((cos(ang), sin(ang), 0))
    velocities[i] = v
    f.location = vec((uniform(-wall,wall), uniform(-wall,wall), 0))
    f.keyframe_insert(data_path="location", frame=1)
    f.rotation_euler[2] = atan2(v.y, v.x) + pi
    f.keyframe_insert(data_path="rotation_euler", frame=1)

def calculate_wall_repulsion(position):
    repulsion = vec([0, 0, 0])
    for axis in range(2):
        if position[axis] > wall - repulsion_zone_width:
            repulsion[axis] -= wall_repulsion_strength * (wall - position[axis]) / repulsion_zone_width
        elif position[axis] < -wall + repulsion_zone_width:
            repulsion[axis] += wall_repulsion_strength * (-wall - position[axis]) / repulsion_zone_width
    return repulsion

c = 1 

def gaussian_repulsion(distance):
    return exp(-((distance)**2) / (2 * c**2))

def calculate_separation(i, Fish):
    repulsion = vec([0, 0, 0])
    closest_fish = None
    closest_fish_distance = 99
    for j, other_fish in enumerate(Fish):
        if i != j:
            distance_vector = Fish[i].location - other_fish.location
            distance = distance_vector.length
            if distance < repulsion_zone_width and distance > 0:
                repulsion_direction = distance_vector.normalized()
                repulsion += (repulsion_zone_width - distance) / repulsion_zone_width * repulsion_direction * repulsion_strength * gaussian_repulsion(distance)
                print('repulsion', repulsion)
    return repulsion

for frame_no in range(1,int(frames/keyframe_interval)):
    for i,f in enumerate(Fish):
        p = f.location
        v = velocities[i]
        
        wall_repulsion = calculate_wall_repulsion(p)
        
        fish_separation = calculate_separation(i, Fish)
        
        v += wall_repulsion + fish_separation
        

        
        p += v
        f.keyframe_insert(data_path="location", frame=frame_no*keyframe_interval)
        f.rotation_euler[2] = atan2(v.y, v.x) + pi
        f.keyframe_insert(data_path="rotation_euler", frame=frame_no*keyframe_interval)