From 333651ffad0462c3e5ffaafdef474c9d23f4354f Mon Sep 17 00:00:00 2001
From: Jarrett Ye <jarrett.ye@outlook.com>
Date: Thu, 2 Jan 2025 15:43:49 +0800
Subject: [PATCH] Fix/S0 dataset initialization fallback (#156)

---
 pyproject.toml                       | 2 +-
 src/fsrs_optimizer/fsrs_optimizer.py | 9 +++++++++
 2 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index eafb07d..8837ecf 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "FSRS-Optimizer"
-version = "5.6.2"
+version = "5.6.3"
 readme = "README.md"
 dependencies = [
     "matplotlib>=3.7.0",
diff --git a/src/fsrs_optimizer/fsrs_optimizer.py b/src/fsrs_optimizer/fsrs_optimizer.py
index c9fe57a..50d8ee8 100644
--- a/src/fsrs_optimizer/fsrs_optimizer.py
+++ b/src/fsrs_optimizer/fsrs_optimizer.py
@@ -1002,6 +1002,15 @@ def pretrain(self, dataset=None, verbose=True):
             )
         else:
             self.dataset = dataset
+            self.dataset["first_rating"] = self.dataset["r_history"].map(
+                lambda x: x[0] if len(x) > 0 else ""
+            )
+            self.S0_dataset_group = (
+                self.dataset[self.dataset["i"] == 2]
+                .groupby(by=["first_rating", "delta_t"], group_keys=False)
+                .agg({"y": ["mean", "count"]})
+                .reset_index()
+            )
         self.dataset = self.dataset[
             (self.dataset["i"] > 1) & (self.dataset["delta_t"] > 0)
         ]