Anno 1800 island optimizer using OR-Tools CP-SAT constraint solver

anno.py

#!/usr/bin/env python3

from collections import namedtuple
from ortools.sat.python import cp_model
import itertools
import sys

MAX_INT = 1000000000000
RATE_SCALING = 1000000

class EntityType:
    def __init__(self, name):
        self.name = name
        self.model = None

    def set_model(self, model):
        if self.model is not None:
            raise IllegalStateException('Already have a model')
        self.model = model
        self.variables = self._get_variables(model)

    def add_to_model(self, model):
        raise NotImplementedError

    def _get_variables(self, model):
        raise NotImplementedError

    def get_profit_variable(self, model):
        raise NotImplementedError

    def print_solution(self, solver):
        raise NotImplementedError


RESIDENTS = []

ResidentVariables = namedtuple('ResidentVariables', ['number', 'houses', 'income', 'happiness', 'needs_enabled'])
ResidentNeed = namedtuple('ResidentNeed', ['resource', 'houses_supplied_per_unit', 'influx', 'income', 'happiness', 'pop_requirement'])

class ResidentType(EntityType):
    def __init__(self, name, max_per_house, needs):
        super().__init__(name)
        self.max_per_house = max_per_house
        self.needs = needs
        self.employers = []
        RESIDENTS.append(self)

    def __repr__(self):
        return f'<Resident: {self.name}>'

    def _get_variables(self, model):
        return ResidentVariables(
            number=model.NewIntVar(0, MAX_INT, f'{self.name}Count'),
            houses=model.NewIntVar(0, MAX_INT, f'{self.name}Houses'),
            income=model.NewIntVar(0, MAX_INT, f'{self.name}Income'),
            happiness=model.NewIntVar(-20, 20, f'{self.name}Happiness'),
            needs_enabled=dict((need.resource, model.NewBoolVar(f'{self.name}{need.resource.name}ConsumptionEnabled')) for need in self.needs),
        )

    def add_employer(self, employer):
        self.employers.append(employer)

    @property
    def min_per_house(self):
        return self.max_per_house - sum(need.influx for need in self.needs)
    
    @property
    def happiness_from_buildings(self):
        return 20 - sum(need.happiness for need in self.needs)

    def add_to_model(self, model):
        # Max residents per house
        model.Add(self.variables.number <= self.variables.houses * self.max_per_house)
        # Min residents per house
        model.Add(self.variables.number >= self.variables.houses * self.min_per_house)
        # Employers (factories)
        model.Add(self.variables.number >= cp_model.LinearExpr.Sum(self.employers))
        # Needs
        influxes = [self.variables.houses * self.min_per_house]
        incomes = []
        luxuries = [self.happiness_from_buildings]
        self.need_variables = []
        for need in self.needs:
            consumption_enabled = self.variables.needs_enabled[need.resource]
            # Add the consumption rate of these residents to the resource
            consumption_rate = model.NewIntVar(0, MAX_INT, f'{self.name}ConsumptionOf{need.resource.name}')
            need.resource.add_consumption(consumption_rate)
            # Consumption, if disabled, is zero
            model.Add(consumption_rate == 0).OnlyEnforceIf(consumption_enabled.Not())
            # Consumption, if enabled, is proportional to number of houses
            model.Add(consumption_rate == self.variables.houses * int(RATE_SCALING / need.houses_supplied_per_unit)).OnlyEnforceIf(consumption_enabled)

            # Needs with a population requirement only consume resources when number of residents meets threshold
            pop_requirement_met = model.NewBoolVar(f'{self.name}Needs{need.resource.name}')
            if need.pop_requirement:
                model.Add(pop_requirement_met == self.variables.number >= need.pop_requirement)
            else:
                model.Add(pop_requirement_met == True)
            # Do not enable consumption if population requirement is not met.
            model.Add(consumption_enabled == 0).OnlyEnforceIf(pop_requirement_met.Not())

            # If need has an influx effect, it adds N workers per house
            influx = model.NewIntVar(0, MAX_INT, f'{self.name}InfluxFrom{need.resource.name}')
            influxes.append(influx)
            if need.influx:
                model.Add(influx == self.variables.houses * need.influx).OnlyEnforceIf(consumption_enabled)
                model.Add(influx == 0).OnlyEnforceIf(consumption_enabled.Not())
            else:
                model.Add(influx == 0)
            
            # If need has an income effect, it adds N coins per house
            income = model.NewIntVar(0, MAX_INT, f'{self.name}IncomeFrom{need.resource.name}')
            incomes.append(income)
            if need.income:
                model.Add(income == self.variables.houses * need.income).OnlyEnforceIf(consumption_enabled)
                model.Add(income == 0).OnlyEnforceIf(consumption_enabled.Not())
            else:
                model.Add(income == 0)
            
            # If need is a luxury (i.e. has a happiness effect), it adds/removes N happiness per house
            happiness = model.NewIntVar(-need.happiness, need.happiness, f'{self.name}HappinessFrom{need.resource.name}')
            luxuries.append(happiness)
            if need.happiness:
                model.Add(happiness == need.happiness).OnlyEnforceIf(consumption_enabled)
                model.Add(happiness == -need.happiness).OnlyEnforceIf(consumption_enabled.Not())
            else:
                model.Add(happiness == 0)
            
            self.need_variables.append((need, {
                'consumption_enabled': consumption_enabled,
                'consumption_rate': consumption_rate,
                'pop_requirement_met': pop_requirement_met,
                'influx': influx,
                'income': income,
                'happiness': happiness,
            }))
        model.Add(self.variables.number == cp_model.LinearExpr.Sum(influxes))
        model.Add(self.variables.income == cp_model.LinearExpr.Sum(incomes))
        model.Add(self.variables.happiness == cp_model.LinearExpr.Sum(luxuries))

    def get_profit_variable(self):
        return self.variables.income

    def print_solution(self, solver):
        print(f'Worker Type: {self.name}')
        print(f'  labor force: {solver.Value(self.variables.number)}')
        print(f'  houses: {solver.Value(self.variables.houses)}')
        print(f'  income: {solver.Value(self.variables.income)}')
        print(f'  happiness: {solver.Value(self.variables.happiness)} (from buildings: {self.happiness_from_buildings})')
        print(f'  needs:')
        for need, need_vars in self.need_variables:
            print(f'    need: {need.resource.name}')
            for need_var_name, variable in need_vars.items():
                if need_var_name == 'consumption_rate':
                    print(f'      {need_var_name}: {solver.Value(variable)/RATE_SCALING}')
                else:
                    print(f'      {need_var_name}: {solver.Value(variable)}')

RESOURCES = []

ResourceVariables = namedtuple('ResourceVariables', ['production', 'consumption', 'balance'])

class ResourceType(EntityType):
    def __init__(self, name):
        super().__init__(name)
        self.consumers = []
        self.producers = []
        RESOURCES.append(self)
    
    def __repr__(self):
        return f'<Resource: {self.name}>'

    def _get_variables(self, model):
        return ResourceVariables(
            production=model.NewIntVar(0, MAX_INT, f'{self.name}Production'),
            consumption=model.NewIntVar(0, MAX_INT, f'{self.name}Consumption'),
            balance=model.NewIntVar(-MAX_INT, MAX_INT, f'{self.name}Balance'),
        )

    def add_consumption(self, expr):
        self.consumers.append(expr)

    def add_production(self, expr):
        self.producers.append(expr)

    def add_to_model(self, model):
        # Balance is production - consumption
        model.Add(self.variables.balance == self.variables.production - self.variables.consumption)
        # Production should always be greater than or equal to consumption
        # TODO: Revisit this, we might want to allow a negative balance sometimes, e.g. if modeling an island that relies on trade.
        model.Add(self.variables.production >= self.variables.consumption)
        # Consumption rate is the sum of all consumers
        model.Add(self.variables.consumption == cp_model.LinearExpr.Sum(self.consumers))
        # Production rate is the sum of all producers
        model.Add(self.variables.production == cp_model.LinearExpr.Sum(self.producers))

    def get_profit_variable(self):
        return 0
    
    def print_solution(self, solver):
        print(f'Resource Type: {self.name}')
        print(f'  production: {solver.Value(self.variables.production) / RATE_SCALING}')
        print(f'  consumption: {solver.Value(self.variables.consumption) / RATE_SCALING}')
        print(f'  balance: {solver.Value(self.variables.balance) / RATE_SCALING}')

FACTORIES = []

FactoryVariables = namedtuple('FactoryVariables', ['number'])

class FactoryType(EntityType):
    def __init__(self, name, maint_cost, employee_type, employee_num, output_resource, input_resources, production_time_seconds):
        super().__init__(name)
        self.maint_cost = maint_cost
        self.employee_type = employee_type
        self.employee_num = employee_num
        self.output_resource = output_resource
        self.input_resources = input_resources
        self.production_time_seconds = production_time_seconds
        FACTORIES.append(self)

    def __repr__(self):
        return f'<Factory: {self.name}>'

    def _get_variables(self, model):
        return FactoryVariables(number=model.NewIntVar(0, MAX_INT, f'{self.name}Count'))

    def add_to_model(self, model):
        self.employee_type.add_employer(self.variables.number * self.employee_num)
        for input_resource in self.input_resources:
            input_resource.add_consumption(self.variables.number * (RATE_SCALING * 60 // self.production_time_seconds))
        self.output_resource.add_production(self.variables.number * (RATE_SCALING * 60 // self.production_time_seconds))
    
    def get_profit_variable(self):
        return self.variables.number * -self.maint_cost

    def print_solution(self, solver):
        print(f'Factory Type: {self.name}')
        print(f'  number: {solver.Value(self.variables.number)}')
        print(f'  cost: {solver.Value(self.variables.number * -self.maint_cost)}')

Fish = ResourceType('Fish')

Wool = ResourceType('Wool')
WorkClothes = ResourceType('WorkClothes')

Potatoes = ResourceType('Potatoes')
Schnapps = ResourceType('Schnapps')

Pigs = ResourceType('Pigs')
Sausages = ResourceType('Sausages')
Tallow = ResourceType('Tallow')
Soap = ResourceType('Soap')

Grain = ResourceType('Grain')
Malt = ResourceType('Malt')
Hops = ResourceType('Hops')
Beer = ResourceType('Beer')

Flour = ResourceType('Flour')
Bread = ResourceType('Bread')

Farmer = ResidentType('Farmer',
    max_per_house=10,
    needs=[
        ResidentNeed(Fish, houses_supplied_per_unit=40, influx=3, income=1, happiness=0, pop_requirement=50),
        ResidentNeed(WorkClothes, houses_supplied_per_unit=33, influx=2, income=3, happiness=0, pop_requirement=150),
        ResidentNeed(Schnapps, houses_supplied_per_unit=30, influx=0, income=3, happiness=8, pop_requirement=100),
    ])

Worker = ResidentType('Worker',
max_per_house=20,
needs=[
    ResidentNeed(Fish, houses_supplied_per_unit=20, influx=3, income=1, happiness=0, pop_requirement=0),
    ResidentNeed(WorkClothes, houses_supplied_per_unit=16.25, influx=2, income=7, happiness=0, pop_requirement=0),
    ResidentNeed(Sausages, houses_supplied_per_unit=50, influx=3, income=5, happiness=0, pop_requirement=1),
    ResidentNeed(Bread, houses_supplied_per_unit=55, influx=3, income=5, happiness=0, pop_requirement=150),
    ResidentNeed(Soap, houses_supplied_per_unit=120, influx=2, income=5, happiness=0, pop_requirement=300),
    ResidentNeed(Schnapps, houses_supplied_per_unit=15, influx=0, income=7, happiness=4, pop_requirement=0),
    ResidentNeed(Beer, houses_supplied_per_unit=65, influx=0, income=12, happiness=3, pop_requirement=500),
])

Fishery = FactoryType('Fishery', maint_cost=40, employee_type=Farmer, employee_num=25, output_resource=Fish, input_resources=[], production_time_seconds=30)

SheepFarm = FactoryType('SheepFarm', maint_cost=20, employee_type=Farmer, employee_num=10, output_resource=Wool, input_resources=[], production_time_seconds=30)
Knittery = FactoryType('Knittery', maint_cost=50, employee_type=Farmer, employee_num=50, output_resource=WorkClothes, input_resources=[Wool], production_time_seconds=30)

PotatoFarm = FactoryType('PotatoFarm', maint_cost=20, employee_type=Farmer, employee_num=20, output_resource=Potatoes, input_resources=[], production_time_seconds=30)
SchnappsDistillery = FactoryType('SchnappsDistillery', maint_cost=40, employee_type=Farmer, employee_num=50, output_resource=Schnapps, input_resources=[Potatoes], production_time_seconds=30)

PigFarm = FactoryType('PigFarm', maint_cost=40, employee_type=Farmer, employee_num=30, output_resource=Pigs, input_resources=[], production_time_seconds=60)
Slaughterhouse = FactoryType('Slaughterhouse', maint_cost=80, employee_type=Worker, employee_num=50, output_resource=Sausages, input_resources=[Pigs], production_time_seconds=60)
RenderingWorks = FactoryType('RenderingWorks', maint_cost=40, employee_type=Worker, employee_num=40, output_resource=Tallow, input_resources=[Pigs], production_time_seconds=60)
SoapFactory = FactoryType('SoapFactory', maint_cost=50, employee_type=Worker, employee_num=50, output_resource=Soap, input_resources=[Tallow], production_time_seconds=30)

HopFarm = FactoryType('HopFarm', maint_cost=20, employee_type=Farmer, employee_num=20, output_resource=Hops, input_resources=[], production_time_seconds=90)
GrainFarm = FactoryType('GrainFarm', maint_cost=20, employee_type=Farmer, employee_num=20, output_resource=Grain, input_resources=[], production_time_seconds=60)
Malthouse = FactoryType('Malthouse', maint_cost=150, employee_type=Worker, employee_num=25, output_resource=Malt, input_resources=[Grain], production_time_seconds=30)
Brewery = FactoryType('Brewery', maint_cost=200, employee_type=Worker, employee_num=75, output_resource=Beer, input_resources=[Hops, Malt], production_time_seconds=60)

FlourMill = FactoryType('FlourMill', maint_cost=50, employee_type=Farmer, employee_num=10, output_resource=Flour, input_resources=[Grain], production_time_seconds=30)
Bakery = FactoryType('Bakery', maint_cost=60, employee_type=Worker, employee_num=50, output_resource=Bread, input_resources=[Flour], production_time_seconds=60)

def main():
    model = cp_model.CpModel()
    profit = []
    for entities in [RESOURCES, RESIDENTS, FACTORIES]:
        for entity in entities:
            entity.set_model(model)
            profit.append(entity.get_profit_variable())
    for factory in FACTORIES:
        factory.add_to_model(model)
    for resident in RESIDENTS:
        resident.add_to_model(model)
    for resource in RESOURCES:
        resource.add_to_model(model)

    # 96 farmer houses max, all full and happy
    model.Add(Farmer.variables.houses <= 64)
    model.Add(Farmer.variables.number == Farmer.variables.houses * Farmer.max_per_house)
    model.Add(Farmer.variables.happiness == 20)

    # 96 worker houses max, all full and happy
    model.Add(Worker.variables.houses <= 32)
    model.Add(Worker.variables.number == Worker.variables.houses * Worker.max_per_house)
    model.Add(Worker.variables.happiness == 20)

    # Maximize profit
    model.Maximize(sum(profit))

    solver = cp_model.CpSolver()
    status = solver.Solve(model)

    if status == cp_model.INFEASIBLE:
        print(f'NO SOLUTION')
        print(solver.ResponseStats())
        sys.exit(1)
    elif status == cp_model.OPTIMAL:
        print(f'OPTIMAL SOLUTION FOUND\n')
    elif status == cp_model.FEASIBLE:
        print(f'FEASIBLE SOLUTION FOUND\n')

    print(f'Farmers: {solver.Value(Farmer.variables.number)} in {solver.Value(Farmer.variables.houses)} houses')
    print(f'Workers: {solver.Value(Worker.variables.number)} in {solver.Value(Worker.variables.houses)} houses')
    print(f'Resident revenue: {solver.Value(sum(resident.get_profit_variable() for resident in RESIDENTS))}')
    print(f'Factories: {solver.Value(sum(factory.variables.number for factory in FACTORIES))}')
    print(f'Factory costs: {solver.Value(sum(factory.get_profit_variable() for factory in FACTORIES))}')
    print(f'Profit: {solver.Value(sum(profit))}\n')
    print(f'')
    print('===== Residents =====\n')
    for resident in RESIDENTS:
        resident.print_solution(solver)
        print()

    print('===== Factories =====\n')
    for factory in FACTORIES:
        factory.print_solution(solver)
        print()

    print('===== Resources =====\n')
    for resource in RESOURCES:
        resource.print_solution(solver)
        print()


if __name__ == '__main__':
    main()