mirror of
https://github.com/brandons209/Red-bot-Cogs.git
synced 2024-04-28 17:42:59 +12:00
559 lines
23 KiB
Python
559 lines
23 KiB
Python
import asyncio
|
|
import discord
|
|
import random
|
|
import matplotlib.pyplot as plt
|
|
import pandas as pd
|
|
from typing import Optional, Literal
|
|
from io import BytesIO
|
|
|
|
from redbot.core import Config, checks, commands
|
|
from redbot.core.utils.chat_formatting import *
|
|
from redbot.core.utils.predicates import MessagePredicate
|
|
|
|
from .discord_thread_feature import send_thread_message, create_thread
|
|
from .time_utils import parse_time, parse_timedelta
|
|
|
|
from datetime import datetime
|
|
from datetime import timedelta
|
|
|
|
|
|
class ThreadRotate(commands.Cog):
|
|
"""
|
|
Rotate threads for events, roleplay, etc
|
|
"""
|
|
|
|
def __init__(self, bot):
|
|
self.bot = bot
|
|
self.config = Config.get_conf(self, identifier=45612361654894681623, force_registration=True)
|
|
self.methods = ["replacement", "unique"]
|
|
default_channel = {
|
|
"topics": {},
|
|
"topic_threads": {},
|
|
"ping_roles": [],
|
|
"rotation_interval": 10080,
|
|
"rotate_on": None,
|
|
"method": "replacement",
|
|
"prev_topics": [],
|
|
"last_topic": None,
|
|
}
|
|
self.config.register_channel(**default_channel)
|
|
|
|
self.task = asyncio.create_task(self.thread_rotation_task())
|
|
|
|
def cog_unload(self):
|
|
if self.task is not None:
|
|
self.task.cancel()
|
|
return super().cog_unload()()
|
|
|
|
async def thread_rotation_task(self):
|
|
await self.bot.wait_until_ready()
|
|
|
|
while True:
|
|
for guild in self.bot.guilds:
|
|
for channel in guild.text_channels:
|
|
topics = await self.config.channel(channel).topics()
|
|
if not topics:
|
|
continue
|
|
rotate_on = datetime.fromtimestamp(await self.config.channel(channel).rotate_on())
|
|
if datetime.now() > rotate_on:
|
|
await self.rotate_thread(channel)
|
|
await asyncio.sleep(1)
|
|
|
|
await asyncio.sleep(60)
|
|
|
|
async def rotate_thread(self, channel: discord.TextChannel):
|
|
topics = await self.config.channel(channel).topics()
|
|
topic_threads = await self.config.channel(channel).topic_threads()
|
|
ping_roles = await self.config.channel(channel).ping_roles()
|
|
rotation = timedelta(seconds=await self.config.channel(channel).rotation_interval())
|
|
rotate_on = datetime.now()
|
|
last_topic = await self.config.channel(channel).last_topic()
|
|
method = await self.config.channel(channel).method()
|
|
prev_topics = await self.config.channel(channel).prev_topics()
|
|
|
|
# choose new topic
|
|
# don't want to choose the last topic, so set it's weight to 0 so it is not choosen
|
|
if last_topic is not None:
|
|
topics[last_topic] = 0
|
|
|
|
# select without replacement if method is "unique"
|
|
if method == "unique":
|
|
# first check if the list of topics have been exhausted
|
|
topic_names = list(topics.keys())
|
|
cnt = 0
|
|
for t in topic_names:
|
|
if t in prev_topics:
|
|
cnt += 1
|
|
|
|
# reset list
|
|
if cnt == len(topic_names):
|
|
prev_topics = []
|
|
|
|
for t in prev_topics:
|
|
topics[t] = 0
|
|
|
|
new_topic = random.choices(list(topics.keys()), weights=list(topics.values()), k=1)[0]
|
|
|
|
# rotate the thread, create new thread, ping roles, etc
|
|
# ping roles
|
|
roles = [channel.guild.get_role(r) for r in ping_roles]
|
|
roles = [r for r in roles if r is not None]
|
|
|
|
role_msg = " ".join([r.mention for r in roles])
|
|
role_msg += f"\n\nHello, a new topic has been set for {channel.mention}: `{new_topic}`"
|
|
|
|
# if a thread already exists for the topic, try to send a message to it to unarchive it
|
|
new_thread_id = None
|
|
if new_topic in topic_threads:
|
|
try:
|
|
await send_thread_message(self.bot, topic_threads[new_topic], role_msg, mention_roles=ping_roles)
|
|
new_thread_id = topic_threads[new_topic]
|
|
except (
|
|
discord.HTTPException
|
|
): # may occur if bot cant unarchive manually archived threads or thread is deleted
|
|
try:
|
|
new_thread_id = await create_thread(self.bot, channel, new_topic, archive=10080)
|
|
except discord.HTTPException:
|
|
return
|
|
await send_thread_message(self.bot, new_thread_id, role_msg, mention_roles=ping_roles)
|
|
else:
|
|
try:
|
|
new_thread_id = await create_thread(self.bot, channel, new_topic, archive=10080)
|
|
except discord.HTTPException:
|
|
return
|
|
await send_thread_message(self.bot, new_thread_id, role_msg, mention_roles=ping_roles)
|
|
|
|
# update next rotation
|
|
async with self.config.channel(channel).topic_threads() as topic_threads:
|
|
topic_threads[new_topic] = new_thread_id
|
|
|
|
await self.config.channel(channel).rotate_on.set(int((rotate_on + rotation).timestamp()))
|
|
await self.config.channel(channel).last_topic.set(new_topic)
|
|
if method == "unique":
|
|
prev_topics.append(new_topic)
|
|
await self.config.channel(channel).prev_topics.set(prev_topics)
|
|
|
|
@commands.group(name="rotate")
|
|
@commands.guild_only()
|
|
@checks.admin()
|
|
async def thread_rotate(self, ctx):
|
|
"""
|
|
Manage thread rotations per channel
|
|
"""
|
|
pass
|
|
|
|
@thread_rotate.command(name="simulation")
|
|
async def thread_rotate_simulation(self, ctx, channel: discord.TextChannel):
|
|
"""
|
|
Run a simulation using the settings for the channel to see how often topics are chosen.
|
|
"""
|
|
topics = await self.config.channel(channel).topics()
|
|
if not topics:
|
|
await ctx.send(error("That channel has not been setup for thread rotation!"), delete_after=30)
|
|
return
|
|
|
|
topic_names = [k + f": {v}" for k, v in topics.items()]
|
|
dist = random.choices(topic_names, weights=list(topics.values()), k=10000 * len(topics))
|
|
|
|
# make graph and send it
|
|
fontsize = 30
|
|
fig = plt.figure(figsize=(50, 20 + 10 * (len(topics) % 10)))
|
|
|
|
# define graph and table save paths
|
|
save_path = BytesIO()
|
|
|
|
pd.Series(dist).value_counts(sort=False).plot(kind="barh")
|
|
|
|
# make graph look nice
|
|
plt.title(
|
|
f"Simulation for {channel} with {len(topics)} unique topics and {10000 * len(topics)} rotations",
|
|
fontsize=fontsize,
|
|
)
|
|
plt.xlabel("# of times chosen", fontsize=fontsize)
|
|
plt.ylabel("Topics and Weights", fontsize=fontsize)
|
|
plt.xticks(fontsize=fontsize)
|
|
plt.yticks(fontsize=fontsize)
|
|
plt.grid(True)
|
|
|
|
fig.tight_layout()
|
|
|
|
fig.savefig(save_path, dpi=fig.dpi)
|
|
plt.close()
|
|
|
|
save_path.seek(0)
|
|
files = [discord.File(save_path, filename="graph.png")]
|
|
await ctx.send(files=files)
|
|
save_path.close()
|
|
|
|
@thread_rotate.command(name="manual")
|
|
async def thread_rotate_manual(self, ctx, channel: discord.TextChannel):
|
|
"""
|
|
Manually rotate a thread topic
|
|
"""
|
|
current = await self.config.channel(channel).topics()
|
|
if not current:
|
|
await ctx.send(error("That channel has not been setup for thread rotation!"), delete_after=30)
|
|
return
|
|
|
|
await self.rotate_thread(channel)
|
|
await ctx.tick()
|
|
|
|
@thread_rotate.command(name="method")
|
|
async def thread_rotate_method(self, ctx, channel: discord.TextChannel, method: str):
|
|
"""
|
|
Set the method for random selection of topics.
|
|
|
|
Possible methods are:
|
|
- replacement: select a channel with replacement, meaning the same channel can be selected again for every rotation.
|
|
- unique: select a channel without replacement, meaning for each rotation a new topic will be selected that was not selected previously, until all topics are exhausted.
|
|
"""
|
|
current = await self.config.channel(channel).topics()
|
|
if not current:
|
|
await ctx.send(error("That channel has not been setup for thread rotation!"), delete_after=30)
|
|
return
|
|
|
|
method = method.lower()
|
|
if method not in self.methods:
|
|
await ctx.send(
|
|
error(f"Unknown method {method}, please choose from this list: {humanize_list(self.methods)}."),
|
|
delete_after=60,
|
|
)
|
|
return
|
|
|
|
prev_method = await self.config.channel(channel).method()
|
|
|
|
if method == prev_method:
|
|
await ctx.send(info(f"Selection method for {channel.mention} is already set to {method}!"), delete_after=60)
|
|
return
|
|
|
|
if prev_method == "replacement":
|
|
await self.config.channel(channel).prev_topics.clear()
|
|
|
|
await self.config.channel(channel).method.set(method)
|
|
await ctx.tick()
|
|
|
|
@thread_rotate.command(name="interval")
|
|
async def thread_rotate_interval(self, ctx, channel: discord.TextChannel, interval: str):
|
|
"""
|
|
Modify the rotation interval for a thread rotation
|
|
|
|
The channel must of already been setup for thread rotation
|
|
|
|
This will apply on the next thread rotation for the channel!
|
|
"""
|
|
current = await self.config.channel(channel).topics()
|
|
if not current:
|
|
await ctx.send(error("That channel has not been setup for thread rotation!"), delete_after=30)
|
|
return
|
|
|
|
interval = parse_timedelta(interval.strip())
|
|
if interval is None:
|
|
await ctx.send(error("Invalid time interval, please try again!"), delete_after=60)
|
|
return
|
|
|
|
await self.config.channel(channel).rotation_interval.set(interval.total_seconds())
|
|
await ctx.tick()
|
|
|
|
@thread_rotate.command(name="roles")
|
|
async def thread_rotate_roles(self, ctx, channel: discord.TextChannel, *roles: discord.Role):
|
|
"""
|
|
Modify the ping roles for a thread rotation
|
|
|
|
The channel must of already been setup for thread rotation
|
|
"""
|
|
current = await self.config.channel(channel).topics()
|
|
if not current:
|
|
await ctx.send(error("That channel has not been setup for thread rotation!"), delete_after=30)
|
|
return
|
|
|
|
await self.config.channel(channel).ping_roles.set([r.id for r in roles])
|
|
await ctx.tick()
|
|
|
|
@thread_rotate.command(name="topics")
|
|
async def thread_rotate_topics(self, ctx, channel: discord.TextChannel):
|
|
"""
|
|
Modify topics for thread rotation.
|
|
|
|
The channel must of already been setup for thread rotation
|
|
"""
|
|
current = await self.config.channel(channel).topics()
|
|
if not current:
|
|
await ctx.send(error("That channel has not been setup for thread rotation!"), delete_after=30)
|
|
return
|
|
|
|
await ctx.send(info(f"{channel.mention}'s topics:"))
|
|
topic_msg = "Topics:\n"
|
|
for topic, weight in current.items():
|
|
topic_msg += f"{topic}: {weight}\n"
|
|
|
|
for page in pagify(topic_msg):
|
|
await ctx.send(box(page), delete_after=300)
|
|
|
|
await ctx.send(
|
|
info(
|
|
"Please list the thread topics and their selection weights.\nThe weight is how likely the topic will be choosen.\nA weight of `1` means it will not be choosen more or less than other topics.\nA weight between 0 and 1 means it is that weight times less likely to be choosen, with a weight of 0 meaning it will never be choosen.\nA weight greater than 1 means it will be that times more likely to be choosen.\n\nFor example, a weight of 1.5 means that topic is 1.5 more likely to be choose over the others. A weight of 0.5 means that topic is half as likely to be choosen over others.\n\nPlease use this format for listing the weights:\n"
|
|
),
|
|
delete_after=300,
|
|
)
|
|
msg = await ctx.send(
|
|
box("topic name: weight_value\ntopic 2 name: weight_value\ntopic 3 name: weight_value")
|
|
+ "\n\nYou can send as many messages as needed, when you are done, type `done`."
|
|
)
|
|
|
|
topic_msg = ""
|
|
while msg.content.lower() != "done":
|
|
pred = MessagePredicate.same_context(ctx)
|
|
try:
|
|
msg = await self.bot.wait_for("message", check=pred, timeout=301)
|
|
except asyncio.TimeoutError:
|
|
await ctx.send(error("Took too long, cancelling setup!"), delete_after=30)
|
|
return
|
|
topic_msg += msg.content + "\n"
|
|
|
|
topics = topic_msg.strip().split("\n")[:-1] # remove done from end
|
|
parsed_topics = {}
|
|
for topic in topics:
|
|
topic = topic.split(":")
|
|
try:
|
|
if len(topic) > 2:
|
|
parsed_topics[":".join(topic[0:-1])] = float(topic[-1])
|
|
else:
|
|
parsed_topics[topic[0]] = float(topic[-1])
|
|
|
|
if float(topic[-1]) < 0:
|
|
raise ValueError()
|
|
except:
|
|
await ctx.send(
|
|
error(
|
|
"Please make sure to use the correct format, every topic and weight should be split by a `:` and the weight should be a single decimal value greater than or equal to 0. Topic {topic} caused this error."
|
|
),
|
|
delete_after=60,
|
|
)
|
|
return
|
|
|
|
await self.config.channel(channel).topics.set(parsed_topics)
|
|
await ctx.send(info("Topics changed successfully!"), delete_after=60)
|
|
|
|
@thread_rotate.command(name="clear")
|
|
async def thread_rotate_clear(self, ctx, channel: discord.TextChannel):
|
|
"""
|
|
Clear a channel's thread rotation settings
|
|
"""
|
|
await ctx.send(
|
|
warning(f"Are you sure you want to delete all settings for {channel.mention}? This cannot be reversed."),
|
|
delete_after=31,
|
|
)
|
|
pred = MessagePredicate.yes_or_no(ctx)
|
|
try:
|
|
await self.bot.wait_for("message", check=pred, timeout=30)
|
|
except asyncio.TimeoutError:
|
|
await ctx.send(error("Took too long, cancelling clear!"), delete_after=30)
|
|
return
|
|
|
|
if not pred.result:
|
|
await ctx.send(info("Cancelling clear."), delete_after=30)
|
|
return
|
|
|
|
await self.config.channel(channel).clear()
|
|
await ctx.send(info(f"Settings for {channel.mention} cleared."), delete_after=30)
|
|
|
|
@thread_rotate.command(name="setup")
|
|
async def thread_rotate_setup(self, ctx, channel: discord.TextChannel):
|
|
"""
|
|
Interactively setup a thread rotation for a channel
|
|
"""
|
|
guild = ctx.guild
|
|
now = datetime.now()
|
|
|
|
await ctx.send(
|
|
info(
|
|
"Welcome to the thread rotation setup wizard!\n\nFirst, please specifiy the rotation interval. Rotation intervals can be formatted as follows:\n\t5 minutes\n\t1 minute 30 seconds\n\t1 hour\n\t2 days\n\t30 days\n\t5h30m\n\t(etc)"
|
|
),
|
|
delete_after=300,
|
|
)
|
|
|
|
pred = MessagePredicate.same_context(ctx)
|
|
try:
|
|
msg = await self.bot.wait_for("message", check=pred, timeout=121)
|
|
except asyncio.TimeoutError:
|
|
await ctx.send(error("Took too long, cancelling setup!"), delete_after=30)
|
|
return
|
|
|
|
interval = parse_timedelta(msg.content.strip())
|
|
if interval is None:
|
|
await ctx.send(error("Invalid time interval, please run setup again!"), delete_after=60)
|
|
return
|
|
|
|
await ctx.send(
|
|
info(
|
|
"Thank you.\n\nNow, please specify the date and time to start rotation. You can say `now` to start rotation as soon as setup is complete.\n\nValid date formats are:\n\tFebruary 14 at 6pm EDT\n\t2019-04-13 06:43:00 PST\n\t01/20/18 at 21:00:43\n\t(etc)"
|
|
),
|
|
delete_after=300,
|
|
)
|
|
|
|
pred = MessagePredicate.same_context(ctx)
|
|
try:
|
|
msg = await self.bot.wait_for("message", check=pred, timeout=121)
|
|
except asyncio.TimeoutError:
|
|
await ctx.send(error("Took too long, cancelling setup!"), delete_after=30)
|
|
return
|
|
|
|
if msg.content.strip().lower() == "now":
|
|
date = datetime.now()
|
|
else:
|
|
date = parse_time(msg.content.strip())
|
|
|
|
if date is None:
|
|
await ctx.send(error("Invalid date, please run setup again!"), delete_after=60)
|
|
return
|
|
|
|
if date < now:
|
|
await ctx.send(
|
|
error("Invalid date, the date must be in the future! Please run the setup again."), delete_after=60
|
|
)
|
|
|
|
await ctx.send(
|
|
info(
|
|
"Great, next step is to list all roles that should be pinged and added to each thread when it rotates.\n\nList each role **seperated by a comma `,`**.\nYou can use role IDs, role mentions, or role names. If you do not want to ping any roles type `next` or `no`."
|
|
),
|
|
delete_after=300,
|
|
)
|
|
|
|
pred = MessagePredicate.same_context(ctx)
|
|
try:
|
|
msg = await self.bot.wait_for("message", check=pred, timeout=241)
|
|
except asyncio.TimeoutError:
|
|
await ctx.send(error("Took too long, cancelling setup!"), delete_after=30)
|
|
return
|
|
|
|
if msg.content.strip().lower() != "no" and msg.content.strip().lower() != "next":
|
|
roles = [m.strip().strip("<").strip(">").strip("@").strip("&") for m in msg.content.split(",")]
|
|
role_objs = []
|
|
for r in roles:
|
|
try:
|
|
role = guild.get_role(int(r))
|
|
except:
|
|
role = discord.utils.find(lambda c: c.name == r, guild.roles)
|
|
if role is None:
|
|
await ctx.send(
|
|
error(
|
|
f"Unknown role: `{r}`, please run the command again. (make sure to seperate roles by commas!)"
|
|
),
|
|
delete_after=60,
|
|
)
|
|
return
|
|
|
|
role_objs.append(role)
|
|
else:
|
|
role_objs = []
|
|
|
|
await ctx.send(
|
|
info(
|
|
"Next step: select the method I will use to rotate topics with.\nAvailable methods are:\n\t- `replacement`: the same channel can be selected again for every rotation.\n\t- `unique`: for each rotation a new topic will be selected that was not selected previously, until all topics are exhausted."
|
|
),
|
|
delete_after=300,
|
|
)
|
|
|
|
pred = MessagePredicate.same_context(ctx)
|
|
try:
|
|
msg = await self.bot.wait_for("message", check=pred, timeout=241)
|
|
except asyncio.TimeoutError:
|
|
await ctx.send(error("Took too long, cancelling setup!"), delete_after=30)
|
|
return
|
|
|
|
method = msg.content.lower().strip()
|
|
if method not in self.methods:
|
|
await ctx.send(
|
|
error(f"Unknown method {method}, please choose from this list: {humanize_list(self.methods)}."),
|
|
delete_after=30,
|
|
)
|
|
return
|
|
|
|
await ctx.send(
|
|
info(
|
|
"Final step is to list the thread topics and their selection weights.\nThe weight is how likely the topic will be choosen.\nA weight of `1` means it will not be choosen more or less than other topics.\nA weight between 0 and 1 means it is that weight times less likely to be choosen, with a weight of 0 meaning it will never be choosen.\nA weight greater than 1 means it will be that times more likely to be choosen.\n\nFor example, a weight of 1.5 means that topic is 1.5 more likely to be choose over the others. A weight of 0.5 means that topic is half as likely to be choosen over others.\n\nPlease use this format for listing the weights:\n"
|
|
),
|
|
delete_after=300,
|
|
)
|
|
msg = await ctx.send(
|
|
box("topic name: weight_value\ntopic 2 name: weight_value\ntopic 3 name: weight_value")
|
|
+ "\n\nYou can send as many messages as needed, when you are done, type `done`."
|
|
)
|
|
|
|
topic_msg = ""
|
|
while msg.content.lower() != "done":
|
|
pred = MessagePredicate.same_context(ctx)
|
|
try:
|
|
msg = await self.bot.wait_for("message", check=pred, timeout=301)
|
|
except asyncio.TimeoutError:
|
|
await ctx.send(error("Took too long, cancelling setup!"), delete_after=30)
|
|
return
|
|
topic_msg += msg.content + "\n"
|
|
|
|
topics = topic_msg.strip().split("\n")[:-1] # remove done from end
|
|
parsed_topics = {}
|
|
for topic in topics:
|
|
topic = topic.split(":")
|
|
try:
|
|
if len(topic) > 2:
|
|
parsed_topics[":".join(topic[0:-1])] = float(topic[-1])
|
|
else:
|
|
parsed_topics[topic[0]] = float(topic[-1])
|
|
|
|
if float(topic[-1]) < 0:
|
|
raise ValueError()
|
|
except:
|
|
await ctx.send(
|
|
error(
|
|
"Please make sure to use the correct format, every topic and weight should be split by a `:` and the weight should be a single decimal value greater than or equal to 0. Topic {topic} caused this error."
|
|
),
|
|
delete_after=60,
|
|
)
|
|
return
|
|
|
|
topic_msg = "Topics:\n"
|
|
for topic, weight in parsed_topics.items():
|
|
topic_msg += f"{topic}: {weight}\n"
|
|
|
|
await ctx.send(
|
|
info(f"Please review the settings for thread rotation on channel {channel.mention}:"),
|
|
delete_after=300,
|
|
)
|
|
await ctx.send(
|
|
box(
|
|
f"Rotation interval: {humanize_timedelta(seconds=interval.total_seconds())}\n\nRotation Start: {date}\n\nRotation Method: `{method}`\n\nPing roles: {humanize_list([r.name for r in role_objs])}"
|
|
),
|
|
delete_after=300,
|
|
)
|
|
for page in pagify(topic_msg):
|
|
await ctx.send(
|
|
box(page),
|
|
delete_after=300,
|
|
)
|
|
|
|
await ctx.send("Type yes to confirm the thread rotation, type no to cancel thread rotation setup.")
|
|
|
|
pred = MessagePredicate.yes_or_no(ctx)
|
|
|
|
try:
|
|
msg = await self.bot.wait_for("message", check=pred, timeout=240)
|
|
except asyncio.TimeoutError:
|
|
await ctx.send(error("Took too long, cancelling setup!"), delete_after=30)
|
|
return
|
|
|
|
if not pred.result:
|
|
await ctx.send(info("Cancelled setup."), delete_after=60)
|
|
return
|
|
|
|
# setup the channel
|
|
await self.config.channel(channel).topics.set(parsed_topics)
|
|
await self.config.channel(channel).ping_roles.set([r.id for r in role_objs])
|
|
await self.config.channel(channel).rotation_interval.set(interval.total_seconds())
|
|
await self.config.channel(channel).rotate_on.set(int(date.timestamp()))
|
|
await self.config.channel(channel).method.set(method)
|
|
|
|
await ctx.send(
|
|
info(f"Thread rotation setup! The first rotation will start at <t:{int(date.timestamp())}>"),
|
|
delete_after=60,
|
|
)
|