#! /usr/bin/python3
#
# Copyright (c) 2021 - Adjacent Link LLC, Bridgewater, New Jersey
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#
# * Redistributions of source code must retain the above copyright
#   notice, this list of conditions and the following disclaimer.
# * Redistributions in binary form must reproduce the above copyright
#   notice, this list of conditions and the following disclaimer in
#   the documentation and/or other materials provided with the
#   distribution.
# * Neither the name of Adjacent Link LLC nor the names of its
#   contributors may be used to endorse or promote products derived
#   from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
# COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
# ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
#

from __future__ import absolute_import, division, print_function
import argparse
import os

import matplotlib.pyplot as plt
import pandas as pd
import sqlite3


usagestr = 'Generate a line plot of the UE rnti values when' \
           'a UE is in the RRC connected state to an ENB.'

suffix = 'connected_rnti_vs_time'

parser = argparse.ArgumentParser(description=usagestr)

parser.add_argument('--graphfile',
                    metavar='GRAPHFILE',
                    action='store',
                    default=None,
                    help='The output graph file. Default is to rename ' \
                    'SQLITEFILE with a "%s.png" suffix.' % suffix)
parser.add_argument('--plotstyle',
                    metavar='PLOTSTYLE',
                    action='store',
                    default='ggplot',
                    help='Any valid matplotlib.pyplot plot style. ' \
                    'Default: ggplot.')
parser.add_argument('--height',
                    type=float,
                    action='store',
                    default=None,
                    help='Override graph height.')
parser.add_argument('--width',
                    type=float,
                    action='store',
                    default=None,
                    help='Override graph width.')
parser.add_argument('--margins',
                    metavar='MARGINS',
                    default='0.85:0.15:0.15:0.9',
                    help='Colon separated value of top:left:bottom:right graph ' \
                    'margins. Margins are expressed as a decimal from 0.0 to 1.0 ' \
                    'where the horizontal margin runs from left to right and the ' \
                    'vertical margin from bottom to top. So a bottom margin of 0.15 ' \
                    'is 15%% from the bottom of the page and a top margin of 0.8 is ' \
                    '20%% from the top of the page. Default: 0.85:0.15:0.15:0.9')
parser.add_argument('sqlitefile',
                    metavar='SQLITEFILE',
                    action='store',
                    help='An ETCE sqlite database.')

args = parser.parse_args()

plt.style.use(args.plotstyle)

if not os.path.isfile(args.sqlitefile):
    print('Could not find file "%s". Quitting.' % args.sqlitefile, file=sys.stderr)
    exit(1)

con = sqlite3.connect(args.sqlitefile)

df = pd.read_sql('select * from srslte_ue_tables_upper_nodeinfotable', con)

# connected_rnti contains the current rnti value only during
# reports where thee rrcstate is not idle
df['connected_rnti'] = df.rnti
df.loc[df.rrcstate=='IDLE', ['connected_rnti']] = 0

ymax = (df.connected_rnti.max() // 5 + 1) * 5

trials = sorted(df.trial.unique())

height = max(ymax // 72, 6)
if args.height:
    height = args.height

width = 8.0
if args.width:
    width = args.width

top_margin, left_margin, bottom_margin, right_margin = \
    map(float, args.margins.split(':'))

for trial in trials:
    fig = plt.figure(figsize=(width, height))

    ax = fig.add_subplot(111)

    trial_rntis = df.loc[df.trial == trial]

    print(trial_rntis)

    dfrnti = trial_rntis.pivot_table(values='connected_rnti',
                                     index='timestamp',
                                     columns='nodename')

    dfrnti.plot(ax=ax,
                marker='.')

    ax.set_xlabel('time (sec)')

    ax.set_ylabel('rnti')

    ax.set_ylim(0, ymax)

    ax.set_title('RNTI vs. Time')

    fig.subplots_adjust(top=top_margin,
                        left=left_margin,
                        bottom=bottom_margin,
                        right=right_margin)

    graphfile = '.'.join(os.path.basename(args.sqlitefile).split('.')[:-1]) + '-' + \
        suffix + \
        ('-trial_%02d.png' % trial)

    if args.graphfile:
        graphfile = args.graphfile % trial

    fig.savefig(graphfile)
