import os
import sys
import time
import argparse
import numpy as np
sys.path.append('../decode_rf_data')
import matplotlib.pyplot as plt
from colormap import spectral
import rf_data as rf

parser = argparse.ArgumentParser(description="""
        Will plot waterfall plots
        """)

parser.add_argument('--rf_dir', metavar='\b', default='./../../data/', help='Path to rf data directory. Default=./../../data/')
parser.add_argument('--ref_name', metavar='\b', default='rf0XX_2019-10-10-02:30', help='Name of ref data file. Default=rf0XX_2019-10-10-02:30')
parser.add_argument('--tile_name', metavar='\b', default='S10XX_2019-10-10-02:30', help='Name of tile data file. Default=S10XX_2019-10-10-02:30')
parser.add_argument('--out_dir', metavar='\b', default='./../../outputs/paper_plots/', help='Output dir. Default=./../../outputs/paper_plots/')

args = parser.parse_args()
rf_dir = args.rf_dir
ref_name = args.ref_name
tile_name = args.tile_name
out_dir = args.out_dir

# Make output dir if it doesn't exist
os.makedirs(os.path.dirname(out_dir), exist_ok=True)




# read in raw data
ref_p, ref_t = rf.read_data(f'{rf_dir}/{ref_name}.txt')
tile_p, tile_t = rf.read_data(f'{rf_dir}/{tile_name}.txt')

# scale median to zero
ref_p_median = np.median(ref_p)
tile_p_median = np.median(tile_p)
r_image = ref_p - ref_p_median
t_image = tile_p - tile_p_median

# setting dynamic range of waterfall to be 30 dB above the median
vmin = 0
vmax = 30

# Custom spectral colormap
cmap = spectral()

nice_fonts = {
        # Use LaTeX to write all text
        "text.usetex": True,
        "font.family": "sans-serif",
        # Use 10pt font in plots, to match 10pt font in document
        "axes.labelsize": 10,
        "font.size": 10,
        # Make the legend/label fonts a little smaller
        "legend.fontsize": 8,
        "xtick.labelsize": 8,
        "ytick.labelsize": 8,
        #"ytick.color" : "w",
        #"xtick.color" : "w",
        #"axes.labelcolor" : "w",
        #"axes.edgecolor" : "w"
        }

plt.rcParams.update(nice_fonts)

fig = plt.figure(figsize=(7, 4.2))
#fig, axs = plt.subplots(1,2, figsize=(7,5))
#axs = axs.ravel()

ax1 = fig.add_axes([0.00, 0.0, 0.46, 1])
ax2 = fig.add_axes([0.50, 0.0, 0.46, 1])
cax = fig.add_axes([0.98, 0.0, 0.015,1])

ax1.imshow(t_image, vmin=vmin, vmax=vmax, interpolation='none', cmap=cmap)
ax1.set_aspect('auto')
image = ax1.get_images()[0]
cbar = fig.colorbar(image, cax=cax, label='Power [dB]')

# Number of time steps on y-axis
number_t = 5
t_step = int(len(tile_t)/(number_t-1))
times = list(tile_t)
times = times[::t_step]

t_tz = []

# Convert UNIX time to local HH:MM time
for i in range(len(times)):
    
    perth_t = float(times[i])+28800 #28800=+8GMT @ PERTH
    hms = time.strftime('%H:%M', time.gmtime(perth_t))
    t_tz.append(hms)

# Frequency: x-axis
start_freq = 137.15
stop_freq = 138.55

# X-axis stuff
x_ax = t_image.shape[1]
freqs = np.arange(start_freq, stop_freq, 0.25)
x_ticks = np.arange(0, x_ax, (0.25/0.0125)) #.0125MHz/ch
ax1.set_xticks(x_ticks)
ax1.set_xticklabels(freqs)
ax1.set_xlabel('Freqency [MHz]')

# Y-axis stuff
y_ax = t_image.shape[0]
y_ticks = np.arange(0, y_ax, t_step)
ax1.set_yticks(y_ticks)
ax1.set_yticklabels(t_tz)
ax1.set_ylabel('MWA local time [HH:MM]')


ax2.imshow(r_image, vmin=vmin, vmax=vmax, interpolation='none', cmap=cmap)
ax2.set_aspect('auto')
# Number of time steps on y-axis
number_t = 5
t_step = int(len(ref_t)/(number_t-1))
times = list(ref_t)
times = times[::t_step]

t_tz = []

# Convert UNIX time to local HH:MM time
for i in range(len(times)):
    
    perth_t = float(times[i])+28800 #28800=+8GMT @ PERTH
    hms = time.strftime('%H:%M', time.gmtime(perth_t))
    t_tz.append(hms)

# Frequency: x-axis
start_freq = 137.15
stop_freq = 138.55

# X-axis stuff
x_ax = r_image.shape[1]
freqs = np.arange(start_freq, stop_freq, 0.25)
x_ticks = np.arange(0, x_ax, (0.25/0.0125)) #.0125MHz/ch
ax2.set_xticks(x_ticks)
ax2.set_xticklabels(freqs)
ax2.set_xlabel('Freqency [MHz]')

# Y-axis stuff
ax2.set_yticklabels([])
ax2.set_yticks([])


plt.savefig(f'{out_dir}/waterfall.png', dpi=600, bbox_inches='tight')
#plt.savefig(f'waterfall.png', dpi=600, transparent=True, bbox_inches='tight')

