#!/usr/bin/env python
"""

 Plot tracking & cycle slip data and RAIM flags.  Either
 processes a single T0x file or else all of the T04 files
 in a specified directory using wildcard.

   Interface to plot_tracking() and plot_raim() in mutils.py

 Example:
  ./plt_tracking.py /home/ddelore/DataDir/RX6-2209
  Defaults to tracking plots.
"""


import os
import argparse
import matplotlib
from mutils import *


def plt_tracking(args):

    path_name = args.path_name
    subchan_raim_figs = []  # Initialize before any conditionals
    
    if path_name.lower().endswith('.t02') or path_name.lower().endswith('t04'):
        # Process a single T02 or T04 file
        title_txt = os.path.basename(path_name)
    else:
        # Process all of the T04 files in a directory
        title_txt = path_name
        path_name = os.path.join(path_name, '*.T04')

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

    sat_types = None
    if args.sat_types:
        sat_types = [int(x) for x in args.sat_types.split(',')]

    if args.track and args.raim:
        # Call plot_raim() and plot_tracking()
        d = vd2cls(path_name, rec='-d35:19')
    elif args.raim or args.sub_chan_raim:
        # Just call plot_raim()
        d = vd2cls(path_name, rec='-d35:19 -t')
    else:
        # Just call plot_tracking() [default]
        d = vd2cls(path_name, rec='-d27,35:19')

    if args.antenna_num < 0:
        antenna_list = unique(d.ANTENNA)
    else:
        antenna_list = array([args.antenna_num])
    raim_figs = {}
    trk_figs = {}
    subchan_raim_figs = {}
    for antenna in antenna_list:
        # Select data for this antenna
        d0 = d[d.ANTENNA==antenna]

        curr_txt = title_txt + f' (ant{antenna:.0f})'
        if args.raim:
            raim_figs[antenna]= plot_raim( d0, d0.kf, sat_types=sat_types, title_txt=curr_txt )

        if args.track:
            trk_figs[antenna] = plot_tracking( d0, d0.k, sat_types=sat_types, title_txt=curr_txt )

        if args.sub_chan_raim:
            subchan_raim_figs[antenna] = plot_subchan_raim( d0, d0.kf,
                                                            sat_types=sat_types, title_txt=curr_txt )

    systemLUT = { 0:'GPS',
                  1:'SBAS',
                  2:'GLONASS',
                  3:'Galileo',
                  4:'QZSS',
                  9:'NavIC',
                  10:'BeiDou'} 

    if args.png:
        if args.raim:
            for antenna in antenna_list:
                for fig,sat_type in raim_figs[antenna]:
                    if len(antenna_list) > 1:
                        filename = 'raim-%s-%d.png' % (systemLUT[int(sat_type)],antenna)
                    else:
                        filename = 'raim-%s.png' % (systemLUT[int(sat_type)],)
                    fig.savefig(filename)
                    print('saved %s' %filename)
        if args.track:
            for antenna in antenna_list:
                for fig,sat_type in trk_figs[antenna]:
                    if len(antenna_list) > 1:
                        filename = 'track-%s-%d.png' % (systemLUT[int(sat_type)],antenna)
                    else:
                        filename = 'track-%s.png' % (systemLUT[int(sat_type)])
                    fig.savefig(filename)
                    print('saved %s' %filename)
        if args.sub_chan_raim:
            for antenna in antenna_list:
                for fig,sat_type in subchan_raim_figs[antenna]:
                    if len(antenna_list) > 1:
                        filename = 'subchan-raim-%s-%d.png' % (systemLUT[int(sat_type)],antenna)
                    else:
                        filename = 'subchan-raim-%s.png' % (systemLUT[int(sat_type)])
                    fig.savefig(filename)
                    print('saved %s' %filename)

    show()


if __name__ == '__main__':
    parser = argparse.ArgumentParser(formatter_class=argparse.RawTextHelpFormatter,
                                     description=__doc__)

    parser.add_argument('path_name', help='name of or path to your T0x file(s)')
    parser.add_argument('-r', '--raim', action="store_true",
                        help='Plot RAIM info - only for T04 files')
    parser.add_argument('-t', '--track', action="store_true",
                        help='Plot tracking info - for T02 or T04 files')
    parser.add_argument('-p', '--png', action="store_true",
                        help='Generate PNG files in headless operation')
    parser.add_argument('--sat_types',
                        help='List of RT satellite types to plot')
    parser.add_argument('--antenna_num', type=int, default=0,
                        help='Antenna number [default 0. Note: -1 is all antennas]')
    parser.add_argument('--sub_chan_raim', action="store_true",
                    help='Plot sub-channel RAIM info')
    
    args = parser.parse_args()
    if not args.raim and not args.track:
        args.track = True
    plt_tracking(args)