#
# Copyright Trimble Inc 2023
#
# Provided under NDA 
#
# Extracts IonoGuard messages and plots a "traffic" light indicator similar to the
# web GUI. Only GPS data is extracted, each level (red, yellow, orange, red) is offset
# by the (PRN - 16) / 160
#
# Requires:
#   t04Extract (at least version 2.05)
#   awk
#

import matplotlib
# Allow running headless from the command line
matplotlib.use("agg")
from pylab import *

import numpy as np
import os
import io 
import pandas as pd
from PIL import Image

# Override pylab find() to avoid deprecation warning
# pylint: disable=function-redefined
def find(x):
    return where(x)[0]

def doload(filename, verbose=True, error_bad_lines=True):
    """
    filename=text file or viewdat --mat output.
    [verbose=True - print status info when loading file?]
    [error_bad_lines=False - pass error_bad_lines to pandas read_csv?  Skips lines with errors.]
    Accepts gzip'd files ending in '.gz' and bzip'd files ending in '.bz2'
    Returns numpy array with data.
    """
    import os

    filename = os.path.expanduser(filename)
    if verbose:
        print("Loading %s..." % filename)

    def myopen( fname, mode ):
        if filename.endswith('.gz'):
            import gzip
            return gzip.open(filename,mode)
        elif filename.endswith('.bz2'):
            import bz2
            return bz2.BZ2File(filename,mode)
        else:
            return open(filename,mode)


    # Test for a Matlab file (& then use loadmat())
    f = myopen(filename,'rb')
    bytes = f.read(512)
    if b'\0' in bytes:
        from scipy.io import loadmat
        f = myopen(filename,'rb')
        raw_d = loadmat(f)
        # raw_d may have several data arrays if the size is > 2GB.
        # first raw_d['data'], then raw_d['data1'], raw_d['data2'], etc.
        all_d = []
        curr_i = 0

        while True:
            if curr_i == 0:
                var_name = 'data'
            else:
                var_name = 'data%d' % curr_i
            if not var_name in raw_d:
                break
            if b'C library' in bytes:
                all_d.append( raw_d[var_name].T )
            else:
                all_d.append( raw_d[var_name] )
            curr_i += 1
        
        return concatenate( all_d )

    # Strip header from file
    def is_string_a_float(s):
        try:
            float(s)
        except ValueError:
            return False
        return True
    f = myopen(filename,'r')
    n = 0
    for l in f.readlines():
        w = l.strip().split()
        if len(w) > 0 and is_string_a_float(w[0]):
            break
        n += 1

    # pandas reads text data really quickly
    import pandas as pd
    from packaging.version import parse as parse_version
    f = myopen(filename,'rb')
    try:
        csv_args = {'skiprows':n,
                    'header':None,
                    'delim_whitespace':True}
        if parse_version(pd.__version__)>=parse_version("1.3.0"):
            if not error_bad_lines:
                csv_args['on_bad_lines']='skip'
        else:
            if not error_bad_lines:
                csv_args['error_bad_lines']=False
        return pd.read_csv(f,**csv_args)\
                  .apply(pd.to_numeric,errors='coerce')\
                  .values
    except pd.errors.EmptyDataError:
        return array([])

def save_compressed_png(filename, dpi=None, make_dir=True):
    """Save current plot as a PNG file. Matplotlib isn't very good at
    compressing PNG image files. Therefore, use PIL functions to
    compress a memory version of the file. The compressed images are
    close to half the size of the Matplotlib generated version.

    Note: this reduces the image to 256 colors, so in images with
    gradients or lots of colors you may notice the difference.

    Inputs: filename = output PNG name
    Optional inputs: dpi = resolution of image (default None)
                     make_dir = create directory if needed? (default True)
    """

    if make_dir:
        dirname = os.path.abspath(os.path.dirname(filename))
        if not os.path.exists(dirname):
            os.makedirs(dirname)

    # Save the png to memory
    ram = io.BytesIO()
    savefig(ram,format='png',dpi=dpi)
    ram.seek(0)
    # Compress the png data before saving to the file system, the
    # matplotlib png's are not well compressed
    im = Image.open(ram)
    if hasattr(Image,'ADAPTIVE'):
        im2 = im.convert('RGB').convert('P', palette=Image.ADAPTIVE)
    else:
        im2 = im.convert('RGB').convert('P', palette=Image.Palette.ADAPTIVE)
    im2.save(filename,format='PNG')

def intersect( *args, **kwargs ):
    """
    intersect( a, b, c, ..., make_unique=False )
    Find values that appear in both a and b, c, d, ...
    a = array
    b = array
    etc. [optional]
    ret[0] = common values in a,b,c, ... Will be sorted!
    ret[1] = index into a[] (not sorted)
    ret[2] = index into b[] (not sorted)
    etc. [optional]
    Note that any repeated input values will be in the output too
    unless you set make_unique=True.
    If any of c,d,etc. are None, then it won't be used for comparison
    and ret[3],ret[4],etc. will be [].
    """
    from functools import reduce
    make_unique = False
    if 'make_unique' in kwargs:
        make_unique=kwargs['make_unique']
    make_quick_unique = False
    duplicted_values = []
    if 'make_quick_unique' in kwargs:
        make_quick_unique=kwargs['make_quick_unique']

    common = reduce( lambda l,r: intersect1d(l,r,True), (i for i in args if i is not None) )
    if make_unique:
        common = unique(common)
    if make_quick_unique:
        rst = np.ediff1d(common)
        index = np.where(rst[:]==0)[0]
        duplicted_values = common[index]
        common = np.delete(common, index)

    def get_arg_idx(arg):
        if arg is None:
            return []
        if make_unique:
            arg = array(arg).astype(object)
            m = zeros_like(arg, dtype=bool)
            m[unique(arg,return_index=True)[1]] = True
            arg[~m] = None
        if make_quick_unique:
            for x in duplicted_values:
                index = np.where(arg[:]==x)[0]
                arg[index[:-1]] = None            
        common_arg = in1d(arg,common).nonzero()[0]
        if len(common) != len(common_arg):
            print_("intersect length mismatch - perhaps there is repeated data?")
        return common_arg

    # the first 2 args[] are mandatory
    ret = [common] + [get_arg_idx(arg) for arg in [args[0],args[1]]]
    # further args[] are optional
    ret += [get_arg_idx(arg) for arg in args[2:]]
    return ret

if(__name__ == "__main__"):

  if(len(sys.argv) != 2):
    print('Usage:\n  python plotIonoGuard.py YourFile.T04')
    sys.exit()

  filename = sys.argv[1]

  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])

  # t04Extract -d outputs a *lot* of data. We need to extract the elevation angle, so decimate
  # the data to 1Hz AND use awk in the system call to only extract the columns we need.

  obsKeyTIME     = 0 # GPS Time [Sec]
  obsKeySV       = 1 # SV ID
  obsKeySAT_TYPE = 2 # Satellite Type
  obsKeyEL       = 3 # Elevation [deg]
  obsKeyFREQ     = 4 # Frequency Band
  obsKeyTRACK    = 5 # Track Type
  obsKeyANTENNA  = 6 # Antenna #
 
  # Get the observable data for the elevation angle - use a system call and awk
  os.system('t04Extract -o -d 1000 ' + filename + ' | awk \'{print $2,$5,$8,$11,$12,$13,$20}\' > obs.txt')

  ionoKeyGps_week   =  0
  ionoKeyTIME       =  1
  ionoKeySAT_TYPE   =  2
  ionoKeySV         =  3
  ionoKeyFlags      =  4
  ionoKeyAmplitude  =  5
  ionoKeyPhase      =  6
  ionoKeyGradient   =  7
  
  # Get the iono activity index data
  os.system('t04Extract -g ' + filename + ' > iono.txt')
  
  obs  = doload('obs.txt')
  iono = doload('iono.txt')

  figure()
  
  # This script only plots GPS data (satellite type of 0). Update for other satellite
  # types.
  for sv in range(1,33):
    i_iono = find( (iono[:,ionoKeySV] == sv) & (iono[:,ionoKeySAT_TYPE] == 0) )
    if(len(i_iono) > 0):
      # GPS L1 C/A on the primary antenna
      i_obs = find(   (obs[:,obsKeySV] == sv) 
                    & (obs[:,obsKeySAT_TYPE] == 0)
                    & (obs[:,obsKeyANTENNA] == 0)
                    & (obs[:,obsKeyFREQ] == 0)
                    & (obs[:,obsKeyTRACK] == 0) )

      thisObs = obs[i_obs,:]
      u,index = np.unique(thisObs[:,obsKeyTIME],return_index=True)
      if(len(thisObs) != len(index)):
          thisObs = thisObs[index,:]
      
      thisIono = iono[i_iono,:]
      u,index = np.unique(thisIono[:,ionoKeyTIME],return_index=True)
      if(len(thisIono) != len(index)):
          thisIono = thisObs[index,:]

      _, IA, IB = intersect( thisObs[:,obsKeyTIME], thisIono[:, ionoKeyTIME] )

      # The data is time aligned. We only need the elevation angle from the 
      # obs object
      thisElev = thisObs[IA, obsKeyEL]

      # Get the iono data that matches obs/thisElev.
      thisIono = thisIono[IB,:]

      amp   = thisIono[:,ionoKeyAmplitude].astype(int)
      phase = thisIono[:,ionoKeyPhase].astype(int)
      grad  = thisIono[:,ionoKeyGradient].astype(int)

      thisIndex = find(   (amp >= 0) & (amp <=7)
                        & (phase >= 0) & (phase <=7)
                        & (grad >= 0) & (grad <=7) )

      thisElev = thisElev[thisIndex]
      thisIono = thisIono[thisIndex,:]
      amp   = amp[thisIndex]
      phase = phase[thisIndex]
      grad  = grad[thisIndex]

      sinElev = np.sin(thisElev * np.pi / 180.0)
      
      # Convert the enumeration to a normalized scale
      ampScaled   = np.sqrt(sinElev) * ampSF[amp]/0.158
      phaseScaled = sinElev * (phaseSF[phase]/0.003)
      
      # We only use the normalized scale if the enumeration
      # is greater than 2. Zero any data <= 2. 
      i = find( amp <= 2)
      ampScaled[i] = 0
      i = find( phase <= 2)
      phaseScaled[i] = 0

      # Find the max phase/amplitude data
      metric = np.maximum(ampScaled, phaseScaled)

      # Get the indicator from the gradient information
      green  = find(grad< 5)
      yellow = find(grad==5)
      orange = find(grad==6)
      red    = find(grad==7)

      metricGrad = np.copy(metric)
      metricGrad[green]  = 0
      metricGrad[yellow] = 5
      metricGrad[orange] = 10
      metricGrad[red]    = 20
      metric = np.maximum(metric, metricGrad)
      
      # Color according to the normalized scale
      green  = find(metric < 5)
      yellow = find( (metric >= 5) & (metric < 10) )
      orange = find( (metric >= 10) & (metric < 15) )
      red    = find( metric >= 20)
  
      Time = thisIono[:,ionoKeyTIME]
      dayNum = np.floor(np.median(Time)/86400)
      Time -= dayNum*86400
      
      if(len(green) > 0):
        plot(Time[green]/3600, np.zeros(len(green)) + (sv-16)/160, 'g.')
      
      if(len(yellow) > 0):
        plot(Time[yellow]/3600, np.ones(len(yellow)) + (sv-16)/160, '.', color=(1,1,0))

      if(len(orange) > 0):
        plot(Time[orange]/3600, 2*np.ones(len(orange)) + (sv-16)/160, '.', color=(1,165.0/255.0,0))
      
      if(len(red) > 0):
        plot(Time[red]/3600, 3*np.ones(len(red)) + (sv-16)/160, 'r.')
  
  grid(True)
  xlabel('Time [Hour of Day]')
  ylabel('Iono Activity Web Indicator')
  title('GPS Iono Activity data')
  xlim([0,24])
  ylim([-0.2,3.2])
  tight_layout()
  pngFilename = filename + '-IonoLevel.png'
  tokens = pngFilename.split('/') # Linux 
  print('Saving ... ' + tokens[-1])
  save_compressed_png( tokens[-1], dpi=300)
  close()