#!/usr/bin/env python3
import matplotlib.pyplot as plt
import aiscat
import statistics
import argparse
import time

parser = argparse.ArgumentParser()
parser.add_argument('file')
parser.add_argument('-t', nargs='*', type=int, help='Filter by msg type(s). For example -t 4 21 for base stations and AtoN')
parser.add_argument('-m', nargs='*', type=int, help='Filter by MMSI(s). For example -m 3669154 3669152.')
parser.add_argument('-mb', action='store_true', help='Filter base station MMSIs.')
parser.add_argument('-g', choices=['mmsi', 'type', 'channel'], default='mmsi', help='AIS field to use for grouping. Default: mmsi')
parser.add_argument('-x', choices=['time', 'freq', 'power', 'distance'], default='time', help="X-axis. Default: time")
parser.add_argument('-y', choices=['time', 'freq', 'power', 'distance'], default='time', help="Y-axis. Default: time")
parser.add_argument('-lon', type=float, help='Longitude in degrees of receiver.')
parser.add_argument('-lat', type=float, help='Latitude in degrees of receiver.')
parser.add_argument('-kts', type=float, help='Filter targets by minimum speed in kts.')
parser.add_argument('-st', type=float, help='Filter targets by maximum slot time in microseconds.')
parser.add_argument('-c', type=int, help='Filter groups with fewer counts.')
parser.add_argument('-v', action='store_true', help='Validate and discard entries that fail sanity check (distance>1000km, slot timing>5ms).')
args = parser.parse_args()

if (args.x == 'distance' or args.y == 'distance') and (not args.lon or not args.lat):
    print("Error: Plotting distance requires longitude and latitude of receiver (-lon, -lat)!")
    exit()

def is_base_station(mmsi):
    # Base stations have MMSIs in 002000000-009999999 (first two digits 00)
    return 2000000 <= mmsi <= 9999999

from math import radians, cos, sin, asin, sqrt
def haversine(lon1, lat1, lon2, lat2):
    lon1, lat1, lon2, lat2 = map(radians, [lon1, lat1, lon2, lat2])
    dlon = lon2 - lon1
    dlat = lat2 - lat1
    a = sin(dlat/2)**2 + cos(lat1) * cos(lat2) * sin(dlon/2)**2
    c = 2 * asin(sqrt(a))
    return 6371 * c

slot_timestamps = {}
rxux_timestamps = {}
power_dbm = {}
freq_ppm = {}
distances = {}
meta = {}
total = 0

start_time = time.time()
row_count = 0

for m in aiscat.from_file(args.file):
    row_count += 1
    # Prefer toa for timing accuracy; fall back to rxuxtime.
    rxuxtime = m.get('toa')
    if rxuxtime is None:
        continue

    msg_type = m.get('type')
    mmsi = m.get('mmsi')

    if args.t or args.m or args.mb:
        match = False
        if args.m and mmsi in args.m:
            match = True
        if args.mb and is_base_station(mmsi):
            match = True
        if args.t:
            if args.m or args.mb:
                if msg_type not in args.t:
                    match = False
            elif msg_type in args.t:
                match = True
        if not match:
            continue

    lat = m.get('lat')
    lon = m.get('lon')
    if (args.x == 'distance' or args.y == 'distance') and (lat is None or lon is None):
        continue

    if lat is not None and lon is not None and args.lon and args.lat:
        distance = haversine(args.lon, args.lat, lon, lat)
    else:
        distance = 0

    speed = m.get('speed', 0) or 0

    # Slot timing from rxuxtime, reference at expected ToA (T2 = 3.333ms)
    slot_time_us = int((((rxuxtime % 60) * (9600.0/256.0)) % 1.0) * (256.0/9600.0) * 1000000 - 3333.33 + 0.5)

    if args.st and slot_time_us > args.st:
        continue
    if args.kts and speed < args.kts:
        continue
    if args.v:
        if distance > 1000:
            continue
        if slot_time_us > 5000 or slot_time_us < -100:
            continue

    if args.g == 'type':
        g = msg_type
    elif args.g == 'channel':
        g = m.get('channel', '?')
    else:
        g = mmsi

    if g not in slot_timestamps:
        slot_timestamps[g] = []
        rxux_timestamps[g] = []
        power_dbm[g] = []
        freq_ppm[g] = []
        distances[g] = []
        meta[g] = []
    slot_timestamps[g].append(slot_time_us)
    rxux_timestamps[g].append(rxuxtime)
    power_dbm[g].append(m.get('signalpower', 0))
    freq_ppm[g].append(m.get('ppm', 0))
    distances[g].append(distance)
    meta[g].append(mmsi)
    total += 1

if args.g == 'type':
    group_label = 'msg type'
elif args.g == 'channel':
    group_label = 'channel'
else:
    group_label = 'MMSI'

print(f"Processed {row_count} lines in {time.time() - start_time:.2f} seconds")
print(f"Plotted {total} data points for {len(slot_timestamps)} unique {group_label}s")
if len(slot_timestamps) > 0:
    print(f"{group_label:<9}  {'count':>5}  {'mean':>7}  {'median':>7}  {'stdev':>8}  {'power min':>9}  {'max':>9}  {'mean':>9}  {'freq min':>9} {'max':>9} {'mean':>9}")
for m in slot_timestamps:
    if args.c and len(slot_timestamps[m]) < args.c:
        continue
    mean = statistics.mean(slot_timestamps[m])
    median = statistics.median(slot_timestamps[m])
    stdev = statistics.stdev(slot_timestamps[m]) if len(slot_timestamps[m]) > 1 else 0
    min_power = min(power_dbm[m])
    max_power = max(power_dbm[m])
    mean_power = statistics.mean(power_dbm[m])
    min_freq = min(freq_ppm[m])
    max_freq = max(freq_ppm[m])
    mean_freq = statistics.mean(freq_ppm[m])
    print(f"{str(m):<9}  {len(slot_timestamps[m]):>5}  {mean:>5.0f}us  {median:>5.0f}us  {stdev:>6.1f}us  {min_power:>6.1f}dBm  {max_power:>6.1f}dBm  {mean_power:>6.1f}dBm  {min_freq:>6.1f}ppm  {max_freq:>6.1f}ppm  {mean_freq:>6.1f}ppm")

if len(slot_timestamps) > 0:
    fig = plt.figure(figsize=(12, 6))
    ax = plt.gca()
    lines = []
    for m in slot_timestamps:
        if args.c and len(slot_timestamps[m]) < args.c:
            continue
        if args.x == 'power':
            xdata = power_dbm[m]
            xlabel = 'signal power (dBm)'
        elif args.x == 'freq':
            xdata = freq_ppm[m]
            xlabel = 'frequency offset (ppm)'
        elif args.x == 'distance':
            xdata = distances[m]
            xlabel = 'distance (km)'
        else:
            xdata = rxux_timestamps[m]
            xlabel = 'Unix timestamp (seconds)'

        if args.y == 'power':
            ydata = power_dbm[m]
            ylabel = 'signal power (dBm)'
        elif args.y == 'freq':
            ydata = freq_ppm[m]
            ylabel = 'frequency offset (ppm)'
        elif args.y == 'distance':
            ydata = distances[m]
            ylabel = 'distance (km)'
        else:
            ydata = slot_timestamps[m]
            ylabel = 'AIS slot timing offset (microseconds)'

        lines.append(plt.plot(xdata, ydata, '.', markersize=5, label=m, alpha=0.6))

    # Speed-of-light reference: 1/c = 3.336 µs/km, offset by receiver pipeline latency.
    T0 = 0.0  # µs: fixed receiver-side delay (intercept)
    if {args.x, args.y} == {'distance', 'time'}:
        xmax = max((max(distances[m]) for m in slot_timestamps if distances[m]), default=0)
        xmin = min((min(distances[m]) for m in slot_timestamps if distances[m]), default=100000)
        if xmax > 0:
            label = f'1/c (3.336 µs/km) + {T0:.0f} µs'
            plt.plot([xmin, xmax], [T0 + xmin * 3.336, T0 + xmax * 3.336], 'k:', alpha=0.5, label=label)

    ax.ticklabel_format(style='plain', axis='x', useOffset=False)
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)
    plt.title(f'AIS messages grouped by {group_label}')
    if len(slot_timestamps) <= 20:
        plt.legend()
    plt.grid(True, alpha=0.3)
    plt.tight_layout()

    def hover(event):
        vis = annot.get_visible()
        if event.inaxes == ax:
            for l in lines:
                cont, ind = l[0].contains(event)
                if cont is True:
                    x, y = l[0].get_data()
                    i = l[0].get_label()
                    if args.g == 'mmsi':
                        i = int(i)
                    elif args.g == 'type':
                        i = int(i)
                    annot.xy = (x[ind["ind"][0]], y[ind["ind"][0]])
                    if distances[i][ind['ind'][0]] > 0:
                        annot.set_text(f"{meta[i][ind['ind'][0]]}\n{power_dbm[i][ind['ind'][0]]} dBm\n{distances[i][ind['ind'][0]]:.1f} km")
                    else:
                        annot.set_text(f"{meta[i][ind['ind'][0]]}\n{power_dbm[i][ind['ind'][0]]} dBm")
                    annot.get_bbox_patch().set_alpha(0.4)
                    annot.set_visible(True)
                    fig.canvas.draw_idle()
                    return
            if vis:
                annot.set_visible(False)
                fig.canvas.draw_idle()

    annot = ax.annotate("", xy=(0, 0), xytext=(-20, 20), textcoords="offset points",
                        bbox=dict(boxstyle="round", fc="w"),
                        arrowprops=dict(arrowstyle="->"))
    annot.set_visible(False)
    fig.canvas.mpl_connect("motion_notify_event", hover)

    plt.show()