#!/usr/bin/env python
usage = """\
Compare CDF between RF replay runs or local files.

Examples:
 1) Use "-r" to include stuff from the RF replay system
  ./compare_runs_cdf.py -r RX15-3790,RX16-3790,RX22-3790 -z 80

 2) Use "-l" to include stuff locally (outside the RF replay system)
  ./compare_runs_cdf.py -r RX15-3790,RX16-3790,RX22-3790 -l File1.T04,File2.T04

"""

if __name__ == "__main__":
    import matplotlib
    # Allow running headless from the command line
    matplotlib.use("agg")

from mutils import *
from mutils.PosPacConst import *
from ProcessResults import parse_sampleConfig, fill_in_auto_config_spans
from glob import glob
import argparse
import os
import shutil
from collections import defaultdict

def main():
    """Stand-alone operation.  See usage at top of file"""
    parser = argparse.ArgumentParser(formatter_class=argparse.RawTextHelpFormatter,
                                     description=usage)
    parser.add_argument('-l','--local',
                        help="comma-separated local T04 file/directory (dir1,dir2)")
    parser.add_argument('-r','--replay',
                        help="comma-separated RF replay run name (RX15-3790,RX16-3790)")
    parser.add_argument('-z','--zoom_y',
                        type=float,
                        default=75.,
                        help='Zoom CDF Y-axis to only given percent and above [75]')
    parser.add_argument('-x','--x_limit',
                        type=float,
                        default=None,
                        help='Limit CDF X-axis maximum value [None]')
    parser.add_argument('-o','--out_dir',
                        default=".",
                        help='Output directory for plots ["."].  Will get overwritten')
    parser.add_argument('-c','--config',
                        help='Optionally specify .xml config for time ranges.')
    args = parser.parse_args()

    # Create output dir?
    if args.out_dir != '.':
        if os.path.isdir( args.out_dir ):
            shutil.rmtree( args.out_dir )
        os.makedirs( args.out_dir )

    # Get path to all files
    all_files = []
    if args.replay:
        rf_replay_loc = '/mnt/data_drive/SpirentTest/DataDir/'
        if not os.path.isdir(rf_replay_loc):
            rf_replay_loc = '/net/fermion' + rf_replay_loc
        all_files.extend( [rf_replay_loc+x for x in args.replay.split(',')] )
    if args.local:
        all_files.extend( args.local.split(',') )
    if len(all_files) == 0:
        print("Error: need either -r or -l option.  Use -h for more info.")
        sys.exit(1)

    # Convert directories to T04 patterns
    for i in range(len(all_files)):
        if not all_files[i].lower().endswith('.t04'):
            all_files[i] = all_files[i] + '/*.T04'

    # Use given config, or try to find it
    if args.config:
        config_filename = args.config
    else:
        dir_name = os.path.dirname(all_files[0]) + '/'
        config_filename = None
        for tmp in glob(dir_name + '*.xml'):
            if tmp.endswith('ALM.xml') or tmp.endswith('config.xml'):
                continue
            config_filename = tmp
            break
        if config_filename is None:
            print("Can't find scenario XML.  Use -c?")
            sys.exit(1)
    config = parse_sampleConfig( config_filename )
    if config.truth_format != "POSPAC_ASCII":
        raise RuntimeError("Currently only works with dynamic data")

    # Get truth data (Assume we only need 10Hz)
    pp = doload(config.truth_file)
    pp=pp[(around(pp[:,dPP_TIME]*1000)%100)==0,:]

    # get and combine timespans (segments to analyze):
    #   spans[timespan_desc] = [(start1,end1), (start2,end2), ...]
    fill_in_auto_config_spans( config, pp )
    spans = defaultdict(list)
    for desc,start_stops in config.span:
        for start,stop in start_stops:
            spans[desc].append( (start,stop) )

    # Load data files
    all_d = []
    for in_file in all_files:
        d = vd2cls(in_file,'-d35:2')
        short_file = in_file.split('/')
        if len(short_file) == 1:
            short_file = short_file[-1]
        else:
            short_file = short_file[-2]
        all_d.append((short_file,d))

    # Create CDF plots
    for name,start_stops in spans.items():
        print("Processing segment %s"%name)
        i = []
        for start,stop in start_stops:
            i.extend( find((pp[:,dPP_TIME]>=start)&(pp[:,dPP_TIME]<=stop)).tolist() )
        i = unique(i)
        pp0 = pp[i,:]
        fig_3d,ax_3d=subplots()
        fig_2d,ax_2d=subplots()
        for d_name,d in all_d:
            t, i_pp0, i_d = intersect( pp0[:,dPP_TIME], d.TIME )
            (dE, dN, dU) = llh2enu( d[i_d,d.k.LAT:d.k.LAT+3], pp0[i_pp0,dPP_LAT:dPP_LAT+3] )
            err_3d = sqrt( dE**2 + dN**2 + dU**2 )
            err_2d = sqrt( dE**2 + dN**2 )
            cx3,cy3 = docdf( err_3d )
            cy3 *= 100
            cx2,cy2 = docdf( err_2d )
            cy2 *= 100
            if args.x_limit is not None:
                i3=find((cy3>=args.zoom_y)&(cx3<=args.x_limit))
                i2=find((cy2>=args.zoom_y)&(cx2<=args.x_limit))
            else:
                i3=find(cy3>=args.zoom_y)
                i2=find(cy2>=args.zoom_y)
            ax_3d.plot( cx3[i3], cy3[i3], label=d_name )
            ax_2d.plot( cx2[i2], cy2[i2], label=d_name )
        for ax in [ax_3d,ax_2d]:
            ax.grid()
            ax.legend()
            ax.set_ylabel('CDF %')
        ax_3d.set_xlabel('3D err[m]')
        ax_2d.set_xlabel('2D err[m]')
        ax_3d.set_title('3D %s'%name)
        ax_2d.set_title('2D %s'%name)
        fig_3d.savefig(args.out_dir+'/%s_3d.png'%name)
        fig_2d.savefig(args.out_dir+'/%s_2d.png'%name)
        close('all')

if __name__ == '__main__':
    main()