import argparse
import json
import sys
from datetime import datetime
import numpy as np
import matplotlib.pyplot as plt

from geometry.Geometry import Geometry

def parse_posix_time(value):
    try:
        return int(value)
    except ValueError:
        raise argparse.ArgumentTypeError("Invalid POSIX time format")

def parse_command_line_arguments():
    parser = argparse.ArgumentParser(description="Process command line arguments.")
    
    parser.add_argument("json_file", type=str, help="Input JSON file path")
    parser.add_argument("target_name", type=str, help="Target name")
    parser.add_argument("--start_time", type=parse_posix_time, help="Optional start time in POSIX seconds")
    parser.add_argument("--stop_time", type=parse_posix_time, help="Optional stop time in POSIX seconds")

    return parser.parse_args()

def main():
  
    # input handling
    args = parse_command_line_arguments()
    json_data = []
    with open(args.json_file, 'r') as json_file:
        for line in json_file:
            try:
                json_object = json.loads(line)
                json_data.append(json_object)
            except json.JSONDecodeError:
                print(f"Error decoding JSON from line: {line}")
    json_data = [item for item in json_data if item]
    start_time = args.start_time if args.start_time else None
    stop_time = args.stop_time if args.stop_time else None
    print("JSON String (Last Non-Empty Data):", json_data[-1])
    print("Target Name:", args.target_name)
    print("Start Time:", start_time)
    print("Stop Time:", stop_time)

    # extract data of interest
    server = json_data[0][0]["server"]
    timestamp = []
    associated = {}
    for item in json_data:

        first_result = item[0]

        if first_result["server"] != server:
            print('error')
            sys.exit(-1)

        if start_time and first_result["timestamp_event"]/1000 < start_time:
            continue

        if stop_time and first_result["timestamp_event"]/1000 > stop_time:
            continue

        # store association data
        if "detections_associated" in first_result:

            if args.target_name in first_result["detections_associated"]:

                for radar in first_result["detections_associated"][args.target_name]:

                    if radar['radar'] not in associated:
                        associated[radar['radar']] = []
                    else:
                        associated[radar['radar']].append(first_result["timestamp_event"])

        timestamp.append(first_result["timestamp_event"])

    # data massaging
    timestamp = list(dict.fromkeys(timestamp))
    associated = dict(sorted(associated.items(), key=lambda x: x[0]))
    radars = list(associated.keys())
    radar_label = []
    for radar in radars:
        radar_label.append(radar.split('.', 1)[0])

    # get start and stop times from data
    start_time = min(min(arr) for arr in associated.values())
    stop_time = max(max(arr) for arr in associated.values())
    timestamp = [value for value in timestamp if value >= start_time]
    timestamp = [value for value in timestamp if value <= stop_time]

    print(associated)
    
    data = []
    for radar in radars:
        result = [1 if value in associated[radar] else 0 for value in timestamp]
        data.append(result)
    
    print(data)

    # plot x, y, z
    plt.figure(figsize=(8,4))
    img = plt.imshow(data, aspect='auto', interpolation='none')
    y_extent = plt.gca().get_ylim()
    img.set_extent([start_time/1000, stop_time/1000, y_extent[1], y_extent[0]])
    plt.yticks(np.arange(len(radar_label)), radar_label, rotation='vertical')
    plt.xlabel('Timestamp')
    plt.ylabel('Radar')
    plt.tight_layout()
    filename = 'plot_associate_' + args.target_name + '.png'
    plt.savefig('save/' + filename, bbox_inches='tight', pad_inches=0.01)

if __name__ == "__main__":
    main()