From 3a43722317b7ec7ba89a078a2f95b220ef04dad1 Mon Sep 17 00:00:00 2001
From: RussellDash332 <russellsaerang@gmail.com>
Date: Sun, 6 Aug 2023 18:28:36 +0800
Subject: [PATCH] Fix colormapping of problem plot

---
 README.md                  |  1 +
 autokattis/api/__init__.py | 13 ++++++++++---
 setup.py                   |  2 +-
 test/main.py               |  3 ++-
 4 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/README.md b/README.md
index 812f8a0..8aa3203 100644
--- a/README.md
+++ b/README.md
@@ -29,6 +29,7 @@ where `'username'` is your Kattis username/email and `'password'` is your Kattis
 kt.problems()                               # problems you have solved so far
 kt.problems(show_partial=False)             # exclude partial submissions
 kt.problems(*[True]*4)                      # literally all problems on Kattis
+
 kt.list_unsolved()                          # let's grind!
 
 kt.plot_problems()                          # plot the points distribution
diff --git a/autokattis/api/__init__.py b/autokattis/api/__init__.py
index 74acebd..50ccd69 100644
--- a/autokattis/api/__init__.py
+++ b/autokattis/api/__init__.py
@@ -106,7 +106,8 @@ def problems(self, show_solved=True, show_partial=True, show_tried=False, show_u
                                 # for example:
                                 # - difficulty 9.1-9.6 -> [9.1, 9.6]
                                 # - difficulty 5.0 -> [5.0]
-                            category = re.findall('[A-Za-z]+', columns[7].text)[0]
+                            try: category = re.findall('[A-Za-z]+', columns[7].text)[0]
+                            except: category = 'N/A'
                             data.append({
                                 'name': name,
                                 'fastest': fastest,
@@ -127,8 +128,14 @@ def plot_problems(self, filepath=None, show_solved=True, show_partial=True, show
         '''
 
         df = pd.DataFrame(self.problems(show_solved, show_partial, show_tried, show_untried))
-        hist = sns.histplot(data=df, x='difficulty', hue='category', multiple='stack', binwidth=0.1)
-        hist.set(title=f'Solved Kattis Problems ({df.shape[0]})', xlabel='Difficulty')
+        categories = set(df.category)
+        hue_order = [c for c in ['Easy', 'Medium', 'Hard', 'N/A'][::-1] if c in categories]
+        palette = {'Easy': '#39a137', 'Medium': '#ffbe00', 'Hard': '#ff411a', 'N/A': 'gray'}
+        for c in [*palette]:
+            if c not in categories: del palette[c]
+        hist = sns.histplot(data=df, x='difficulty', hue='category', multiple='stack', binwidth=0.1, hue_order=hue_order, palette=palette)
+        hist.set(title=f'Solved Kattis Problems by {self.user} ({df.shape[0]})', xlabel='Difficulty')
+        plt.legend(title='Category', loc='upper right', labels=hue_order[::-1])
         plt.xticks([*range(math.floor(min(df.difficulty)), math.ceil(max(df.difficulty))+1)])
         if filepath != None:
             plt.savefig(filepath)
diff --git a/setup.py b/setup.py
index c0c7cef..8747ca2 100644
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@
 
 setup_args = dict(
     name='autokattis',
-    version='1.4.5',
+    version='1.4.6',
     description='Updated Kattis API wrapper',
     long_description_content_type="text/markdown",
     long_description=README,
diff --git a/test/main.py b/test/main.py
index b0fb116..3a6994d 100644
--- a/test/main.py
+++ b/test/main.py
@@ -9,7 +9,8 @@
     kt.problems(),                              # problems you have solved so far
     kt.problems(show_partial=False),            # exclude partial submissions
     kt.problems(*[True]*4),                     # literally all problems on Kattis
-    kt.list_unsolved(),                          # let's grind!
+
+    kt.list_unsolved(),                         # let's grind!
 
     kt.plot_problems(),                         # plot the points distribution
     kt.plot_problems(filepath='plot.png'),      # save to a filepath