#!/usr/bin/env python3

import random
import math

DIM = 80 # dimensionality of problem
ME = 100 # number of epochs
AXIS = 500 # -AXIS, AXIS is range for search space

class Molecule:
    def __init__(self):
        self.pos = [random.uniform(-AXIS, AXIS) for _ in range(DIM)]
        self.fit = self.get_fit()
        self.vel = [0.00 for _ in range(DIM)]
        self.mass = 0.00

    def get_fit(self):
        # schwefel function, minimum is 0
        summ = 0
        for d in self.pos:
            summ += (d * math.sin(math.sqrt(abs(d))))
        return (418.9829 * DIM) - summ

    def __lt__(self, other):
        return self.fit < other.fit

def fun():
    eps = 0.001 # epsilon
    pop = []

    for i in range(1000):
        new_molecule = Molecule()
        pop.append(new_molecule)
    temp = 1000 # temperature of the system initially
    k = 1000 # boltzman constant
    a, b = 0.5, 0.2
    for e in range(ME):
        pop.sort() # sort molecules by fitness

        avg_fit = sum(a.fit for a in pop)/len(pop) # get avg fitness

        best = pop[0].fit # best fitness is top of pop
        worst = pop[-1].fit # worst is bottom

        print(f'{best:.4}')

        temp = temp - (1 / avg_fit)

        for agent in pop:
            # for each agent, init its mass
            agent.mass = (agent.fit - worst) / (best - worst)
            for d in range(DIM):
                # for each dimension, calculate new velocity and position
                agent.vel[d] += math.sqrt((3*k*temp)/(agent.mass + eps))
                agent.pos[d] += agent.vel[d]
                agent.pos[d] += (b - (a / (2*math.pi)) * \
                    math.sin(2*math.pi * agent.pos[d])) % 1
                if agent.pos[d] < -AXIS or agent.pos[d] > AXIS:
                    # if molecule outside of range, place randomly
                    agent.pos[d] = random.uniform(-AXIS, AXIS)
            agent.fit = agent.get_fit() # update molecule's fitness

if __name__ == '__main__':
    fun()