import copy
import os
from transformers import Qwen2VLForConditionalGeneration, AutoTokenizer, AutoProcessor
from qwen_vl_utils import process_vision_info
from tqdm import tqdm
import h5py
import torch
import numpy as np
import cv2
from collections import Counter
import json
RED = '\033[31m'
GREEN = '\033[32m'
YELLOW = '\033[33m'
BLUE = '\033[34m'
RESET = '\033[0m'  # Reset to default color
def load_hdf5(dataset_dir, dataset_name):
    dataset_path = os.path.join(dataset_dir, dataset_name)
    if not os.path.isfile(dataset_path):
        print(f'Dataset does not exist at \n{dataset_path}\n')
        exit()

    with h5py.File(dataset_path, 'r') as root:
        is_sim = root.attrs['sim']
        # qpos = root['/observations/qpos'][()]
        # qvel = root['/observations/qvel'][()]
        # effort = root['/observations/effort'][()]
        # action = root['/action'][()]
        subtask = root['/subtask'][()]

        image_dict = dict()
        for cam_name in root[f'/observations/images/'].keys():
            image_dict[cam_name] = root[f'/observations/images/{cam_name}'][()]

    return image_dict, subtask
def load_model(model_path='/media/rl/HDD/data/weights/Qwen2-VL-7B-Instruct'):
    #"/gpfs/private/tzb/wjj/model_param/Qwen2-VL-7B-Instruct/"

    model = Qwen2VLForConditionalGeneration.from_pretrained(
        model_path, torch_dtype="auto", device_map="auto"
    )

    # We recommend enabling flash_attention_2 for better acceleration and memory saving, especially in multi-image and video scenarios.
    # model = Qwen2VLForConditionalGeneration.from_pretrained(
    #     model_path,
    #     torch_dtype=torch.bfloat16,
    #     attn_implementation="flash_attention_2",
    #     device_map="auto",
    # )

    # default processer
    processor = AutoProcessor.from_pretrained(model_path)

    # The default range for the number of visual tokens per image in the model is 4-16384.
    # You can set min_pixels and max_pixels according to your needs, such as a token range of 256-1280, to balance performance and cost.
    # min_pixels = 256*28*28
    # max_pixels = 1280*28*28
    # processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-7B-Instruct", min_pixels=min_pixels, max_pixels=max_pixels)
    return model, processor

chat_template = [
    {
        "role": "user",
        "content": [
        ],
    }
]
prompt = """There are four images. Please detect the objects on the table and return the objects in a list. The object names can only be one of the predefined list: [<objects>]. The first image contains all objects in predefined list and the first list equals to predefined list.
Notice that the first image contains 4 objects, the second image contains 3 objects, the third image contains 2 objects and the last image only contains 1 object. So the length of answer lists must be 4,3,2,1.
Your answer must be four lists corresponding to the chosen objects for each image. 
Answer example:['a','b','c','d']; ['b','c','a']; ['b','c']; ['c']
"""
# prompt = ("There are four images and the objects in images are following [<objects>]. The objects on the image is grandually picked away one by one. Please find out the order in which the objects are taken away."
#           "Your answer must be a list such as [a,b,c,d].")
def model_inference(model, processor, messages):


    # Preparation for inference
    text = processor.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    image_inputs, video_inputs = process_vision_info(messages)
    inputs = processor(
        text=[text],
        images=image_inputs,
        videos=video_inputs,
        padding=True,
        return_tensors="pt",
    )
    inputs = inputs.to("cuda")

    # Inference: Generation of the output
    generated_ids = model.generate(**inputs, max_new_tokens=128)
    generated_ids_trimmed = [
        out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
    ]
    output_text = processor.batch_decode(
        generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
    )
    print(output_text)
    results = output_text[0].split(';')
    results = [eval(each.strip()) for each in results]
    return results

def filter_images_by_subtask(image_dict, subtask, OUTPUT_DIR, episode):
    idxs = np.where(subtask != 0)[0]

    temp_idxs =[0] + idxs[:-1].tolist()
    key_frames = []

    for i, idx in enumerate(temp_idxs):
        img = image_dict['cam_high'][idx][180:480, 200:480]
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        save_name = os.path.join(OUTPUT_DIR, f'{episode}_{i}.png')
        cv2.imwrite(save_name, img)
        key_frames.append(save_name)
    return key_frames, idxs

def find_missing_names_counter(a,b):
    count_a = Counter(a)
    count_b = Counter(b)

    missing_names = []
    for name, freq_a in count_a.items():
        freq_b = count_b.get(name, 0)
        if freq_a > freq_b:
            missing_count = freq_a - freq_b
            missing_names.extend([name] * missing_count)
    return missing_names

def label_clean_tables(DATA_DIR, model, processor, task):

    OUTPUT_DIR = os.path.join(DATA_DIR, task, 'annotations_qwen2vl')
    os.makedirs(OUTPUT_DIR, exist_ok=True)
    task_path = os.path.join(DATA_DIR, task)
    objs = []
    try:
        with open(os.path.join(OUTPUT_DIR, 'annotations.json'), 'r') as f:
            anno = json.load(f)
    except Exception as e:
        print(e)
        anno = {}
    ##########################for debug#########################
    # objs = ['empty bottle', 'empty bottle', 'cup', 'mug']
    ############################################################
    with open(os.path.join(task_path, "meta.txt"), 'r', encoding='utf-8') as f:
        lines = f.readlines()
        for each in lines:
            objs.extend(each.strip().split(','))
    # os.makedirs(os.path.join(OUTPUT_DIR, task), exist_ok=True)
    episodes = os.listdir(task_path)
    episodes = [episode for episode in episodes if episode.endswith('.hdf5')]
    episodes = sorted(episodes, key=lambda x: int(x.split('.')[0].split('_')[-1]))

    for episode in tqdm(episodes[:10]):
        if episode in anno.keys() and anno[episode]['status']:
            print(f"Already processed {episode}")
            continue
        episode_path = os.path.join(task_path, episode)
        image_dict, subtask = load_hdf5(task_path, episode)
        key_frames, idxs = filter_images_by_subtask(image_dict, subtask, OUTPUT_DIR, episode.split(".")[0])

        messages = copy.deepcopy(chat_template)
        for i in range(4):
            messages[0]['content'].append({
                "type": "image",
                "image": os.path.join(OUTPUT_DIR, f'{episode.split(".")[0]}_{i}.png'),
            })
        messages[0]['content'].append({"type": "text", "text": f""})
        messages[0]['content'][-1]['text'] = prompt.replace("[<objects>]", f"[{(','.join(objs))}]")

        results = model_inference(model, processor, messages)

        print("<<<<<<<<<<<<<<<<<<Processing missing objects>>>>>>>>>>>>>>>>>>")
        objects = []
        status = True
        for i in range(0, len(results) - 1, 1):
            res = find_missing_names_counter(results[i], results[i + 1])
            objects.append(res)
            if len(res) > 1 or len(res) == 0:
                print(f"{YELLOW} Detected error in {episode}: {res} {RESET}")
                status = False

        objects.append(results[-1])
        print(f"The order of objects in {RED} {episode} is {objects} {RESET}")
        anno[episode] = {
            'path': episode_path,
            'objects_order': objects,
            'status': status,
        }

    with open(os.path.join(OUTPUT_DIR, 'annotations.json'), 'w', encoding='utf-8') as f:
        json.dump(anno, f, indent=4)

if __name__ == '__main__':
    model, processor = load_model("/home/jovyan/tzb/wjj/model_param/Qwen2-VL-7B-Instruct/")
    tasks = [
        # 'fold_shirt_wjj1213_meeting_room',
        # 'clean_table_ljm_1217',
        'clean_table_zmj_1217_green_plate_coke_can_brown_mug_bottle',
    ]
    DATA_DIR = "/home/jovyan/tzb/wjj/data/aloha_bimanual/aloha_4views/"
    for task in tasks:
        label_clean_tables(DATA_DIR=DATA_DIR, task=task, model=model, processor=processor)