cocostrat/app/economy.py

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


class Resources:
    def __init__(self, rows=10, cols=10):
        self.shape = shape = (rows, cols)
        self.lc_pops = np.zeros(shape, dtype=np.float32)
        self.mc_pops = np.zeros(shape, dtype=np.float32)
        self.uc_pops = np.zeros(shape, dtype=np.float32)

        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_capacity = np.zeros(shape, dtype=np.float32)
        self.food_production = 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_distance = np.zeros(shape, dtype=np.float32)
        self.trade_range = np.full(shape, 5)
        self.trade_value = np.zeros(shape, dtype=np.float32)

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

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

        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)

        # Lets calculate the base food production
        terrain = verkolst["terrain"]

        food_rate = np.full(self.shape, 0, dtype=np.float32)
        food_capacity = np.full(self.shape, 0, dtype=np.float32)
        trade_distance = 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_capacity[i][j] = terrain.cells[i][j].tile.properties["food_capacity"]
            trade_distance[i][j] = terrain.cells[i][j].tile.properties["trade_distance"]

        self.food_rate = food_rate
        self.food_production_capacity = food_capacity
        self.food_production = 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_distance = trade_distance
        self.trade_range = np.full_like(self.lc_pops, 5)
        self.update_trade_value()

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

    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 * np.floor(self.lc_pops)).clip(
            min=0, max=self.food_production_capacity, out=self.food_production
        )

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

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

    def eat_food(self):
        # First, add any produced food to the stockpile
        self.food_stored += self.food_production

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

        self.food_stored -= np.floor(self.mc_pops) * 2
        self.mc_starvation = self.food_stored <= 0

        self.food_stored -= np.floor(self.lc_pops) * 1
        self.lc_starvation = self.food_stored <= 0

        self.any_starvation = np.logical_or(
            self.lc_starvation, np.logical_or(self.mc_starvation, self.uc_starvation)
        )

        # 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.stack([self.lc_pops, self.mc_pops, self.uc_pops]) / (
            total_pops
        )
        target_proportions = np.stack([
            np.full_like(self.lc_pops, 0.8),
            np.full_like(self.lc_pops, 0.15),
            np.full_like(self.lc_pops, 0.05),
        ])
        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 = total_pops * new_proportions[0]
        self.mc_pops = total_pops * new_proportions[1]
        self.uc_pops = total_pops * new_proportions[2]

    def make_babies(self):

        food_stored_proportion = self.food_stored / self.food_stored_capacity

        repro_rate = 1 + 0.10 * food_stored_proportion
        repro_rate[self.any_starvation] = 1.0
        repro_rate[self.lc_starvation] = 0.95
        np.multiply(self.lc_pops, repro_rate, out=self.lc_pops, casting="unsafe")

        repro_rate = 1 + 0.08 * food_stored_proportion
        repro_rate[self.any_starvation] = 1.0
        repro_rate[self.mc_starvation] = 0.95
        np.multiply(self.mc_pops, repro_rate, out=self.mc_pops, casting="unsafe")

        repro_rate = 1 + 0.06 * food_stored_proportion
        repro_rate[self.any_starvation] = 1.0
        repro_rate[self.uc_starvation] = 0.95
        np.multiply(self.uc_pops, repro_rate, out=self.uc_pops, casting="unsafe")

        self.food_consumption = (
            5 * np.floor(self.uc_pops)
            + 2 * np.floor(self.mc_pops)
            + 1 * np.floor(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()