import ast
import os
import json
from matplotlib.patches import Polygon
from matplotlib.collections import PatchCollection
import matplotlib.pyplot as plt
import numpy as np
import cv2
import inflect

p = inflect.engine()

img_dir = "imgs"
bg_prompt_text = "Background prompt: "
# h, w
box_scale = (512, 512)
size = box_scale
size_h, size_w = size
print(f"Using box scale: {box_scale}")

def parse_input(text=None, no_input=False):
    if not text:
        if no_input:
            return
        
        text = input("Enter the response: ")
    if "Objects: " in text:
        text = text.split("Objects: ")[1]
        
    text_split = text.split(bg_prompt_text)
    if len(text_split) == 2:
        gen_boxes, bg_prompt = text_split
    elif len(text_split) == 1:
        if no_input:
            return
        gen_boxes = text
        bg_prompt = ""
        while not bg_prompt:
            # Ignore the empty lines in the response
            bg_prompt = input("Enter the background prompt: ").strip()
        if bg_prompt_text in bg_prompt:
            bg_prompt = bg_prompt.split(bg_prompt_text)[1]
    else:
        raise ValueError(f"text: {text}")
    try:
        gen_boxes = ast.literal_eval(gen_boxes)    
    except SyntaxError as e:
        # Sometimes the response is in plain text
        if "No objects" in gen_boxes:
            gen_boxes = []
        else:
            raise e
    bg_prompt = bg_prompt.strip()
    
    return gen_boxes, bg_prompt

def filter_boxes(gen_boxes, scale_boxes=True, ignore_background=True, max_scale=3):
    if len(gen_boxes) == 0:
        return []
    
    box_dict_format = False
    gen_boxes_new = []
    for gen_box in gen_boxes:
        if isinstance(gen_box, dict):
            name, [bbox_x, bbox_y, bbox_w, bbox_h] = gen_box['name'], gen_box['bounding_box']
            box_dict_format = True
        else:
            name, [bbox_x, bbox_y, bbox_w, bbox_h] = gen_box
        if bbox_w <= 0 or bbox_h <= 0:
            # Empty boxes
            continue
        if ignore_background:
            if (bbox_w >= size[1] and bbox_h >= size[0]) or bbox_x > size[1] or bbox_y > size[0]:
                # Ignore the background boxes
                continue
        gen_boxes_new.append(gen_box)
    
    gen_boxes = gen_boxes_new
    
    if len(gen_boxes) == 0:
        return []
    
    filtered_gen_boxes = []
    if box_dict_format:
        # For compatibility
        bbox_left_x_min = min([gen_box['bounding_box'][0] for gen_box in gen_boxes])
        bbox_right_x_max = max([gen_box['bounding_box'][0] + gen_box['bounding_box'][2] for gen_box in gen_boxes])
        bbox_top_y_min = min([gen_box['bounding_box'][1] for gen_box in gen_boxes])
        bbox_bottom_y_max = max([gen_box['bounding_box'][1] + gen_box['bounding_box'][3] for gen_box in gen_boxes])
    else:
        bbox_left_x_min = min([gen_box[1][0] for gen_box in gen_boxes])
        bbox_right_x_max = max([gen_box[1][0] + gen_box[1][2] for gen_box in gen_boxes])
        bbox_top_y_min = min([gen_box[1][1] for gen_box in gen_boxes])
        bbox_bottom_y_max = max([gen_box[1][1] + gen_box[1][3] for gen_box in gen_boxes])
    
    # All boxes are empty
    if (bbox_right_x_max - bbox_left_x_min) == 0:
        return []
    
    # Used if scale_boxes is True
    shift = -bbox_left_x_min
    scale = size_w / (bbox_right_x_max - bbox_left_x_min)
    
    scale = min(scale, max_scale)
    
    for gen_box in gen_boxes:
        if box_dict_format:
            name, [bbox_x, bbox_y, bbox_w, bbox_h] = gen_box['name'], gen_box['bounding_box']
        else:
            name, [bbox_x, bbox_y, bbox_w, bbox_h] = gen_box
            
        if scale_boxes:
            # Vertical: move the boxes if out of bound
            # Horizontal: move and scale the boxes so it spans the horizontal line
            
            bbox_x = (bbox_x + shift) * scale
            bbox_y = bbox_y * scale
            bbox_w, bbox_h = bbox_w * scale, bbox_h * scale
            # TODO: verify this makes the y center not moving
            bbox_y_offset = 0
            if bbox_top_y_min * scale + bbox_y_offset < 0:
                bbox_y_offset -= bbox_top_y_min * scale
            if bbox_bottom_y_max * scale + bbox_y_offset >= size_h:
                bbox_y_offset -= bbox_bottom_y_max * scale - size_h
            bbox_y += bbox_y_offset
            
            if bbox_y < 0:
                bbox_y, bbox_h = 0, bbox_h - bbox_y
                
        name = name.rstrip(".")
        bounding_box = (int(np.round(bbox_x)), int(np.round(bbox_y)), int(np.round(bbox_w)), int(np.round(bbox_h)))
        if box_dict_format:
            gen_box = {
                'name': name,
                'bounding_box': bounding_box
            }
        else:
            gen_box = (name, bounding_box)
        
        filtered_gen_boxes.append(gen_box)
        
    return filtered_gen_boxes

def draw_boxes(anns):
    ax = plt.gca()
    ax.set_autoscale_on(False)
    polygons = []
    color = []
    for ann in anns:
        c = (np.random.random((1, 3))*0.6+0.4)
        [bbox_x, bbox_y, bbox_w, bbox_h] = ann['bbox']
        poly = [[bbox_x, bbox_y], [bbox_x, bbox_y+bbox_h],
                [bbox_x+bbox_w, bbox_y+bbox_h], [bbox_x+bbox_w, bbox_y]]
        np_poly = np.array(poly).reshape((4, 2))
        polygons.append(Polygon(np_poly))
        color.append(c)

        # print(ann)
        name = ann['name'] if 'name' in ann else str(ann['category_id'])
        ax.text(bbox_x, bbox_y, name, style='italic',
                bbox={'facecolor': 'white', 'alpha': 0.7, 'pad': 5})

    p = PatchCollection(polygons, facecolor='none',
                        edgecolors=color, linewidths=2)
    ax.add_collection(p)


def show_boxes(gen_boxes, bg_prompt=None, ind=None, show=False):
    if len(gen_boxes) == 0:
        return
    
    if isinstance(gen_boxes[0], dict):
        anns = [{'name': gen_box['name'], 'bbox': gen_box['bounding_box']}
                for gen_box in gen_boxes]
    else:
        anns = [{'name': gen_box[0], 'bbox': gen_box[1]} for gen_box in gen_boxes]

    # White background (to allow line to show on the edge)
    I = np.ones((size[0]+4, size[1]+4, 3), dtype=np.uint8) * 255

    plt.imshow(I)
    plt.axis('off')

    if bg_prompt is not None:
        ax = plt.gca()
        ax.text(0, 0, bg_prompt, style='italic',
                bbox={'facecolor': 'white', 'alpha': 0.7, 'pad': 5})

        c = (np.zeros((1, 3)))
        [bbox_x, bbox_y, bbox_w, bbox_h] = (0, 0, size[1], size[0])
        poly = [[bbox_x, bbox_y], [bbox_x, bbox_y+bbox_h],
                [bbox_x+bbox_w, bbox_y+bbox_h], [bbox_x+bbox_w, bbox_y]]
        np_poly = np.array(poly).reshape((4, 2))
        polygons = [Polygon(np_poly)]
        color = [c]
        p = PatchCollection(polygons, facecolor='none',
                            edgecolors=color, linewidths=2)
        ax.add_collection(p)

    draw_boxes(anns)
    if show:
        plt.show()
    else:
        print("Saved to", f"{img_dir}/boxes.png", f"ind: {ind}")
        if ind is not None:
            plt.savefig(f"{img_dir}/boxes_{ind}.png")
        plt.savefig(f"{img_dir}/boxes.png")


def show_masks(masks):
    masks_to_show = np.zeros((*size, 3), dtype=np.float32)
    for mask in masks:
        c = (np.random.random((3,))*0.6+0.4)

        masks_to_show += mask[..., None] * c[None, None, :]
    plt.imshow(masks_to_show)
    plt.savefig(f"{img_dir}/masks.png")
    plt.show()
    plt.clf()

def convert_box(box, height, width):
    # box: x, y, w, h (in 512 format) -> x_min, y_min, x_max, y_max
    x_min, y_min = box[0] / width, box[1] / height
    w_box, h_box = box[2] / width, box[3] / height
    
    x_max, y_max = x_min + w_box, y_min + h_box
    
    return x_min, y_min, x_max, y_max

def convert_spec(spec, height, width, include_counts=True, verbose=False):
    # Infer from spec
    prompt, gen_boxes, bg_prompt = spec['prompt'], spec['gen_boxes'], spec['bg_prompt']
    
    # This ensures the same objects appear together because flattened `overall_phrases_bboxes` should EXACTLY correspond to `so_prompt_phrase_box_list`. 
    gen_boxes = sorted(gen_boxes, key=lambda gen_box: gen_box[0])
    
    gen_boxes = [(name, convert_box(box, height=height, width=width)) for name, box in gen_boxes]
    
    # NOTE: so phrase should include all the words associated to the object (otherwise "an orange dog" may be recognized as "an orange" by the model generating the background).
    # so word should have one token that includes the word to transfer cross attention (the object name).
    # Currently using the last word of the object name as word.
    if bg_prompt:
        so_prompt_phrase_word_box_list = [(f"{bg_prompt} with {name}", name, name.split(" ")[-1], box) for name, box in gen_boxes]
    else:
        so_prompt_phrase_word_box_list = [(f"{name}", name, name.split(" ")[-1], box) for name, box in gen_boxes]
    
    objects = [gen_box[0] for gen_box in gen_boxes]
    
    objects_unique, objects_count = np.unique(objects, return_counts=True)

    num_total_matched_boxes = 0
    overall_phrases_words_bboxes = []
    for ind, object_name in enumerate(objects_unique):
        bboxes = [box for name, box in gen_boxes if name == object_name]
        
        if objects_count[ind] > 1:
            phrase = p.plural_noun(object_name.replace("an ", "").replace("a ", ""))
            if include_counts:
                phrase = p.number_to_words(objects_count[ind]) + " " + phrase
        else:
            phrase = object_name
        # Currently using the last word of the phrase as word.
        word = phrase.split(' ')[-1]
        
        num_total_matched_boxes += len(bboxes)
        overall_phrases_words_bboxes.append((phrase, word, bboxes))
        
    assert num_total_matched_boxes == len(gen_boxes), f"{num_total_matched_boxes} != {len(gen_boxes)}"

    objects_str = ", ".join([phrase for phrase, _, _ in overall_phrases_words_bboxes])
    if objects_str:
        if bg_prompt:
            overall_prompt = f"{bg_prompt} with {objects_str}"
        else:
            overall_prompt = objects_str
    else:
        overall_prompt = bg_prompt
        
    if verbose:
        print("so_prompt_phrase_word_box_list:", so_prompt_phrase_word_box_list)
        print("overall_prompt:", overall_prompt)
        print("overall_phrases_words_bboxes:", overall_phrases_words_bboxes)
    
    return so_prompt_phrase_word_box_list, overall_prompt, overall_phrases_words_bboxes