#!/usr/bin/env python3

#Cunren Liang, JPL/Caltech, 28-NOV-2016

#https://matplotlib.org/3.1.1/gallery/text_labels_and_annotations/date.html

import os
import sys
import glob
import datetime
import argparse
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.dates as mdates


def read_alosstack_baseline(baseline_file):
    '''read baseline file generated by alosStack
    '''
    baseline_dict = {}
    with open(baseline_file, 'r') as f:
        lines = [line for line in f if line.strip() != '']
        for x in lines[2:]:
            blist = x.split()
            #to fit into the format of other processors, all alos satellites are after 2000
            #blist[0] = '20' + blist[0]
            #blist[1] = '20' + blist[1]
            baseline_dict[blist[1]] = float(blist[3])
        baseline_dict[blist[0]] = 0

    return baseline_dict


def cmdLineParse():
    '''
    Command line parser.
    '''
    parser = argparse.ArgumentParser(description='plot baselines')
    parser.add_argument('-baseline', dest='baseline', type=str, required=True,
            help = 'baseline file')
    parser.add_argument('-pairs_dir', dest='pairs_dir', type=str, required=True,
            help = 'pairs directory containing YYMMDD-YYMMDD folders. Only folders are recognized.')
    parser.add_argument('-pairs_exc', dest='pairs_exc', type=str, nargs='+', default=None,
            help = 'a number of pairs seperated by blanks. format: YYMMDD-YYMMDD YYMMDD-YYMMDD... If provided, these pairs will be excluded from plotting')
    parser.add_argument('-output', dest='output', type=str, default='baseline.pdf',
            help = 'output file name')

    if len(sys.argv) <= 1:
        print('')
        parser.print_help()
        sys.exit(1)
    else:
        return parser.parse_args()


if __name__ == '__main__':

    inps = cmdLineParse()

    baseline = inps.baseline
    pairs_dir = inps.pairs_dir
    pairs_exc = inps.pairs_exc
    output = inps.output

    baseline_dict = read_alosstack_baseline(baseline)
    pairs = [os.path.basename(x) for x in sorted(glob.glob(os.path.join(pairs_dir, '*-*'))) if os.path.isdir(x)]
    if pairs_exc != None:
        for x in pairs_exc:
            if x in pairs:
                pairs.remove(x)

    #start plot
    plt.rcParams['font.family'] = 'Times New Roman'
    plt.rcParams['font.size'] = 12
    fig, ax = plt.subplots()

    time = [datetime.datetime.strptime(x, "%y%m%d") for x in baseline_dict]
    baseline = [baseline_dict[x] for x in baseline_dict]
    ax.plot(time, baseline, 'o', alpha=0.7, c='g')
    
    year_min = datetime.datetime(min(time).year, 1, 1)
    year_max = datetime.datetime(max(time).year+1, 1, 1)

    for x in pairs:
        rdate, sdate = x.split('-')
        rtime = datetime.datetime.strptime(rdate, "%y%m%d")
        stime = datetime.datetime.strptime(sdate, "%y%m%d")
        time = [rtime, stime]
        baseline = [baseline_dict[rdate], baseline_dict[sdate]]
        ax.plot(time, baseline, '-', lw=.5, c='b')

    ax.xaxis.set_major_locator(mdates.YearLocator())
    ax.xaxis.set_major_formatter(mdates.DateFormatter('%Y'))
    ax.xaxis.set_minor_locator(mdates.MonthLocator())

    ax.minorticks_on()
    ax.tick_params('both', length=7, which='major', width=1)
    ax.tick_params('both', length=4, which='minor', width=0.5)
    ax.set_xlim(year_min, year_max)

    ax.format_xdata = mdates.DateFormatter('%Y-%m-%d')

    # rotates and right aligns the x labels, and moves the bottom of the
    # axes up to make room for them
    #fig.autofmt_xdate()


    ax.set_xlabel('Time [years]')
    ax.set_ylabel('Perpendicular Baseline [meters]')


    plt.savefig(os.path.splitext(output)[0]+'.pdf')