#!/usr/bin/env python

###############################################################################
# Copyright (c) 2017-2022 Trimble Inc
# $Id: small_DI.py,v 1.21 2024/02/11 20:58:51 wlentz Exp $
###############################################################################
#
# Small utility to plot NoPi data quickly.
# Run "./small_DI.py -h" for help.
#
###############################################################################
from __future__ import print_function
import argparse, textwrap
from builtins import input, range  # install "future" package if needed

import pandas as pd
#from pylab import *
from matplotlib.figure import Figure
from pylab import where, grid, unique, insert, diff, nan, mean, std, np
from glob import glob
from collections import namedtuple
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg, NavigationToolbar2Tk
from functools import partial
import os.path
import subprocess
import sys
from six.moves import tkinter as Tk

def find( x ) :
    return where(x)[0]

def parse_cli():
    """Parse the command line interface"""
    parser = argparse.ArgumentParser(formatter_class=argparse.RawTextHelpFormatter,
    description=textwrap.dedent('''\
    Plot NoPi data.

    Usage: small_DI.py [file1.T04 file2.T04]
    If diffs_summary.txt is present and newer than the T04 inputs, then
    the program will load and plot the results.  Otherwise, the program
    will try to run NoPi and generate the data.  If file1/file2 are missing,
    then the program will just read the NoPi MTB files.
    NOTE - before running NoPi:
    - NoPi executable must be in your path
    - you must create piline.ini manually or with "../NoPi/gen_ini*.py"
    ''' ))
    parser.add_argument('file1', nargs='?', default=[], help='Base T0x file')
    parser.add_argument('file2', nargs='?', default=[], help='Rover T0x file')
    parser.add_argument('--combo', help=textwrap.dedent('''\
    Produce PNG files for given combos # and exit. Example:
    --combo GAL:E1,GPS:L1CA  (or --combo all)
    Also writes all text statistics to "summary.txt".'''))
    parser.add_argument('--nonint_amb', action="store_true",help=textwrap.dedent('''\
    Use this for Android data with accumulated doppler'''))
    parser.add_argument('--yes', action="store_true",
                        help="For scripting, answer any inputs w/'yes'")
    parser.add_argument('--dual_ant', action="store_true",
                        help="Baseline antenna 0 to antenna 1")
    args = parser.parse_args()

    if args.combo:
        import matplotlib as mpl
        mpl.use('pdf') # 'pdf' (and 'svg') is almost twice as fast as 'Agg'

    return args

class State(object):
    """Used to hold data on NoPi results and graphics window"""
    def __init__(self, inc_doppler, func_extra_init=None):
        self.data = {}
        DataIdx = namedtuple('DataIdx',['ms','sv','az','el','snr','dSnr','dCode','dCarr','dDopp','status'])
        if ( inc_doppler ) :
          self.k = DataIdx(0,1,2,3,4,5,6,7,8,9)
        else :
          self.k = DataIdx(0,1,2,3,4,5,6,7,-1,8) # dDopp is invalid
        self.files = []
        self.filetypes = []
        self.inc_doppler = inc_doppler
        self.animating = False
        self.fig = None
        self.axis = None
        self.lines = []
        self.last_sv_list = np.array([])
        self.last_col = -100
        self.last_xlim_max = -1
        self.canvas = None
        self.check_frame = None
        self.check_list = []
        self.console = True
        self.outfile = None

        self.func_make_menubar = make_menubar
        self.func_get_data = get_data
        self.func_checkboxes = sv_checkboxes
        self.func_init_animate = None
        self.ref_sig = None
        self.tst_sig = None

        if func_extra_init is not None:
            func_extra_init(self)

def print_summary( console, outfile, txt ):
    """Print summary info to screen and file
       console = print to console
       outfile = file handle or None
       txt = what to print
    """
    if console:
        print(txt)
    if outfile is not None:
        print(txt,file=outfile)

def get_mtb(n):
    """Get name of mtb file.  Normally it is just "diffs_combo_0.mtb",
       but sometimes it is compressed (e.g. "diffs_combo_0.mtb.bz2").
    """
    try:
        return glob('diffs_combo_%d.mtb*'%n)[0]
    except:
        raise RuntimeError("Can't find diffs_combo_%d.mtb"%n)

def run_nopi_if_needed( file1, file2, auto_answer_yes ):
    """Determine if NoPi results are up to date (and potentially run NoPi)"""
    need_update = False
    if not os.path.exists('diffs_summary.txt'):
        need_update = True
    elif os.path.getmtime(file1) > os.path.getmtime('diffs_summary.txt'):
        need_update = True
    elif os.path.getmtime(file2) > os.path.getmtime('diffs_summary.txt'):
        need_update = True
    elif os.stat(get_mtb(0)).st_size < 1000:
        need_update = True

    if need_update:
        if not os.path.exists('piline.ini'):
            print('Please create piline.ini')
            sys.exit(1)
        if not os.path.exists('NoPi_Default.cfg'):
            print('Generating NoPi_Default.cfg')
            df = open('NoPi_Default.cfg','w')
            df.write('-is*\n')
            df.close()
        if auto_answer_yes:
            yn = 'y'
        else:
            yn = input("Run NoPI? ([y]/n) ")
        if yn == 'y' or yn == '':
            other_opts = ""
            if args.nonint_amb:
                other_opts += " --nonint_amb"
            if args.dual_ant:
                other_opts += " -ab0 -ar1"
            subprocess.check_call("NoPi%s %s %s" % (other_opts,file1,file2),shell=True)
        else:
            print('Please run NoPi manually')
            sys.exit(1)

    if os.stat(get_mtb(0)).st_size < 1000:
        print('diffs_combo_0.mtb seems to small - is there an error?')
        sys.exit(1)


def prep_state():
    """Return "State" with info on NoPi data"""
    # Create list of filenames, data types, and space for data
    summary = open('diffs_summary.txt').readlines()
    ncombos = int(summary[17])
    inc_doppler = False
    # Only supports file format 1, extended, and 5, extended with Doppler
    if ( int(summary[18 + ncombos]) == 5 ) :
        inc_doppler = True

    s = State(inc_doppler)
    for i in range(ncombos):
        s.files.append( get_mtb(i) )
        sat_type = summary[18+i].split()[0].split('_')[-1]
        signal = summary[18+i].split()[1].split('_')[-1]
        s.filetypes.append( '%s:%s' % (sat_type, signal ) )

    return s

def show_slips( s ):
    """Show data for all cycle slips on plot"""
    for cb in s.check_list:
        cb.set_points_visible(True)
    s.fig.canvas.draw()

def hide_slips( s ):
    """Hide data for all cycle slips on plot"""
    for cb in s.check_list:
        cb.set_points_visible(False)
    s.fig.canvas.draw()

def show_all_svs( s ):
    """Show data for all satellites on plot"""
    for cb in s.check_list:
        cb.set_visible(True)
    s.fig.canvas.draw()

def hide_all_svs( s ):
    """Hide data for all satellites on plot"""
    for cb in s.check_list:
        cb.set_visible(False)
    s.fig.canvas.draw()

def load_file(filename):
    """Load NoPi text data file 'filename'"""
    return pd.read_csv(filename,header=None,delim_whitespace=True).values

class LineCheckbox(object):
    def __init__(self,master,fig,sv,line,points):
        self.var = Tk.IntVar()
        self.fig = fig
        self.line = line
        self.points = points
        self.points_visible = False
        line.set_visible(False)
        if points is not None:
            points.set_visible(False)
        self.widget = Tk.Checkbutton(master,
                                     text="%2d"%sv,
                                     selectcolor=line.get_color(),
                                     variable=self.var,
                                     command=self.cb)
        self.widget.pack(side=Tk.TOP)
    def destroy(self):
        self.widget.destroy()
    def set_points_visible(self,vis):
        self.points_visible = vis
        self.set_visible( self.var.get() )
    def set_visible(self,vis):
        if vis:
            self.var.set(1)
            self.line.set_visible(True)
            if self.points_visible and self.points is not None:
                self.points.set_visible(True)
        else:
            self.var.set(0)
            self.line.set_visible(False)
            if self.points is not None:
                self.points.set_visible(False)
    def cb(self):
        self.set_visible( self.var.get() )
        self.fig.canvas.draw()

def sv_checkboxes(s, clear_list, sv, l, slip_l):
    """Append a new SV enable/disablecheckbox to the list
       Clear the old list first based on clear_list
    """
    if clear_list:
        for cb in s.check_list:
            cb.destroy()
        s.check_list = []
    cb = LineCheckbox(s.check_frame, s.fig, sv, l, slip_l)
    s.check_list.append( cb )

def get_data(s, file_idx):
    """Get the NoPi diffs_combo data for file_idx
       Loads the data file as needed
    """
    if not file_idx in s.data:
        print('loading data %s' % s.files[file_idx])
        s.data[file_idx] = load_file(s.files[file_idx])
    return s.data[file_idx]

def update_plot(lines, s, sv, x, y, **kwargs):
    """If s.animating is False, just plot normally.
    If s.animating is True, use set_data() for speed.
    Returns current plotted "line".
    """
    do_fresh = True
    if s.animating:
        if sv in lines.keys():
            do_fresh = False
            lines[sv].set_data( x, y )
            line = lines[sv]
    if do_fresh:
        line, = s.axis.plot( x, y, **kwargs )
        lines[sv] = line
    return line

def draw_col(s,col,file_idx,scale,label,units,ddiff,hide_ref=True):
    """Draw data from a single NoPi file:
      s = State for all data/graphics
      col = column in diffs_combo_x.mtb to plot
      file_idx = file # (0 = diffs_combo_0.mtb, etc.)
      scale = scale data by this
      label = text name for data
      units = units for data (e.g., 'm' for meters)
      ddiff = diff. type label e.g. 'ddiff' or 'sdiff'
      hide_ref = if True, ignore data from reference satellite
    """
    do_redraw = True
    d = s.func_get_data(s, file_idx)
    if d is None:
        return
    if hide_ref:
        i = find( d[:,s.k.status] != 17 )
    else:
        i = range(len(d))
    if len(i) > 0:
        sv_list = unique( d[:,1] )
    else:
        sv_list = []

    if s.animating:
        # avoid redrawing everything if possible
        if s.axis is not None and np.array_equal(sv_list,s.last_sv_list) and col==s.last_col:
            if s.axis.get_xlim()[-1] >= s.last_xlim_max:
                do_redraw = False
            s.last_xlim_max = s.axis.get_xlim()[-1]

        if do_redraw:
            s.fig.clf()
            s.axis = s.fig.add_subplot(111)
            s.axis.grid( True )
            s.lines = {}
            s.slip_lines = {}
            s.last_sv_list = sv_list
            s.last_col = col
            s.last_xlim_max = -1
    else:
        s.fig.clf()
        s.axis = s.fig.add_subplot(111)
        s.axis.grid( True )
        s.lines = {}
        s.slip_lines = {}
    if len(i) > 0:
        y = d[:,col]*scale
        allmean = mean(y[i])
        allstd = std(y[i])
        allmax = max(y[i])
        allmin = min(y[i])
        clear_cbs = True
        for sv in sv_list:
            i = find( d[:,s.k.sv] == sv )
            if hide_ref:
                i2 = find( (d[:,s.k.sv] == sv) & (d[:,s.k.status] != 17) )
            else:
                i2 = find( d[:,s.k.sv] == sv )
            tdiff_ms = abs(diff(d[i,0]))
            if len(tdiff_ms) == 0:
                min_tdiff_ms = 1e9
            else:
                min_tdiff_ms = min(tdiff_ms)
            jump_i = find( tdiff_ms > min_tdiff_ms+2 ) + 1
            l = update_plot(s.lines, s, sv,
                            insert(d[i,0]*1e-3,jump_i,nan), insert(y[i],jump_i,nan) )
            slip_l = None
            if col == s.k.dCarr:
                islip = i[find( d[i,s.k.status].astype(int)&2 )]
                slip_l = update_plot(s.slip_lines, s, sv,
                                     d[islip,0]*1e-3, y[islip],
                                     color=l.get_color(),
                                     marker='v', markeredgecolor='k',
                                     linestyle="" )
            if len(i2) > 0:
                non_ref_mean = mean(y[i2])
                non_ref_std = std(y[i2])
                non_ref_max = max(y[i2])
                non_ref_min = min(y[i2])
            else:
                non_ref_mean = non_ref_std = non_ref_max = non_ref_min = nan
            print_summary(s.console, s.outfile,
                          '%s %s sv %3d mean %6.3f std %5.3f max %.1f min %.1f [%s], len %d (ref %d)' % \
                          (label,\
                           s.filetypes[file_idx],
                           sv,
                           non_ref_mean,
                           non_ref_std,
                           non_ref_max,
                           non_ref_min,
                           units,
                           len(i2),
                           len(i)-len(i2)
                          ) )
            if s.canvas is not None:
                s.func_checkboxes(s, clear_cbs and do_redraw, sv, l, slip_l)
            clear_cbs = False

        info = '%s %s - mean %.3f std %.3f max %.1f min %.1f [%s]' % \
               (label,s.filetypes[file_idx],allmean,allstd,allmax,allmin,units)
        print_summary(s.console, s.outfile, 'Total: '+info)
        s.axis.set_title(info)
        if s.animating:
            s.axis.relim()
            s.axis.autoscale_view(True,True,True)
        if do_redraw:
            s.axis.set_xlabel('GPS time[s]')
            s.axis.set_ylabel('%s %s [%s]' % (label,ddiff,units))
    if not s.animating and s.canvas is not None:
        s.canvas.draw()

def draw_base_elev(s,file_idx):
    """Display SV elevations at base for file #file_idx."""
    draw_col( s, s.k.el, file_idx, 1.0, 'Elev', 'deg', 'base', hide_ref=False )

def draw_base_cno(s,file_idx):
    """Display raw CNo at base for file #file_idx."""
    draw_col( s, s.k.snr, file_idx, 1.0, 'CNo', 'dBHz', 'base', hide_ref=False )

def draw_dd_code(s,file_idx):
    """Display code double-diff residuals for file #file_idx."""
    draw_col( s, s.k.dCode, file_idx, 1.0, 'code', 'm', 'ddiff' )

def draw_sd_cno(s,file_idx,scaling=0.1):
    """Display CNo single-diff residuals for file #file_idx."""
    draw_col( s, s.k.dSnr, file_idx, scaling, 'CNo', 'dBHz', 'sdiff', hide_ref=False )

def draw_dd_carr(s,file_idx):
    """Display carrier double-diff residuals for file #file_idx."""
    #LIGHT = 2.99792458e8
    subtype = s.filetypes[file_idx].split(':')[-1]
    is_FDMA = False
    if subtype == 'G1C':
        is_FDMA = True
    elif subtype == 'G1P':
        is_FDMA = True
    elif subtype == 'G2C':
        is_FDMA = True
    elif subtype == 'G2P':
        is_FDMA = True
    if is_FDMA:
        scale = 1000.
        units = 'mm'
    else:
        scale = 1000.
        units = 'mcyc'
    draw_col( s, s.k.dCarr, file_idx, scale, 'carr', units, 'ddiff' )

def draw_dd_dopp(s,file_idx):
    """Display Doppler double-diff residuals for file #file_idx."""
    draw_col( s, s.k.dDopp, file_idx, 1.0, 'Doppler', 'm/s', 'ddiff', hide_ref=False )

def make_menubar(s, menubar, cfg=None):
    """Create the menubar for the UI
       cfg is used by NoPiLive to set the default menu state
    """
    filemenu = Tk.Menu(menubar, tearoff=0)
    filemenu.add_command(label="Quit", command=s.root.destroy)
    menubar.add_cascade(label="File", menu=filemenu, underline=0)

    codemenu = Tk.Menu(menubar)
    carrmenu = Tk.Menu(menubar)
    if ( s.inc_doppler ) :
        doppmenu = Tk.Menu(menubar)
    cnomenu = Tk.Menu(menubar)
    for i in range(len(s.filetypes)):
        codemenu.add_command(label=s.filetypes[i], command=partial(draw_dd_code,s,i))
        carrmenu.add_command(label=s.filetypes[i], command=partial(draw_dd_carr,s,i))
        if ( s.inc_doppler ) :
            doppmenu.add_command(label=s.filetypes[i], command=partial(draw_dd_dopp,s,i))
        cnomenu.add_command(label=s.filetypes[i], command=partial(draw_sd_cno,s,i))
    menubar.add_cascade(label="Code", menu=codemenu, underline=2)
    menubar.add_cascade(label="Carr", menu=carrmenu, underline=2)
    if ( s.inc_doppler ) :
        menubar.add_cascade(label="Dopp", menu=doppmenu, underline=2)
    menubar.add_cascade(label="CNo", menu=cnomenu, underline=2)

    togglemenu = Tk.Menu(menubar, tearoff=0)
    togglemenu.add_command(label="Show all SVs", command=partial(show_all_svs,s))
    togglemenu.add_command(label="Hide all SVs", command=partial(hide_all_svs,s))
    togglemenu.add_command(label="Show slips", command=partial(show_slips,s))
    togglemenu.add_command(label="Hide slips", command=partial(hide_slips,s))
    menubar.add_cascade(label="Toggle", menu=togglemenu, underline=0)

def make_ui(s, ui_title="Small NoPi Display", figsize=(8,6), cfg=None):
    """Set up basic UI with menu items/etc."""
    s.root = Tk.Tk()
    s.root.wm_title(ui_title)
    s.fig = Figure(figsize=figsize)

    mainFrame = Tk.Frame(s.root)
    mainFrame.pack()
    s.canvas = FigureCanvasTkAgg(s.fig, mainFrame)
    toolbar = NavigationToolbar2Tk(s.canvas, mainFrame)
    toolbar.update()
    toolbar.pack(side=Tk.BOTTOM)
    s.check_frame = Tk.Frame(mainFrame)
    s.check_frame.pack(side=Tk.RIGHT, fill=Tk.BOTH, expand=1)
    s.canvas._tkcanvas.pack(side=Tk.LEFT, fill=Tk.BOTH, expand=1)
    s.canvas.get_tk_widget().pack(side=Tk.LEFT, fill=Tk.BOTH, expand=1)

    menubar = Tk.Menu(s.root)
    if s.func_make_menubar is not None:
        s.func_make_menubar(s, menubar, cfg=cfg)
    s.root.config(menu=menubar)

    if s.func_init_animate is not None:
        s.func_init_animate(s)
    else:
        draw_dd_code(s,0)

    s.root.mainloop()

######################################################################

if __name__ == "__main__":
    args = parse_cli()

    if args.file1!=[] and args.file2!=[]:
        run_nopi_if_needed( args.file1, args.file2, args.yes )
    s = prep_state()
    if args.combo:
        s.fig = Figure()
        s.outfile = open('summary.txt','w')
        if args.combo == 'all':
            combo_list = s.filetypes
        else:
            combo_list = args.combo.split(',')
        for combo in combo_list:
            combo_num = -1
            for i,name in enumerate(s.filetypes):
                if name == combo:
                    combo_num = i
                    break
            if combo_num < 0:
                print("WARNING: invalid combo '%s'" % combo)
                continue
            draw_dd_code(s,combo_num)
            s.fig.savefig('code_' + s.filetypes[combo_num].replace(':','_') + '.png')
            draw_dd_carr(s,combo_num)
            s.fig.savefig('carr_' + s.filetypes[combo_num].replace(':','_') + '.png')
            if ( s.inc_doppler ) :
                draw_dd_dopp(s,combo_num)
                s.fig.savefig('dopp_' + s.filetypes[combo_num].replace(':','_') + '.png')
            draw_sd_cno(s,combo_num)
            s.fig.savefig('cno_' + s.filetypes[combo_num].replace(':','_') + '.png')
    else:
        make_ui(s)