#!/usr/bin/env python

import matplotlib
matplotlib.use("Agg")
import matplotlib.animation as manimation
from mutils import *
import argparse
import time, sys

parser = argparse.ArgumentParser(description='Create a movie of rec 35:261 FFT correlator.')
parser.add_argument('filename', help='T04 filename with rec35:261 data')
parser.add_argument('-s', '--satellite', help='Satellite ID/sat_type, e.g. -s 25/0 or -s */0 or -s 3/0,5/0', required=True)
parser.add_argument('-r', '--range_m', help='+/- range of plot in meters (default 1200)')
parser.add_argument('-b', '--begin', help='Optional start time in GPS week seconds')
parser.add_argument('-e', '--end', help='Optional end time in GPS week seconds')
args = parser.parse_args()

plot_sv_list = []
for combo in args.satellite.split(','):
    sv,sat_type = combo.split('/')
    if sv == '*':
        break
    plot_sv_list.append( (int(sv), int(sat_type)) )
if args.range_m:
    range_m = float(args.range_m)
else:
    range_m = 1200.


if(args.begin):
  start = int(args.begin)
else:
  start = 0

if(args.end):
  stop = int(args.end)
else:
  stop = 7 * 86400

if not 'd' in globals():
    d,k=vd2arr(args.filename,rec='-d35:261')
    print(start)
    print(stop)
    i = find( (d[:,k.TIME] >= start) & (d[:,k.TIME] <= stop) )
    d = d[i,:]
    # Scale to dB (RX units are dB x 10)
    d[:,k.cno] = d[:,k.cno] * 0.1
    sv_list = get_sv_list(d,k)
if len(plot_sv_list) == 0:
    plot_sv_list = sv_list

FFMpegWriter = manimation.writers['ffmpeg']
metadata = dict(title='rec35:261', artist='Matplotlib')
writer = FFMpegWriter(fps=10, metadata=metadata)

# Create all figure lines in advance because updating data with
# set_ydata() is much faster than calling plot()
close('all')

fig = figure()
ax0 = subplot(211)
ax1 = subplot(212)
fft_mid = len(d[0,k.corr_start:])/2
x = (r_[0:2*fft_mid] - fft_mid)*GnssConst.LIGHT/4.096e6

l_corr_fn = []
l_all_cno = []
l_curr_cno = []
for sv,sat_type in plot_sv_list:
    l0, = ax0.plot( x, x ) #, '.-' )
    l_corr_fn.append(l0)

    i = find( (d[:,k.SV]==sv) & (d[:,k.SAT_TYPE]==sat_type) )
    if len(i) == 0:
        print("Can't find sv %d/%d" % (sv,sat_type))
        print("Try something from:")
        print(sv_list)
        sys.exit(1)
    l1_0, = ax1.plot( d[i,k.TIME], d[i,k.cno] )
    l1_1, = ax1.plot( [], [], 'r.' )
    l_all_cno.append( l1_0 )
    l_curr_cno.append( l1_1 )

ax0.set_xlim([-range_m,range_m])
ax0.set_ylim([0,1.])
ax0.set_ylabel('FFT corr')
ax0.set_xlabel('corr offset [m]')
ax1.set_ylabel('CNo [dBHz]')
ax1.set_xlabel('GPS secs')
tight_layout()

# Line up data in a pandas dataframe
import pandas as pd
all_d = {}
for sv,sat_type in plot_sv_list:
    i = find( (d[:,k.SV]==sv) & (d[:,k.SAT_TYPE]==sat_type) )
    s = pd.Series( i, index=d[i,k.TIME] )
    dups = s.index.duplicated()
    all_d[(sv,sat_type)] = s[~dups]
data_i = pd.DataFrame( all_d )

# Create move
filename = "sv.mp4"
DPI=150
with writer.saving(fig, filename, DPI):
    t0 = time.time()
    tot_jump = 0
    for pos in range(len(data_i)):
        if time.time() - t0 > 10.:
            t0 = time.time()
            print('%.1f%% done' % (100.*pos/len(data_i)))
        data_i_pos = data_i.iloc[pos]
        has_title = False
        for n in range(len(plot_sv_list)):
            idx = data_i_pos.iloc[n]
            sv,sat_type = data_i_pos.index[n]
            if isnan(idx):
                l_corr_fn[n].set_visible(False)
                l_curr_cno[n].set_visible(False)
                continue
            else:
                l_corr_fn[n].set_visible(True)
                l_curr_cno[n].set_visible(True)
            idx = int(idx)

            ri = d[idx,k.track_idx]
            scale = 1. / max(d[idx,k.corr_start:])
            l_corr_fn[n].set_ydata( roll(d[idx,k.corr_start:] * scale, int(fft_mid-ri-1) ) )

            l_curr_cno[n].set_xdata( d[idx,k.TIME] )
            l_curr_cno[n].set_ydata( d[idx,k.cno] )
            if not has_title:
                if len(plot_sv_list) == 1:
                    if (int(d[idx,k.info]) & k.dINFO_FORCE_JUMP)!=0:
                        tot_jump += 1
                    res_ok = (int(d[idx,k.info]) & k.dINFO_RES_VALID)!=0
                    res_trk = d[idx,k.res_track]
                    ax0.set_title('time %.1f tot_jump=%d res_ok=%d res_trk=%.1f' %
                                  (d[idx,k.TIME],tot_jump,res_ok,res_trk))
                else:
                    ax0.set_title('time %.1f' % d[idx,k.TIME])
                has_title = True

        writer.grab_frame()
print('generated %s' % filename)