cocostrat/app/economy.py

import json
import pickle
from os.path import join
from itertools import product
import numpy as np


class Resources:

    base_trade_range = 2

    def __init__(self, rows=10, cols=10):
        self.shape = shape = (rows, cols)
        self.lc_pops = np.zeros(shape, dtype=np.int32)
        self.mc_pops = np.zeros(shape, dtype=np.int32)
        self.uc_pops = np.zeros(shape, dtype=np.int32)
        self.pops_cap = np.zeros(shape, dtype=np.int32)

        self.lc_starvation = np.zeros(shape, dtype=np.bool)
        self.mc_starvation = np.zeros(shape, dtype=np.bool)
        self.uc_starvation = np.zeros(shape, dtype=np.bool)
        self.any_starvation = np.zeros(shape, dtype=np.bool)

        self.food_rate = np.zeros(shape, dtype=np.float32)
        self.food_production_cap = np.zeros(shape, dtype=np.float32)
        self.food_production = np.zeros(shape, dtype=np.float32)
        self.food_imports = np.zeros(shape, dtype=np.float32)
        self.food_consumption = np.zeros(shape, dtype=np.float32)
        self.food_stored = np.zeros(shape, dtype=np.float32)
        self.food_stored_capacity = np.zeros(shape, dtype=np.float32)

        self.trade_impedance = np.zeros(shape, dtype=np.float32)
        self.trade_range = np.full(shape, self.base_trade_range, dtype=np.float32)
        self.trade_value = np.zeros(shape, dtype=np.float32)
        self.trade_modifier = np.zeros(shape, dtype=np.float32)

        self.export_partner = np.full((*shape, 2), 0, dtype=np.int32)

        self.goods_production = {
            'stone': np.full(shape, 0, dtype=np.float32),
            'wood': np.full(shape, 0, dtype=np.float32),
        }

    def init(self, verkolst):
        self.shape = (verkolst.width, verkolst.height)
        self.lc_pops = np.full(self.shape, 25, dtype=np.int32)
        self.mc_pops = np.full(self.shape, 0, dtype=np.int32)
        self.uc_pops = np.full(self.shape, 0, dtype=np.int32)

        self.lc_starvation = np.full_like(self.lc_pops, False, dtype=np.bool)
        self.mc_starvation = np.full_like(self.lc_pops, False, dtype=np.bool)
        self.uc_starvation = np.full_like(self.lc_pops, False, dtype=np.bool)
        self.any_starvation = np.full_like(self.lc_pops, False, dtype=np.bool)

        terrain = verkolst["terrain"]
        buildings = verkolst["buildings"]

        food_rate = np.full(self.shape, 0, dtype=np.float32)
        food_production_cap = np.full(self.shape, 0, dtype=np.float32)
        trade_impedance = np.full(self.shape, 0, dtype=np.float32)
        stone_production = np.full(self.shape, 0, dtype=np.float32)
        wood_production = np.full(self.shape, 0, dtype=np.float32)

        trade_modifier = np.full(self.shape, 0, dtype=np.float32)
        pops_cap = np.full(self.shape, 0, dtype=np.float32)

        for i, j in product(range(self.shape[0]), range(self.shape[1])):
            food_rate[i][j] = terrain.cells[i][j].tile.properties["food_rate"]
            food_production_cap[i][j] = terrain.cells[i][j].tile.properties["food_production_cap"]
            trade_impedance[i][j] = terrain.cells[i][j].tile.properties["trade_impedance"]
            trade_modifier[i][j] = buildings.cells[i][j].tile.properties["trade_modifier"]
            pops_cap[i][j] = terrain.cells[i][j].tile.properties["pops_cap"]

            stone_production[i][j] = terrain.cells[i][j].tile.properties["stone_production"]
            wood_production[i][j] = terrain.cells[i][j].tile.properties["wood_production"]

        self.goods_production = {
            'stone': stone_production,
            'wood': wood_production,
        }

        self.food_rate = food_rate
        self.food_production_cap = food_production_cap
        self.food_production = np.full_like(self.lc_pops, 0)
        self.food_imports = np.full_like(self.lc_pops, 0)
        self.update_food_production()
        self.food_consumption = 5 * self.uc_pops + 2 * self.mc_pops + 1 * self.lc_pops
        self.food_stored = 2 * self.food_production
        self.food_stored_capacity = np.full_like(self.lc_pops, 5000)

        self.trade_impedance = trade_impedance
        self.trade_range = np.full_like(self.lc_pops, self.base_trade_range, dtype=np.float32)
        self.trade_modifier = trade_modifier

        self.export_partner = np.full((*self.shape, 2), 0)
        # self.update_export_partners()

        self.pops_cap = pops_cap

        # Now process a single turn to sync up everything
        self.process_turn()

    def load(self, save_dir):
        with open(join(save_dir, "pops.json"), "r") as f:
            obs = json.load(f)
            self.lc_pops = pickle.loads(obs["lc_pops"])
            self.mc_pops = pickle.loads(obs["mc_pops"])
            self.uc_pops = pickle.loads(obs["uc_pops"])

    def dump(self, save_dir):
        with open(join(save_dir, "pops.json"), "w") as f:
            json.dump(
                {
                    "lc_pops": self.lc_pops.dumps(),
                    "mc_pops": self.mc_pops.dumps(),
                    "uc_pops": self.uc_pops.dumps(),
                },
                f,
            )

    def update_food_production(self):
        (self.food_rate * self.lc_pops).clip(
            min=0, max=self.food_production_cap, out=self.food_production
        )

    def update_trade_value(self):
        self.trade_value = (self.lc_pops*1 + self.mc_pops*5 + self.uc_pops*20).astype(np.float32) * self.trade_modifier

    def update_export_partners(self):
        from app.trade.trade_utils import update_export_partner
        update_export_partner(self.trade_range, self.trade_impedance, self.trade_value, self.export_partner)

    def eat_food(self):
        # from app.trade.trade_utils import share_food
        # First, add any produced food to the stockpile
        self.food_stored += self.food_production
        # share_food(self.food_production, self.food_consumption, self.food_stored, self.food_stored_capacity,
        #            self.export_partner, 0.1)

        # Next eat the food from uc->mc->lc marking starvation as it occurs
        self.food_stored -= self.uc_pops * 5
        self.uc_starvation = self.food_stored <= 0

        self.food_stored -= self.mc_pops * 2
        self.mc_starvation = self.food_stored <= 0

        self.food_stored -= self.lc_pops * 1
        self.lc_starvation = self.food_stored <= 0

        self.any_starvation = np.any([self.lc_starvation, self.mc_starvation, self.uc_starvation], axis=0)

        # Finally, clip the amount of food stored to between 0 and the food storage capacity
        self.food_stored.clip(
            min=0, max=self.food_stored_capacity, out=self.food_stored
        )

    def redistribute_classes(self):
        total_pops = self.lc_pops + self.mc_pops + self.uc_pops
        current_proportions = np.divide(np.stack([self.lc_pops, self.mc_pops, self.uc_pops]), total_pops)
        np.nan_to_num(current_proportions, copy=False, nan=0, posinf=0)

        target_proportions = np.stack([
            np.full_like(self.lc_pops, 0.8, dtype=np.float32),
            np.full_like(self.lc_pops, 0.15, dtype=np.float32),
            np.full_like(self.lc_pops, 0.05, dtype=np.float32),
        ])

        target_proportions[0, self.any_starvation] = 1
        target_proportions[1, self.any_starvation] = 0
        target_proportions[2, self.any_starvation] = 0

        new_proportions = current_proportions + (target_proportions - current_proportions) * 0.1
        self.lc_pops = np.ceil(np.multiply(total_pops, new_proportions[0])).astype(np.int32)
        self.mc_pops = np.ceil(np.multiply(total_pops, new_proportions[1])).astype(np.int32)
        self.uc_pops = np.ceil(np.multiply(total_pops, new_proportions[2])).astype(np.int32)

    def make_babies(self):

        food_stored_proportion = self.food_stored / self.food_stored_capacity

        repro_rate = 0.05 * food_stored_proportion
        repro_rate[self.any_starvation] = 0.0
        repro_rate[self.lc_starvation] = -0.05
        np.add(self.lc_pops, np.sign(repro_rate) * np.random.binomial(self.lc_pops,  np.abs(repro_rate)),
               out=self.lc_pops, casting='unsafe')

        repro_rate = 0.02 * food_stored_proportion
        repro_rate[self.any_starvation] = 0.0
        repro_rate[self.mc_starvation] = -0.05
        np.add(self.mc_pops, np.sign(repro_rate) * np.random.binomial(self.mc_pops,  np.abs(repro_rate)),
               out=self.mc_pops, casting='unsafe')

        repro_rate = 0.01 * food_stored_proportion
        repro_rate[self.any_starvation] = 0.0
        repro_rate[self.uc_starvation] = -0.05
        np.add(self.uc_pops, np.sign(repro_rate) * np.random.binomial(self.uc_pops,  np.abs(repro_rate)),
               out=self.uc_pops, casting='unsafe')

        overpop_fraction = np.divide(self.pops_cap, self.lc_pops+self.mc_pops+self.uc_pops)
        np.nan_to_num(overpop_fraction, copy=False, nan=1.0, posinf=1.0, neginf=1.0)
        overpop_fraction.clip(max=1.0, out=overpop_fraction)

        self.lc_pops = np.multiply(self.lc_pops, overpop_fraction).astype(np.int32)
        self.mc_pops = np.multiply(self.mc_pops, overpop_fraction).astype(np.int32)
        self.uc_pops = np.multiply(self.uc_pops, overpop_fraction).astype(np.int32)

        self.food_consumption = 5 * self.uc_pops + 2 * self.mc_pops + 1 * self.lc_pops

    def process_turn(self):
        self.eat_food()
        self.redistribute_classes()
        self.make_babies()
        self.update_food_production()
        self.update_trade_value()
        self.update_export_partners()