import datetime

import django
from django.conf import settings
from django.db import connection, models
from django.db.models import Avg, F, Window

if not settings.configured:
    settings.configure(
        DATABASES={
            "default": {"ENGINE": "django.db.backends.sqlite3", "NAME": ":memory:"}
        },
        INSTALLED_APPS=["__main__"],
    )
    django.setup()


class Movie(models.Model):
    title = models.CharField(max_length=100)
    studio = models.CharField(max_length=50)
    genre = models.CharField(max_length=50)
    released = models.DateField()
    rating = models.IntegerField()

    class Meta:
        app_label = "__main__"


with connection.schema_editor() as schema_editor:
    schema_editor.create_model(Movie)

Movie.objects.create(
    title="M1", studio="A", genre="Action", released=datetime.date(2020, 1, 1), rating=8
)
Movie.objects.create(
    title="M2",
    studio="A",
    genre="Action",
    released=datetime.date(2020, 6, 1),
    rating=10,
)
Movie.objects.create(
    title="M3", studio="A", genre="Action", released=datetime.date(2021, 1, 1), rating=2
)

docs_qs = Movie.objects.annotate(
    avg=Window(
        expression=Avg("rating"),
        partition_by=[F("studio"), F("genre")],
        order_by="released__year",
    )
)
print("Output of documentation query:")
for m in docs_qs:
    print(f"{m.title} ({m.released.year}): {m.avg:.2f}")


fixed_qs = Movie.objects.annotate(
    avg=Window(
        expression=Avg("rating"),
        partition_by=[F("studio"), F("genre"), F("released__year")],
    )
)

print("\nOutput when partitioning by year:")
for m in fixed_qs:
    print(f"{m.title} ({m.released.year}): {m.avg:.2f}")