from dataloader import create_dataloader_aris
from aris import BEAM_WIDTH_DIR
import json
import cv2
import base64

VIDEO_HEIGHT = 700

def load_frames(video_path, preds_path):
    """Load frames for annotation editing
    """

    dataloader, dataset = create_dataloader_aris(video_path, BEAM_WIDTH_DIR, None)

    didson = dataset.didson
    
    frames = didson.load_frames(start_frame=0)
    frame_info, h, w = get_frame_info(frames, preds_path)
    
    return frame_info

def get_frame_info(frames, preds_path):
    """Get visualized video frames ready for output, given raw ARIS/DIDSON frames.
    Warning: all frames in frames will be stored in memory - careful of OOM errors. Consider processing large files
    in batches, such as in generate_video_batches()
    
    Returns:
        list(np.ndarray), height (int), width (int)
    """
    preds = json.load(open(preds_path, 'r'))
    color_map = { fish['id'] : fish['color'] for fish in preds['fish'] }
    
    frame_info = []
    if len(frames):
        # assumes all frames the same size
        h, w = frames[0].shape
        
        # enforce a standard size so that text/box thickness is consistent
        scale_factor = VIDEO_HEIGHT / h
        h = VIDEO_HEIGHT
        w = int(scale_factor*w)

        num_frames = min(len(frames), len(preds['frames']))
        
        for i, frame_raw in enumerate(frames[:num_frames]):
            image = cv2.resize(cv2.cvtColor(frame_raw, cv2.COLOR_GRAY2BGR), (w,h))
            retval, buffer = cv2.imencode('.jpg', image)
            jpg_as_text = base64.b64encode(buffer).decode("utf-8")

            frame = {
                'annotations': [],
                'frame': jpg_as_text
            }
            for fish in preds['frames'][i]['fish']:
                xmin, ymin, xmax, ymax = fish['bbox']
                hexx = color_map[fish['fish_id']].lstrip('#')
                frame['annotations'].append({
                    'bbox': {
                        'left': int(round(xmin * w)),
                        'right': int(round(xmax * w)),
                        'top': int(round(ymin * h)),
                        'bottom': int(round(ymax * h)),
                    },
                    'id': str(fish['fish_id']),
                    'conf': fish['conf']
                })
            frame_info.append(frame)
            
    return frame_info, h, w