join_datasets

import jsonlines
import re
from odin_to_gqfcn import Robotv2ToGQFCNv2
from dateutil import parser

class DebugParser:
    def __init__(self):
        self.start_marker = 'Acquired FSM lock in state S4(PNP - Picking)'
        self.end_marker = 'Released FSM lock [<function FSM.place_result'

    def extract_mission_data(self, lines):
        mission_id = None
        length = None
        width = None
        height = None
        hostname = None

        for line in lines:
            if mission_id is None and 'mission_id' in line:
                mission_id = line['mission_id']
            
            if hostname is None and 'hostname' in line:
                hostname = line['hostname']

            if 'item' in line:
                if 'length' in line['item']:
                    if length is None:
                        length = line['item']['length']
                
                if 'width' in line['item']:
                    if width is None:
                        width = line['item']['width']
                
                if 'height' in line['item']:
                    if height is None:
                        height = line['item']['height']

        return {
            'mission_id': mission_id,
            'min_time': parser.parse(lines[0]['timestamp']),
            'max_time': parser.parse(lines[-1]['timestamp']),
            'length': length,
            'width': width,
            'height': height,
            'hostname': hostname,
            #'lines': lines
        }

    def get_missions(self, lines):
        missions = []

        current_pick = []
        for line in lines:
            event = line['event']
            timestamp = line['timestamp']

            if self.start_marker in event:
                current_pick = []

            if current_pick is not None:
                current_pick.append(line)

            if self.end_marker in event:
                pick_data = self.extract_mission_data(current_pick)
                if pick_data['mission_id']:
                    missions.append(pick_data)
                else:
                    print("Could not extract mission_id")

                current_pick = None

        return missions

    def parse(self, filename):
        with jsonlines.open(filename) as f:
            missions = self.get_missions(f.iter())

        return missions

def remove_blacklisted_logs(lines):
    blacklist = ['Sending bits', 'IncomingTelegramHandlers', 'Received bits', 'Handling inbound telegram', 'Enqueueing execution', 'Added acknowledge', 'exceeded timeout']

    result = []
    for line in lines:
        event = line['event']
        black = False
        for item in blacklist:
            if item in event:
                black = True

        if not black:
            result.append(line)
    return result


def match_missions_and_outputs(missions, outputs):
    mission_count = len(missions)
    output_count = len(outputs)
    matches = []

    for output in outputs:
        candidate_time = None
        candidate = None
        ts = re.sub(r'\+.*$', 'Z', output['timestamp'])
        timestamp = parser.parse(ts)

        for mission in missions:
            time_delta = (mission['max_time'] - timestamp).total_seconds()
            if candidate is None:
                candidate_time = time_delta
                candidate = mission
            elif abs(time_delta) < abs(candidate_time):
                candidate_time = time_delta
                candidate = mission

        matches.append({
            'mission': candidate, 
            'output': output,
        })

    return matches
            

cache_path = '/Users/brian/.geri/cache'

#2 element
dataset = '6d8eb230-63e6-4a1e-a05d-fea279fbdf19'

#early large dataset
dataset = 'c606d679-b64f-49a2-90a8-712f366262f9'

debug_filename = f"{cache_path}/{dataset}/logs/debug.jsonl"
debug_parser = DebugParser()
missions = debug_parser.parse(debug_filename)

output_filename = f"{cache_path}/{dataset}/grasp_pose_generator/output.json"
with jsonlines.open(output_filename) as f:
    outputs = list(f.iter())

matches = match_missions_and_outputs(missions, outputs)

import pprint
pprint.pprint(matches)