#!/usr/bin/env python3

import numpy as np
import matplotlib.pyplot as plt
import argparse, sys, math
from astropy.time import Time
from astropy.coordinates import SkyCoord, EarthLocation, AltAz

#import astropy.units as u

parser = argparse.ArgumentParser()

parser.add_argument('-rate', '--rate', help="Plot Rate", action="store_true")
parser.add_argument('-U', '--U', help="Plot U", action="store_true")
parser.add_argument('-V', '--V', help="Plot V", action="store_true")
parser.add_argument('-W', '--W', help="Plot W", action="store_true")
parser.add_argument('-uv', '--uv', help="Plot UV", action="store_true")
parser.add_argument('-f', '-freq', '--freq', help="Observing Frequency (MHz)", type=float)
parser.add_argument('-r', '-refant', '--refant', help="Reference antenna")
parser.add_argument('-el', '--el', help="Plot vs elevation", action="store_true")
parser.add_argument('-t', '-time', '--time', help="Start time of re")
parser.add_argument('file')

args = parser.parse_args()

if args.el and args.time is None:
    print("Error: Must pass time if plotting elevation")
    sys.exit(1)

askap = EarthLocation.from_geocentric(-2556088.333, 5097405.644, -2848428.244, unit='m')
source = SkyCoord(0, 0.5, unit='rad')

c = 1
scale = 1
ylabel = 'Delay (usec)'

if args.freq is not None:
    scale *= args.freq*1e6*360/1e6
    ylabel = 'Phase (deg)'

if args.rate:
    c = 2
    if args.freq is not None:
        ylabel = 'Fringe Rate (Deg/sec)'
    else:
        ylabel= 'Delay Rate (psec/sec)'
        scale = 1e6

if args.uv:
    subRef = True

subRef = False
    
if args.U:
    scale = 1
    c = 3
    ylabel= 'U (m)'
    subRef = True

if args.V:
    scale = 1
    c = 4
    ylabel= 'V (m)'
    subRef = True

if args.W:
    scale = 1
    c = 5
    ylabel= 'W (m)'
    subRef = True

if subRef and args.refant is None:
    print("Error: Must pass reference Antenna")
    sys.exit(1)

if args.refant is not None:
    subRef = True

refAnt = args.refant

if args.el:
    startTime = Time(args.time)
    
def readCalc(file):

    calc = {}

    data=np.genfromtxt(file, dtype=(float,'S4',float,float,float,float,float), comments='#', usecols=(0,2,4,5,6,7,8))
    

    for d in data:
        ant = str(d[1],'utf-8')
        if not ant in calc:
            calc[ant] = [[] for i in range(6)]

        calc[ant][0].append(d[0])
        calc[ant][1].append(d[2])
        calc[ant][2].append(d[3])
        calc[ant][3].append(d[4])
        calc[ant][4].append(d[5])
        calc[ant][5].append(d[6])

    for a in calc.keys():
        for i in range(6):
            calc[a][i] = np.array(calc[a][i])
        
    return(calc)


calc = readCalc(args.file)


for a in calc.keys():
#for a in ['AK18', 'AK32','AK34', 'AK35']:
    if args.uv:
        x = calc[a][3] - calc[refAnt][3]
        y = calc[a][4] - calc[refAnt][4]
    else:
        x = calc[a][0]/60.0/60.0
        y = calc[a][c]*scale
        if subRef:
            y  -= calc[refAnt][c]*scale

    if args.el:
        times = startTime + (x/24.0)
        aa = AltAz(location=askap, obstime=times)
        azel = source.transform_to(aa)
        x = azel.alt.degree
            
    plt.plot(x,y,label=a)

if args.uv:
    plt.ylabel('U (m)')
    plt.xlabel('V (m)')
else:
    plt.ylabel(ylabel)
    if args.el:
        plt.xlabel('Elevation (Degrees)')
    else:
        plt.xlabel('Time (hrs)')

if len(calc.keys()) > 18:
    ncol = 2
else:
    ncol = 1
    
plt.legend(loc='center left', bbox_to_anchor=(1, 0.5), prop={"size":9}, ncol=ncol)
plt.tight_layout()
plt.show()