From 7288fc1c7c1c098d9d06847ece8c9253c8493e4d Mon Sep 17 00:00:00 2001
From: chanomkaimuk <22185824+chanomkaimuk@users.noreply.github.com>
Date: Fri, 11 Aug 2023 03:14:49 +0200
Subject: [PATCH] improved regex and added Lithuanian keywords for Baltic heart

---
 commands/react_baltics/consts.py        | 27 +++++++++++++++++++++++++
 commands/react_baltics/react_baltics.py | 25 +++++++++++++++++++++++
 commands/reacttw/consts.py              |  8 ++++++--
 commands/reacttw/react_tw.py            |  2 +-
 main.py                                 |  4 ++++
 5 files changed, 63 insertions(+), 3 deletions(-)
 create mode 100644 commands/react_baltics/consts.py
 create mode 100644 commands/react_baltics/react_baltics.py

diff --git a/commands/react_baltics/consts.py b/commands/react_baltics/consts.py
new file mode 100644
index 0000000..fbc4012
--- /dev/null
+++ b/commands/react_baltics/consts.py
@@ -0,0 +1,27 @@
+from re import compile, IGNORECASE, UNICODE
+
+POSSIBLE_REACTS = (
+    "<:tw_baltics_heart:1132524940587962388>",
+    "<:bubblemilktea:1132632348651966596>",
+)
+
+KEYWORDS = (
+    # Lithuanian
+    "Taivanas",
+    "Taivane",
+    "Taivano",
+    "Taivanui",
+    "Taivanietis",
+    "Taivanietė",
+    "Taivaniečiai",
+    "Taivaniečių",
+    "Taivaniečiui",
+    "Taivaniečiams",
+    "Taivanietiškas",
+    "Taivana"
+    # Latvian
+    # Estonian
+)
+
+
+BALTIC_REGEX = compile(rf"({'|'.join(KEYWORDS)})", flags=IGNORECASE | UNICODE)
diff --git a/commands/react_baltics/react_baltics.py b/commands/react_baltics/react_baltics.py
new file mode 100644
index 0000000..dee00e8
--- /dev/null
+++ b/commands/react_baltics/react_baltics.py
@@ -0,0 +1,25 @@
+import discord
+import random
+from .consts import POSSIBLE_REACTS, BALTIC_REGEX
+
+
+def is_baltic_message(message: discord.Message):
+    return BALTIC_REGEX.search(message.content)
+
+
+def mock_bernoulli(p: float) -> bool:
+    """Returns True with probability p.
+
+    Args:
+        p (float): a float between 0 and 1.
+
+    Returns:
+        bool: True or False.
+    """
+    return random.random() < p
+
+
+async def send_react_baltic(message: discord.Message):
+    for react in POSSIBLE_REACTS:
+        if mock_bernoulli(0.69):
+            await message.add_reaction(react)
diff --git a/commands/reacttw/consts.py b/commands/reacttw/consts.py
index b2ccf49..bbc9734 100644
--- a/commands/reacttw/consts.py
+++ b/commands/reacttw/consts.py
@@ -1,4 +1,4 @@
-from re import compile
+from re import compile, IGNORECASE, UNICODE
 
 POSSIBLE_REACTS = (
     "<:flag_twi:1133045891780071436>",
@@ -44,10 +44,13 @@
     "台灣",
     "臺灣",
     "臺北",
+    "台北",
     "新北",
     "桃園",
     "臺中",
+    "台中",
     "臺南",
+    "台南",
     "高雄",
     "新竹",
     "苗栗",
@@ -59,6 +62,7 @@
     "宜蘭",
     "花蓮",
     "臺東",
+    "台東",
     "澎湖",
     "金門",
     "連江",
@@ -69,4 +73,4 @@
 )
 
 
-TW_REGEX = compile(rf"({'|'.join(KEYWORDS)})")
+TW_REGEX = compile(rf"({'|'.join(KEYWORDS)})", flags=IGNORECASE | UNICODE)
diff --git a/commands/reacttw/react_tw.py b/commands/reacttw/react_tw.py
index 75ae938..a091023 100644
--- a/commands/reacttw/react_tw.py
+++ b/commands/reacttw/react_tw.py
@@ -4,7 +4,7 @@
 
 
 def is_TW_message(message: discord.Message):
-    return TW_REGEX.search(message.content.upper())
+    return TW_REGEX.search(message.content)
 
 
 def mock_bernoulli(p: float) -> bool:
diff --git a/main.py b/main.py
index 223004a..a4c8727 100644
--- a/main.py
+++ b/main.py
@@ -15,6 +15,7 @@
 from commands.one_o_one import one_o_one
 from commands import hgs
 from commands.reacttw import react_tw
+from commands.react_baltics import react_baltics
 from commands.shiba import random_shiba
 import sys
 
@@ -101,5 +102,8 @@ async def on_message(message: discord.Message):
     if react_tw.is_TW_message(message):
         await react_tw.send_react_tw(message)
 
+    if react_baltics.is_baltic_message(message):
+        await react_baltics.send_react_baltic(message)
+
 
 client.run(TOKEN)