import os
from data.prefix_instruction import get_image_prompt, get_task_instruction, get_layout_instruction, get_content_instruction
import random
from PIL import Image
from .gradio_tasks import dense_prediction_data


style_transfer = [
    dict(
        name='Style Transfer', 
        image_type=["target", "style_source", "style_target"]),
]
style_transfer_text = [[x['name']] for x in style_transfer]


style_condition_fusion = [
    dict(
        name='Canny+Style to Image', 
        image_type=["canny", "style_source", "style_target"]),
    dict(
        name='Depth+Style to Image', 
        image_type=["depth", "style_source", "style_target"]),
    dict(
        name='Hed+Style to Image', 
        image_type=["hed", "style_source", "style_target"]),
    dict(
        name='Normal+Style to Image', 
        image_type=["normal", "style_source", "style_target"]),
    dict(
        name='Pose+Style to Image', 
        image_type=["openpose", "style_source", "style_target"]),
    dict(
        name='SAM2+Style to Image', 
        image_type=["sam2_mask", "style_source", "style_target"]),
    dict(
        name='Mask+Style to Image', 
        image_type=["mask", "style_source", "style_target"]),
]
style_condition_fusion_text = [[x['name']] for x in style_condition_fusion]


def process_style_transfer_tasks(x):
    for task in style_transfer:
        if task['name'] == x[0]:
            image_type = task['image_type']
            image_prompt_list = [get_image_prompt(x)[0] for x in image_type]
            image_prompt_list = [f"[IMAGE{idx+1}] {image_prompt}" for idx, image_prompt in enumerate(image_prompt_list)]
            condition_prompt = ", ".join(image_prompt_list[:-1])
            target_prompt = image_prompt_list[-1]
            task_prompt = get_task_instruction(condition_prompt, target_prompt)

            # sample examples
            valid_data = [x for x in dense_prediction_data if all([(x.get(t, None) is not None and os.path.exists(x[t])) for t in image_type])]
            n_samples = random.randint(2, min(len(valid_data), 3))
            images = random.sample(valid_data, k=n_samples)
            rets = []
            for image in images:
                for t in image_type:
                    if t == "style_source":
                        target = Image.open(image["style_target"])
                        source = Image.open(image[t])
                        source = source.resize(target.size)
                        rets.append(source)
                    else:
                        rets.append(Image.open(image[t]))
            
            content_prompt = ""

            grid_h = n_samples
            grid_w = len(image_type)
            mask = task.get('mask', [0 for _ in range(grid_w - 1)] + [1])
            layout_prompt = get_layout_instruction(grid_w, grid_h)

            upsampling_noise = None
            steps = None
            outputs = [mask, grid_h, grid_w, layout_prompt, task_prompt, content_prompt, upsampling_noise, steps] + rets
            break

    return outputs


def process_style_condition_fusion_tasks(x):
    for task in style_condition_fusion:
        if task['name'] == x[0]:
            image_type = task['image_type']
            image_prompt_list = [get_image_prompt(x)[0] for x in image_type]
            image_prompt_list = [f"[IMAGE{idx+1}] {image_prompt}" for idx, image_prompt in enumerate(image_prompt_list)]
            condition_prompt = ", ".join(image_prompt_list[:-1])
            target_prompt = image_prompt_list[-1]
            task_prompt = get_task_instruction(condition_prompt, target_prompt)

            # sample examples
            valid_data = [x for x in dense_prediction_data if all([(x.get(t, None) is not None and os.path.exists(x[t])) for t in image_type])]
            x = dense_prediction_data[0]
            n_samples = random.randint(2, min(len(valid_data), 3))
            images = random.sample(valid_data, k=n_samples)
            rets = []
            for image in images:
                for t in image_type:
                    if t == "style_source":
                        target = Image.open(image["style_target"])
                        source = Image.open(image[t])
                        source = source.resize(target.size)
                        rets.append(source)
                    else:
                        rets.append(Image.open(image[t]))   

            content_prompt = ""

            grid_h = n_samples
            grid_w = len(image_type)
            mask = task.get('mask', [0 for _ in range(grid_w - 1)] + [1])
            layout_prompt = get_layout_instruction(grid_w, grid_h)  

            upsampling_noise = None
            steps = None
            outputs = [mask, grid_h, grid_w, layout_prompt, task_prompt, content_prompt, upsampling_noise, steps] + rets
            break

    return outputs