add heatmap

ai/urls/oj.py

@@ -4,10 +4,12 @@ from ..views.oj import (
     AIAnalysisAPI,
     AIDetailDataAPI,
     AIWeeklyDataAPI,
+    AIHeatmapDataAPI,
 )
 
 urlpatterns = [
     path("ai/detail", AIDetailDataAPI.as_view()),
     path("ai/weekly", AIWeeklyDataAPI.as_view()),
     path("ai/analysis", AIAnalysisAPI.as_view()),
+    path("ai/heatmap", AIHeatmapDataAPI.as_view()),
 ]

ai/views/oj.py

@@ -5,7 +5,8 @@ import json
 
 from dateutil.relativedelta import relativedelta
 from django.core.cache import cache
-from django.db.models import Min
+from django.db.models import Min, Count
+from django.db.models.functions import TruncDate
 from django.http import StreamingHttpResponse
 from django.utils import timezone
 from openai import OpenAI
@@ -412,3 +413,48 @@ class AIAnalysisAPI(APIView):
         )
         response["Cache-Control"] = "no-cache"
         return response
+
+
+class AIHeatmapDataAPI(APIView):
+    @login_required
+    def get(self, request):
+        user = request.user
+        cache_key = get_cache_key("ai_heatmap", user.id, user.class_name or "")
+        cached_result = cache.get(cache_key)
+        if cached_result:
+            return self.success(cached_result)
+
+        end = datetime.now()
+        start = end - timedelta(days=365)
+
+        # 使用单次查询获取所有数据，按日期分组统计
+        submission_counts = (
+            Submission.objects.filter(
+                user_id=user.id, create_time__gte=start, create_time__lte=end
+            )
+            .annotate(date=TruncDate("create_time"))
+            .values("date")
+            .annotate(count=Count("id"))
+            .order_by("date")
+        )
+
+        # 将查询结果转换为字典，便于快速查找
+        submission_dict = {item["date"]: item["count"] for item in submission_counts}
+
+        # 生成365天的热力图数据
+        heatmap_data = []
+        current_date = start.date()
+        for i in range(365):
+            day_date = current_date + timedelta(days=i)
+            submission_count = submission_dict.get(day_date, 0)
+            heatmap_data.append(
+                {
+                    "timestamp": int(datetime.combine(
+                        day_date, datetime.min.time()
+                    ).timestamp() * 1000),
+                    "value": submission_count,
+                }
+            )
+
+        cache.set(cache_key, heatmap_data, CACHE_TIMEOUT)
+        return self.success(heatmap_data)

contest/models.py

@@ -46,10 +46,13 @@ class Contest(models.Model):
 
     # 是否有权查看problem 的一些统计信息 诸如submission_number, accepted_number 等
     def problem_details_permission(self, user):
-        return self.rule_type == ContestRuleType.ACM or \
-               self.status == ContestStatus.CONTEST_ENDED or \
-               user.is_authenticated and user.is_contest_admin(self) or \
-               self.real_time_rank
+        return (
+            self.rule_type == ContestRuleType.ACM
+            or self.status == ContestStatus.CONTEST_ENDED
+            or user.is_authenticated
+            and user.is_contest_admin(self)
+            or self.real_time_rank
+        )
 
     class Meta:
         db_table = "contest"

submission/migrations/0002_submission_user_create_time_idx.py

@@ -0,0 +1,19 @@
+# Generated by Django 5.2.3 on 2025-09-25 07:03
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('contest', '0001_initial'),
+        ('problem', '0001_initial'),
+        ('submission', '0001_initial'),
+    ]
+
+    operations = [
+        migrations.AddIndex(
+            model_name='submission',
+            index=models.Index(fields=['user_id', 'create_time'], name='user_create_time_idx'),
+        ),
+    ]

submission/models.py

@@ -41,7 +41,12 @@ class Submission(models.Model):
     ip = models.TextField(null=True)
 
     def check_user_permission(self, user, check_share=True):
-        if self.user_id == user.id or user.is_super_admin() or user.can_mgmt_all_problem() or self.problem.created_by_id == user.id:
+        if (
+            self.user_id == user.id
+            or user.is_super_admin()
+            or user.can_mgmt_all_problem()
+            or self.problem.created_by_id == user.id
+        ):
             return True
 
         if check_share:
@@ -54,6 +59,11 @@ class Submission(models.Model):
     class Meta:
         db_table = "submission"
         ordering = ("-create_time",)
+        indexes = [
+            models.Index(
+                fields=["user_id", "create_time"], name="user_create_time_idx"
+            ),
+        ]
 
     def __str__(self):
         return self.id