#!/usr/bin/python
#
# email_stats.py
#
# Time vs date plot for an email mbox file.
#
# Copyright (C) 2015 George C. Privon
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.


import mailbox
import datetime
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
import argparse
import re
import numpy as np
import cubehelix


parser = argparse.ArgumentParser()
parser.add_argument('mbox', help='Mailbox to analyze.')
parser.add_argument('--plotfile', '-p', default=False, action='store',
                    help='Name of output plotting file.')
parser.add_argument('--title', '-t', default='Email Send Times',
                    action='store', help='Plot title.')
parser.add_argument('--sendercolors', '-s', default=False, action='store',
                    help='Comma separated list of search strings for the \
                          sender field. Each will be displayed with a \
                          different color.')
args = parser.parse_args()

plt.figure()
plt.ylim([0, 24])
plt.yticks(4*np.arange(7))
plt.ylabel('Hour')
plt.xlabel('Date')
plt.minorticks_on()
plt.title(args.title)
a = mailbox.mbox(args.mbox)

scolor = cubehelix.cmap(startHue=240, endHue=-300,
                        minSat=1, maxSat=2.5,
                        minLight=.3, maxLight=.8,
                        gamma=.9)

pldata = {}
if args.sendercolors:
    slist = args.sendercolors.split(',')
    nsend = len(slist)
    for item in slist:
        pldata[item] = []
pldata['unknown'] = []

for msg in a:
    cid = None
    label = None
    if msg['date'] is not None:
        try:
            z = datetime.datetime.strptime(msg['date'], '%a, %d %b %Y %H:%M:%S %z')
        except ValueError:
            try:
                z = datetime.datetime.strptime(msg['date'], '%a, %d %b %Y %H:%M:%S %Z')
            except ValueError:
                print("Skipping message from " + msg['date'])
                continue
        if args.sendercolors:
            for search in enumerate(slist):
                if re.search(search[1], msg['From'], re.IGNORECASE):
                    cid, label = search
                    pldata[label].append([z.date(), z.hour + z.minute/60.])
                    break
        if cid is None:
            pldata['unknown'].append([z.date(), z.hour + z.minute/60.])

if args.sendercolors:
    for plid in enumerate(slist):
        plt.plot_date(np.array(pldata[plid[1]])[:,0],
                      np.array(pldata[plid[1]])[:,1],
                      color=scolor(plid[0] / (nsend + 1)),
                      marker='.',
                      #tz=z.tzname(),
                      label=plid[1],
                      xdate=True)

if len(pldata['unknown']) > 0:
    plt.plot_date(np.array(pldata['unknown'])[:,0],
                  np.array(pldata['unknown'])[:,1],
                  color=scolor(1),
                  marker='.',
                  #tz=z.tzname(),
                  label='unknown',
                  xdate=True)

if args.sendercolors:
    plt.legend(loc='upper left', ncol=nsend + 1)

if args.plotfile:
    plt.savefig(args.plotfile)
else:
    plt.savefig('email_times.png')