busplan.py

import itertools
import math
import random
import time


DEPENDENCIES = {
    'wires': {'copper'},
    'green circuits': {'iron', 'wires'},
    'red circuits': {'plastic', 'wires', 'green circuits'},
    'plastic': {'coal', 'gas'},
    'sulphur': {'gas', 'water'},
    'acid': {'iron', 'sulphur'},
    'blue circuits': {'acid', 'red circuits', 'green circuits'},
    'batteries': {'acid', 'iron', 'copper'},
    'cogs': {'iron'},
    'red science': {'cogs', 'copper'},
    'arms': {'green circuits', 'iron', 'cogs'},
    'belts': {'iron', 'cogs'},
    'green science': {'arms', 'belts'},
    'ammo': {'iron'},
    'grenades': {'iron', 'coal'},
    'piercing ammo': {'steel', 'copper', 'ammo'},
    'bricks': {'stone'},
    'walls': {'bricks'},
    'grey science': {'piercing ammo', 'grenades', 'walls'},
    'furnaces': {'bricks', 'steel', 'red circuits'},
    'prod modules': {'red circuits', 'green circuits'},
    'iron sticks': {'iron'},
    'rails': {'stone', 'steel', 'iron sticks'},
    'purple science': {'furnaces', 'prod modules', 'rails'},
    'pipes': {'iron'},
    'engines': {'steel', 'cogs', 'pipes'},
    'blue science': {'sulphur', 'red circuits', 'engines'},
    'eletric motor': {'engines', 'lube', 'green circuits'},
    'robot frame': {'electric motor', 'battery', 'green circuit', 'steel'},
    'low density structure': {'copper', 'steel', 'plastic'},
    'yellow science': {'red circuits', 'robot frame', 'low density structure'},
    'research': {'red science', 'green science', 'grey science', 'purple science', 'blue science'},
}


def powerset(iterable):
    "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
    s = list(iterable)
    return itertools.chain.from_iterable(itertools.combinations(s, r) for r in range(len(s) + 1))


def random_subset(iterable):
    subset = []
    for x in iterable:
        if random.randrange(2):
            subset.append(x)
    return subset


class Problem:
    def __init__(self, bus_items, factories):
        self.bus_items = bus_items
        self.factories = factories
        self.bus_size = len(bus_items)
        self.num_factories = len(factories)

    def size(self):
        num_partitions = 2 ** self.num_factories
        num_orders = math.factorial(self.bus_size)
        return num_partitions * num_orders

    def generate_partitions(self):
        remaining_items = self.factories.copy()

        ## Without loss of generality, put the first item on one side
        #remaining_items = items.copy()
        #remaining_items.pop()

        for subset in powerset(remaining_items):
            subset = set(subset)
            yield self.factories - subset, subset

    def generate_bus_orders(self):
        for order in itertools.permutations(self.bus_items):
            yield list(order)

    def generate_solutions(self):
        for left, right in self.generate_partitions():
            for bus_order in self.generate_bus_orders():
                yield Solution(self, bus_order, left, right)

    def get_random_solution(self):
        bus_order = list(self.bus_items)
        random.shuffle(bus_order)

        left = set(random_subset(self.factories))
        right = self.factories - left

        return Solution(self, bus_order, left, right)

    def get_neighbours(self, solution):
        # Two kinds of neighbour:
        #   - Move factory from one side to the other
        #   - Swap two bus lanes

        for item in self.factories:
            left = solution.left.copy()
            right = solution.right.copy()
            if item in solution.left:
                left.discard(item)
                right.add(item)
            else:
                right.discard(item)
                left.add(item)
            yield Solution(self, solution.bus_order, left, right)

        for i in range(self.bus_size-1):
            bus_order = solution.bus_order[:]
            bus_order[i], bus_order[i+1] = bus_order[i+1], bus_order[i]
            yield Solution(self, bus_order, solution.left, solution.right)


class Solution:
    def __init__(self, problem, bus_order, left, right):
        self.problem = problem
        self.bus_order = bus_order
        self.left = left
        self.right = right
        self.bus_cost = {}
        self.bus_cost_reverse = {}
        for item in self.problem.bus_items:
            self.bus_cost[item] = self.bus_order.index(item)
            self.bus_cost_reverse[item] = self.problem.bus_size - self.bus_order.index(item)

    def __hash__(self):
        return hash((tuple(self.bus_order), frozenset(self.left), frozenset(self.right)))

    def __eq__(self, other):
        return self.bus_order == other.bus_order and\
            self.left == other.left and\
            self.right == other.right

    def __str__(self):
        return '%s | %s | %s' % (self.left, self.bus_order, self.right)

    def __repr__(self):
        return 'Solution(<Problem>, %s, %s, %s)' % (repr(self.bus_order), repr(self.left), repr(self.right))

    def evaluate(self, trace=None):
        if trace:
            trace('Evaluating %s', self)
        score = 0
        score += self.evaluate_side(self.bus_cost, self.left, 'left', self.right, trace)
        score += self.evaluate_side(self.bus_cost_reverse, self.right, 'right', self.left, trace)
        return score

    def evaluate_side(self, bus_cost, side, side_name, other_side, trace):
        score = 0
        for name in side:
            score += self.evaluate_item(bus_cost, name, side, side_name, other_side, trace)
        return score

    def evaluate_item(self, bus_cost, name, side, side_name, other_side, trace):
        score = 0
        if trace:
            trace('For %s on %s', name, side_name)
        item_dependencies = DEPENDENCIES[name]
        for dep in item_dependencies:
            if dep not in self.problem.bus_items:
                if dep in side:
                    if trace:
                        trace('    %s is on the same side already', dep)
                    continue
                elif dep in other_side:
                    unds = self.problem.bus_size
                    if trace:
                        trace('    %s is on the far side and needs %d undergrounds', dep, unds)
                else:
                    if trace:
                        trace('    %s needs an unknown number of undergrounds', dep)
                    continue
            else:
                unds = bus_cost[dep]
                if trace:
                    trace('    %s needs %d undergrounds', dep, unds)
            score += unds

        if name in self.problem.bus_items:
            unds = bus_cost[name]
            if trace:
                trace('    output needs %d undergrounds', unds)
            score += unds

        return score


def print_dot(problem):
    print('graph {')
    solution_to_label = {}
    for i, solution in enumerate(problem.generate_solutions()):
        solution_to_label[solution] = i
        score = solution.evaluate()
        print('    %d[label=%d]' % (i, score))

    for sol in solution_to_label:
        label1 = solution_to_label[sol]
        score = sol.evaluate()
        for n in problem.get_neighbours(sol):
            label2 = solution_to_label[n]
            if label1 < label2:
                if sol.bus_order == n.bus_order:
                    colour = 'red'
                else:
                    colour = 'blue'
                diff = abs(score - n.evaluate())
                print('    %d -- %d[color=%s, label=%d, fontcolor=%s]' % (label1, label2, colour, diff, colour))
    print('}')


def anneal(problem):
    step = 0
    max_steps = 1000
    solution = problem.get_random_solution()
    while step < max_steps:
        temperature = 1 - step / float(max_steps)
        neighbours = list(problem.get_neighbours(solution))
        neighbour = random.choice(neighbours)
        score = solution.evaluate()
        new_score = neighbour.evaluate()

        value = math.exp((score - new_score)/temperature)
        random_value = random.random()
        if random_value < value:
            solution = neighbour

        print('New solution with score %d: %s' % (score,solution))

        step += 1

    def trace(msg, *args):
        print(msg % args)
    solution.evaluate(trace=trace)


def brute_force(problem):
    best_score = None
    best_solution = None
    for solution in problem.generate_solutions():
        score = solution.evaluate()
        if best_score is None or score < best_score:
            best_score = score
            best_solution = solution
            print('New solution with score %d: %s' % (score,solution))

    def trace(msg, *args):
        print(msg % args)
    solution.evaluate(trace=trace)

def main():
    problem = Problem(
        bus_items={'iron', 'copper', 'green circuits', 'red circuits', 'steel', 'cogs', 'stone', 'blue circuits'},
        factories={'wires', 'green circuits', 'red circuits', 'blue circuits', 'cogs', 'red science', 'arms', 'belts', 'green science',
                   'purple science', 'furnaces', 'rails', 'sulphur', 'iron sticks', 'blue science', 'engines'}
    )

    print('Number of solutions', problem.size())

    t0 = time.time()
    anneal(problem)
    print('Annealing took %0.3f seconds' % (time.time() - t0))

    t0 = time.time()
    brute_force(problem)
    print('Brute forcing took %0.3f seconds' % (time.time() - t0))

    #print_dot(problem)


if __name__ == '__main__':
    main()