#!/usr/bin/env python3

from argparse import ArgumentParser


import pandas as pd
import matplotlib.pyplot as plt

if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument("-f", "--file", required=True, help="Serial CSV")
    parser.add_argument("--save", default=None, help="Location to save pdf file.")

    args = parser.parse_args()

    df = pd.read_csv(args.file)
    df["time_rel"] = df["time"] - df["time"].iloc[0]
    df.index = df["time_rel"] / 60

    for i in range(1, 5):
        df["LTE_SCC{}_effective_bw".format(i)] = df["LTE_SCC{}_bw".format(i)]

        mask = df["LTE_SCC{}_state".format(i)].isin(["ACTIVE"])
        df["LTE_SCC{}_effective_bw".format(i)] = df[
            "LTE_SCC{}_effective_bw".format(i)
        ].where(mask, other=0)

    df["SCC1_NR5G_effective_bw"] = df["SCC1_NR5G_bw"].fillna(0)
    df["effective_bw_sum"] = (
        df["SCC1_NR5G_effective_bw"]
        + df["LTE_SCC1_effective_bw"]
        + df["LTE_SCC2_effective_bw"]
        + df["LTE_SCC3_effective_bw"]
        + df["LTE_SCC4_effective_bw"]
        + df["LTE_bw"]
    )
    bw_cols = [
        "SCC1_NR5G_effective_bw",
        "LTE_bw",
        "LTE_SCC1_effective_bw",
        "LTE_SCC2_effective_bw",
        "LTE_SCC3_effective_bw",
        "LTE_SCC4_effective_bw",
    ]

    ax = df[bw_cols].plot.area(stacked=True)
    ax.set_ylabel("bandwidth [MHz]")
    ax.set_xlabel("time [minutes]")

    L = plt.legend()
    L.get_texts()[0].set_text("5G main")
    L.get_texts()[1].set_text("4G main")
    L.get_texts()[2].set_text("4G SCC 1")
    L.get_texts()[3].set_text("4G SCC 2")
    L.get_texts()[4].set_text("4G SCC 3")
    L.get_texts()[5].set_text("4G SCC 4")

    if args.save:
        plt.savefig("{}-used_bandwidth".format(args.save))
    else:
        plt.show()