#!/usr/bin/env python

""" A script to plot the performance of the stinger pulsed jammer mitigation feature

plot_pjam.py:
 For usage details "plot_pjam.py -h"

Desc:
 Given a T02 or T04 file the script will generate
 some simple plots to show the performance of the pulsed jammer mitigation
 stinger feature.
 Works with the JAMTH, JAMSU 16:199 kvp's documented here:
 https://wiki.eng.trimble.com/index.php/Diagnostic_Key_Value_Records

 Performance note: This script makes multiple passes over the t0x file, one pass per key.

 Copyright Trimble Inc 2021
"""
# pylint: disable=line-too-long

import argparse
import matplotlib
import matplotlib.pyplot as plt
import mutils as m
import numpy as np

if __name__ == "__main__":
    # Parse arguments
    parser = argparse.ArgumentParser(description='Parse JAMTH and JAMSU diagnostics from the stinger pulsed jammer mitigation feature')
    parser.add_argument('filename', help='T02 or T04 file containing 16:199 keys JAMTH and JAMSU')
    parser.add_argument('-s', '--show_figures', help='Show the figures and summary data, instead of saving to a file.', action="store_true")
    parser.add_argument('-o', '--output_prefix', help='By default use "filename".png as the output for graphs/summary.  With this option, choose a different output prefix.')
    args = parser.parse_args()

    filename = args.filename

    if not args.show_figures:
        # Allow running headless from the command line
        matplotlib.use("agg")

    if args.output_prefix:
        out_basename = args.output_prefix
    else:
        out_basename = filename

    # load the summary of the jamming detector operation over the previous second
    jamsu = m.kv_info2arr(filename, 'JAMSU', convert_float=True)
    k_jamsu = m.dotdict({})
    k_jamsu.tow_msec = 0         # Time of week (ms) at start of data
    k_jamsu.epochs_total = 1     # Number of epochs (ms) the detector ran. This is just that data was collected, the jammer detector
                                # was not necessarily run. Typically 666 or 667 when not jammer present. Typically 1000 when
                                # a jammer is present.
    k_jamsu.epochs_detected = 2  # Number of epochs with a jammer present.
    k_jamsu.epochs_invalid = 3   # Number of epochs with one or more detectors in an invalid state.
                                # epochs_total - epochs_invalid = # epochs the jammer detector was run.
                                # The value will be non-zero at start-up or, for example, under continuous jamming conditions.

    d = m.kv_info2arr(filename, 'JAMTH').astype(int)
    k = m.dotdict({})
    k.tow_msec = 0 # Approx. time of week (ms) when threshold calculated.
    k.nominal = 1  # Nominal detector value.
    k.sigma = 2    # Detector sigma
    k.frq_ind = 3  # Detector Number / frequency index, 0 or 1.
                # Typically running two detectors with different nominal and sigma values.
    k.err_cnt = 4  # Detector error count.
    k.state = 5    # Detector state :
                # 0 = Unknown.
                # 1 = Initialising thresholds (start-up).
                # 2 = Checking thresholds (after initialisation).
                # 3 = Thresholds valid, detector active.
                # 4 = Threshold error, re-initialise.

    idx_f0 = m.find((d[:, k.frq_ind] == 0))
    idx_f1 = m.find((d[:, k.frq_ind] == 1))

    fig, ax = m.subplots(6, 1, sharex=True, tight_layout=True)

    ax[0].plot(d[idx_f0, k.tow_msec]*.001, d[idx_f0, k.nominal], label='F#0', color='blue', linestyle="None", marker=".")
    ax[0].plot(d[idx_f1, k.tow_msec]*.001, d[idx_f1, k.nominal], label='F#1', color='red', linestyle="None", marker=".")
    ax[0].set_ylabel('Mean Detector')
    ax[0].grid(True)
    ax[0].legend(loc='best')

    ax[1].plot(d[idx_f0, k.tow_msec]*.001, d[idx_f0, k.sigma], label='F#0 mean={mean:.0f}'.format(mean=np.mean(d[idx_f0, k.sigma])), color='blue', linestyle="None", marker=".")
    ax[1].plot(d[idx_f1, k.tow_msec]*.001, d[idx_f1, k.sigma], label='F#1 mean={mean:.0f}'.format(mean=np.mean(d[idx_f1, k.sigma])), color='red', linestyle="None", marker=".")
    ax[1].set_ylabel('Sigma Detector')
    ax[1].grid(True)
    ax[1].legend(loc='best')

    ax[2].plot(d[idx_f0, k.tow_msec]*.001, d[idx_f0, k.err_cnt], label='F#0', color='blue', linestyle="None", marker=".")
    ax[2].plot(d[idx_f1, k.tow_msec]*.001, d[idx_f1, k.err_cnt], label='F#1', color='red', linestyle="None", marker=".")
    ax[2].set_ylabel('Error Count')
    ax[2].grid(True)
    ax[2].legend(loc='best')

    ax[3].set_ylim(0, 4)
    ax[3].plot(d[idx_f0, k.tow_msec]*.001, d[idx_f0, k.state], label='F#0', color='blue', linestyle="None", marker=".")
    ax[3].plot(d[idx_f1, k.tow_msec]*.001, d[idx_f1, k.state], label='F#1', color='red', linestyle="None", marker=".")
    ax[3].set_title('Detector Status: 1-Init, 2-Checking, 3-Valid, Detector active, 4-Error, re-initialize',fontsize=8)
    ax[3].grid(True)
    ax[3].legend(loc='best')

    # for JAMSU, on startup, the first entry sometimes has time zero, followed by the time being set. Ignore the first entry
    ax[4].plot(jamsu[1:, k_jamsu.tow_msec]*.001, jamsu[1:, k_jamsu.epochs_detected], label='Detected Epochs', color='blue', linestyle="None", marker=".")
    ax[4].plot(jamsu[1:, k_jamsu.tow_msec]*.001, jamsu[1:, k_jamsu.epochs_invalid], label='Invalid Epochs', color='red', linestyle="None", marker=".")
    ax[4].set_ylabel('Epochs')
    ax[4].grid(True)
    ax[4].legend(loc='best')

    ax[5].plot(jamsu[1:, k_jamsu.tow_msec]*.001, jamsu[1:,k_jamsu.epochs_total], color='blue', linestyle="None", marker=".")
    ax[5].set_ylabel('Epochs Total')
    ax[5].set_ylim(0, 1000)
    ax[5].grid(True)

    fig.set_size_inches(18.5, 10.5, forward=True)
    if args.show_figures:
        m.show()
    else:
        plt.savefig(out_basename + '_pulsed_jammer_summary.png', dpi=150)
        plt.close()