import os
import re
import json
import pickle


import numpy as np
from PIL import Image
from torchvision import transforms

from configs.config import ConfigManager
from datasets.dataset import AdvancedDataset
from utils.omni_util import OmniUtil
from utils.pcl_util import PclUtil
from utils.pose_util import PoseUtil


class NextOneBestViewDataset(AdvancedDataset):
    

    def __init__(self, dataset_config):
        super(NextOneBestViewDataset, self).__init__(dataset_config)
        self.data_type = dataset_config["data_type"]
        self.source = dataset_config["source"]
        self.gsnet_label_name = dataset_config["gsnet_label"]
        #self.foundation_pose_label_name = dataset_config["foundation_pose_label"]
        self.data_dir = ConfigManager.get("datasets", "general", "data_dir")
        self.score_limit = ConfigManager.get("datasets", "general", "score_limit")
        self.target_pts_num = ConfigManager.get("datasets", "general", "target_pts_num")
        self.scene_pts_num = ConfigManager.get("datasets", "general", "scene_pts_num")
        self.image_size = ConfigManager.get("datasets", "general", "image_size")
        self.rgb_feat_cache = ConfigManager.get("datasets", "general", "rgb_feat_cache")
        self.canonical = ConfigManager.get("datasets", "general", "canonical")
        self.small_batch_overfit = ConfigManager.get("settings", "experiment", "small_batch_overfit")
        self.container_path = str(os.path.join(self.data_dir, self.source, "container_set.pickle"))
        self.data_path = str(os.path.join(self.data_dir, self.source, self.data_type))
        self.gsnet_label_path = str(os.path.join(self.data_dir, self.source, self.gsnet_label_name))
        #self.foundation_pose_label_path = str(os.path.join(self.data_dir, self.source, self.foundation_pose_label_name))
        self.scene_list = os.listdir(self.data_path)
        self.task = ConfigManager.get("settings", "experiment", "task")
        self.container_set = self.load_container_set()
        self.data_list = self.get_datalist()
        
        
        if self.small_batch_overfit:
            small_batch_size = ConfigManager.get("settings", "experiment", "small_batch_size")
            small_batch_times = ConfigManager.get("settings", "experiment", "small_batch_times")
            self.data_list = self.data_list[:small_batch_size] * small_batch_times
            
        self.transform = transforms.Compose([           
                        transforms.Resize(self.image_size),
                        transforms.CenterCrop(int(self.image_size//14)*14),              
                        transforms.ToTensor(),                    
                        transforms.Normalize(mean=0.5, std=0.2)
                        ])


    def __len__(self):
        return len(self.data_list)

    def getitem(self, index) -> dict:  
        data_pair = self.data_list[index]
        src_path = data_pair[0]["frame_path"]
        dst_path = data_pair[1]["frame_path"]
        target_name = data_pair[0]["object_name"]
        scene_name = data_pair[0]["scene"]
        src_data = self.load_src_data(src_path, target_name, canonical=self.canonical)
        dst_data = self.load_dst_data(dst_path)
        src_rot = src_data["cam_transform"][:3,:3]
        dst_rot = dst_data["cam_transform"][:3,:3]
        delta_rot = np.dot(dst_rot.T, src_rot)
        delta_rot_6d = PoseUtil.matrix_to_rotation_6d_numpy(delta_rot)
            
        item_data = {
            "src_path": src_path,
            "target_name": target_name,
            "scene_name": scene_name,
            "data_type": self.data_type,
            "source": self.source,
            "target_pts": src_data["target_pts"].astype(np.float32),
            "scene_pts": src_data["scene_pts"].astype(np.float32),
            "delta_rot_6d": delta_rot_6d.astype(np.float32),
            "src_rot_mat": src_rot.astype(np.float32),
            "dst_rot_mat": dst_rot.astype(np.float32),
            "src_transform": src_data["cam_transform"].astype(np.float32),
            "dst_transform": dst_data["cam_transform"].astype(np.float32),
        } 
        # if self.rgb_feat_cache:
        #     item_data["rgb_feat"] = src_data["rgb_feat"].astype(np.float32)
        # else:
        #     item_data["rgb"] = src_data["rgb"]
        return item_data
    
    def load_dst_data(self, frame_path):
        """ camera params """
        cam_transform = OmniUtil.get_transform_mat(frame_path)
        frame_data = {'cam_transform': cam_transform}
        return frame_data

    def load_src_data(self, frame_path, target_object_name, canonical = False):
        """ pts """
        scene_pts = OmniUtil.get_points(path=frame_path, object_name=OmniUtil.FOREGROUND)
        target_pts = OmniUtil.get_points(
            path=frame_path, object_name=target_object_name
        )
        scene_pts = PclUtil.sample_pcl(scene_pts, self.scene_pts_num)
        target_pts = PclUtil.sample_pcl(target_pts, self.target_pts_num)
        

        """ camera params """
        cam_transform = OmniUtil.get_transform_mat(frame_path)
        if canonical:
            target_pts = PclUtil.cam2canonical(target_pts, cam_transform)
            scene_pts = PclUtil.cam2canonical(scene_pts, cam_transform)

        frame_data = {
            "target_pts": target_pts,
            "scene_pts": scene_pts,
            "cam_transform": cam_transform,
        }
        
        """ rgb """
        # if self.rgb_feat_cache:
        #     rgb_feat = OmniUtil.get_rgb_feat(frame_path)
        #     frame_data["rgb_feat"] = rgb_feat
        # else:
        #     rgb = OmniUtil.get_rgb(frame_path)
        #     rgb = Image.fromarray(rgb)
        #     rgb = self.transform(rgb)
        #     frame_data["rgb"] = rgb

        return frame_data
    
    def load_container_set(self):
        container_list = ['chair_028', 'chair_029', 'chair_026', 'chair_027', 'table_025', 'table_027', 'table_026', 'table_028', 'sofa_014', 'sofa_013', 'picnic_basket_010', 'picnic_basket_011', 'cabinet_009', 'flower_pot_023', 'flower_pot_022', 'flower_pot_021', 'chair_017', 'chair_020', 'chair_012', 'chair_010', 'chair_018', 'chair_025', 'chair_024', 'chair_011', 'chair_001', 'chair_013', 'chair_004', 'chair_021', 'chair_023', 'chair_006', 'chair_014', 'chair_007', 'chair_003', 'chair_009', 'chair_022', 'chair_015', 'chair_016', 'chair_008', 'chair_005', 'chair_019', 'chair_002', 'table_004', 'table_023', 'table_014', 'table_024', 'table_019', 'table_022', 'table_007', 'table_017', 'table_013', 'table_002', 'table_016', 'table_009', 'table_008', 'table_003', 'table_015', 'table_001', 'table_018', 'table_005', 'table_020', 'table_021', 'sofa_001', 'sofa_005', 'sofa_012', 'sofa_009', 'sofa_006', 'sofa_008', 'sofa_011', 'sofa_004', 'sofa_003', 'sofa_002', 'sofa_007', 'sofa_010', 'picnic_basket_005', 'picnic_basket_004', 'picnic_basket_001', 'picnic_basket_008', 'picnic_basket_002', 'picnic_basket_009', 'picnic_basket_006', 'picnic_basket_003', 'picnic_basket_007', 'cabinet_006', 'cabinet_008', 'cabinet_002', 'cabinet_001', 'cabinet_005', 'cabinet_007', 'flower_pot_013', 'flower_pot_005', 'flower_pot_008', 'flower_pot_001', 'flower_pot_003', 'flower_pot_020', 'flower_pot_006', 'flower_pot_012', 'flower_pot_018', 'flower_pot_007', 'flower_pot_002', 'flower_pot_011', 'flower_pot_010', 'flower_pot_016', 'flower_pot_004', 'flower_pot_014', 'flower_pot_017', 'flower_pot_019']
        container_set = set(container_list)
        return container_set

    def get_ground_object_set(self, scene_name):
        fall_path  = os.path.join(self.data_path, scene_name, "fall_objects.pickle")
        with open(fall_path, 'rb') as f:
            fall_objects = pickle.load(f)
        return fall_objects
    
    def get_datalist(self):
        if self.task == "object_pose":
            raise NotImplementedError("object_pose task is not supported now.")
            #return self.get_foundation_pose_datalist()
        elif self.task == "grasp_pose":
            return self.get_grasp_pose_datalist()
        else:
            raise ValueError("task must be 'object_pose' or 'grasp_pose'.")
        
    '''
    def get_foundation_pose_datalist(self):
        data_list = []
        for scene in self.scene_list:
            scene_path = os.path.join(self.data_path, scene)
            gsnet_label_scene_path = os.path.join(self.foundation_pose_label_path, scene)
            file_list = os.listdir(scene_path)
            scene_frame_list = []
            target_object_set = self.get_target_object_set(scene)
            unseen_object_set = self.get_target_object_set(scene)
            cnt_under = 0
            cnt_above = 0
            limit = 0.002
            for file in file_list:
                if file.startswith("camera_params"):
                    frame_index = re.findall(r"\d+", file)[0]
                    frame_path = os.path.join(scene_path, frame_index)
                    
                    score_label_path = os.path.join(
                        gsnet_label_scene_path, OmniUtil.SCORE_LABEL_TEMPLATE.format(frame_index)
                    )
                    with open(score_label_path, "r") as f:
                        score_label = json.load(f)
                    for obj_name in score_label.keys():
                        if obj_name in target_object_set:
                            scene_frame_list.append(
                                {
                                    "frame_path": frame_path,
                                    "object_name": obj_name,
                                    "score": score_label[obj_name]["eval_result"]["ADD-S"],
                                    "scene": scene
                                }
                            )
                        if score_label[obj_name]["eval_result"]["ADD-S"] <= limit:
                            cnt_under += 1
                        else:
                            cnt_above += 1
            print(f"under {limit}: {cnt_under}, above {limit}: {cnt_above}")
            for i in range(len(scene_frame_list)):
                for j in range(i, len(scene_frame_list)):
                    fm_i, fm_j = scene_frame_list[i], scene_frame_list[j]
                    if fm_i["object_name"] == fm_j["object_name"]:
                        bad_view, good_view = None, None
                        if fm_i["score"] <= limit < fm_j["score"]:
                            good_view, bad_view  = fm_i, fm_j
                        elif fm_i["score"] > limit >= fm_j["score"]:
                            good_view, bad_view  = fm_j, fm_i
                        if bad_view is None or good_view is None:
                            continue
                        data_list.append((bad_view, good_view))
                        if bad_view["object_name"] in unseen_object_set:
                            unseen_object_set.remove(bad_view["object_name"])
        return data_list
    '''
        
    def get_grasp_pose_datalist(self):
        data_list = []
        for scene in self.scene_list:
            
            scene_path = os.path.join(self.data_path, scene)
            gsnet_label_scene_path = os.path.join(self.gsnet_label_path, scene)
            file_list = os.listdir(scene_path)
            scene_frame_list = []
            ground_object_set = self.get_ground_object_set(scene)
            unseen_object_set = set()
            for file in file_list:
                if file.startswith("camera_params"):
                    frame_index = re.findall(r"\d+", file)[0]
                    frame_path = os.path.join(scene_path, frame_index)
                    
                    score_label_path = os.path.join(
                        gsnet_label_scene_path, OmniUtil.SCORE_LABEL_TEMPLATE.format(frame_index)
                    )
                    with open(score_label_path, "r") as f:
                        score_label = json.load(f)
                    for obj_name in score_label["avg_score"].keys():
                        if obj_name not in ground_object_set and obj_name not in self.container_set:
                            scene_frame_list.append(
                                {
                                    "frame_path": frame_path,
                                    "object_name": obj_name,
                                    "score": score_label["avg_score"][obj_name],
                                    "scene": scene
                                }
                            )
                            unseen_object_set.add(obj_name)
            for i in range(len(scene_frame_list)):
                for j in range(i, len(scene_frame_list)):
                    fm_i, fm_j = scene_frame_list[i], scene_frame_list[j]
                    if fm_i["object_name"] == fm_j["object_name"]:
                        bad_view, good_view = None, None
                        if fm_i["score"] <= self.score_limit < fm_j["score"]:
                            bad_view, good_view = fm_i, fm_j
                        elif fm_i["score"] > self.score_limit >= fm_j["score"]:
                            bad_view, good_view = fm_j, fm_i
                        if bad_view is None or good_view is None:
                            continue
                        sample_prob = ((max(0,good_view["score"] - bad_view["score"]))/0.3)**2
                        if np.random.rand() > sample_prob:
                            continue
                        data_list.append((bad_view, good_view))
                        if bad_view["object_name"] in unseen_object_set:
                            unseen_object_set.remove(bad_view["object_name"])
            for obj_name in unseen_object_set:
                views = []
                for frame in scene_frame_list:
                    if frame["object_name"] == obj_name:
                        views.append(frame)
                sorted_views = sorted(views, key=lambda x: x["score"], reverse=True)
                total_view_num = len(sorted_views)
                good_view_num = int(total_view_num * 0.1)
                good_views = sorted_views[:good_view_num]
                bad_views = sorted_views[good_view_num:]
                filtered_good_view = []
                filtered_bad_view = bad_views
                for good_view in good_views:
                    if good_view["score"] >= 0.01:
                        filtered_good_view.append(good_view)
                    else:
                        filtered_bad_view.append(good_view)
                for good_view in filtered_good_view:
                    for bad_view in filtered_bad_view:
                        data_list.append((bad_view, good_view))
        return data_list
