#!/usr/bin/python3

import logging
import logging.handlers

import os
from subprocess import call
from pathlib import Path

import csv
from datetime import datetime, timedelta
import time
import dateutil.parser as dp

import pandas as pd
import numpy as np

import matplotlib.pyplot as plt
import matplotlib.cbook as cbook
import matplotlib.dates as mdates
from matplotlib.ticker import FuncFormatter
from matplotlib import rcParams

rcParams["font.family"] = "sans-serif"
rcParams["font.sans-serif"] = ["DejaVu Sans"]
rcParams["font.weight"] = "medium"
rcParams["font.size"] = 10.0
rcParams["axes.linewidth"] = 0.4
rcParams["legend.fontsize"] = 8.0

from scipy.interpolate import spline
from scipy.interpolate import interp1d
from scipy.interpolate import UnivariateSpline

from twitter import *

CSV_FILE = "/home/ubuntu/bndwdth.csv"

SPEED_MAX_DOWN = 367001600          # 350 Mbps
SPEED_MIN_DOWN = SPEED_MAX_DOWN / 2 # 175 Mbps
SPEED_MAX_UP   = 20971520           #  20 Mbps
SPEED_MIN_UP   = SPEED_MAX_UP / 2   #  10 Mbps

DOWN_LINE_COLOUR  = "blue"
UP_LINE_COLOUR    = "magenta"
PING_LINE_COLOUR  = "orange"
TTFB_LINE_COLOUR = "red"

my_logger = logging.getLogger('bndwdth-log')
my_logger.setLevel(logging.DEBUG)
my_logger.addHandler(logging.handlers.SysLogHandler(address = '/dev/log'))

def notice(s):
    my_logger.debug(s)
    print(s)

def mean(numbers):
    return float(sum(numbers)) / max(len(numbers), 3)

#
## Get bandwidth in human readable form
#

def ghr(b, u=True):
    return get_human_readable(b, u)

def get_human_readable(bps, units=True):
    K = 1024
    M = K * K
    G = M * K
    Bps = bps / 8
    if bps < K:
        return "%d%s" % (bps, "bps" if units else "")
    if bps < M:
        return "%d%s" % ((bps / K), "Kbps" if units else "")
    if bps < G:
        return "%d%s" % ((bps / M), "Mbps" if units else "")
    return "%d%s" % ((bps / G), "Gbps" if units else "")

def sufficiently_fast_up(speed):
    return speed["up"] > SPEED_MIN_UP

def sufficiently_fast_down(speed):
    return speed["down"] > SPEED_MIN_DOWN

def host_missing(speed):
    return speed["up"] == 0 and speed["down"] == 0

def insufficiently_fast(speed):
    return speed["up"] < SPEED_MIN_UP and speed["down"] < SPEED_MIN_DOWN

def sufficiently_fast(speed):
    return sufficiently_fast_up(speed) and sufficiently_fast_down(speed)

def extra_fast(speed):
    return speed["up"] > SPEED_MAX_UP and speed["down"] > SPEED_MAX_DOWN

#
## Generate a graph of the previous 24 hours bandwidth
#

def parse_bandwidth(duration=24):
    date = []
    down = []
    up = []
    ping = []
    ttfb = []
    with open(CSV_FILE, "r") as c:
        reader = csv.reader(c)
        now = datetime.now()
        for row in reader:
            d = dp.parse(row[0])
            if now - d < timedelta(hours=duration):
                date.append(d)
                if len(row) > 1:
                    down.append(int(row[1]))
                    up.append(int(row[2]))
                    if len(row) > 3:
                        ping.append(float(row[3]))
                        if len(row) > 4:
                            ttfb.append(float(row[4]))
                        else:
                            ttfb.append(0)
                    else:
                        ping.append(0)
                        ttfb.append(0)
                else:
                    down.append(0)
                    up.append(0)
                    ping.append(0)
                    ttfb.append(0)

    return { "date": date, "down": down, "up": up, "ping": ping, "ttfb": ttfb }

def average_bandwidth(data, avg=2):
    date = []
    down = []
    up = []
    ping = []
    ttfb = []
    for d in range(int(len(data["date"]) / avg)):
        t = d * avg
        date.append(data["date"][t])

        ad = 0
        for i in range(avg):
            ad = ad + data["down"][t + i]
        ad = ad / avg
        down.append(ad)

        au = 0
        for i in range(avg):
            au = au + data["up"][t + i]
        au = au / avg
        up.append(au)

        ap = 0
        for i in range(avg):
            ap = ap + data["ping"][t + i]
        ap = ap / avg
        ping.append(ap)

        ac = 0
        for i in range(avg):
            ac = ac + data["ttfb"][t + i]
        ac = ac / avg
        ttfb.append(ac)

    return { "date": date, "down": down, "up": up, "ping": ping, "ttfb": ttfb }

def generate_graph(data, scale="log", xkcd=False):
    if xkcd:
        plt.xkcd()
    fig, speed_plot = plt.subplots()
    duration_plot = speed_plot.twinx()

    time = []
    timestamp = []
    down = []
    up = []
    ping = []
    ttfb = []

    data_points = len(data["date"])

    for i in range(data_points):
        time.append(data["date"][i])
        timestamp.append(float(data["date"][i].timestamp()))
        down.append(float(data["down"][i]) / 1000000)
        up.append(float(data["up"][i]) / 1000000)
        ping.append(float(data["ping"][i]))
        ttfb.append(float(data["ttfb"][i]))

    now = datetime.now()
    plt.suptitle("Bandwidth [" + now.strftime("%Y-%m-%d %-I%p") + "]")
    speed_plot.set_xlabel("Date / Time")
    speed_plot.set_ylabel("Bandwidth (Mbps)")
    duration_plot.set_ylabel("Delay (ms)")

    # plot the original data points
    #speed_plot.plot(time, down, DOWN_LINE_COLOUR, linewidth=0, marker=".")
    #speed_plot.plot(time, up, UP_LINE_COLOUR, linewidth=0, marker=".")
    #duration_plot.plot(time, ping, PING_LINE_COLOUR, linewidth=0, marker=".")
    #duration_plot.plot(time, ttfb, TTFB_LINE_COLOUR, linewidth=0, marker=".")

    # calculate (and draw) a pretty looking line
    timestamp_min = np.array(timestamp).min()
    timestamp_max = np.array(timestamp).max()

    time_continuous = np.linspace(timestamp_min, timestamp_max, (timestamp_max - timestamp_min) / 60)
    down_continuous = interp1d(timestamp, down, kind="cubic")(time_continuous)
    up_continuous = interp1d(timestamp, up, kind="cubic")(time_continuous)
    ping_continuous = interp1d(timestamp, ping, kind="cubic")(time_continuous)
    #ttfb_continuous_i = interp1d(timestamp, ttfb, kind="cubic")(time_continuous)

    ttfb_weight = []
    tm = max(ttfb)
    for i in range(len(ttfb)):
        ttfb_weight.append(2.75 / ttfb[i])
    ttfb_continuous = UnivariateSpline(timestamp, ttfb, w=ttfb_weight, k=3)(time_continuous)

    time_extra = []
    down_expected = []
    up_expected = []
    for i in time_continuous:
        time_extra.append(datetime.fromtimestamp(i))
        down_expected.append(float(ghr(SPEED_MAX_DOWN, False)))
        up_expected.append(float(ghr(SPEED_MAX_UP, False)))

    for i in range(len(ping_continuous)):
        if ping_continuous[i] < 1:
            ping_continuous[i] = np.NaN

    for i in range(len(ttfb_continuous)):
        #print(datetime.fromtimestamp(time_continuous[i]).strftime('%Y-%m-%d %H:%M:%S') + "," + str(ttfb_continuous[i]))
        if ttfb_continuous[i] < 1:
            ttfb_continuous[i] = np.NaN

    ln = speed_plot.plot(time_extra, down_continuous, DOWN_LINE_COLOUR, linewidth=1, label="Down")
    ln = ln + speed_plot.plot(time_extra, down_expected, DOWN_LINE_COLOUR, linewidth=1, linestyle="dashed", label="Expected: 350Mbps")
    ln = ln + speed_plot.plot(time_extra, up_continuous, UP_LINE_COLOUR, linewidth=1, label="Up")
    ln = ln + speed_plot.plot(time_extra, up_expected, UP_LINE_COLOUR, linewidth=1, linestyle="dashed", label="Expected: 20Mbps")

    ln = ln + duration_plot.plot(time_extra, ttfb_continuous, TTFB_LINE_COLOUR, linewidth=1, label="TTFB")
    ln = ln + duration_plot.plot(time_extra, ping_continuous, PING_LINE_COLOUR, linewidth=1, label="Ping")

    # finish up
    plt.legend(ln, [l.get_label() for l in ln], loc= "upper center", ncol=4, bbox_to_anchor=(0.5, 1.06), framealpha=1).get_frame().set_linewidth(0.4)

    speed_plot.xaxis.set_major_formatter(mdates.DateFormatter("%m/%d %H:%M"))

    plt.xlim([time_extra[0], datetime.now()])

    speed_plot.set_yscale(scale)
    speed_plot.set_ylim(1, 1000)
    speed_plot.yaxis.set_major_formatter(FuncFormatter(lambda y, p: format(int(y), ',')))

    duration_plot.set_yscale(scale)
    duration_plot.set_ylim(1, max(data["ping"] + data["ttfb"]) * 1.5)
    duration_plot.yaxis.set_major_formatter(FuncFormatter(lambda y, p: format(int(y), ',')))
    fig.autofmt_xdate()
    fig.set_size_inches(10, 8)

    out = "%s-bw-%s.png" % (datetime.now().strftime("%Y%m%d%H%M"), scale)
    plt.savefig(out, dpi=320)
    #plt.show()
    return out

#
## Send the latest results to Twitter
#

def get_twitter(up=False):
    # token, token_secret, consumer_key, consumer_secret
    oa = OAuth("", "", "", "")
    if up:
        return Twitter(domain="upload.twitter.com", auth=oa)
    else:
        return Twitter(auth=oa)

def alert_network_down():
    t = get_twitter()
    s = "@virginmedia it looks like there might be network trouble in #Watford! My #awscloud EC2 instance has lost its connection with my home PC."
    t.statuses.update(status=s);
    call(["/home/ubuntu/bin/bndwdth-notify", "down"])

def send_tweet(speed, graphs):
    # upload graph to twitter
    img_ids = []
    for graph in graphs:
        img = None
        with open(graph, "rb") as i:
            img = i.read()
        t = get_twitter(True)
        img_ids.append(t.media.upload(media=img)["media_id_string"])

    # send tweet
    t = get_twitter()
    s = "Average #bandwidth in #Watford for the last hour was %s down and %s up." % (ghr(speed["down"]), ghr(speed["up"]))
    if insufficiently_fast(speed):
        s = "@virginmedia Why's my hourly average #bandwidth in #Watford only %s down / %s up? Not even close to %s/%s!" % (ghr(speed["down"]), ghr(speed["up"]), ghr(SPEED_MAX_DOWN, False), ghr(SPEED_MAX_UP, False))
        #call(["/home/ubuntu/bin/bndwdth-notify", "slow"])
    elif not sufficiently_fast_down(speed):
        s = "@virginmedia Why's my hourly average download #bandwidth in #Watford only %s? Nowhere close to %s! (%s up)" % (ghr(speed["down"]), ghr(SPEED_MAX_DOWN, False), ghr(speed["up"]))
        #call(["/home/ubuntu/bin/bndwdth-notify", "slow/down"])
    elif not sufficiently_fast_up(speed):
        s = "@virginmedia Why's my hourly average upload #bandwidth in #Watford only %s? Not even close to %s! (%s down)" % (ghr(speed["up"]), ghr(SPEED_MAX_UP, False), ghr(speed["down"]))
        #call(["/home/ubuntu/bin/bndwdth-notify", "slow/up"])
    elif extra_fast(speed):
        s = "@virginmedia Excellent: my broadband #bandwidth in #Watford is %s down and %s up! #VIVID350 #superfast" % (ghr(speed["down"]), ghr(speed["up"]))
        call(["/home/ubuntu/bin/bndwdth-notify", "fast"])
    t.statuses.update(status=s, media_ids=",".join(img_ids));
    notice(s)

if __name__ == "__main__":
    s = parse_bandwidth(duration=1)
    s = average_bandwidth(s, len(s["date"]))
    s["up"] = s["up"][0]
    s["down"] = s["down"][0]

    alert = Path("bndwdth-alert")
    if int(s["up"]) > 0 and int(s["down"]) > 0:
        if alert.is_file():
            alert.unlink()
        d = parse_bandwidth()
        d = average_bandwidth(d)
        g = []
        g.append(generate_graph(d, scale="log"))
        send_tweet(s, g)
    else:
        if not alert.is_file():
            alert.open("a")
            alert_network_down()

    for f in g:
        os.remove(f)