import os
import torch
from zipfile import ZipFile

from aris import create_manual_marking, BEAM_WIDTH_DIR, add_metadata_to_result, prep_for_mm
from dataloader import create_dataloader_aris
from inference import do_full_inference, json_dump_round_float
from visualizer import generate_video_batches

WEIGHTS = 'models/v5m_896_300best.pt'

def predict_task(filepath, weights=WEIGHTS, gradio_progress=None):
    """
    Main processing task to be run in gradio
        - Writes aris frames to dirname(filepath)/frames/{i}.jpg
        - Writes json output to dirname(filepath)/{filename}_results.json
        - Writes manual marking to dirname(filepath)/{filename}_marking.txt
        - Writes video output to dirname(filepath)/{filename}_results.mp4
        - Zips all results to dirname(filepath)/{filename}_results.zip
    Args:
        filepath (str): path to aris file
        
    TODO: Separate into subtasks in different queues; have a GPU-only queue.
    """
    if (gradio_progress): gradio_progress(0, "In task...")
    print("Cuda available in task?", torch.cuda.is_available())
    
    print(filepath)
    dirname = os.path.dirname(filepath)
    filename = os.path.basename(filepath).replace(".aris","").replace(".ddf","")
    results_filepath = os.path.join(dirname, f"{filename}_results.json")
    marking_filepath = os.path.join(dirname, f"{filename}_marking.txt")
    video_filepath = os.path.join(dirname, f"{filename}_results.mp4")
    zip_filepath = os.path.join(dirname, f"{filename}_results.zip")
    os.makedirs(dirname, exist_ok=True)

    # create dataloader
    if (gradio_progress): gradio_progress(0, "Initializing Dataloader...")
    dataloader, dataset = create_dataloader_aris(filepath, BEAM_WIDTH_DIR, None)

    # extract aris/didson info. didson does not yet have pixel-meter info
    if ".ddf" in filepath:
        image_meter_width = -1
        image_meter_height = -1
    else:
        image_meter_width = dataset.didson.info['xdim'] * dataset.didson.info['pixel_meter_width']
        image_meter_height = dataset.didson.info['ydim'] * dataset.didson.info['pixel_meter_height']
    frame_rate = dataset.didson.info['framerate']
    
    # run detection + tracking
    results = do_full_inference(dataloader, image_meter_width, image_meter_height, gp=gradio_progress, weights=weights)

    # re-index results if desired - this should be done before writing the file
    results = prep_for_mm(results)
    results = add_metadata_to_result(filepath, results)
    
    # write output to disk
    json_dump_round_float(results, results_filepath)
    
    if dataset.didson.info['version'][3] == 5: # ARIS only
        create_manual_marking(results, out_path=marking_filepath)
    
    # generate a video with tracking results
    generate_video_batches(dataset.didson, results, frame_rate, video_filepath, 
                   image_meter_width=image_meter_width, image_meter_height=image_meter_height, gp=gradio_progress)
    
    # zip up the results
    with ZipFile(zip_filepath, 'w') as z:
        for file in [results_filepath, marking_filepath, video_filepath, os.path.join(dirname, 'bg_start.jpg')]:
            if os.path.exists(file):
                z.write(file, arcname=os.path.basename(file))
    
    # release GPU memory 
    torch.cuda.empty_cache()
    
    return results, results_filepath, zip_filepath, video_filepath, marking_filepath