<commit_before>#!/usr/bin/env python

# Copyright (C) 2014 SignalFuse, Inc.

# Maestro extension that can be used for wrapping the execution of a
# long-running service with log management scaffolding, potentially sending the
# log output to a file, or through Pipestash, or both, depending on the use
# case and parameters.

import os
import random
import signal
import subprocess

from ...guestutils import get_container_name, get_service_name, get_node_list

def run_service(cmd, logbase=None, logtarget=None):
    """Wrap the execution of a service with the necessary logging nets.

    If logbase is provided (it is by default), log output will be redirected
    (or teed) to a file named after the container executing the service inside
    the logbase directory.

    If Redis nodes are available in the environment as referenced by the given
    logtarget, log output will be streamed via pipestash to one of the
    available node containers, chosen at random when the service starts.

    The way this is accomplished varied on whether logbase is provided or not,
    and whether Redis nodes are available:

        - if neither, log output flows to stdout and will be captured by
          Docker;
        - if logbase is provided, but no Redis nodes are available, the
          output of the service is directly redirected to the log file;
        - if logbase is not provided, but Redis nodes are available, the
          output of the service is piped to pipestash;
        - if logbase is provided and Redis nodes are available, the output
          of the service is piped to a tee that will write the log file, and
          the output of the tee is piped to pipestash.

    The whole pipeline, whatever its construct is, waits for the service to
    terminate. SIGTERM is also redirected from the parent to the service.
    """
    if type(cmd) == str:
        cmd = cmd.split(' ')

    log = logbase \
        and os.path.join(logbase, '{}.log'.format(get_container_name())) \
        or None
    if logbase and not os.path.exists(logbase):
        os.makedirs(logbase)

    redis = logtarget \
        and get_node_list(logtarget, ports=['redis'], minimum=0) \
        or None
    stdout = redis and subprocess.PIPE or (log and open(log, 'w+') or None)

    # Start the service with the provided command.
    service = subprocess.Popen(cmd, stdout=stdout,
                               stderr=subprocess.STDOUT)
    last = service

    # Connect SIGTERM to the service process.
    signal.signal(signal.SIGTERM, lambda signum, frame: service.terminate())

    if redis:
        if log:
            # Tee to a local log file.
            tee = subprocess.Popen(['tee', log], stdin=last.stdout,
                                   stdout=subprocess.PIPE)
            last.stdout.close()
            last = tee

        pipestash = subprocess.Popen(
            ['pipestash', '-t', 'log',
             '-r', 'redis://{}/0'.format(random.choice(redis)),
             '-R', 'logstash',
             '-f', 'service={}'.format(get_service_name()),
             '-S', get_container_name()],
            stdin=last.stdout)
        last.stdout.close()
        last = pipestash

    # Wait for the service to exit and return its return code.
    last.communicate()
    return service.wait()
<commit_msg>Allow for the logstash logtype to be specified in run_service()<commit_after>#!/usr/bin/env python

# Copyright (C) 2014 SignalFuse, Inc.

# Maestro extension that can be used for wrapping the execution of a
# long-running service with log management scaffolding, potentially sending the
# log output to a file, or through Pipestash, or both, depending on the use
# case and parameters.

import os
import random
import signal
import subprocess

from ...guestutils import get_container_name, get_service_name, get_node_list

def run_service(cmd, logtype='log', logbase=None, logtarget=None):
    """Wrap the execution of a service with the necessary logging nets.

    If logbase is provided (it is by default), log output will be redirected
    (or teed) to a file named after the container executing the service inside
    the logbase directory.

    If Redis nodes are available in the environment as referenced by the given
    logtarget, log output will be streamed via pipestash to one of the
    available node containers, chosen at random when the service starts.

    The way this is accomplished varied on whether logbase is provided or not,
    and whether Redis nodes are available:

        - if neither, log output flows to stdout and will be captured by
          Docker;
        - if logbase is provided, but no Redis nodes are available, the
          output of the service is directly redirected to the log file;
        - if logbase is not provided, but Redis nodes are available, the
          output of the service is piped to pipestash;
        - if logbase is provided and Redis nodes are available, the output
          of the service is piped to a tee that will write the log file, and
          the output of the tee is piped to pipestash.

    The whole pipeline, whatever its construct is, waits for the service to
    terminate. SIGTERM is also redirected from the parent to the service.
    """
    if type(cmd) == str:
        cmd = cmd.split(' ')

    log = logbase \
        and os.path.join(logbase, '{}.log'.format(get_container_name())) \
        or None
    if logbase and not os.path.exists(logbase):
        os.makedirs(logbase)

    redis = logtarget \
        and get_node_list(logtarget, ports=['redis'], minimum=0) \
        or None
    stdout = redis and subprocess.PIPE or (log and open(log, 'w+') or None)

    # Start the service with the provided command.
    service = subprocess.Popen(cmd, stdout=stdout,
                               stderr=subprocess.STDOUT)
    last = service

    # Connect SIGTERM to the service process.
    signal.signal(signal.SIGTERM, lambda signum, frame: service.terminate())

    if redis:
        if log:
            # Tee to a local log file.
            tee = subprocess.Popen(['tee', log], stdin=last.stdout,
                                   stdout=subprocess.PIPE)
            last.stdout.close()
            last = tee

        pipestash = subprocess.Popen(
            ['pipestash', '-t', logtype,
             '-r', 'redis://{}/0'.format(random.choice(redis)),
             '-R', 'logstash',
             '-f', 'service={}'.format(get_service_name()),
             '-S', get_container_name()],
            stdin=last.stdout)
        last.stdout.close()
        last = pipestash

    # Wait for the service to exit and return its return code.
    last.communicate()
    return service.wait()
<|endoftext|>


<commit_before>import collections
import datetime
import itertools
import random

from django.conf import settings
from django.db import models

from django.contrib.auth.models import User

from manoria.utils import weighted_choices

class Player(models.Model):

    user = models.OneToOneField(User)
    name = models.CharField(max_length=20, unique=True)

    # @@@ points

    def __unicode__(self):
        return self.name

    def resource_counts(self):
        # @@@ instance cache
        counts = []
        for row in self.playerresourcecount_set.distinct().values("kind"):
            kind = ResourceKind.objects.get(id=row["kind"])
            current = PlayerResourceCount.current(kind, player=self)
            counts.append(current)
        return counts

class Continent(models.Model):

    name = models.CharField(max_length=20)
    allocation = models.TextField()

    def __unicode__(self):
        return self.name

    @property
    def size(self):
        return settings.CONTINENT_SIZE

    def cells(self):
        for cell in self.settlement_set.all():
            yield cell

class Settlement(models.Model):

    name = models.CharField(max_length=20)
    kind = models.CharField(
        max_length=15,
        choices=[
            ("homestead", "Homestead"),
            ("hamlet", "Hamlet"),
            ("village", "Village"),
            ("town", "Town"),
        ],
        default="homestead",
    )
    player = models.ForeignKey(Player, related_name="settlements")
    continent = models.ForeignKey(Continent)

    # location on continent
    x = models.IntegerField()
    y = models.IntegerField()

    allocation = models.TextField()

    # @@@ points

    def __unicode__(self):
        return u"%s (%s)" % (self.name, self.player)

    @property
    def size(self):
        return settings.SETTLEMENT_SIZE

    def place(self):
        CX, CY = settings.CONTINENT_SIZE
        SX, SY = settings.SETTLEMENT_SIZE
        # @@@ need to test if continent is full otherwise an infinite loop
        # will occur
        y = None
        while y is None:
            x = random.randint(1, CX)
            S = set(range(1, CY+1)) - set([s.y for s in Settlement.objects.filter(x=x)])
            if S:
                y = random.choice(list(S))
        self.x = x
        self.y = y
        self.save()

        # mark x,y used on the continent
        self.continent.allocation += "%s%d,%d" % (" ", x, y)
        self.continent.save()

        for resource_kind in ResourceKind.objects.filter(player=False):
            self.settlementresourcecount_set.create(
                kind=resource_kind,
                count=0,
                natural_rate=0,
                rate_adjustment=0,
                timestamp=datetime.datetime.now(),
                limit=0,
            )

        def check_cell(settlement, x, y):
            if not 1 <= x <= SX or not 1 <= y <= SY:
                terrain = None
            else:
                try:
                    terrain = settlement.terrain.select_related("kind").get(x=x, y=y)
                except SettlementTerrain.DoesNotExist:
                    terrain = None
            return terrain

        allocation = []

        for i in range(settings.SETTLEMENT_RESOURCE_COUNT):
            occupied = True
            while occupied:
                x = random.randint(1, SX)
                y = random.randint(1, SY)
                # check if the randomly chosen x, y is not already occupied
                if not (x, y) in allocation:
                    break
            allocation.append((x, y))
            # check surrounding cells
            neighbors = [
                check_cell(self, x+1, y),
                check_cell(self, x-1, y),
                check_cell(self, x, y+1),
                check_cell(self, x, y-1),
                check_cell(self, x+1, y+1),
                check_cell(self, x-1, y-1),
                check_cell(self, x+1, y-1),
                check_cell(self, x-1, y+1),
            ]
            counts = collections.defaultdict(int)
            for neighbor in filter(bool, neighbors):
                counts[neighbor.kind.slug] += 1
            population = []
            for kind in SettlementTerrainKind.objects.all():
                population.append((kind, counts.get(kind.slug, 0) + 1))
            kind = weighted_choices(population, 1)[0]
            terrain = self.terrain.create(kind=kind, x=x, y=y)
            # create resource counts for what the terrain kind produces
            for resource_kind in kind.produces.all():
                count = random.randint(1, 50000)
                terrain.settlementterrainresourcecount_set.create(
                    kind=resource_kind,
                    count=count,
                    natural_rate=count / 100,
                    rate_adjustment=0,
                    timestamp=datetime.datetime.now(),
                    limit=count
                )

        self.allocation = " ".join(("%d,%d" % (x, y) for x, y in allocation))
        # for updating the allocation table
        self.save()

    def cells(self):
        cells = itertools.chain(
            self.build_queue(), self.buildings(), self.terrain.all()
        )
        for cell in cells:
            yield cell

    def build_queue(self):
        queue = SettlementBuilding.objects.filter(
            settlement=self,
            construction_end__gt=datetime.datetime.now()
        )
        queue = queue.order_by("construction_start")
        return queue

    def buildings(self):
        return SettlementBuilding.objects.filter(
            settlement=self,
            construction_end__lte=datetime.datetime.now()
        )

    def resource_counts(self):
        # @@@ instance cache
        counts = []
        for row in self.settlementresourcecount_set.distinct().values("kind"):
            kind = ResourceKind.objects.get(id=row["kind"])
            current = SettlementResourceCount.current(kind, settlement=self)
            counts.append(current)
        return counts

class ResourceKind(models.Model):

    name = models.CharField(max_length=25)
    slug = models.SlugField()
    player = models.BooleanField()

    def __unicode__(self):
        return self.name

def pairwise(iterable):
    """
    pulled from itertools recipes, but modified to return last item and None
    s -> (s0,s1), (s1,s2), (s2, s3), (s3, None)
    """
    a, b = itertools.tee(iterable)
    next(b, None)
    for a, b in itertools.izip(a, b):
        yield a, b
    yield b, None

class BaseResourceCount(models.Model):

    count = models.IntegerField(default=0)
    timestamp = models.DateTimeField(default=datetime.datetime.now)
    natural_rate = models.DecimalField(max_digits=7, decimal_places=1)
    rate_adjustment = models.DecimalField(max_digits=7, decimal_places=1)
    limit = models.IntegerField(default=0)

    class Meta:
        abstract = True

    @classmethod
    def current(cls, kind, **kwargs):
        when = kwargs.pop("when", datetime.datetime.now())
        lookup_params = {
            "kind": kind,
            "timestamp__lt": when,
        }
        lookup_params.update(kwargs)
        past = cls._default_manager.filter(**lookup_params).order_by("-timestamp")
        return past[0]

    @classmethod
    def calculate_extremum(cls, kind, **kwargs):
        current = cls.current(kind, **kwargs)
        del kwargs["when"]
        lookup_params = {
            "kind": kind,
            "timestamp__gte": current.timestamp,
        }
        lookup_params.update(kwargs)
        future_counts = cls._default_manager.filter(**lookup_params).order_by("timestamp")
        for a, b in pairwise(future_counts):
            if b is None:
                start, end = a.timestamp, None
            else:
                start, end = a.timestamp, b.timestamp
            count = a.count
            rate = a.rate
            limit = a.limit
            if a.rate < 0:
                # when the count will hit zero
                timestamp = start + datetime.timedelta(hours=(count / -float(rate)))
                if end is not None and timestamp > end:
                    continue
                return timestamp, False
            elif a.rate > 0:
                # when the count will hit the limit
                timestamp = start + datetime.timedelta(hours=((limit - count) / float(rate)))
                if end is not None and timestamp > end:
                    continue
                return timestamp, True
            else:
                continue
        return None, None

    @property
    def rate(self):
        return self.natural_rate + self.rate_adjustment

    def amount(self, when=None):
        if when is None:
            when = datetime.datetime.now()
        change = when - self.timestamp
        amt = int(self.count + float(self.rate) * (change.days * 86400 + change.seconds) / 3600.0)
        if self.limit == 0:
            return max(0, amt)
        else:
            return min(max(0, amt), self.limit)

class PlayerResourceCount(BaseResourceCount):

    kind = models.ForeignKey(ResourceKind)
    player = models.ForeignKey(Player)

    def __unicode__(self):
        return u"%s (%s)" % (self.kind, self.player)

class SettlementResourceCount(BaseResourceCount):

    kind = models.ForeignKey(ResourceKind)
    settlement = models.ForeignKey(Settlement)

    def __unicode__(self):
        return u"%s (%s)" % (self.kind, self.settlement)

class BuildingKind(models.Model):

    name = models.CharField(max_length=30)
    slug = models.SlugField()

    def __unicode__(self):
        return self.name

class BuildingKindProduct(models.Model):

    building_kind = models.ForeignKey(BuildingKind, related_name="products")
    resource_kind = models.ForeignKey(ResourceKind, related_name="produced_by")
    source_terrain_kind = models.ForeignKey("SettlementTerrainKind", null=True)
    base_rate = models.IntegerField()

    def __unicode__(self):
        bits = []
        bits.append("%s produces %s" % (self.building_kind, self.resource_kind))
        if self.source_terrain_kind:
            bits.append("from %s" % self.source_terrain_kind)
        bits.append("at %d/hr" % self.base_rate)
        return " ".join(bits)

class SettlementBuilding(models.Model):

    kind = models.ForeignKey(BuildingKind)
    settlement = models.ForeignKey(Settlement)

    # location in settlement
    x = models.IntegerField()
    y = models.IntegerField()

    # build queue
    construction_start = models.DateTimeField(default=datetime.datetime.now)
    construction_end = models.DateTimeField(default=datetime.datetime.now)

    class Meta:
        unique_together = [("settlement", "x", "y")]

    def __unicode__(self):
        return u"%s on %s" % (self.kind, self.settlement)

    def queue(self, commit=True):
        # look for most recently added building to queue (None if none)
        try:
            oldest = self.settlement.build_queue().reverse()[0]
        except IndexError:
            oldest = None

        # set when construction starts based on the most recently added
        # building to the queue (if None use now)
        # @@@ hard-coded two minute build times
        if oldest:
            self.construction_start = oldest.construction_end
        else:
            self.construction_start = datetime.datetime.now()
        self.construction_end = self.construction_start + datetime.timedelta(minutes=2)

        # allocate space on the map using settlement allocation table
        self.settlement.allocation += "%s%d,%d" % (" ", self.x, self.y)
        self.settlement.save()

        # iterate over what the building kind produces setting up the state
        # of resource counts when the building will be finished building
        for product in self.kind.products.all():
            # establish some common params used while creating resource counts
            common_params = {}
            if product.resource_kind.player:
                ResourceCount = PlayerResourceCount
                common_params["player"] = self.settlement.player
            else:
                ResourceCount = SettlementResourceCount
                common_params["settlement"] = self.settlement

            # adjust the settlement/player resource count based on what the building
            # will produce in the best case scenario.
            lookup_params = {
                "when": self.construction_end,
            }
            lookup_params.update(common_params)
            current = ResourceCount.current(product.resource_kind, **lookup_params)
            create_kwargs = {
                "kind": product.resource_kind,
                "count": current.amount(self.construction_end),
                "timestamp": self.construction_end,
                "natural_rate": current.natural_rate,
                "rate_adjustment": current.rate_adjustment + product.base_rate,
                "limit": 0, # @@@ storage
            }
            create_kwargs.update(common_params)
            ResourceCount.objects.create(**create_kwargs)

            # find the closest terrain which produces what the building needs
            # to be productive.
            # @@@ change to work based on aadjacency instead of a single
            # closest terrain
            terrain = SettlementTerrain.objects.filter(
                settlement=self.settlement, kind=product.source_terrain_kind
            )
            closest = sorted([
                (((self.x - t.x) ** 2) + ((self.y - t.y) ** 2), t)
                for t in terrain
            ])[0][1]

            # when the building finishes set what affect it will have on the
            # terrain resource count
            current = SettlementTerrainResourceCount.current(
                product.resource_kind,
                terrain=closest, when=self.construction_end
            )
            SettlementTerrainResourceCount.objects.create(
                kind=product.resource_kind,
                terrain=closest,
                count=current.amount(self.construction_end),
                timestamp=self.construction_end,
                natural_rate=current.natural_rate,
                rate_adjustment=current.rate_adjustment - product.base_rate,
                limit=current.limit,
            )

            # determine *when* the terrain will either run out or hit its
            # limit based on its current rate and if the terrain will hit
            # zero by the time the building finishes we need to adjust how
            # it affects the settlement resource counts
            when, hit_limit = SettlementTerrainResourceCount.calculate_extremum(
                product.resource_kind,
                terrain=closest, when=self.construction_end
            )
            # if when is None the terrain will never run out or hit a limit
            if when and not hit_limit:
                lookup_params = {
                    "when": when,
                }
                lookup_params.update(common_params)
                current = ResourceCount.current(product.resource_kind, **lookup_params)
                create_kwargs = {
                    "kind": product.resource_kind,
                    "settlement": self.settlement,
                    "count": current.amount(when),
                    "timestamp": when,
                    "natural_rate": current.natural_rate,
                    "rate_adjustment": current.rate_adjustment - product.base_rate,
                    "limit": 0, # @@@ storage
                }
                create_kwargs.update(common_params)
                ResourceCount.objects.create(**create_kwargs)

        if commit:
            self.save()

    def status(self):
        now = datetime.datetime.now()
        if self.construction_start > now:
            return "queued"
        elif self.construction_end > now:
            return "under construction"
        else:
            return "built"

class SettlementBuildingResourceCount(BaseResourceCount):

    building = models.ForeignKey(SettlementBuilding)

class SettlementTerrainKind(models.Model):

    name = models.CharField(max_length=50)
    slug = models.SlugField()
    buildable_on = models.BooleanField(default=True)
    produces = models.ManyToManyField(ResourceKind)

    def __unicode__(self):
        return self.name

class SettlementTerrain(models.Model):

    kind = models.ForeignKey(SettlementTerrainKind)
    settlement = models.ForeignKey(Settlement, related_name="terrain")

    # location in settlement
    x = models.IntegerField()
    y = models.IntegerField()

    def __unicode__(self):
        return u"%s on %s" % (self.kind, self.settlement)

    def resource_counts(self):
        # @@@ instance cache
        counts = []
        for row in self.settlementterrainresourcecount_set.distinct().values("kind"):
            kind = ResourceKind.objects.get(id=row["kind"])
            current = SettlementTerrainResourceCount.current(kind, terrain=self)
            counts.append(current)
        return counts

class SettlementTerrainResourceCount(BaseResourceCount):

    kind = models.ForeignKey(ResourceKind)
    terrain = models.ForeignKey(SettlementTerrain)
<commit_msg>fixed farm creation (ala buildings which produce by themselves)<commit_after>import collections
import datetime
import itertools
import random

from django.conf import settings
from django.db import models

from django.contrib.auth.models import User

from manoria.utils import weighted_choices

class Player(models.Model):

    user = models.OneToOneField(User)
    name = models.CharField(max_length=20, unique=True)

    # @@@ points

    def __unicode__(self):
        return self.name

    def resource_counts(self):
        # @@@ instance cache
        counts = []
        for row in self.playerresourcecount_set.distinct().values("kind"):
            kind = ResourceKind.objects.get(id=row["kind"])
            current = PlayerResourceCount.current(kind, player=self)
            counts.append(current)
        return counts

class Continent(models.Model):

    name = models.CharField(max_length=20)
    allocation = models.TextField()

    def __unicode__(self):
        return self.name

    @property
    def size(self):
        return settings.CONTINENT_SIZE

    def cells(self):
        for cell in self.settlement_set.all():
            yield cell

class Settlement(models.Model):

    name = models.CharField(max_length=20)
    kind = models.CharField(
        max_length=15,
        choices=[
            ("homestead", "Homestead"),
            ("hamlet", "Hamlet"),
            ("village", "Village"),
            ("town", "Town"),
        ],
        default="homestead",
    )
    player = models.ForeignKey(Player, related_name="settlements")
    continent = models.ForeignKey(Continent)

    # location on continent
    x = models.IntegerField()
    y = models.IntegerField()

    allocation = models.TextField()

    # @@@ points

    def __unicode__(self):
        return u"%s (%s)" % (self.name, self.player)

    @property
    def size(self):
        return settings.SETTLEMENT_SIZE

    def place(self):
        CX, CY = settings.CONTINENT_SIZE
        SX, SY = settings.SETTLEMENT_SIZE
        # @@@ need to test if continent is full otherwise an infinite loop
        # will occur
        y = None
        while y is None:
            x = random.randint(1, CX)
            S = set(range(1, CY+1)) - set([s.y for s in Settlement.objects.filter(x=x)])
            if S:
                y = random.choice(list(S))
        self.x = x
        self.y = y
        self.save()

        # mark x,y used on the continent
        self.continent.allocation += "%s%d,%d" % (" ", x, y)
        self.continent.save()

        for resource_kind in ResourceKind.objects.filter(player=False):
            self.settlementresourcecount_set.create(
                kind=resource_kind,
                count=0,
                natural_rate=0,
                rate_adjustment=0,
                timestamp=datetime.datetime.now(),
                limit=0,
            )

        def check_cell(settlement, x, y):
            if not 1 <= x <= SX or not 1 <= y <= SY:
                terrain = None
            else:
                try:
                    terrain = settlement.terrain.select_related("kind").get(x=x, y=y)
                except SettlementTerrain.DoesNotExist:
                    terrain = None
            return terrain

        allocation = []

        for i in range(settings.SETTLEMENT_RESOURCE_COUNT):
            occupied = True
            while occupied:
                x = random.randint(1, SX)
                y = random.randint(1, SY)
                # check if the randomly chosen x, y is not already occupied
                if not (x, y) in allocation:
                    break
            allocation.append((x, y))
            # check surrounding cells
            neighbors = [
                check_cell(self, x+1, y),
                check_cell(self, x-1, y),
                check_cell(self, x, y+1),
                check_cell(self, x, y-1),
                check_cell(self, x+1, y+1),
                check_cell(self, x-1, y-1),
                check_cell(self, x+1, y-1),
                check_cell(self, x-1, y+1),
            ]
            counts = collections.defaultdict(int)
            for neighbor in filter(bool, neighbors):
                counts[neighbor.kind.slug] += 1
            population = []
            for kind in SettlementTerrainKind.objects.all():
                population.append((kind, counts.get(kind.slug, 0) + 1))
            kind = weighted_choices(population, 1)[0]
            terrain = self.terrain.create(kind=kind, x=x, y=y)
            # create resource counts for what the terrain kind produces
            for resource_kind in kind.produces.all():
                count = random.randint(1, 50000)
                terrain.settlementterrainresourcecount_set.create(
                    kind=resource_kind,
                    count=count,
                    natural_rate=count / 100,
                    rate_adjustment=0,
                    timestamp=datetime.datetime.now(),
                    limit=count
                )

        self.allocation = " ".join(("%d,%d" % (x, y) for x, y in allocation))
        # for updating the allocation table
        self.save()

    def cells(self):
        cells = itertools.chain(
            self.build_queue(), self.buildings(), self.terrain.all()
        )
        for cell in cells:
            yield cell

    def build_queue(self):
        queue = SettlementBuilding.objects.filter(
            settlement=self,
            construction_end__gt=datetime.datetime.now()
        )
        queue = queue.order_by("construction_start")
        return queue

    def buildings(self):
        return SettlementBuilding.objects.filter(
            settlement=self,
            construction_end__lte=datetime.datetime.now()
        )

    def resource_counts(self):
        # @@@ instance cache
        counts = []
        for row in self.settlementresourcecount_set.distinct().values("kind"):
            kind = ResourceKind.objects.get(id=row["kind"])
            current = SettlementResourceCount.current(kind, settlement=self)
            counts.append(current)
        return counts

class ResourceKind(models.Model):

    name = models.CharField(max_length=25)
    slug = models.SlugField()
    player = models.BooleanField()

    def __unicode__(self):
        return self.name

def pairwise(iterable):
    """
    pulled from itertools recipes, but modified to return last item and None
    s -> (s0,s1), (s1,s2), (s2, s3), (s3, None)
    """
    a, b = itertools.tee(iterable)
    next(b, None)
    for a, b in itertools.izip(a, b):
        yield a, b
    yield b, None

class BaseResourceCount(models.Model):

    count = models.IntegerField(default=0)
    timestamp = models.DateTimeField(default=datetime.datetime.now)
    natural_rate = models.DecimalField(max_digits=7, decimal_places=1)
    rate_adjustment = models.DecimalField(max_digits=7, decimal_places=1)
    limit = models.IntegerField(default=0)

    class Meta:
        abstract = True

    @classmethod
    def current(cls, kind, **kwargs):
        when = kwargs.pop("when", datetime.datetime.now())
        lookup_params = {
            "kind": kind,
            "timestamp__lt": when,
        }
        lookup_params.update(kwargs)
        past = cls._default_manager.filter(**lookup_params).order_by("-timestamp")
        return past[0]

    @classmethod
    def calculate_extremum(cls, kind, **kwargs):
        current = cls.current(kind, **kwargs)
        del kwargs["when"]
        lookup_params = {
            "kind": kind,
            "timestamp__gte": current.timestamp,
        }
        lookup_params.update(kwargs)
        future_counts = cls._default_manager.filter(**lookup_params).order_by("timestamp")
        for a, b in pairwise(future_counts):
            if b is None:
                start, end = a.timestamp, None
            else:
                start, end = a.timestamp, b.timestamp
            count = a.count
            rate = a.rate
            limit = a.limit
            if a.rate < 0:
                # when the count will hit zero
                timestamp = start + datetime.timedelta(hours=(count / -float(rate)))
                if end is not None and timestamp > end:
                    continue
                return timestamp, False
            elif a.rate > 0:
                # when the count will hit the limit
                timestamp = start + datetime.timedelta(hours=((limit - count) / float(rate)))
                if end is not None and timestamp > end:
                    continue
                return timestamp, True
            else:
                continue
        return None, None

    @property
    def rate(self):
        return self.natural_rate + self.rate_adjustment

    def amount(self, when=None):
        if when is None:
            when = datetime.datetime.now()
        change = when - self.timestamp
        amt = int(self.count + float(self.rate) * (change.days * 86400 + change.seconds) / 3600.0)
        if self.limit == 0:
            return max(0, amt)
        else:
            return min(max(0, amt), self.limit)

class PlayerResourceCount(BaseResourceCount):

    kind = models.ForeignKey(ResourceKind)
    player = models.ForeignKey(Player)

    def __unicode__(self):
        return u"%s (%s)" % (self.kind, self.player)

class SettlementResourceCount(BaseResourceCount):

    kind = models.ForeignKey(ResourceKind)
    settlement = models.ForeignKey(Settlement)

    def __unicode__(self):
        return u"%s (%s)" % (self.kind, self.settlement)

class BuildingKind(models.Model):

    name = models.CharField(max_length=30)
    slug = models.SlugField()

    def __unicode__(self):
        return self.name

class BuildingKindProduct(models.Model):

    building_kind = models.ForeignKey(BuildingKind, related_name="products")
    resource_kind = models.ForeignKey(ResourceKind, related_name="produced_by")
    source_terrain_kind = models.ForeignKey("SettlementTerrainKind", null=True)
    base_rate = models.IntegerField()

    def __unicode__(self):
        bits = []
        bits.append("%s produces %s" % (self.building_kind, self.resource_kind))
        if self.source_terrain_kind:
            bits.append("from %s" % self.source_terrain_kind)
        bits.append("at %d/hr" % self.base_rate)
        return " ".join(bits)

class SettlementBuilding(models.Model):

    kind = models.ForeignKey(BuildingKind)
    settlement = models.ForeignKey(Settlement)

    # location in settlement
    x = models.IntegerField()
    y = models.IntegerField()

    # build queue
    construction_start = models.DateTimeField(default=datetime.datetime.now)
    construction_end = models.DateTimeField(default=datetime.datetime.now)

    class Meta:
        unique_together = [("settlement", "x", "y")]

    def __unicode__(self):
        return u"%s on %s" % (self.kind, self.settlement)

    def queue(self, commit=True):
        # look for most recently added building to queue (None if none)
        try:
            oldest = self.settlement.build_queue().reverse()[0]
        except IndexError:
            oldest = None

        # set when construction starts based on the most recently added
        # building to the queue (if None use now)
        # @@@ hard-coded two minute build times
        if oldest:
            self.construction_start = oldest.construction_end
        else:
            self.construction_start = datetime.datetime.now()
        self.construction_end = self.construction_start + datetime.timedelta(minutes=2)

        # allocate space on the map using settlement allocation table
        self.settlement.allocation += "%s%d,%d" % (" ", self.x, self.y)
        self.settlement.save()

        # iterate over what the building kind produces setting up the state
        # of resource counts when the building will be finished building
        for product in self.kind.products.all():
            # establish some common params used while creating resource counts
            common_params = {}
            if product.resource_kind.player:
                ResourceCount = PlayerResourceCount
                common_params["player"] = self.settlement.player
            else:
                ResourceCount = SettlementResourceCount
                common_params["settlement"] = self.settlement

            # adjust the settlement/player resource count based on what the building
            # will produce in the best case scenario.
            lookup_params = {
                "when": self.construction_end,
            }
            lookup_params.update(common_params)
            current = ResourceCount.current(product.resource_kind, **lookup_params)
            create_kwargs = {
                "kind": product.resource_kind,
                "count": current.amount(self.construction_end),
                "timestamp": self.construction_end,
                "natural_rate": current.natural_rate,
                "rate_adjustment": current.rate_adjustment + product.base_rate,
                "limit": 0, # @@@ storage
            }
            create_kwargs.update(common_params)
            ResourceCount.objects.create(**create_kwargs)

            if product.source_terrain_kind:
                # find the closest terrain which produces what the building needs
                # to be productive.
                # @@@ change to work based on aadjacency instead of a single
                # closest terrain
                terrain = SettlementTerrain.objects.filter(
                    settlement=self.settlement, kind=product.source_terrain_kind
                )
                closest = sorted([
                    (((self.x - t.x) ** 2) + ((self.y - t.y) ** 2), t)
                    for t in terrain
                ])[0][1]

                # when the building finishes set what affect it will have on the
                # terrain resource count
                current = SettlementTerrainResourceCount.current(
                    product.resource_kind,
                    terrain=closest, when=self.construction_end
                )
                SettlementTerrainResourceCount.objects.create(
                    kind=product.resource_kind,
                    terrain=closest,
                    count=current.amount(self.construction_end),
                    timestamp=self.construction_end,
                    natural_rate=current.natural_rate,
                    rate_adjustment=current.rate_adjustment - product.base_rate,
                    limit=current.limit,
                )

                # determine *when* the terrain will either run out or hit its
                # limit based on its current rate and if the terrain will hit
                # zero by the time the building finishes we need to adjust how
                # it affects the settlement resource counts
                when, hit_limit = SettlementTerrainResourceCount.calculate_extremum(
                    product.resource_kind,
                    terrain=closest, when=self.construction_end
                )
                # if when is None the terrain will never run out or hit a limit
                if when and not hit_limit:
                    lookup_params = {
                        "when": when,
                    }
                    lookup_params.update(common_params)
                    current = ResourceCount.current(product.resource_kind, **lookup_params)
                    create_kwargs = {
                        "kind": product.resource_kind,
                        "settlement": self.settlement,
                        "count": current.amount(when),
                        "timestamp": when,
                        "natural_rate": current.natural_rate,
                        "rate_adjustment": current.rate_adjustment - product.base_rate,
                        "limit": 0, # @@@ storage
                    }
                    create_kwargs.update(common_params)
                    ResourceCount.objects.create(**create_kwargs)

        if commit:
            self.save()

    def status(self):
        now = datetime.datetime.now()
        if self.construction_start > now:
            return "queued"
        elif self.construction_end > now:
            return "under construction"
        else:
            return "built"

class SettlementBuildingResourceCount(BaseResourceCount):

    building = models.ForeignKey(SettlementBuilding)

class SettlementTerrainKind(models.Model):

    name = models.CharField(max_length=50)
    slug = models.SlugField()
    buildable_on = models.BooleanField(default=True)
    produces = models.ManyToManyField(ResourceKind)

    def __unicode__(self):
        return self.name

class SettlementTerrain(models.Model):

    kind = models.ForeignKey(SettlementTerrainKind)
    settlement = models.ForeignKey(Settlement, related_name="terrain")

    # location in settlement
    x = models.IntegerField()
    y = models.IntegerField()

    def __unicode__(self):
        return u"%s on %s" % (self.kind, self.settlement)

    def resource_counts(self):
        # @@@ instance cache
        counts = []
        for row in self.settlementterrainresourcecount_set.distinct().values("kind"):
            kind = ResourceKind.objects.get(id=row["kind"])
            current = SettlementTerrainResourceCount.current(kind, terrain=self)
            counts.append(current)
        return counts

class SettlementTerrainResourceCount(BaseResourceCount):

    kind = models.ForeignKey(ResourceKind)
    terrain = models.ForeignKey(SettlementTerrain)
<|endoftext|>


<commit_before>import sys, os
import json
import argparse
import glob2 as glob

MAP_WORKING_DIR = os.path.dirname(os.path.realpath(__file__))
sys.path.insert(0, os.path.abspath(os.path.join(MAP_WORKING_DIR, '..')))

import common, pa2checker, frontend_pa_inference

JAIF_FILE_ONTOLOGY_HEADER ="\
package ontology.qual:\n\
  annotation @Ontology:\n\
    enum ontology.qual.OntologyValue[] values\n"

def field_mappings_to_annotation(json_file):
    with open(json_file) as data_file:
        data = json.load(data_file)

    mappings = data["mappings"]
    ontology_set = set()
    project_list = common.get_project_list()

    for mapping in mappings:
        ontology_set.add(mapping['label'][0])

    convert_2_ontology_value(ontology_set)

    corpus_jaif_file = create_corpus_jaif(mappings)

    # for project in project_list:
    #     insert_anno_to_project(project, corpus_jaif_file)

    # for project in project_list:
    #     frontend_pa_inference.run_inference(project)

def type_mappings_to_rules(json_file):
    with open(json_file) as data_file:
        data = json.load(data_file)

    for mapping in data["type_mappings"]:
        ontology_value = mapping["ontology_value"]
        java_types = mapping["java_types"]
        pa2checker.insert_ontology_value(ontology_value)
        pa2checker.update_ontology_utils(ontology_value, java_types)

    pa2checker.recompile_checker_framework()

def insert_anno_to_project(project, jaif_file):
    """ Insert annotation info in the ${jaif_file} to ${project}.
    """
    project_dir = common.get_project_dir(project)
    with common.cd(project_dir):
        common.setup_checker_framework_env()
        insert_cmd = ['insert-annotations-to-source', '-i', jaif_file]
        # using glob2.glob to recursive get java files under project dir
        java_files = glob.glob('{}/**/*.java'.format(project_dir))
        insert_cmd.extend(java_files)
        common.run_cmd(insert_cmd, print_output=True)

def create_corpus_jaif(mappings):
    return create_jaif_file("corpus", mappings)

def create_jaif_file(project, mappings):
    """ create a {project_name}.jaif file under project_dir
        this jair file contains the insertted annotations info for this project
        Note: if ${project} value is "corpus", then it will create a "corpus.jaif" under corpus dir
    """
    if project == "corpus":
        project_dir = common.CORPUS_DIR
    else:
        project_dir = common.get_project_dir(project)

    jaif_file = os.path.join(project_dir, "{}.jaif".format(project))

    print ("Writing project {} annotated info to file {}".format(project, jaif_file))

    with open(jaif_file, 'w') as out_file:
        # write ontology package info
        out_file.write(JAIF_FILE_ONTOLOGY_HEADER)
        jaif_dict = dict()
        for mapping in mappings:
            for qualified_field in mapping['fields']:
                (package, clazz, field) = parse_field(qualified_field)
                if not package in jaif_dict:
                    jaif_dict[package] = dict()
                if not clazz in jaif_dict[package]:
                    jaif_dict[package][clazz] = dict()
                if not field in jaif_dict[package][clazz]:
                    jaif_dict[package][clazz][field] = set()

                # assume 'label' at least has one element
                jaif_dict[package][clazz][field].add(mapping['label'][0].upper())

        for package, classes in jaif_dict.items():
            out_file.write("\npackage {}:\n".format(package))
            for clazz, fields in classes.items():
                out_file.write("  class {}:\n".format(clazz))
                for field, value_set in fields.items():
                    out_file.write("    field {}:\n".format(field))
                    out_file.write("    @Ontology(values={{{value_name}}})\n".format(value_name=', '.join(value_set)))
    return jaif_file

def parse_field(qualified_field):
    """ given a ${qualified_field} which describes a full qualified path to a class field with value like
        "xxx.xxx.xxx.Class.field", parse it to a tuple of package path, class name, and field name as
        (xxx.xxx.xxx, Class, field)
    """
    return tuple(qualified_field.rsplit('.', 2))

def convert_2_ontology_value(ontology_set):
    """ 1. clean up the ontology source code back to unchanged git version
        2. takes the ontology_set, and insert elements into ontology source as new ontologies
        3. re-compile ontology to make sure the insertions doesn't break the compilation
    """
    pa2checker.revert_checker_source()
    for new_ontology in ontology_set:
        pa2checker.insert_ontology_value(new_ontology)
    pa2checker.recompile_checker_framework()

def main():
    parser = argparse.ArgumentParser(description='command line interface for map2annotation')
    parser.add_argument('--type-mapping',dest='type_mapping_file')
    parser.add_argument('--field-mapping', dest='field_mapping_file')
    args = parser.parse_args()
    if args.type_mapping_file is None and args.field_mapping_file is None:
        print "error, required at least one mapping file to be indicated."
        parser.print_help()

    pa2checker.revert_checker_source()

    if not args.type_mapping_file is None:
        type_mappings_to_rules(args.type_mapping_file)

    if not args.field_mapping_file is None:
        field_mappings_to_annotation(args.field_mapping_file)

if __name__ == '__main__':
    main()
<commit_msg>in command line mode, run inference on projects after updating Ontology system.<commit_after>import sys, os
import json
import argparse
import glob2 as glob

MAP_WORKING_DIR = os.path.dirname(os.path.realpath(__file__))
sys.path.insert(0, os.path.abspath(os.path.join(MAP_WORKING_DIR, '..')))

import common, pa2checker, frontend_pa_inference

JAIF_FILE_ONTOLOGY_HEADER ="\
package ontology.qual:\n\
  annotation @Ontology:\n\
    enum ontology.qual.OntologyValue[] values\n"

def field_mappings_to_annotation(json_file):
    with open(json_file) as data_file:
        data = json.load(data_file)

    mappings = data["mappings"]
    ontology_set = set()
    project_list = common.get_project_list()

    for mapping in mappings:
        ontology_set.add(mapping['label'][0])

    convert_2_ontology_value(ontology_set)

    corpus_jaif_file = create_corpus_jaif(mappings)

    # for project in project_list:
    #     insert_anno_to_project(project, corpus_jaif_file)

    # for project in project_list:
    #     frontend_pa_inference.run_inference(project)

def type_mappings_to_rules(json_file):
    with open(json_file) as data_file:
        data = json.load(data_file)

    for mapping in data["type_mappings"]:
        ontology_value = mapping["ontology_value"]
        java_types = mapping["java_types"]
        pa2checker.insert_ontology_value(ontology_value)
        pa2checker.update_ontology_utils(ontology_value, java_types)

    pa2checker.recompile_checker_framework()

def insert_anno_to_project(project, jaif_file):
    """ Insert annotation info in the ${jaif_file} to ${project}.
    """
    project_dir = common.get_project_dir(project)
    with common.cd(project_dir):
        common.setup_checker_framework_env()
        insert_cmd = ['insert-annotations-to-source', '-i', jaif_file]
        # using glob2.glob to recursive get java files under project dir
        java_files = glob.glob('{}/**/*.java'.format(project_dir))
        insert_cmd.extend(java_files)
        common.run_cmd(insert_cmd, print_output=True)

def create_corpus_jaif(mappings):
    return create_jaif_file("corpus", mappings)

def create_jaif_file(project, mappings):
    """ create a {project_name}.jaif file under project_dir
        this jair file contains the insertted annotations info for this project
        Note: if ${project} value is "corpus", then it will create a "corpus.jaif" under corpus dir
    """
    if project == "corpus":
        project_dir = common.CORPUS_DIR
    else:
        project_dir = common.get_project_dir(project)

    jaif_file = os.path.join(project_dir, "{}.jaif".format(project))

    print ("Writing project {} annotated info to file {}".format(project, jaif_file))

    with open(jaif_file, 'w') as out_file:
        # write ontology package info
        out_file.write(JAIF_FILE_ONTOLOGY_HEADER)
        jaif_dict = dict()
        for mapping in mappings:
            for qualified_field in mapping['fields']:
                (package, clazz, field) = parse_field(qualified_field)
                if not package in jaif_dict:
                    jaif_dict[package] = dict()
                if not clazz in jaif_dict[package]:
                    jaif_dict[package][clazz] = dict()
                if not field in jaif_dict[package][clazz]:
                    jaif_dict[package][clazz][field] = set()

                # assume 'label' at least has one element
                jaif_dict[package][clazz][field].add(mapping['label'][0].upper())

        for package, classes in jaif_dict.items():
            out_file.write("\npackage {}:\n".format(package))
            for clazz, fields in classes.items():
                out_file.write("  class {}:\n".format(clazz))
                for field, value_set in fields.items():
                    out_file.write("    field {}:\n".format(field))
                    out_file.write("    @Ontology(values={{{value_name}}})\n".format(value_name=', '.join(value_set)))
    return jaif_file

def parse_field(qualified_field):
    """ given a ${qualified_field} which describes a full qualified path to a class field with value like
        "xxx.xxx.xxx.Class.field", parse it to a tuple of package path, class name, and field name as
        (xxx.xxx.xxx, Class, field)
    """
    return tuple(qualified_field.rsplit('.', 2))

def convert_2_ontology_value(ontology_set):
    """ 1. clean up the ontology source code back to unchanged git version
        2. takes the ontology_set, and insert elements into ontology source as new ontologies
        3. re-compile ontology to make sure the insertions doesn't break the compilation
    """
    pa2checker.revert_checker_source()
    for new_ontology in ontology_set:
        pa2checker.insert_ontology_value(new_ontology)
    pa2checker.recompile_checker_framework()

def main():
    parser = argparse.ArgumentParser(description='command line interface for map2annotation')
    parser.add_argument('--type-mapping',dest='type_mapping_file')
    parser.add_argument('--field-mapping', dest='field_mapping_file')
    args = parser.parse_args()
    if args.type_mapping_file is None and args.field_mapping_file is None:
        print "error, required at least one mapping file to be indicated."
        parser.print_help()

    pa2checker.revert_checker_source()

    if not args.type_mapping_file is None:
        type_mappings_to_rules(args.type_mapping_file)

    if not args.field_mapping_file is None:
        field_mappings_to_annotation(args.field_mapping_file)

    project_list = common.get_project_list()

    for project in project_list:
        frontend_pa_inference.run_inference(project)

if __name__ == '__main__':
    main()
<|endoftext|>


<commit_before># vim:fileencoding=utf-8

__version__ = '0.6.1'
<commit_msg>Bumping version<commit_after># vim:fileencoding=utf-8

__version__ = '0.7.0'
<|endoftext|>


<commit_before># see LICENSE file for copyright and license details

"""
Package merging functions and helpers
"""

import os
from os.path import join
from gzip import open as gzip_open
from lzma import open as lzma_open
from shutil import copyfile

from lib.config import mergedir, packages_keys, sources_keys, spooldir
from lib.parse import parse_dependencies, parse_packages

def write_packages(packages, filename, sort=True, sources=False):
    """
    Writes `packages` to a file (per debian Packages format)
    If sort=True, the packages are sorted by name.
    """
    os.makedirs(os.path.dirname(filename), exist_ok=True)

    # Copy the arch-specific Release file from devuan if it's not there
    bsnm = 'Packages.gz'
    if sources:
        bsnm = 'Sources.gz'
    rl = filename.replace(bsnm, 'Release')
    sprl = rl.replace(mergedir, join(spooldir, 'devuan'))
    if not os.path.isfile(rl) and os.path.isfile(sprl):
        copyfile(sprl, rl)

    gzf = gzip_open(filename, 'w')
    xzf = lzma_open(filename.replace('.gz', '.xz'), 'w')
    f = open(filename.replace('.gz', ''), 'wb')

    pkg_items = packages.items()
    if sort:
        pkg_items = sorted(pkg_items, key=lambda x: x[0])

    if sources:
        keylist = sources_keys
    else:
        keylist = packages_keys

    for pkg_name, pkg_contents in pkg_items:
        for key in keylist:
            if key in pkg_contents:
                s = '%s: %s\n' % (key, pkg_contents[key])
                gzf.write(s.encode('utf-8'))
                xzf.write(s.encode('utf-8'))
                f.write(s.encode('utf-8'))
        gzf.write(b'\n')
        xzf.write(b'\n')
        f.write(b'\n')

    gzf.close()
    xzf.close()
    f.close()

def load_packages_file(filename):
    """ Load a gzip'd packages file.
    Returns a dictionary of package name and package key-values.
    """
    # TODO: should we skip files like this if they don't exist?
    if filename is not None and os.path.isfile(filename):
        packages_contents = gzip_open(filename).read()
        packages_contents = packages_contents.decode('utf-8')
        return parse_packages(packages_contents)

    return None

def package_banned(pkg, banned_pkgs):
    """
    Returns True is the package contains a banned dependency.
    Currently checks and parses both the 'Depends:' and the 'Pre-Depends'
    fields of the package.
    """
    if pkg.get('Package') in banned_pkgs:
        return True

    depends = parse_dependencies(pkg.get('Depends', ''))
    pre_depends = parse_dependencies(pkg.get('Pre-Depends', ''))

    depends = [v[0] for v in depends]
    pre_depends = [v[0] for v in pre_depends]

    deps = set(depends).union(set(pre_depends))

    return bool(deps.intersection(banned_pkgs))

def merge_packages(pkg1, pkg2, name1, name2, banned_packages=set(),
                   rewriter=None):
    """
    Merges two previously loaded/parsed (using load_packages_file) packages
    dictionaries, preferring `pkg1` over `pkg2`, and optionally discarding any
    banned packages.
    """
    new_pkgs = {}
    package_names = set(pkg1.keys()).union(set(pkg2.keys()))

    for pkg in package_names:
        pkg1_pkg = pkg1.get(pkg)
        pkg2_pkg = pkg2.get(pkg)

        if pkg1_pkg and pkg2_pkg:
            if rewriter:
                pkg1_pkg = rewriter(pkg1_pkg, name1)
            new_pkgs[pkg] = pkg1_pkg
        elif pkg1_pkg:
            if not package_banned(pkg1_pkg, banned_packages):
                if rewriter:
                    pkg1_pkg = rewriter(pkg1_pkg, name1)
                new_pkgs[pkg] = pkg1_pkg
        elif pkg2_pkg:
            if not package_banned(pkg2_pkg, banned_packages):
                if rewriter:
                    pkg2_pkg = rewriter(pkg2_pkg, name2)
                new_pkgs[pkg] = pkg2_pkg
        else:
            assert False, 'Impossibru'

    return new_pkgs

def merge_packages_many(packages, banned_packages=set(), rewriter=None):
    """
    Merges two (or more) previously loaded/parsed (using load_packages_file)
    packages dictionaries, priority is defined by the order of the `packages`
    list, optionally discarding any banned packages.
    """
    if not len(packages) > 1:
        while not len(packages) > 1:
            packages.append({'name': '', 'packages': {}})

    new_pkgs = {}

    pkg1 = packages[0]
    pkg2 = packages[1]

    new_pkgs = merge_packages(pkg1['packages'], pkg2['packages'],
                              pkg1['name'], pkg2['name'],
                              banned_packages=banned_packages,
                              rewriter=rewriter)

    for pkg in packages[2:]:
        new_pkgs = merge_packages(new_pkgs, pkg['packages'], '', pkg['name'],
                                  banned_packages=banned_packages)

    return new_pkgs
<commit_msg>implement checking if a package with a lower priority is newer<commit_after># see LICENSE file for copyright and license details

"""
Package merging functions and helpers
"""

import os
from os.path import join
from gzip import open as gzip_open
from lzma import open as lzma_open
from shutil import copyfile

from lib.config import mergedir, packages_keys, sources_keys, spooldir
from lib.parse import parse_dependencies, parse_packages

def write_packages(packages, filename, sort=True, sources=False):
    """
    Writes `packages` to a file (per debian Packages format)
    If sort=True, the packages are sorted by name.
    """
    os.makedirs(os.path.dirname(filename), exist_ok=True)

    # Copy the arch-specific Release file from devuan if it's not there
    bsnm = 'Packages.gz'
    if sources:
        bsnm = 'Sources.gz'
    rl = filename.replace(bsnm, 'Release')
    sprl = rl.replace(mergedir, join(spooldir, 'devuan'))
    if not os.path.isfile(rl) and os.path.isfile(sprl):
        copyfile(sprl, rl)

    gzf = gzip_open(filename, 'w')
    xzf = lzma_open(filename.replace('.gz', '.xz'), 'w')
    f = open(filename.replace('.gz', ''), 'wb')

    pkg_items = packages.items()
    if sort:
        pkg_items = sorted(pkg_items, key=lambda x: x[0])

    if sources:
        keylist = sources_keys
    else:
        keylist = packages_keys

    for pkg_name, pkg_contents in pkg_items:
        for key in keylist:
            if key in pkg_contents:
                s = '%s: %s\n' % (key, pkg_contents[key])
                gzf.write(s.encode('utf-8'))
                xzf.write(s.encode('utf-8'))
                f.write(s.encode('utf-8'))
        gzf.write(b'\n')
        xzf.write(b'\n')
        f.write(b'\n')

    gzf.close()
    xzf.close()
    f.close()

def load_packages_file(filename):
    """ Load a gzip'd packages file.
    Returns a dictionary of package name and package key-values.
    """
    # TODO: should we skip files like this if they don't exist?
    if filename is not None and os.path.isfile(filename):
        packages_contents = gzip_open(filename).read()
        packages_contents = packages_contents.decode('utf-8')
        return parse_packages(packages_contents)

    return None

def package_banned(pkg, banned_pkgs):
    """
    Returns True is the package contains a banned dependency.
    Currently checks and parses both the 'Depends:' and the 'Pre-Depends'
    fields of the package.
    """
    if pkg.get('Package') in banned_pkgs:
        return True

    depends = parse_dependencies(pkg.get('Depends', ''))
    pre_depends = parse_dependencies(pkg.get('Pre-Depends', ''))

    depends = [v[0] for v in depends]
    pre_depends = [v[0] for v in pre_depends]

    deps = set(depends).union(set(pre_depends))

    return bool(deps.intersection(banned_pkgs))

def package_newer(pkg1, pkg2):
    """
    Checks whether the package of a lower priority is newer than the package
    of the higher priority and returns True if so.
    """
    higher_prio = pkg1.get('Version').split('+')
    lower_prio = pkg2.get('Version').split('+')

    if lower_prio[0] > higher_prio[0]:
        return True

    return False

def merge_packages(pkg1, pkg2, name1, name2, banned_packages=set(),
                   rewriter=None):
    """
    Merges two previously loaded/parsed (using load_packages_file) packages
    dictionaries, preferring `pkg1` over `pkg2`, and optionally discarding any
    banned packages.
    """
    new_pkgs = {}
    package_names = set(pkg1.keys()).union(set(pkg2.keys()))

    for pkg in package_names:
        pkg1_pkg = pkg1.get(pkg)
        pkg2_pkg = pkg2.get(pkg)

        if pkg1_pkg and pkg2_pkg:
            if rewriter:
                pkg1_pkg = rewriter(pkg1_pkg, name1)
            new_pkgs[pkg] = pkg1_pkg
            if package_newer(pkg1_pkg, pkg2_pkg):
                # TODO: proper logfile
                print('There is a newer version of "%s"' %
                      pkg1_pkg.get('Package'))
                print('Old version: %s' % pkg1_pkg.get('Version'))
                print('New version: %s' % pkg2_pkg.get('Version'))
        elif pkg1_pkg:
            if not package_banned(pkg1_pkg, banned_packages):
                if rewriter:
                    pkg1_pkg = rewriter(pkg1_pkg, name1)
                new_pkgs[pkg] = pkg1_pkg
        elif pkg2_pkg:
            if not package_banned(pkg2_pkg, banned_packages):
                if rewriter:
                    pkg2_pkg = rewriter(pkg2_pkg, name2)
                new_pkgs[pkg] = pkg2_pkg
        else:
            assert False, 'Impossibru'

    return new_pkgs

def merge_packages_many(packages, banned_packages=set(), rewriter=None):
    """
    Merges two (or more) previously loaded/parsed (using load_packages_file)
    packages dictionaries, priority is defined by the order of the `packages`
    list, optionally discarding any banned packages.
    """
    if not len(packages) > 1:
        while not len(packages) > 1:
            packages.append({'name': '', 'packages': {}})

    new_pkgs = {}

    pkg1 = packages[0]
    pkg2 = packages[1]

    new_pkgs = merge_packages(pkg1['packages'], pkg2['packages'],
                              pkg1['name'], pkg2['name'],
                              banned_packages=banned_packages,
                              rewriter=rewriter)

    for pkg in packages[2:]:
        new_pkgs = merge_packages(new_pkgs, pkg['packages'], '', pkg['name'],
                                  banned_packages=banned_packages)

    return new_pkgs
<|endoftext|>


<commit_before>import numpy as np

from menpo.feature import gradient as fast_gradient
from menpo.image import Image
from menpo.visualize import print_dynamic

from menpo3d.rasterize import rasterize_barycentric_coordinates

from .derivatives import (d_orthographic_projection_d_shape_parameters,
                          d_perspective_projection_d_shape_parameters,
                          d_orthographic_projection_d_warp_parameters,
                          d_perspective_projection_d_warp_parameters)
from ..projection import compute_rotation_matrices
from ..result import MMAlgorithmResult

class LucasKanade(object):
    def __init__(self, model, n_samples, projection_type, eps=1e-3):
        self.model = model
        self.eps = eps
        self.n_samples = n_samples
        self.projection_type = projection_type

        # Call precomputation
        self._precompute()

    @property
    def n(self):
        r"""
        Returns the number of active components of the shape model.

        :type: `int`
        """
        return self.model.shape_model.n_active_components

    @property
    def m(self):
        r"""
        Returns the number of active components of the texture model.

        :type: `int`
        """
        return self.model.texture_model.n_active_components

    @property
    def n_vertices(self):
        r"""
        Returns the number of vertices of the shape model's trimesh.

        :type: `int`
        """
        return self.model.shape_model.template_instance.n_points

    @property
    def n_channels(self):
        r"""
        Returns the number of channels of the texture model.

        :type: `int`
        """
        return self.model.n_channels

    def visible_sample_points(self, instance_in_img, image_shape):
        r"""
        Computes the warping map.

        Parameters
        ----------
        instance : `menpo.shape.ColouredTriMesh`
            The input mesh instance.

        Returns
        -------
        vertex_indices : ``(n_samples, 3)`` `ndarray`
            The vertices indices per sample.
        b_coords : ``(n_samples, 3)`` `ndarray`
            The barycentric coordinates per sample.
        yx : ``(n_samples,)`` `ndarray`
            The indices of the true points.
        """
        # Inverse rendering
        yx, tri_indices, b_coords = rasterize_barycentric_coordinates(
            instance_in_img, image_shape)

        # Select triangles randomly
        rand = np.random.permutation(b_coords.shape[0])
        b_coords = b_coords[rand[:self.n_samples]]
        yx = yx[rand[:self.n_samples]]
        tri_indices = tri_indices[rand[:self.n_samples]]

        # Build the vertex indices (3 per pixel) for the visible triangles
        vertex_indices = instance_in_img.trilist[tri_indices]

        return vertex_indices, tri_indices, b_coords, yx

    def sample(self, x, vertex_indices, b_coords):
        r"""
        Method that samples an object at locations specified by the
        barycentric coordinates.

        Parameters
        ----------
        x : `ndarray`
            The object to sample.
        vertex_indices : ``(n_samples, 3)`` `ndarray`
            The vertices indices per sample.
        b_coords : ``(n_samples, 3)`` `ndarray`
            The barycentric coordinates per sample.

        Returns
        -------
        sampled_x : ``(n_samples, ...)`` `ndarray`
            The sampled object.
        """
        per_vert_per_pixel = x[vertex_indices]
        return np.sum(per_vert_per_pixel * b_coords[..., None], axis=1)

    def gradient(self, image):
        r"""
        Computes the gradient of an image.

        Parameters
        ----------
        image : `menpo.image.Image` or subclass
            The input image.

        Returns
        -------
        gradient_X : ``(n_channels, height, width)`` `menpo.image.Image`
            The image gradient over the X direction.
        gradient_Y : ``(n_channels, height, width)`` `menpo.image.Image`
            The image gradient over the Y direction.
        """
        # Compute the gradient of the image
        grad = fast_gradient(image)

        # Create gradient image for X and Y
        grad_y = Image(grad.pixels[:self.n_channels])
        grad_x = Image(grad.pixels[self.n_channels:])

        return grad_x, grad_y

    def d_projection_d_shape_parameters(self, warped_uv, shape_pc_uv,
                                        focal_length, rotation_transform):
        if self.projection_type == 'perspective':
            dp_da = d_perspective_projection_d_shape_parameters(
                shape_pc_uv, focal_length, rotation_transform, warped_uv)
        else:
            dp_da = d_orthographic_projection_d_shape_parameters(
                shape_pc_uv, focal_length, rotation_transform)
        return dp_da

    def d_projection_d_warp_parameters(self, shape_uv, warped_uv,
                                       warp_parameters):
        r_phi, r_theta, r_varphi, _= compute_rotation_matrices(
            warp_parameters[1], warp_parameters[2], warp_parameters[3])
        if self.projection_type == 'perspective':
            dp_dr = d_perspective_projection_d_warp_parameters(
                shape_uv, warped_uv, warp_parameters, r_phi, r_theta, r_varphi)
        else:
            dp_dr = d_orthographic_projection_d_warp_parameters(
                shape_uv, warped_uv, warp_parameters, r_phi, r_theta, r_varphi)
        return dp_dr

    def compute_steepest_descent(self, dp_da, grad_x_uv, grad_y_uv):
        permuted_grad_x = np.transpose(grad_x_uv[..., None], (0, 2, 1))
        permuted_grad_y = np.transpose(grad_y_uv[..., None], (0, 2, 1))
        return permuted_grad_x * dp_da[0] + permuted_grad_y * dp_da[1]

    def compute_hessian(self, sd):
        n_channels = sd.shape[0]
        n_params = sd.shape[1]
        h = np.zeros((n_params, n_params))
        sd = sd.T
        for i in range(n_channels):
            h += np.dot(sd[:, :, i].T, sd[:, :, i])
        return h

    def compute_sd_error(self, sd, error_uv):
        n_channels = sd.shape[0]
        n_parameters = sd.shape[1]
        sd_error_product = np.zeros(n_parameters)
        sd = sd.T
        for i in range(n_channels):
            sd_error_product += np.dot(error_uv[i, :], sd[:, :, i])
        return sd_error_product.T

    def _precompute(self):
        # Rescale shape and appearance components to have size:
        # n_vertices x (n_active_components * n_dims)
        shape_pc = self.model.shape_model.components.T
        self.shape_pc = shape_pc.reshape([self.n_vertices, -1])

class Simultaneous(LucasKanade):
    r"""
    Class for defining Simultaneous Morphable Model optimization algorithm.
    """
    def run(self, image, initial_mesh, camera, gt_mesh=None,
            camera_update=False, max_iters=20, return_costs=False):
        # Define cost closure
        def cost_closure(x):
            return x.T.dot(x)

        # Retrieve warp (camera) parameters from the provided view and
        # projection transforms.
        camera_parameters = camera.as_vector()
        shape_parameters = self.model.shape_model.project(initial_mesh)
        texture_parameters = self.model.project_instance_on_texture_model(
            initial_mesh)

        # Reconstruct provided instance
        instance = self.model.instance(shape_weights=shape_parameters,
                                       texture_weights=texture_parameters)

        # Compute input image gradient
        grad_x, grad_y = self.gradient(image)

        # Initialize lists
        a_list = [shape_parameters]
        b_list = [texture_parameters]
        r_list = [camera]
        instances = [instance]
        costs = None
        if return_costs:
            costs = []

        # Initialize iteration counter and epsilon
        k = 0
        eps = np.Inf
        while k < max_iters and eps > self.eps:
            print_dynamic("{}/{}".format(k + 1, max_iters))

            instance_in_image = camera.apply(instance)

            # Compute indices locations for warping
            (vertex_indices, tri_indices,
             b_coords, yx) = self.visible_sample_points(instance_in_image,
                                                        image.shape)

            # Warp the mesh with the view matrix
            W = camera.view_transform.apply(instance.points)

            # Sample all the terms we need at our sample locations.
            shape_uv = self.sample(instance.points, vertex_indices, b_coords)
            warped_uv = self.sample(W, vertex_indices, b_coords)
            texture_uv = instance.sample_texture_with_barycentric_coordinates(
                b_coords, tri_indices)
            texture_pc_uv = self.model.sample_texture_model(b_coords,
                                                            tri_indices)
            shape_pc_uv = self.sample(self.shape_pc, vertex_indices, b_coords)
            # Reshape shape basis after sampling
            shape_pc_uv = shape_pc_uv.reshape([self.n_samples, 3, -1])
            img_uv = image.sample(yx)
            grad_x_uv = grad_x.sample(yx)
            grad_y_uv = grad_y.sample(yx)

            # print('texture_uv.shape: {}, texture_pc_uv.shape: {}'.format(texture_uv.shape, texture_pc_uv.shape))
            # print('shape_uv.shape: {}, shape_pc_uv.shape: {}'.format(shape_uv.shape, shape_pc_uv.shape))

            # Compute derivative of projection wrt shape parameters
            dp_da_dr = self.d_projection_d_shape_parameters(
                warped_uv, shape_pc_uv, camera_parameters[0],
                camera.rotation_transform)

            # Compute derivative of projection wrt warp parameters
            if camera_update:
                dp_dr = self.d_projection_d_warp_parameters(
                    shape_uv, warped_uv, camera_parameters)
                # Concatenate it with the derivative wrt shape parameters
                dp_da_dr = np.hstack((dp_da_dr, dp_dr))

            # Derivative of texture wrt texture parameters
            dt_db = np.rollaxis(texture_pc_uv, 0, 3)

            # Compute steepest descent
            sd_da_dr = self.compute_steepest_descent(dp_da_dr, grad_x_uv,
                                                     grad_y_uv)
            sd = np.hstack((sd_da_dr, -dt_db))

            # Compute hessian
            h = self.compute_hessian(sd)

            # Compute error
            img_error_uv = img_uv - texture_uv.T

            # Compute steepest descent matrix
            sd_error_img = self.compute_sd_error(sd, img_error_uv)

            # Update costs
            if return_costs:
                costs.append(cost_closure(img_error_uv.ravel()))

            # Compute increment
            delta_s = - np.linalg.solve(h, sd_error_img)

            # Update parameters
            shape_parameters += delta_s[:self.n]
            a_list.append(shape_parameters)
            if camera_update:
                camera_parameters += delta_s[self.n:self.n+len(camera_parameters)]
                texture_parameters += delta_s[(self.n+len(camera_parameters)):]
            else:
                texture_parameters += delta_s[self.n:]
            b_list.append(texture_parameters)

            # print("Image UV: {:.5f} - {:.5f}".format(img_uv.min(),
            #                                          img_uv.max()))
            # print("Instance UV: {:.5f} - {:.5f}".format(texture_uv.min(),
            #                                             texture_uv.max()))
            # print("Grad X: {:.5f} - {:.5f}".format(grad_x_uv.min(),
            #                                        grad_x_uv.max()))
            # print("Grad Y: {:.5f} - {:.5f}".format(grad_y_uv.min(),
            #                                        grad_y_uv.max()))
            # print("SD 1.1: {:.5f} - {:.5f}".format(sd_da_dr.min(),
            #                                        sd_da_dr.max()))
            # print("SD 1.2: {:.5f} - {:.5f}".format(dp_da_dr.min(),
            #                                        dp_da_dr.max()))
            # print("SD 2: {:.5f} - {:.5f}".format(dt_db.min(),
            #                                      dt_db.max()))

            # Generate the updated instance
            instance = self.model.instance(
                shape_weights=shape_parameters,
                texture_weights=texture_parameters)
            instances.append(instance)

            if camera_update:
                camera = camera.from_vector(camera_parameters)
            r_list.append(camera)

            # Increase iteration counter
            k += 1

        return MMAlgorithmResult(
            shape_parameters=a_list, texture_parameters=b_list,
            meshes=instances, camera_transforms=r_list, image=image,
            initial_mesh=initial_mesh, initial_camera_transform=r_list[0],
            gt_mesh=gt_mesh, costs=costs)

    def __str__(self):
        return "Simultaneous Lucas-Kanade"
<commit_msg>first version of priors on parameters<commit_after>import numpy as np

from menpo.feature import gradient as fast_gradient
from menpo.image import Image
from menpo.visualize import print_dynamic

from menpo3d.rasterize import rasterize_barycentric_coordinates

from .derivatives import (d_orthographic_projection_d_shape_parameters,
                          d_perspective_projection_d_shape_parameters,
                          d_orthographic_projection_d_warp_parameters,
                          d_perspective_projection_d_warp_parameters)
from ..projection import compute_rotation_matrices
from ..result import MMAlgorithmResult

class LucasKanade(object):
    def __init__(self, model, n_samples, projection_type, eps=1e-3):
        self.model = model
        self.eps = eps
        self.n_samples = n_samples
        self.projection_type = projection_type

        # Call precomputation
        self._precompute()

    @property
    def n(self):
        r"""
        Returns the number of active components of the shape model.

        :type: `int`
        """
        return self.model.shape_model.n_active_components

    @property
    def m(self):
        r"""
        Returns the number of active components of the texture model.

        :type: `int`
        """
        return self.model.texture_model.n_active_components

    @property
    def n_vertices(self):
        r"""
        Returns the number of vertices of the shape model's trimesh.

        :type: `int`
        """
        return self.model.shape_model.template_instance.n_points

    @property
    def n_channels(self):
        r"""
        Returns the number of channels of the texture model.

        :type: `int`
        """
        return self.model.n_channels

    def visible_sample_points(self, instance_in_img, image_shape):
        r"""
        Computes the warping map.

        Parameters
        ----------
        instance : `menpo.shape.ColouredTriMesh`
            The input mesh instance.

        Returns
        -------
        vertex_indices : ``(n_samples, 3)`` `ndarray`
            The vertices indices per sample.
        b_coords : ``(n_samples, 3)`` `ndarray`
            The barycentric coordinates per sample.
        yx : ``(n_samples,)`` `ndarray`
            The indices of the true points.
        """
        # Inverse rendering
        yx, tri_indices, b_coords = rasterize_barycentric_coordinates(
            instance_in_img, image_shape)

        # Select triangles randomly
        rand = np.random.permutation(b_coords.shape[0])
        b_coords = b_coords[rand[:self.n_samples]]
        yx = yx[rand[:self.n_samples]]
        tri_indices = tri_indices[rand[:self.n_samples]]

        # Build the vertex indices (3 per pixel) for the visible triangles
        vertex_indices = instance_in_img.trilist[tri_indices]

        return vertex_indices, tri_indices, b_coords, yx

    def sample(self, x, vertex_indices, b_coords):
        r"""
        Method that samples an object at locations specified by the
        barycentric coordinates.

        Parameters
        ----------
        x : `ndarray`
            The object to sample.
        vertex_indices : ``(n_samples, 3)`` `ndarray`
            The vertices indices per sample.
        b_coords : ``(n_samples, 3)`` `ndarray`
            The barycentric coordinates per sample.

        Returns
        -------
        sampled_x : ``(n_samples, ...)`` `ndarray`
            The sampled object.
        """
        per_vert_per_pixel = x[vertex_indices]
        return np.sum(per_vert_per_pixel * b_coords[..., None], axis=1)

    def gradient(self, image):
        r"""
        Computes the gradient of an image.

        Parameters
        ----------
        image : `menpo.image.Image` or subclass
            The input image.

        Returns
        -------
        gradient_X : ``(n_channels, height, width)`` `menpo.image.Image`
            The image gradient over the X direction.
        gradient_Y : ``(n_channels, height, width)`` `menpo.image.Image`
            The image gradient over the Y direction.
        """
        # Compute the gradient of the image
        grad = fast_gradient(image)

        # Create gradient image for X and Y
        grad_y = Image(grad.pixels[:self.n_channels])
        grad_x = Image(grad.pixels[self.n_channels:])

        return grad_x, grad_y

    def d_projection_d_shape_parameters(self, warped_uv, shape_pc_uv,
                                        focal_length, rotation_transform):
        if self.projection_type == 'perspective':
            dp_da = d_perspective_projection_d_shape_parameters(
                shape_pc_uv, focal_length, rotation_transform, warped_uv)
        else:
            dp_da = d_orthographic_projection_d_shape_parameters(
                shape_pc_uv, focal_length, rotation_transform)
        return dp_da

    def d_projection_d_warp_parameters(self, shape_uv, warped_uv,
                                       warp_parameters):
        r_phi, r_theta, r_varphi, _= compute_rotation_matrices(
            warp_parameters[1], warp_parameters[2], warp_parameters[3])
        if self.projection_type == 'perspective':
            dp_dr = d_perspective_projection_d_warp_parameters(
                shape_uv, warped_uv, warp_parameters, r_phi, r_theta, r_varphi)
        else:
            dp_dr = d_orthographic_projection_d_warp_parameters(
                shape_uv, warped_uv, warp_parameters, r_phi, r_theta, r_varphi)
        return dp_dr

    def compute_steepest_descent(self, dp_da, grad_x_uv, grad_y_uv):
        permuted_grad_x = np.transpose(grad_x_uv[..., None], (0, 2, 1))
        permuted_grad_y = np.transpose(grad_y_uv[..., None], (0, 2, 1))
        return permuted_grad_x * dp_da[0] + permuted_grad_y * dp_da[1]

    def compute_hessian(self, sd):
        n_channels = sd.shape[0]
        n_params = sd.shape[1]
        h = np.zeros((n_params, n_params))
        sd = sd.T
        for i in range(n_channels):
            h += np.dot(sd[:, :, i].T, sd[:, :, i])
        return h

    def compute_sd_error(self, sd, error_uv):
        n_channels = sd.shape[0]
        n_parameters = sd.shape[1]
        sd_error_product = np.zeros(n_parameters)
        sd = sd.T
        for i in range(n_channels):
            sd_error_product += np.dot(error_uv[i, :], sd[:, :, i])
        return sd_error_product.T

    def _precompute(self):
        # Rescale shape and appearance components to have size:
        # n_vertices x (n_active_components * n_dims)
        shape_pc = self.model.shape_model.components.T
        self.shape_pc = shape_pc.reshape([self.n_vertices, -1])

        # Parameters priors
        self.sd_alpha_prior = np.zeros(self.n + self.m)
        self.sd_alpha_prior[:self.n] = 1. / np.sqrt(self.model.shape_model.eigenvalues)
        #self.sd_alpha_prior[:self.n] = 2. / (
        # self.model.shape_model.eigenvalues ** 2)
        self.sd_beta_prior = np.zeros(self.n + self.m)
        self.sd_beta_prior[self.n:] = 1. / np.sqrt(self.model.texture_model.eigenvalues)
        #self.sd_beta_prior[self.n:] = 2. / (
        # self.model.texture_model.eigenvalues ** 2)
        self.H_alpha_prior = np.diag(self.sd_alpha_prior)
        self.H_beta_prior = np.diag(self.sd_beta_prior)

class Simultaneous(LucasKanade):
    r"""
    Class for defining Simultaneous Morphable Model optimization algorithm.
    """
    def run(self, image, initial_mesh, camera, gt_mesh=None,
            camera_update=False, max_iters=20, return_costs=False):
        # Define cost closure
        def cost_closure(x):
            return x.T.dot(x)

        # Retrieve warp (camera) parameters from the provided view and
        # projection transforms.
        camera_parameters = camera.as_vector()
        shape_parameters = self.model.shape_model.project(initial_mesh)
        texture_parameters = self.model.project_instance_on_texture_model(
            initial_mesh)

        # Reconstruct provided instance
        instance = self.model.instance(shape_weights=shape_parameters,
                                       texture_weights=texture_parameters)

        # Compute input image gradient
        grad_x, grad_y = self.gradient(image)

        # Initialize lists
        a_list = [shape_parameters]
        b_list = [texture_parameters]
        r_list = [camera]
        instances = [instance.with_clipped_texture()]
        costs = None
        if return_costs:
            costs = []

        # Initialize iteration counter and epsilon
        k = 0
        eps = np.Inf
        while k < max_iters and eps > self.eps:
            print_dynamic("{}/{}".format(k + 1, max_iters))

            instance_in_image = camera.apply(instance)

            # Compute indices locations for warping
            (vertex_indices, tri_indices,
             b_coords, yx) = self.visible_sample_points(instance_in_image,
                                                        image.shape)

            # Warp the mesh with the view matrix
            W = camera.view_transform.apply(instance.points)

            # Sample all the terms we need at our sample locations.
            shape_uv = self.sample(instance.points, vertex_indices, b_coords)
            warped_uv = self.sample(W, vertex_indices, b_coords)
            texture_uv = instance.sample_texture_with_barycentric_coordinates(
                b_coords, tri_indices)
            texture_pc_uv = self.model.sample_texture_model(b_coords,
                                                            tri_indices)
            shape_pc_uv = self.sample(self.shape_pc, vertex_indices, b_coords)
            # Reshape shape basis after sampling
            shape_pc_uv = shape_pc_uv.reshape([self.n_samples, 3, -1])
            img_uv = image.sample(yx)
            grad_x_uv = grad_x.sample(yx)
            grad_y_uv = grad_y.sample(yx)

            # print('texture_uv.shape: {}, texture_pc_uv.shape: {}'.format(texture_uv.shape, texture_pc_uv.shape))
            # print('shape_uv.shape: {}, shape_pc_uv.shape: {}'.format(shape_uv.shape, shape_pc_uv.shape))

            # Compute derivative of projection wrt shape parameters
            dp_da_dr = self.d_projection_d_shape_parameters(
                warped_uv, shape_pc_uv, camera_parameters[0],
                camera.rotation_transform)

            # Compute derivative of projection wrt warp parameters
            if camera_update:
                dp_dr = self.d_projection_d_warp_parameters(
                    shape_uv, warped_uv, camera_parameters)
                # Concatenate it with the derivative wrt shape parameters
                dp_da_dr = np.hstack((dp_da_dr, dp_dr))

            # Derivative of texture wrt texture parameters
            dt_db = np.rollaxis(texture_pc_uv, 0, 3)

            # Compute steepest descent
            sd_da_dr = self.compute_steepest_descent(dp_da_dr, grad_x_uv,
                                                     grad_y_uv)
            sd = np.hstack((sd_da_dr, -dt_db))

            # Compute hessian
            h = self.compute_hessian(sd) + 1e-1 * self.H_alpha_prior + self.H_beta_prior

            # Compute error
            img_error_uv = img_uv - texture_uv.T

            # Compute steepest descent matrix
            sd_error_img = self.compute_sd_error(sd, img_error_uv)

            # Apply priors
            sd_error_alpha_prior = self.sd_alpha_prior * \
                                   np.concatenate((shape_parameters,
                                                   texture_parameters))
            sd_error_beta_prior = self.sd_beta_prior * \
                                  np.concatenate((shape_parameters,
                                                  texture_parameters))
            sd_error_img = sd_error_img + 1e-1 * sd_error_alpha_prior + sd_error_beta_prior

            # Update costs
            if return_costs:
                costs.append(cost_closure(img_error_uv.ravel()))

            # Compute increment
            delta_s = - np.linalg.solve(h, sd_error_img)

            # Update parameters
            shape_parameters += delta_s[:self.n]
            a_list.append(shape_parameters)
            if camera_update:
                camera_parameters += delta_s[self.n:self.n+len(camera_parameters)]
                texture_parameters += delta_s[(self.n+len(camera_parameters)):]
            else:
                texture_parameters += delta_s[self.n:]
            b_list.append(texture_parameters)

            # print("Image UV: {:.5f} - {:.5f}".format(img_uv.min(),
            #                                          img_uv.max()))
            # print("Instance UV: {:.5f} - {:.5f}".format(texture_uv.min(),
            #                                             texture_uv.max()))
            # print("Grad X: {:.5f} - {:.5f}".format(grad_x_uv.min(),
            #                                        grad_x_uv.max()))
            # print("Grad Y: {:.5f} - {:.5f}".format(grad_y_uv.min(),
            #                                        grad_y_uv.max()))
            # print("SD 1.1: {:.5f} - {:.5f}".format(sd_da_dr.min(),
            #                                        sd_da_dr.max()))
            # print("SD 1.2: {:.5f} - {:.5f}".format(dp_da_dr.min(),
            #                                        dp_da_dr.max()))
            # print("SD 2: {:.5f} - {:.5f}".format(dt_db.min(),
            #                                      dt_db.max()))

            # Generate the updated instance
            instance = self.model.instance(
                shape_weights=shape_parameters,
                texture_weights=texture_parameters)
            instances.append(instance.with_clipped_texture())

            if camera_update:
                camera = camera.from_vector(camera_parameters)
            r_list.append(camera)

            # Increase iteration counter
            k += 1

        return MMAlgorithmResult(
            shape_parameters=a_list, texture_parameters=b_list,
            meshes=instances, camera_transforms=r_list, image=image,
            initial_mesh=initial_mesh.with_clipped_texture(),
            initial_camera_transform=r_list[0], gt_mesh=gt_mesh, costs=costs)

    def __str__(self):
        return "Simultaneous Lucas-Kanade"

def parameters_prior(params, eigenvalues):
    c = params / eigenvalues
    norm = np.sqrt(len(eigenvalues)) / c.dot(c)
    return norm * params
<|endoftext|>


<commit_before>from __future__ import absolute_import
from .core import (State, Transition, Event, EventData, Machine, MachineError,
                   logger)
<commit_msg>Import the version number<commit_after>from __future__ import absolute_import
from .version import __version__
from .core import (State, Transition, Event, EventData, Machine, MachineError,
                   logger)
<|endoftext|>