from test import *

import copy
import gradio as gr

class merge_config_then_run():
    def __init__(self) -> None:
        # Load the tokenizer
        self.pretrained_model_path = '/home/xianyang/Data/code/FateZero/ckpt/stable-diffusion-v1-5'

        # load controlnet
        
    def run(
        self,
        user_input_video,
        num_layouts,
        layout_file1,
        layout_file2,
        layout_file3,
        layout_file4,
        layout_file5,
        prompt,
        model_id,
        n_sample_frame,
        start_sample_frame,
        sampling_rate,
        control_type,
        dwpose_options,
        controlnet_conditioning_scale,
        use_pnp,
        pnp_inject_steps,
        flatten_res,
    ):
        # , ] = inputs
        default_edit_config='config/demo_config.yaml'
        Omegadict_default_edit_config = OmegaConf.load(default_edit_config)
        
        dataset_time_string = get_time_string()
        config_now = copy.deepcopy(Omegadict_default_edit_config)

        config_now['pretrained_model_path'] = self.pretrained_model_path
        print(f"config_now['pretrained_model_path'] = model_id {self.pretrained_model_path}")



        #==========update datset_config===============#

        # 将所有 layout 文件放入列表中
        all_layout_files = [layout_file1, layout_file2, layout_file3, layout_file4, layout_file5]
        # 根据 num_layouts 转换为整数，并只使用前 N 个
        n_layouts = int(num_layouts)
        layout_files = all_layout_files[:n_layouts]


        config_now['dataset_config']['prompt'] = ''
        config_now['dataset_config']['path'] = user_input_video
        config_now['dataset_config']['n_sample_frame'] = n_sample_frame
        config_now['dataset_config']['start_sample_frame'] = start_sample_frame
        config_now['dataset_config']['sampling_rate'] = sampling_rate
        config_now['dataset_config']['layout_files'] = layout_files


        if user_input_video is None:
            raise gr.Error('You need to upload a video or choose a provided video')
        if user_input_video is not None:
            if isinstance(user_input_video, str):
                config_now['dataset_config']['path'] = user_input_video
            elif hasattr(user_input_video, 'name') and user_input_video.name is not None:
                config_now['dataset_config']['path'] = user_input_video.name

        # 检查每个 layout file 是否存在
        layout_files_checked = []
        for idx, lf in enumerate(layout_files):
            if lf is None:
                raise gr.Error(f'Layout file {idx+1} is missing')
            if isinstance(lf, str):
                lf_path = lf
            elif hasattr(lf, 'name') and lf.name is not None:
                lf_path = lf.name
            else:
                raise gr.Error(f'Layout file {idx+1} is invalid')
            if not os.path.exists(lf_path):
                raise gr.Error(f'Layout file "{lf_path}" does not exist')
            layout_files_checked.append(lf_path)
        config_now['dataset_config']['layout_files'] = layout_files_checked
        #==========update datset_config===============#


        #==========update control_config===============#
        config_now['control_config']['control_type'] = control_type
        config_now['control_config']['controlnet_conditioning_scale'] = float(controlnet_conditioning_scale)
        config_now['control_config']['hand'] = 'hand' in dwpose_options
        config_now['control_config']['face'] = 'face' in dwpose_options


        if control_type == "depth_midas":
            pretrained_controlnet_path = "/home/xianyang/Data/code/controlvideo/sd-controlnet-depth"
        elif control_type == "depth_zoe":
            pretrained_controlnet_path = "/home/xianyang/Data/code/FateZero/ckpt/control_v11f1p_sd15_depth"
        elif control_type == "dwpose":
            pretrained_controlnet_path = "/home/xianyang/Data/code/FateZero/ckpt/control_v11p_sd15_openpose"
        #==========update control_config===============#      


        #==========update editing_config===============#
        config_now['editing_config']['use_pnp'] = [use_pnp]
        config_now['editing_config']['inject_step'] = int(pnp_inject_steps)
        config_now['editing_config']['flatten_res'] = [int(x) for x in flatten_res]    
        config_now['editing_config']['editing_prompts'] = [[x.strip() for x in prompt.split(',')]]
        print('editing prompt', prompt)
        #==========update editing_config===============#

        logdir = default_edit_config.replace('config', 'result').replace('.yml', '').replace('.yaml', '')+f'_{dataset_time_string}'      
        config_now['logdir'] = logdir
        print(f'Saving at {logdir}')
        save_path = test(config = config_now,
                        **config_now)
        mp4_path = save_path.replace('_0.gif', '_0_0_0.mp4')
        return mp4_path