#!/usr/bin/env python3

import sys

from datetime import datetime

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

accumulation = sys.argv[1]
resolution = int(sys.argv[2])
logname = sys.argv[3]
outname = sys.argv[4]

print(f"Accumulation: {accumulation}")
print(f"Resolution: {resolution} seconds")
print(f"Input Log: {logname}")
print(f"Output Image: {outname}")

accumulate = {
        "avg": np.mean,
        "peak": np.max,
    }[accumulation]

data = -1000 * np.ones((86400//resolution, 1024), dtype=float)
counts = np.zeros(86400//resolution, dtype=int)

basetime = None
last_lineidx = 0
accumulated_data = []

with open(logname, 'r') as logfile:
    for line in logfile:
        parts = line.split(";", maxsplit=1)
        timestamp = float(parts[0])

        if not basetime:
            basetime = np.floor(timestamp / 86400) * 86400

        lineidx = int((timestamp - basetime) // resolution)
        if lineidx != last_lineidx:
            if len(accumulated_data) != 0:
                # apply the accumulation function
                data[last_lineidx] = accumulate(np.array(accumulated_data).astype(float), 0)

            accumulated_data = []
            last_lineidx = lineidx

        tmp_data = np.fromstring(parts[1], sep=";", dtype=int)
        if tmp_data.size == 1024:
            accumulated_data.append(tmp_data)
        else:
            print(f"Warning: line ignored due to wrong number of elements: has {tmp_data.size}, expected 1024.")

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

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

fig_x_scale = 0.85
fig_y_scale = 0.90
fig_x_off   = 0.05
fig_y_off   = 0.05

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

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

image = ax.imshow(data, 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 vom ' + datetime.fromtimestamp(basetime).strftime('%Y-%m-%d'))

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

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

ax.yaxis_date()

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'))

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')

pp.savefig(outname)