#!/usr/bin/env python3

import sys

import argparse

from datetime import datetime

from lzma import LZMAFile

import numpy as np
import matplotlib.pyplot as pp
from matplotlib.dates import MinuteLocator, HourLocator, DayLocator, DateFormatter, AutoDateLocator, drange, date2num
from matplotlib.ticker import MultipleLocator

parser = argparse.ArgumentParser(description="Plots long term waterfalls from KiwiSDR spectrum recordings.")

parser.add_argument('-a', '--accumulation', choices=['avg_db', 'avg_pwr', 'peak'], default='avg_pwr', help='Accumulation method.')
parser.add_argument('-r', '--resolution', type=int, default=600, help='Spectrogram time resolution in seconds.')
parser.add_argument('-o', '--output', type=argparse.FileType('wb'), required=True, help='The image output name.')
parser.add_argument('input', nargs='+', type=str, help='KiwiSDR spectrum dumps to load (uncompressed or xz-compressed)')

args = parser.parse_args()
print(args)

print(f"Accumulation: {args.accumulation}")
print(f"Resolution: {args.resolution} seconds")
print(f"Input Logs: {args.input}")
print(f"Output Image: {args.output}")

loaded_data = {} # format: base_timestamp => {count, spectrum data as np.array}

for inputfilename in args.input:
    try:
        if inputfilename[-2:].lower() == "xz":
            print(f"Loading '{inputfilename}' as XZ stream...")
            stream = LZMAFile(filename=inputfilename, mode='r')
        else:
            print(f"Loading '{inputfilename}' as uncompressed stream...")
            stream = open(inputfilename, 'rb')

        for line in stream:
            # extract the timestamp
            parts = line.decode('ascii').split(";", maxsplit=1)
            timestamp = float(parts[0])

            # calculate which accumulation block this spectrum fits into
            basetime = (timestamp // args.resolution) * args.resolution

            # split the spectrum data line
            tmp_data = np.fromstring(parts[1], sep=";", dtype=int)
            if tmp_data.size != 1024:
                print(f"Warning: line ignored due to wrong number of elements: has {tmp_data.size}, expected 1024.")
                continue

            # get the block data loaded so far or initialize it
            if basetime not in loaded_data.keys():
                if args.accumulation == 'peak':
                    loaded_data[basetime] = {'data': -1000 * np.ones(1024)}
                else:
                    loaded_data[basetime] = {'count': 0, 'data': np.zeros(1024)}

            block_data = loaded_data[basetime]

            # on the fly accumulation
            if args.accumulation == 'peak':
                block_data['data'] = np.max(np.stack([tmp_data, block_data['data']]), axis=0)
            elif args.accumulation == 'avg_db':
                block_data['data'] += tmp_data
                block_data['count'] += 1
            else: # avg_pwr
                block_data['data'] += 10**(tmp_data/10) # convert dBm to mW
                block_data['count'] += 1

            loaded_data[basetime] = block_data

        stream.close()
    except Exception as e:
        print(f"Error while processing '{inputfilename}': {e}")
        print("This file’s (remaining) data is skipped.")

print("Post-processing accumulation...")

if args.accumulation == 'avg_db':
    for v in loaded_data.values():
        v['data'] /= v['count']
elif args.accumulation == 'avg_pwr':
    for v in loaded_data.values():
        v['count'] = max(1, v['count']) # prevent division by 0
        v['data'] = 10 * np.log10(v['data'] / v['count']) # convert averaged value back to dBm
# peak accumulation does not need finalization

print(f"Loading completed. {len(loaded_data)} blocks accumulated.")

# create a contiguous array from the spectrum lines
timestamps = loaded_data.keys()

first_timestamp = min(timestamps)
last_timestamp = max(timestamps)

num_spectrums = int((last_timestamp - first_timestamp) // args.resolution) + 1
shape = (num_spectrums, loaded_data[first_timestamp]['data'].size)

print(f"Reserving an array with shape {shape}...")

spectrogram = -1000.0 * np.ones(shape, dtype=float)

print(f"Copying the data in...")

for k, v in loaded_data.items():
    idx = int((k - first_timestamp) // args.resolution)

    spectrogram[idx, :] = v['data']

del loaded_data # not needed anymore

print("Plotting image...")

basetime = first_timestamp

x = np.linspace(0, 30, spectrogram.shape[1])
y = np.linspace(basetime, basetime + args.resolution * spectrogram.shape[0])

extent = [
        0, # MHz
        30, # MHz
        date2num(datetime.utcfromtimestamp(basetime + args.resolution * spectrogram.shape[0])),
        date2num(datetime.utcfromtimestamp(basetime))]

pic_w_px = 1024 + 127 + 150

fig_x_scale = 1024 / pic_w_px
fig_x_off   = 127 / pic_w_px # from left side

pic_h_px = spectrogram.shape[0] + 70 + 100

fig_y_scale = spectrogram.shape[0] / pic_h_px
fig_y_off   = 70 / pic_h_px # from bottom

dpi = pp.rcParams['figure.dpi'] #get the default dpi value
fig_size = (spectrogram.shape[1]/dpi/fig_x_scale, spectrogram.shape[0]/dpi/fig_y_scale) # convert pixels to DPI

fix, ax = pp.subplots(figsize=fig_size)

image = ax.imshow(spectrogram, vmin=-100, vmax=0, cmap='inferno',
                  extent=extent)

ax.set_position(pos=[fig_x_off, fig_y_off, fig_x_scale, fig_y_scale])
ax.set_title('Spektrogramm von '
             + datetime.utcfromtimestamp(first_timestamp).strftime('%Y-%m-%d %H:%M:%S')
             + ' bis '
             + datetime.utcfromtimestamp(last_timestamp).strftime('%Y-%m-%d %H:%M:%S')
             + ' UTC')

xrange = extent[1] - extent[0]
yrange = extent[2] - extent[3]

ax.set_aspect((spectrogram.shape[0] / yrange) / (spectrogram.shape[1] / xrange))

print("Formatting...")

ax.yaxis_date()

if len(args.input) == 1:
    # only one argument means that we plot at most 1 day. Use Hour and Minute
    # locators for nicer formatting.
    ax.yaxis.set_major_locator(HourLocator(range(0, 25, 1)))
    ax.yaxis.set_minor_locator(MinuteLocator(range(0, 60, 20)))
    ax.yaxis.set_major_formatter(DateFormatter('%H:%M'))
    ax.yaxis.set_minor_formatter(DateFormatter('%M'))
else:
    # we are plotting multiple files and therefore probably a long time. Use
    # Day and Hour locators.

    ax.yaxis.set_major_locator(DayLocator(interval=max(1, args.resolution // 600)))
    ax.yaxis.set_major_formatter(DateFormatter('%Y-%m-%d'))

    if args.resolution <= 600:
        ax.yaxis.set_minor_locator(HourLocator(interval=max(1, args.resolution // 100)))
        ax.yaxis.set_minor_formatter(DateFormatter('%H'))

ax.xaxis.set_major_locator(MultipleLocator(1.0))
ax.set_xlabel('Frequenz [MHz]')

ax_top = ax.twiny()
ax_top.xaxis.set_major_locator(MultipleLocator(1.0))
ax_top.set_xlabel('Frequenz [MHz]')
ax_top.set_xlim(ax.get_xlim())
ax_top.set_position(pos=[fig_x_off, fig_y_off, fig_x_scale, fig_y_scale])

cax = pp.axes([fig_x_scale+fig_x_off+0.02, fig_y_off, 0.03, fig_y_scale])
pp.colorbar(cax=cax, mappable=image)
cax.set_title('dBm')

print("Saving...")

pp.savefig(args.output)

print("All done!")