#!/usr/bin/python
import argparse
import csv
import datetime as dt
import matplotlib.pyplot as plt
from matplotlib import dates
import os.path

def initdict(dists,distname) :
    dists[distname] = {'date' : [], 'total' : [], 'broken' : [], 'outdated' : [] }
    return dists

def parse(dataset):
    dists = {}
    csv_reader = csv.reader(dataset,delimiter=' ')
    for line in sorted(csv_reader,key=lambda k: k[0]):
        distname = line[1]
        if distname not in dists :
            dists = initdict(dists,distname)

        dists[distname]['date'].append(dt.datetime.strptime(line[0],'%Y%m%dT%H%M%SZ'))
        dists[distname]['total'].append(int(line[2]))
        dists[distname]['broken'].append(int(line[3]))
        dists[distname]['outdated'].append(int(line[4]))

    return dists

def plot(dists,distlist,output) :

    fig = plt.figure()
    fig.suptitle("Evalution during wheezy release cycle")
    ax1 = fig.add_subplot(311,title='Total Packages vs Time')
    ax2 = fig.add_subplot(312,title='Non-Installable Packages vs Time')
    ax3 = fig.add_subplot(313,title='Outdated Packages vs Time')
    for k in distlist :
        ax1.plot(dists[k]['date'],dists[k]['total'],'o-',label=k.capitalize())
        ax2.plot(dists[k]['date'],dists[k]['broken'],'o-',label=k.capitalize())
        ax3.plot(dists[k]['date'],dists[k]['outdated'],'o-',label=k.capitalize())

    if len(distlist) > 1 :
        ax1.legend(loc='upper left')
        ax2.legend()
        ax3.legend(loc='upper left')

    fig.autofmt_xdate()

    plt.savefig(output)

# plot two distribution with different scales
def plotmultiscale(dists,dist1,dist2,output) :

    fig = plt.figure()
    fig.suptitle("Evalution during wheezy release cycle")
    fig.autofmt_xdate()

    ax1 = fig.add_subplot(311,title='Total Packages vs Time')
    ax1.plot(dists[dist1]['date'],dists[dist1]['total'],'o-',label=dist1.capitalize())
    ax1.plot(dists[dist2]['date'],dists[dist2]['total'],'s-',label=dist2.capitalize())
    ax1.legend(loc='upper left')
    ax1.xaxis.set_visible(False)

    ax2 = fig.add_subplot(312,title='Non-Installable Packages vs Time')
    ax2.plot(dists[dist1]['date'],dists[dist1]['broken'],'o-',label=dist1.capitalize())
    ax2.xaxis.set_visible(False)

    ax22 = ax2.twinx() 
    ax22.plot(dists[dist2]['date'],dists[dist2]['broken'],'gs-',label=dist2.capitalize())
    ax22.set_ylim(0, 20)

    ax3 = fig.add_subplot(313,title='Outdated Packages vs Time')
    ax3.plot(dists[dist1]['date'],dists[dist1]['outdated'],'o-',label=dist1.capitalize())

    ax33 = ax3.twinx() 
    ax33.plot(dists[dist2]['date'],dists[dist2]['outdated'],'gs-',label=dist2.capitalize())
    ax33.set_ylim(0, 10)
    # ax33.set_yticks([-1,0,1])

    plt.setp(ax3.xaxis.get_majorticklabels(), rotation=30)

    plt.savefig(output)


def main():
    parser = argparse.ArgumentParser(description='plot outdated/broken')
    parser.add_argument('-v', '--verbose')
    parser.add_argument('-d', '--debug', action='store_true', default=False)
    parser.add_argument('-o', '--output', action='store')
    parser.add_argument('-s', '--split', action='store_true', default=False)
    parser.add_argument('-m', '--multiscale', action='store_true', default=False)
    parser.add_argument('dataset', type=str, nargs=1, help="dataset")
    parser.add_argument('distlist', type=str, nargs='*', help="distribution list")
    args = parser.parse_args()
 
    fname = args.dataset[0]

    dataset = open(fname)
    dists = parse(dataset)
    dataset.close()

    if args.distlist :
        distlist = args.distlist
    else :
        distlist = dists.keys()

    if args.multiscale :
        output = "aggregate-%s.png" % os.path.splitext(fname)[0]
        plotmultiscale(dists, distlist[0], distlist[1], output)
    else :
        if args.split :
            # different graphs, one for distribution
            for d in distlist :
                output = "%s-%s.png" % (d,os.path.splitext(fname)[0])
                plot(dists, [d] , output)
        else :
            # all in one graph
            output = "aggregate-%s.png" % os.path.splitext(fname)[0]
            plot(dists, distlist ,output)
        
 
 
if __name__ == '__main__':
    main()