import torch
from PIL import Image
from conversation import conv_templates
from builder import load_pretrained_model  # Assuming this is your custom model loader
from functools import partial
import numpy as np
DEFAULT_REGION_FEA_TOKEN = "<region_fea>"
DEFAULT_IMAGE_TOKEN = "<image>"
DEFAULT_IM_START_TOKEN = "<im_start>"
DEFAULT_IM_END_TOKEN = "<im_end>"

# define the task categories
box_in_tasks = ['widgetcaptions', 'taperception', 'ocr', 'icon_recognition', 'widget_classification', 'example_0']
box_out_tasks = ['widget_listing', 'find_text', 'find_icons', 'find_widget', 'conversation_interaction']
no_box_tasks = ['screen2words', 'detailed_description', 'conversation_perception', 'gpt4']

# function to generate the mask
def generate_mask_for_feature(coor, raw_w, raw_h, mask=None):
    """
    Generates a region mask based on provided coordinates.
    Handles both point and box input.
    """
    if mask is not None:
        assert mask.shape[0] == raw_w and mask.shape[1] == raw_h
    coor_mask = np.zeros((raw_w, raw_h))

    # if it's a point (2 coordinates)
    if len(coor) == 2:
        span = 5  # Define the span for the point
        x_min = max(0, coor[0] - span)
        x_max = min(raw_w, coor[0] + span + 1)
        y_min = max(0, coor[1] - span)
        y_max = min(raw_h, coor[1] + span + 1)
        coor_mask[int(x_min):int(x_max), int(y_min):int(y_max)] = 1
        assert (coor_mask == 1).any(), f"coor: {coor}, raw_w: {raw_w}, raw_h: {raw_h}"

    # if it's a box (4 coordinates)
    elif len(coor) == 4:
        coor_mask[coor[0]:coor[2]+1, coor[1]:coor[3]+1] = 1
        if mask is not None:
            coor_mask = coor_mask * mask

    # convert to torch tensor and ensure it contains non-zero values
    coor_mask = torch.from_numpy(coor_mask)
    assert len(coor_mask.nonzero()) != 0, "Generated mask is empty :("


    return coor_mask


def infer_single_prompt(image_path, prompt, model_path, region=None, model_name="ferret_llama", conv_mode="ferret_llama_3", add_region_feature=False):
    img = Image.open(image_path).convert('RGB')

    # this loads the model, image processor and tokenizer
    tokenizer, model, image_processor, context_len = load_pretrained_model(model_path, None, model_name)
    # define the image size required by clip
    image_size = {"height": 336, "width": 336}

    # process the image
    image_tensor = image_processor.preprocess(
        img,
        return_tensors='pt',
        do_resize=True,
        do_center_crop=False,
        size=(image_size['height'], image_size['width'])
    )['pixel_values'][0].unsqueeze(0)

    image_tensor = image_tensor.half().cuda()

    # generate the prompt per template requirement
    conv = conv_templates[conv_mode].copy()
    conv.append_message(conv.roles[0], prompt)
    conv.append_message(conv.roles[1], None)
    prompt_input = conv.get_prompt()
    
    # add the special tokens
    prompt_input = DEFAULT_IM_START_TOKEN + DEFAULT_IMAGE_TOKEN + DEFAULT_IM_END_TOKEN + '\n' + prompt_input

    

    # region mask logic (if region is provided)
    region_masks = None
    if add_region_feature and region is not None:
        raw_w, raw_h = img.size
        region_masks = generate_mask_for_feature(region, raw_w, raw_h).unsqueeze(0).cuda().half()
        region_masks = [[region_mask_i.cuda().half() for region_mask_i in region_masks]]
        prompt_input = prompt_input.replace("<bbox_location0>", f"[{region[0]}, {region[1]}, {region[2]}, {region[3]}] {DEFAULT_REGION_FEA_TOKEN}")
        
    # tokenize prompt
    # input_ids = tokenizer(prompt_input, return_tensors='pt')['input_ids'].cuda()

    inputs = tokenizer(prompt_input, return_tensors='pt', padding=True)
    input_ids = inputs['input_ids'].cuda()
    attention_mask = inputs['attention_mask'].cuda() 
        
    # generate model output
    with torch.inference_mode():
        # Use region_masks in model's forward call
        model.orig_forward = model.forward
        model.forward = partial(
            model.orig_forward,
            region_masks=region_masks
        )
        # explcit add of attention mask
        output_ids = model.generate(
            input_ids,
            images=image_tensor,
            attention_mask=attention_mask, 
            max_new_tokens=1024,
            num_beams=1,
            region_masks=region_masks,  # pass the region mask to the model
            image_sizes=[img.size]
        )
        model.forward = model.orig_forward

    # we decode the output
    output_text = tokenizer.batch_decode(output_ids, skip_special_tokens=True)[0]
    return output_text.strip()

# We also define a task-specific inference function
def infer_ui_task(image_path, prompt, model_path, task, region=None, add_region_feature=False):
    # region = torch.tensor(region).cuda()
    """
    Handles task types: box_in_tasks, box_out_tasks, no_box_tasks.
    """
    if region is not None:
        add_region_feature=True
    if task in box_in_tasks and region is None:
        raise ValueError(f"Task {task} requires a bounding box region.")
    
    if task in box_in_tasks:
        print(f"Processing {task} with bounding box region.")
        return infer_single_prompt(image_path, prompt, model_path, region, add_region_feature=add_region_feature)
    
    elif task in box_out_tasks:
        print(f"Processing {task} without bounding box region.")
        return infer_single_prompt(image_path, prompt, model_path)
    
    elif task in no_box_tasks:
        print(f"Processing {task} without image or bounding box.")
        return infer_single_prompt(image_path, prompt, model_path)
    
    else:
        raise ValueError(f"Unknown task type: {task}")