main.py

import random
import math
import sys

CHECKPOINT_RADIUS = 600
POD_RADIUS = 400
MAX_ANGLE_DIFF = 0.31415
FRICTION = 0.85
MAX_THRUST = 100

def circle_collision_time(dp, dv, radius_sum):
    rss = radius_sum * radius_sum
    a = dv * dv
    b = (dv * dp) * 2
    c = dp * dp - rss
    if b >= 0: return None
    if c <= 0: return 0
    d = b * b - 4 * a * c
    if d < 0: return None
    return (-b - math.sqrt(d)) / (2 * a)

class Vec2:
    def __init__(self, x, y):
        self.x = x
        self.y = y

    @classmethod
    def from_angle(cls, angle, size):
        return Vec2(math.cos(angle) * size, math.sin(angle) * size)

    def angle(self):
        return math.atan2(self.y, self.x)

    def len_sq(self):
        return self * self
    
    def len(self):
        return math.sqrt(self.len_sq())

    def norm(self):
        ilen = 1 / self.len()
        return self * ilen
    
    def floor(self):
        return Vec2(math.floor(self.x), math.floor(self.y))
    
    def round(self):
        return Vec2(round(self.x), round(self.y))

    def __add__(self, other):
        return Vec2(self.x + other.x, self.y + other.y)
    
    def __sub__(self, other):
        return Vec2(self.x - other.x, self.y - other.y)
    
    def __mul__(self, other):
        if isinstance(other, Vec2):
            return self.x * other.x + self.y * other.y
        else:
            return Vec2(self.x * other, self.y * other)

    def __repr__(self):
        return "<{}, {}>".format(self.x, self.y)

class DrawingContext:
    def __init__(self, out = None):
        self.out = sys.stdout if out is None else out
        self.scaling = 0.05

    def output(self, *args, **kw):
        print(*args, file = self.out, **kw)

    def start_frame(self):
        self.output("#FRAME_START")

    def no_stroke(self):
        self.output("#NO_STROKE")

    def no_fill(self):
        self.output("#NO_FILL")

    def set_stroke_color(self, r, g, b, a):
        self.output("#STROKE_COLOR", r, g, b, a)

    def set_fill_color(self, r, g, b, a):
        self.output("#FILL_COLOR", r, g, b, a)

    def circle(self, pos, rad):
        self.output("#CIRCLE", pos.x * self.scaling, pos.y * self.scaling, rad * self.scaling)

class Pod:
    def __init__(self, i, pos, vel, heading):
        self.id = i
        self.pos = pos
        self.vel = vel
        self.heading = heading
        self.checkpoint_idx = 1
        self.timeout = 0
        self.taps = 0
        self.dead = False

    def progress(self, game):
        return self.taps * 100000 + (game.get_checkpoint(self.checkpoint_idx) - self.pos).len()

    def simulate(self, game, action):
        cp = game.get_checkpoint(self.checkpoint_idx)
        to_tg = action.target - self.pos
        angle_diff = (to_tg.angle() - self.heading.angle() + math.pi) % (math.pi * 2) - math.pi
        if angle_diff > MAX_ANGLE_DIFF:
            angle_diff = MAX_ANGLE_DIFF
        if angle_diff < -MAX_ANGLE_DIFF:
            angle_diff = -MAX_ANGLE_DIFF
        n = Vec2.from_angle(self.heading.angle() + angle_diff, 1)
        self.heading = n
        thrust = action.thrust
        if thrust > MAX_THRUST:
            thrust = MAX_THRUST
        if thrust < 0:
            thrust = 0
        self.vel += n * action.thrust
        col_t = circle_collision_time(cp - self.pos, Vec2(0, 0) - self.vel, CHECKPOINT_RADIUS)
        self.pos += self.vel
        self.vel *= FRICTION
        self.vel = self.vel.floor()
        self.pos = self.pos.round()
        if col_t is not None and col_t <= 1:
            self.checkpoint_idx = game.next_checkpoint(self.checkpoint_idx)
            self.timeout = 0
            self.taps += 1
        self.timeout += 1
        if self.timeout > 100:
            self.dead = True

    def draw(self, ctx):
        ctx.no_stroke()
        ctx.set_fill_color(255, 0, 0, 255)
        ctx.circle(self.pos, POD_RADIUS)

    def __repr__(self):
        return "<Pod pos={} vel={} heading={} cp_idx={} timeout={} taps={} dead={}>".format(self.pos, self.vel, self.heading, self.checkpoint_idx, self.timeout, self.taps, self.dead)

class GameState:
    def __init__(self, laps, checkpoints):
        self.laps = laps
        self.checkpoints = checkpoints
    
    @classmethod
    def random(cls):
        laps = random.randint(3, 6)
        cp_count = random.randint(24, 36)
        cps = []
        for _ in range(1000):
            new_cp = Vec2(1000 + random.random() * 14000, 1000 + random.random() * 7000)
            reject = False
            for cp in cps:
                if (new_cp - cp).len() < 2000:
                    reject = True
                    break
            if reject:
                continue
            cps.append(new_cp)
            if len(cps) >= cp_count:
                break
        return GameState(laps, cps)

    def get_checkpoint(self, idx):
        return self.checkpoints[idx % len(self.checkpoints)]
    
    def next_checkpoint(self, idx):
        return (idx + 1) % len(self.checkpoints)

    def draw(self, ctx):
        ctx.no_stroke()
        ctx.set_fill_color(99, 99, 99, 255)
        for cp in self.checkpoints:
            ctx.circle(cp, CHECKPOINT_RADIUS)

class TurnState:
    def __init__(self, pods):
        self.pods = pods
        self.leaderboard = None
        self.end = False

    @classmethod
    def init_from(cls, game, pod_count):
        cp = game.get_checkpoint(0)
        to_cp_dir = (game.get_checkpoint(1) - game.get_checkpoint(0)).norm()
        return TurnState([Pod(i, cp, Vec2(0, 0), to_cp_dir) for i in range(pod_count)])

    def draw(self, ctx):
        for pod in self.pods:
            if not pod.dead:
                pod.draw(ctx)

    def simulate(self, game, actions):
        if self.end:
            raise Exception("cannot call simulate on an ended game")
        for i in range(len(self.pods)):
            pod = self.pods[i]
            if pod.dead:
                continue
            act = actions[i]
            pod.simulate(game, act)
        left = len([pod for pod in self.pods if not pod.dead])
        if left == 1:
            for pod in self.pods:
                if not pod.dead:
                    self.update_leaderboard(game)
                    self.end = True
                    return
        elif left == 0:
            self.end = True
            return
        for pod in self.pods:
            if pod.dead:
                continue
            if pod.taps >= len(game.checkpoints) * game.laps:
                self.update_leaderboard(game)
                self.end = True
                # TODO: Ties are a thing that exists. Handle them.
                return

    def update_leaderboard(self, game):
        self.leaderboard = sorted(range(0, len(self.pods)), key = lambda i: self.pods[i].progress(game), reverse = True)

class Action:
    def __init__(self, target, thrust):
        self.target = target
        self.thrust = thrust

    def __repr__(self):
        return "<Action target={}, thrust={}>".format(self.target, self.thrust)

class AI:
    def __init__(self, c):
        self.c = c
        self.fitness = 0

    def decide(self, game, turn, my_pod_id):
        pod = turn.pods[my_pod_id]
        return Action(game.get_checkpoint(pod.checkpoint_idx) - pod.vel * self.c, MAX_THRUST)
    
    def mutate(self, rate):
        c = self.c + (random.random() * 2 - 1) * rate
        return AI(c)

def run_battle(ais, ctx = None):
    if len(ais) == 0:
        raise Exception("empty list of AIs")
    game = GameState.random()
    turn = TurnState.init_from(game, len(ais))
    while True:
        actions = [None if turn.pods[i].dead else ai.decide(game, turn, i) for i, ai in enumerate(ais)]
        turn.simulate(game, actions)
        if ctx is not None:
            ctx.start_frame()
            game.draw(ctx)
            turn.draw(ctx)
        if turn.end:
            return turn.leaderboard

def calculate_fitness(pop):
    leaderboard = run_battle(pop)
    if leaderboard is None:
        ais = [AI(1 + random.random() * 8) for _ in range(50)]
    else:
        scores = list(range(len(leaderboard)))
        for i, s in zip(leaderboard, scores):
            pop[i].fitness = len(leaderboard) - s
        pop.sort(key = lambda p: p.fitness, reverse = True)

pop = [AI(n / 10) for n in range(10, 90)]

calculate_fitness(pop)

new_pop = []

i = 0

while True:
    rate = 1 / (1 + i)
    i += 1
    for i in range(5):
        new_pop.append(pop[i])
    while len(new_pop) < 50:
        baby = random.choice(new_pop[:5]).mutate(rate)
        new_pop.append(baby)
    pop = new_pop
    calculate_fitness(pop)
    print([p.c for p in pop], file = sys.stderr)
    new_pop = []