pretalx-musicrate/pretalx_musicrate/management/commands/computescores.py

import math
import statistics
from operator import attrgetter, itemgetter

from django.core.management.base import BaseCommand, CommandError
from django_scopes import scope
from pretalx.event.models import Event
from pretalx.submission.models import Submission

from ...models import Rating, Score

RATINGS_MIN = 30
SCALE = list(map(float, map(itemgetter(0), Rating.RATING_CHOICES)))

MEAN = statistics.mean(SCALE)
STD = math.sqrt(statistics.mean([(i - MEAN) ** 2 for i in SCALE]))


class Command(BaseCommand):
    help = "Compute submission scores from ratings"

    def add_arguments(self, parser):
        parser.add_argument("event")

    def handle(self, *args, **kwargs):
        try:
            event = Event.objects.get(slug=kwargs["event"])
        except Event.DoesNotExist:
            raise CommandError(f"no event found with slug '{kwargs['event']}'")

        with scope(event=event):
            submissions = {}

            for juror in event.jurors.prefetch_related("ratings__submission").order_by(
                "token"
            ):
                ratings = list(
                    juror.ratings.exclude(rating="").order_by("submission__created")
                )
                values = list(map(int, map(attrgetter("rating"), ratings)))
                ratings = dict(
                    zip(
                        map(attrgetter("code"), map(attrgetter("submission"), ratings)),
                        values,
                    )
                )

                if len(values) < RATINGS_MIN:
                    mean = MEAN
                    std = STD
                else:
                    mean = sum(values) / len(values)
                    std = math.sqrt(
                        sum([(i - mean) ** 2 for i in values]) / len(values)
                    )

                for code in ratings:
                    if code not in submissions:
                        submissions[code] = []
                    submissions[code].append((ratings[code] - mean) / std * STD + MEAN)

            for submission in Submission.objects.select_related("score").filter(
                code__in=submissions.keys()
            ):
                try:
                    score = submission.score
                except Submission.score.RelatedObjectDoesNotExist:
                    score = Score(submission=submission)

                ratings = submissions[submission.code]
                score.value = sum(ratings) / len(ratings)
                score.save()