#!/usr/bin/env python3
import os
from argparse import ArgumentParser
from datetime import datetime

import pandas as pd
import matplotlib.pyplot as plt


def plot_single_transmission(csv_file):
    global direction
    if direction is None:
        direction = "uplink" if "uplink" in csv_file else "downlink"
    bitrate = csv_file.split("bitrate")[1].split("_")[0]
    transmission_df = pd.read_csv(csv_file)
    if len(transmission_df) == 0:
        print("Empty csv.")
        return
    transmission_df["departure_time"] = transmission_df["departure_time"] - transmission_df["departure_time"].iloc[0]
    transmission_df["arrival_time"] = transmission_df["arrival_time"] - transmission_df["arrival_time"].iloc[0]

    transmission_df["data.len_send_cum"] = transmission_df[
        "data.len_send"
    ].cumsum()
    transmission_df["data.len_received_cum"] = transmission_df[
        "data.len_received"
    ].cumsum()

    lossed_packets = transmission_df[transmission_df['delay'].isna()]

    print("Packetloss: {}%".format((len(lossed_packets.departure_time) / len(transmission_df.departure_time)) * 100 ))

    fig, all_axes = plt.subplots(2, 1)
    fig.suptitle("CBR {} {}".format(bitrate, direction))
    fig.subplots_adjust(hspace=0.3)
    axes = all_axes[0]
    axes_boxplot = all_axes[1]

    transmission_df.plot(ax=axes, x="departure_time", y="data.len_send_cum")
    transmission_df.plot(ax=axes, x="arrival_time", y="data.len_received_cum")
    lossed_packets.plot.scatter(ax=axes, x="departure_time", y="data.len_send_cum", color="y", marker="X")

    axes.legend(["D(t)", "A(t)"])
    axes.set_ylabel("Byte")
    axes.set_xlabel("t [s]")

    delay_axes = axes.twinx()
    delay_axes.set_ylabel("delay [s]")
    delay_axes.set_ylim([0, 0.4])
    transmission_df.plot(ax=delay_axes, x="arrival_time", y="delay", color="g", linestyle="dotted")

    transmission_df.boxplot(ax=axes_boxplot, column="delay", vert=False)
    axes_boxplot.set_xlabel("[s]")

    fig.savefig(csv_file.replace(".csv", ".pdf"))
    plt.close("all")

    global all_delays_dict
    if bitrate not in all_delays_dict:
        all_delays_dict[bitrate] = list()
    all_delays_dict[bitrate] += transmission_df["delay"].tolist()


if __name__ == "__main__":
    start_time = datetime.now()
    parser = ArgumentParser()

    parser.add_argument("-f", "--folder", required=True, help="Folder with csv files.")
    parser.add_argument(
        "-o",
        "--output",
        required=False,
        default=None,
        help="Destination for plots as pdf.",
    )
    parser.add_argument(
        "--name_contains", required=False, default=".csv", help="Substring in filename."
    )
    parser.add_argument(
        "--use_subdirs", required=False, default=False, help="Use csvs in subdirs."
    )

    args = parser.parse_args()

    csv_list = list()
    all_delays_dict = dict()
    direction = None

    if args.use_subdirs:
        sub_dirs = os.listdir(args.folder)
        for d in sub_dirs:
            filenames = os.listdir("{}{}".format(args.folder, d))
            for filename in filenames:
                if filename.endswith(".csv") and args.name_contains in filename:
                    csv_list.append("{}/{}".format(d, filename))
    else:
        filenames = os.listdir(args.folder)
        for filename in filenames:
            if filename.endswith(".csv") and args.name_contains in filename:
                csv_list.append(filename)

    csv_list.sort()
    for csv in csv_list:
        print(csv)
        plot_single_transmission("{}/{}".format(args.folder, csv))

    if len(all_delays_dict) > 0:
        #fig, axes = plt.subplots()
        label_list = list()
        unit = None
        datas = list()
        for bitrate, l in all_delays_dict.items():
            if unit is None:
                if "K" in bitrate:
                    unit = "Kbit/s"
                elif "M" in bitrate:
                    unit = "Mbit/s"
                else:
                    unit = "Gbit/s"
            label_list.append(bitrate.replace("K", "").replace("M", "").replace("G", ""))
            datas.append(l)
            print("Bitrate: {}, len: {}".format(bitrate, len(l)))
        #print(datas[1])
        print(len(datas))
        plt.boxplot(datas[1])
        #plt.xticks(range(1, len(label_list) + 1), label_list, rotation=-45)


        #axes.set_xticklabels(label_list)
        plt.xlabel("cbr bitrate [{}]".format(unit))
        plt.ylabel("delay [s]")
        plt.title("CBR {}".format(direction))
        plt.show()
        plt.savefig("{}/all_delays_{}_boxplot.pdf".format(args.folder, direction))