import json
import time

from dataset.tracks_motion import TracksMotion
from GPS import GPS
import gradio as gr

def _synthesis(synthesis_setting, motion_data):
    model = GPS(
        init_mode = f"random_synthesis/{synthesis_setting['frames']}",
        noise_sigma = synthesis_setting['noise_sigma'],
        coarse_ratio = 0.2,
        pyr_factor = synthesis_setting['pyr_factor'],
        num_stages_limit = -1,
        silent=True,
        device='cpu'
    )

    synthesized_motion = model.run(
        motion_data, 
        mode="match_and_blend",
        ext={
            'criteria': {
                'type': 'PatchCoherentLoss',
                'patch_size': synthesis_setting['patch_size'],
                'stride': synthesis_setting['stride'] if 'stride' in synthesis_setting.keys() else 1,
                'loop': synthesis_setting['loop'],
                'coherent_alpha': synthesis_setting['alpha'] if synthesis_setting['completeness'] else None,
                },
            'optimizer': "match_and_blend",
            'num_itrs': synthesis_setting['num_steps'],
        }
    )

    return synthesized_motion

def synthesis(data):
    data = json.loads(data)
    # create track object
    data['setting']['coarse_ratio'] = -1
    motion_data = TracksMotion(data['tracks'], scale=data['scale'])
    start = time.time()
    synthesized_motion = _synthesis(
        data['setting'], 
        [motion_data]
        )
    end = time.time()
    data['time'] = end - start
    data['tracks'] = motion_data.parse(synthesized_motion)

    return data

demo = gr.Interface(fn=synthesis, inputs="json", outputs="json")
demo.launch()