#!/usr/bin/env python3
"""
plotIonoGuardEngineeringMetrics.py
  For usage details "plotIonoGuardEngineeringMetrics.py -h"

  @description: Creates plots of the IonoGuard Station metrics for a given T04 file conating 35:19 and 35:275 records.
  @author: AAA

  Copyright Trimble, Inc. 2024
"""
import argparse
import datetime
import math

import matplotlib
import numpy as np
import pandas as pd
from pylab import close, plt, show

import mutils as m

ampSF = np.array([0, 0.25, 0.75, 1.5, 3.0, 6.0, 12, 32])
phaseSF = np.array([0, 0.005, 0.015, 0.03, 0.06, 0.12, 0.24, 0.64])
gradSF = np.array([0, 0, 0, 0, 0, 5, 10, 20])
realGradSF = np.array([0, 1, 2, 4, 8, 16, 32, 64])

global customerMetrics

# viewdat's UTC conversion does not appear
# to account for leap seconds, so this doesn't either
def _gps_to_utc(gps_weeks, gps_sow) -> datetime.datetime:
    # Define the GPS epoch
    gps_epoch = datetime.datetime(1980, 1, 6, 0, 0, 0)
    delta = datetime.timedelta(weeks=gps_weeks, seconds=gps_sow)
    return gps_epoch + delta


def apply_create_station_metrics_series(data: pd.DataFrame) -> pd.Series:
    """Creates a organized Pandas Series containing the staiton metrics
       for a given week/time combination, also creating a datetime for the time.
       This is expected to be used as an apply function from a resulting groupby operation.

    Args:
        data (pandas.DataFrame): Pandas DataFrame containing the merged 35:19 obs and 35:275 iono data for a single week/time combination. This can be generated from the following example code::

            import pandas
            import mutils
            obs_raw = mutils.vd2cls(filename, "-d35:19 --dec=1000")
            # filter to include primary antenna and exclude SBAS satellites
            obs = obs_raw[(obs_raw.EL >= 0) & (obs_raw.ANTENNA == 0) & (obs_raw.SAT_TYPE != 1)]
            obsDf = pandas.DataFrame(obs, columns=obs.k.keys())

            iono = mutils.vd2cls(filename, "-d35:275")
            ionoDf = pandas.DataFrame(iono, columns=iono.k.keys())

            # Merge this information with the iono data
            ionoAndObsDf = obsDf.merge(
                ionoDf, on=["WN", "TIME", "SAT_TYPE", "SV"], how="inner", suffixes=("_obs", "_iono")
            )

        This function can be invoked from this sample code as follows:
        >>> stationMetrics = ionoAndObsDf.groupby(["WN", "TIME"])[["WN","TIME","EL","stationInfo"]].apply(apply_create_station_metrics_series, include_groups=True)

    Returns:
        pandas.Series: Pandas Series containing the station metrics for the given week/time combination. Can be adjusted by hand in the script to return raw values, or switching the scale factors (specifically for the gradient)
    """
    d = {}
    d["time"] = _gps_to_utc(data["WN"].iloc[0], data["TIME"].iloc[0])

    averageElevation = data["EL"].mean()
    sinAverageElevation = math.sin(math.radians(averageElevation))

    # Since this is a single epoch, the station info is the same for all rows
    stationInfo = int(data["stationInfo"].iloc[0])

    # d["averageElevation"] = averageElevation
    amplitude = stationInfo & 0b111
    # d["amplitude"] = amplitude

    if customerMetrics:
        amplitude = amplitude if amplitude > 2 else 0
    d["amplitudeScaled"] = math.sqrt(sinAverageElevation) * ampSF[amplitude] / 0.158

    phase = (stationInfo >> 3) & 0b111
    # d["phase"] = phase
    if customerMetrics:
        phase = phase if phase > 2 else 0
    d["phaseScaled"] = sinAverageElevation * phaseSF[phase] / 0.003

    gradient = (stationInfo >> 6) & 0b111
    # d["gradient"] = gradient
    if customerMetrics:
        d["gradientScaled"] = gradSF[gradient]
    else:
        d["gradientScaled"] = realGradSF[gradient]

    d["maxMetric"] = max(d["amplitudeScaled"], d["phaseScaled"], d["gradientScaled"])
    return pd.Series(d)

if __name__ == "__main__":
    # Parse arguments
    parser = argparse.ArgumentParser(
        description="Plot IonoGuard Metrics prior to their stoplightization"
    )
    parser.add_argument("filename", help="T04")
    parser.add_argument(
        "-s",
        "--showFigures",
        help="Show the figures and summary data, instead of saving to a file.",
        action="store_true",
    )
    parser.add_argument(
        "-e",
        "--engineeringMetrics",
        help="Plot figures with extra stuff for engineering analysis. Without this flag, plots will appear as we ask customers to interpet them",
        action="store_true",
    )
    args = parser.parse_args()

    filename = args.filename
    customerMetrics = not args.engineeringMetrics
    if not args.showFigures:
        # Allow running headless from the command line
        plt.switch_backend("agg")

    obs_raw = m.vd2cls(filename, "-d35:19 --dec=1000")
    # filter to include primary antenna and exclude SBAS satellites
    obs = obs_raw[(obs_raw.EL >= 0) & (obs_raw.ANTENNA == 0) & (obs_raw.SAT_TYPE != 1)]
    obsDf = pd.DataFrame(obs, columns=obs.k.keys())

    iono = m.vd2cls(filename, "-d35:275")
    ionoDf = pd.DataFrame(iono, columns=iono.k.keys())

    # Merge this information with the iono data
    ionoAndObsDf = obsDf.merge(
        ionoDf, on=["WN", "TIME", "SAT_TYPE", "SV"], how="inner", suffixes=("_obs", "_iono")
    )

    # Build the station metrics.
    # The "[["WN","TIME","EL","stationInfo"]]" section "explicitly selects" the columns that
    # apply_create_station_metrics_series uses. This is necessary to silence a deprecation warning...
    stationMetrics = ionoAndObsDf.groupby(["WN", "TIME"])[["WN","TIME","EL","stationInfo"]].apply(
        apply_create_station_metrics_series, include_groups=True
    )

    # The apply function rewrites the week/seconds into a python datetime, so update the index
    stationMetrics.set_index("time", inplace=True)

    # Make the stoplight plot
    fig, ax = plt.subplots()
    t = matplotlib.dates.date2num(stationMetrics.index)
    ax.plot(
        t, stationMetrics["maxMetric"], linestyle="--", dashes=(1, 6), color="lightgrey"
    )
    ax.plot(
        t,
        stationMetrics["maxMetric"].where(stationMetrics["maxMetric"].between(0, 5)),
        color="green",
    )
    ax.plot(
        t,
        stationMetrics["maxMetric"].where(stationMetrics["maxMetric"].between(5, 10)),
        color="yellow",
    )
    ax.plot(
        t,
        stationMetrics["maxMetric"].where(stationMetrics["maxMetric"].between(10, 20)),
        color="orange",
    )
    ax.plot(
        t,
        stationMetrics["maxMetric"].where(stationMetrics["maxMetric"].between(20, 999)),
        color="red",
    )
    ax.set_title("Ionospheric Activity Station Metric")
    ax.set_xlabel("Time (UTC)")
    if customerMetrics:
        # Hide the tick labels for now
        ax.get_yaxis().set_ticklabels([])

    # Match other plots that limit the x-axis timestamps when plotting 24-hours
    if len(t) / 3600 > 4:
        ax.xaxis.set_major_locator(matplotlib.dates.HourLocator(interval=4))

    ax.xaxis.set_major_formatter(matplotlib.dates.DateFormatter("%Y-%m-%d %H:%M"))

    fig.autofmt_xdate()
    fig.tight_layout()
    ax.grid(True)
    show()

    if not args.showFigures:
        m.save_compressed_png(f"{filename}-IonoGuardCustomerStationMetric.png", dpi=300)
        close()

    # Create the per-metric plot
    fig, ax = plt.subplots()
    t = stationMetrics.index
    ax.plot(t, stationMetrics["amplitudeScaled"], color="blue", label="Amplitude")
    ax.plot(t, stationMetrics["phaseScaled"], color="green", label="Phase")
    ax.plot(t, stationMetrics["gradientScaled"], color="red", label="Gradient")

    # Plot last to ensure it's on top
    ax.plot(
        t, stationMetrics["maxMetric"], linestyle="--", dashes=(1, 6), color="lightgrey"
    )
    ax.grid(True)
    fig.tight_layout()
    ax.legend(loc='best')
    show()

    if not args.showFigures:
        m.save_compressed_png(f"{filename}-IonoGuardStationMetrics.png", dpi=300)
        close()