activity.py

import matplotlib
matplotlib.use('Agg')

import matplotlib.pyplot as plt

from cogs.utils.dataIO import dataIO
from __main__ import send_cmd_help
from discord.ext import commands
import asyncio

import calendar
import discord
import time
import os
import numpy as np

cog_name = 'activity'
data_name = 'activity'


class Activity:
    def __init__(self, bot):
        self.bot = bot
        self.memory = {}
        for f in os.listdir('data/{}'.format(cog_name)):
            d = dataIO.load_json('data/{}/{}'.format(cog_name, f))
            member = f.split('.')[0]
            self.memory[member] = d

    async def listener(self, message):
        if not message.channel.is_private and self.bot.user.id != message.author.id:
            author = message.author
            timestamp = message.timestamp
            filename = 'data/{}/{}.json'.format(cog_name, author.id)
            if not dataIO.is_valid_json(filename):
                data = {}
                data['HOUR'] = {}
                for i in range(1, 25):
                    data['HOUR'][str(i)] = 0
                dataIO.save_json(filename, data)
                self.memory[author.id] = data
            if str(timestamp.year) not in self.memory[author.id]:
                self.memory[author.id][str(timestamp.year)] = {}
                for i in range(1, 13):
                    self.memory[author.id][str(timestamp.year)][str(i)] = 0
            self.memory[author.id]['HOUR'][str(timestamp.hour)] += 1
            self.memory[author.id][str(timestamp.year)][str(timestamp.month)] += 1
            dataIO.save_json(filename, self.memory[author.id])

    @commands.group(pass_context=True, no_pm=True, name='activity', aliases=['ac'])
    async def _activity(self, context):
        if context.invoked_subcommand is None:
            await send_cmd_help(context)

    @_activity.command(pass_context=True, name='hour', aliases=['hr'])
    async def _hour(self, context, *username: discord.Member):
        if username:
            author = username[0]
        else:
            author = context.message.author
        filename = 'data/{}/{}.json'.format(cog_name, author.id)
        if dataIO.is_valid_json(filename):
            data = data = dataIO.load_json(filename)
            t = str(time.time())
            hours = [int(hours) for hours in data['HOUR']]
            messages = [message for message in data['HOUR'].values()]

            plt.bar(hours, messages, width=1.0, facecolor='black', edgecolor='black')
            plt.title('Hourly activity of {} - Total Messages: {}'.format(author.display_name, sum(messages)), y=1.04)
            plt.xlabel('Hour (UTC)')
            plt.xticks(np.arange(24))
            plt.yticks([])
            plt.xlim([0, 23])
            plt.ylim([0, sum(messages)+(sum(messages)/2)])
            plt.savefig('{}.png'.format(t))
            plt.clf()
            try:
                await self.bot.send_file(context.message.channel, '{}.png'.format(t))
            except discord.Forbidden:
                await self.bot.say('I don\'t have any permissions to attach files!')
            os.remove('{}.png'.format(t))

    @_activity.command(pass_context=True, name='year', aliases=['yr'])
    async def _year(self, context, year: str, *username: discord.Member):
        if username:
            author = username[0]
        else:
            author = context.message.author
        server = context.message.server
        filename = 'data/{}/{}.json'.format(cog_name, author.id)
        if dataIO.is_valid_json(filename):
            data = data = dataIO.load_json(filename)
            if year in data:
                t = str(time.time())
                months = [int(months) for months in data[year]]
                messages = [message for message in data[year].values()]
                plt.bar(months, messages, width=1.0, facecolor='black', edgecolor='black')
                plt.title('Monthly activity of {} for {}'.format(author.display_name, server.name, str(year)), y=1.04)
                plt.ylabel('Messages')
                plt.xlabel('Total Messages: {}'.format(sum(messages)))
                plt.xticks(np.arange(13), calendar.month_abbr[0:13])
                plt.xlim([1, 12])
                plt.yticks([])
                plt.ylim([0, sum(messages)+(sum(messages)/2)])
                plt.savefig('{}.png'.format(t))
                plt.autoscale(enable=False, axis="x")
                plt.clf()
                try:
                    await self.bot.send_file(context.message.channel, '{}.png'.format(t))
                except discord.Forbidden:
                    await self.bot.say('I don\'t have any permissions to attach files!')
                os.remove('{}.png'.format(t))


def check_folder():
    if not os.path.exists('data/{}'.format(cog_name)):
        print('Creating data/{}'.format(cog_name))
        os.makedirs('data/{}'.format(cog_name))


def setup(bot):
    check_folder()
    n = Activity(bot)
    bot.add_listener(n.listener, 'on_message')
    bot.add_cog(n)