import random
import numpy as np
from PIL import Image, ImageDraw, ImageFont

colormap = ['blue','orange','green','purple','brown','pink','gray','olive','cyan','red',
            'lime','indigo','violet','aqua','magenta','coral','gold','tan','skyblue']

def plot_bbox(image, data):
    """
    Draws bounding boxes on an image.
    
    Parameters:
    - image: PIL Image object.
    - data: Dictionary containing 'bboxes' and 'labels' keys.
    
    Returns:
    - Image with bounding boxes drawn.
    """
    draw = ImageDraw.Draw(image)
    
    try:
        font = ImageFont.truetype("arial.ttf", 20)
    except IOError:
        font = ImageFont.load_default().font_variant(size=20)
    
    labels = data.get('labels', data.get('bboxes_labels', []))
    for bbox, label in zip(data['bboxes'], labels):
        x1, y1, x2, y2 = bbox
        draw.rectangle([x1, y1, x2, y2], outline="red", width=3)
        
        # Annotate the label
        left, top, right, bottom = font.getbbox(label)
        text_width = right - left
        text_height = bottom - top
        
        padding = 5
        draw.rectangle([x1, y1 - text_height - padding*2, x1 + text_width + padding*2, y1], fill="red")
        draw.text((x1 + padding, y1 - text_height - padding), label, fill="white", font=font)
  
    return image

def draw_polygons(image, prediction, fill_mask=False):  
    """  
    Draws segmentation masks with polygons on an image.  
  
    Parameters:  
    - image: PIL Image object.
    - prediction: Dictionary containing 'polygons' and 'labels' keys.
    - fill_mask: Boolean indicating whether to fill the polygons with color.  

    Returns:
    - Image with polygons drawn.
    """  
    draw = ImageDraw.Draw(image)  
    scale = 1  
      
    for polygons, label in zip(prediction['polygons'], prediction['labels']):  
        color = random.choice(colormap)  
        fill_color = random.choice(colormap) if fill_mask else None  
          
        for _polygon in polygons:  
            _polygon = np.array(_polygon).reshape(-1, 2)  
            if len(_polygon) < 3:  
                print('Invalid polygon:', _polygon)  
                continue  
              
            _polygon = (_polygon * scale).reshape(-1).tolist()  
              
            if fill_mask:  
                draw.polygon(_polygon, outline=color, fill=fill_color)  
            else:  
                draw.polygon(_polygon, outline=color)  
              
            draw.text((_polygon[0] + 8, _polygon[1] + 2), label, fill=color)  
  
    return image

def draw_ocr_bboxes(image, prediction, scale=1):
    """
    Draws OCR bounding boxes on an image.

    Parameters:
    - image: PIL Image object.
    - prediction: Dictionary containing 'quad_boxes' and 'labels' keys.
    - scale: Scale factor for bounding box coordinates.

    Returns:
    - Image with OCR boxes drawn.
    """
    draw = ImageDraw.Draw(image)
        
    try:
        font = ImageFont.truetype("arial.ttf", 18)
    except IOError:
        font = ImageFont.load_default().font_variant(size=18)
    
    if 'quad_boxes' not in prediction or 'labels' not in prediction:
        return image
        
    bboxes, labels = prediction['quad_boxes'], prediction['labels']
    for box, label in zip(bboxes, labels):
        color = random.choice(colormap)
        new_box = (np.array(box) * scale).tolist()
        draw.polygon(new_box, width=3, outline=color)
        draw.text((new_box[0]+8, new_box[1]+2),
                    "{}".format(label),
                    align="right",
                    fill=color,
                    font=font)
       
    return image

def convert_to_od_format(data):  
    """  
    Converts a dictionary with 'bboxes' and 'bboxes_labels' into a standard object detection format.
  
    Parameters:  
    - data: The input dictionary.
  
    Returns:  
    - A dictionary with 'bboxes' and 'labels' keys.
    """  
    return {  
        'bboxes': data.get('bboxes', []),  
        'labels': data.get('bboxes_labels', [])  
    }  

def visualize_results(image, parsed_answer, output_path="result.jpg"):
    """
    Main function to visualize results based on the task.

    Parameters:
    - image: PIL Image object.
    - parsed_answer: The output from the model's post_process_generation.
    - output_path: Path to save the visualized image.

    Returns:
    - The visualized PIL Image object, or None if no visualization is available.
    """
    vis_image = image.copy()
    
    if not parsed_answer or not isinstance(parsed_answer, dict):
        print("Invalid parsed_answer format.")
        return None

    task = list(parsed_answer.keys())[0]
    data = parsed_answer[task]

    if task in ['<OD>', '<DENSE_REGION_CAPTION>', '<REGION_PROPOSAL>', '<CAPTION_TO_PHRASE_GROUNDING>']:
        vis_image = plot_bbox(vis_image, data)
    elif task == '<OPEN_VOCABULARY_DETECTION>':
        bbox_data = convert_to_od_format(data)
        vis_image = plot_bbox(vis_image, bbox_data)
    elif task in ['<REFERRING_EXPRESSION_SEGMENTATION>', '<REGION_TO_SEGMENTATION>']:
        vis_image = draw_polygons(vis_image, data, fill_mask=True)
    elif task == '<OCR_WITH_REGION>':
        vis_image = draw_ocr_bboxes(vis_image, data)
    else:
        print(f"No visualization implemented for task: {task}")
        return None

    if output_path:
        vis_image.save(output_path)
        print(f"Visualization saved to {output_path}")

    return vis_image