import pickle, json
import time
import cv2
import os
import numpy as np
from PIL import Image
import open3d as o3d
import shutil
import math
from utils.io.io_utils import open_file, delete_file, add_to_log, read_py
from utils.transformation_utils import extract_z_axis
from utils.LLM_utils import query_LLM
import glob

prompt_parse_pos = 'prompts/parse_pos.py'
prompt_get_task_pose = read_py('prompts/get_task_pose_str.txt')
prompt_get_pose_from_str = read_py('prompts/get_pose_from_str.py')
prompt_parse_name_str = 'prompts/parse_name.txt'
prompt_change_frame_str = 'prompts/change_frame.txt'
prompt_get_pos_scale = read_py('prompts/get_pos_scale.txt')
prompt_get_obj_name_from_task = read_py('prompts/get_obj_name_from_task.txt')
prompt_update_obj_state_file = 'prompts/update_obj_state.txt'
prompt_get_initial_obj_state = read_py('prompts/get_ini_obj_state.txt')

# global variables
load_detected_objs = None
popped_detected_objs = []
saved_detected_obj = {}
queried_obj = None
reference_frame = 'object'
RADIUS = 0.0135
TASK = None
REAL_ROBOT = None
policy = None
clip_model = None

def get_considered_classes():
    if TASK == 'drawer':
        considered_classes = [["drawer", "drawer handle"], ["pen"], ["cup", "mug"], ["tape"], ["scissors", "scissors handle"], ["apple"]]
        other_classes = [["table cloth", "wooden table", "wooden floor", "gray cloth", "gray background", "gray board"],["tripod"], ["unrecoginized object"]]
    elif TASK == 'coffee':
        considered_classes = [["coffee maker", "coffee machine", "keurig coffee maker"], ["red cup", "red mug"], ["coffee capsule", "cone"], ["pink cup", "pink mug"], ["white spoon"]]
        other_classes = [["unrecognizable object"], ["robot arm"], ["table cloth"], ["work surface"], ["tripod"]]
    elif TASK == 'cup':
        considered_classes = [["pink cup"], ["green cup", "blue cup"], ["gray rack", "gray oval object"]]
        other_classes = [["robot arm"], ["table cloth"], ["work surface"], ["tripod"]]
    elif TASK == 'lego':
        considered_classes = [["white drawer", "white drawer handle"], ["red block"], ["yellow block"], ["blue block"]]
        other_classes = [["robot arm"], ["table cloth"], ["work surface"], ["tripod"]]
    else:
        raise NotImplementedError
    return considered_classes, other_classes

class ToScaledFloat:
    """
    Convert a ``PIL Image`` or ``numpy.ndarray`` to tensor, or keep as is if already a tensor.
    """

    def __init__(self) -> None:
        pass

    def __call__(self, pic):
        # assert pic.dtype == torch.uint8, f"{pic.dtype}"
        pic = pic.float() / 255.0
        return pic

    def __repr__(self) -> str:
        return f"{self.__class__.__name__}()"

class DinoImageTransform:
    def __init__(self):
        self.transforms = pt_transforms.Compose(
            [
                ToScaledFloat(),
                pt_transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
            ]
        )

    def __call__(self, pic):
        return self.transforms(pic)

class detected_object:
    def __init__(self):
        self.pcd = None
        self.prob = None
        self.quat = None
        self.position = None
        self.vec = None
        self.clip_feature = None
        self.dino_feature = None
        self.detected_pose = None

# ---------------------------------- Initialization ----------------------------------

def set_policy_and_task(real_robot, task):
    global TASK, REAL_ROBOT, policy
    global torch, open_clip, F, pt_transforms, device
    global SamAutomaticMaskGenerator, sam_model_registry, clip_with_owl
    REAL_ROBOT = real_robot
    TASK = task
    import torch
    from torchvision import transforms as pt_transforms
    import torch.nn.functional as F
    import open_clip
    from segment_anything import SamAutomaticMaskGenerator, sam_model_registry
    from utils.perception.owl_vit import clip_with_owl
    device = "cuda" if torch.cuda.is_available() else "cpu"
    if REAL_ROBOT:
        import utils.robot.robot_policy as robot_policy
        policy = robot_policy.KptPrimitivePolicy()
    else:
        from utils.robot.dummy_policy import DummyPolicy
        policy = DummyPolicy()

def initialize_detection(first=False, load_image=False, folder_path='cache/image_for_retrieval', image_pattern='*.png', label=['cup', 'drawer']):
    global load_detected_objs, saved_detected_obj, queried_obj, clip_model, clip_preprocess
    queried_obj = detected_object()
    if REAL_ROBOT:
        if first:
            a = input("Can I use last round detection? (y/n)")
        else:
            a = 'y'
        if 'n' in a:
            detect_objs(load_from_cache=False, save_to_cache=False, visualize=False)
            saved_detected_obj = {}
        load_detected_objs = pickle.load(open(f"log/{TASK}/detected_objs.pkl", "rb"))
    else:
        if first:
            a = input("Can I use last round detection? (y/n)")
        else:
            a = 'y'
        if 'n' in a:
            saved_detected_obj = {}
        if first:
            if load_image:
                # Only consider clip feature here
                if clip_model is None:
                    print("Loading CLIP...")
                    clip_model, _, clip_preprocess = open_clip.create_model_and_transforms('ViT-g-14', pretrained='laion2b_s34b_b88k')
                image_files = glob.glob(os.path.join(folder_path, image_pattern))
                img_dict = {}
                for l in label:
                    img_files = [file for file in image_files if l in os.path.basename(file).lower()]
                    img_files = sorted(img_files, key=lambda x: int(x.split(l)[1].split('.png')[0]))
                    img_dict[l] = img_files
                files = {}
                clip_model.to(device=device)
                for l in label:
                    img_files = img_dict[l]
                    all_img = []
                    _files = []
                    for file in img_files:
                        _files.append(file)
                        img = Image.open(file)
                        img = clip_preprocess(img).to(device)
                        all_img.append(img)
                    preprocessed_images = torch.stack(all_img)
                    with torch.no_grad(), torch.cuda.amp.autocast():
                        image_features = clip_model.encode_image(preprocessed_images)
                    image_features = image_features.cpu().numpy()
                    img_dict[l] = image_features
                    files[l] = _files
                load_detected_objs = create_loaded_objs(img_dict)
            else:
                load_detected_objs = create_loaded_objs(None)
        else:
            load_detected_objs = pickle.load(open(f"log/{TASK}/detected_objs.pkl", "rb"))

def create_loaded_objs(img_dict):
    considered_classes, _ = get_considered_classes()
    load_detected_objs = {}
    for cls in considered_classes:
        if img_dict is None:
            obj_name = cls[0]
            load_detected_objs[obj_name] = [(np.ones((10,3)), 0.5, np.random.randn(1024),np.random.randn(512)), (np.ones((10,3)), 0.3, np.random.randn(1024),np.random.randn(512))]
        else:
            if obj_name not in list(img_dict.keys()):
                n = 3
                for i in range(n):
                    obj = obj_name + f'_{i+1}'
                    load_detected_objs[obj] = (np.ones((10,3)), 0.5, np.random.randn(1024),np.random.randn(512))
            else:
                image_feature = img_dict[obj_name]
                n = len(image_feature)
                for i in range(n):
                    obj = obj_name + f'_{i+1}'
                    load_detected_objs[obj] = (np.ones((n+2,3)), 0.5, image_feature[i], np.random.randn(512))
    pickle.dump(load_detected_objs, open(f"log/{TASK}/detected_objs.pkl", "wb"))
    return load_detected_objs

def get_objs():
    global load_detected_objs
    return load_detected_objs

def get_initial_state():
    global load_detected_objs
    # obj_list_str = ''
    # for obj_name in list(load_detected_objs.keys()):
    #     obj_list_str += obj_name + ', '
    # obj_list_str = obj_list_str[:-2]
    # whole_prompt = prompt_get_initial_obj_state + '\n' + 'Object name: ' + obj_list_str +'\n' + 'Object state:'
    # response = query_LLM(whole_prompt, [], "cache/llm_get_initial_state.pkl")
    # initial_state = response.text
    initial_state = 'Object state: N/A'
    return initial_state

def get_extra_classes():
    if TASK == 'drawer':
        return ["scissors", "scissors handle"], 2
    else:
        return None, None

def setup_clip_words():
    considered_classes, other_classes = get_considered_classes()
    clip_candidates = []
    for word in considered_classes + other_classes:
        clip_candidates.extend(word)
    return clip_candidates

def clear_saved_detected_obj():
    global saved_detected_obj
    saved_detected_obj = {}

def set_saved_detected_obj(task_name, task_feature):
    obj_name = get_obj_name_from_task(task_name)
    saved_detected_obj[obj_name] = task_feature

# ---------------------------------- Primitives ----------------------------------

def _change_reference_frame(wrong_direction, right_direction):
    global queried_obj
    with open(prompt_change_frame_str, "r") as template_file:
        template_content = template_file.read()
    values = {"wrong_direction": wrong_direction, "correct_direction": right_direction}
    prompt = template_content.format(**values)
    response = query_LLM(prompt, ['.'], "cache/llm_change_frame.pkl")
    add_to_log(prompt)
    add_to_log(response.text)
    if "left" in response.text:
        correct_forward_vec = get_directional_vec("left")
    if "right" in response.text:
        correct_forward_vec = get_directional_vec("right")
    if "forward" in response.text:
        correct_forward_vec = get_directional_vec("forward")
    if "back" in response.text:
        correct_forward_vec = get_directional_vec("back")
    queried_obj.vec = correct_forward_vec

def get_directional_vec(direction, *kwargs):
    global queried_obj, reference_frame
    if reference_frame == 'absolute':
        if direction == 'forward' or direction == 'along' or direction == 'towards':
            return np.array((1,0,0))
        elif direction == 'backward' or direction == 'back' or direction == 'away':
            return np.array((-1,0,0))
        elif direction == 'left':
            return np.array((0,1,0))
        elif direction == 'right':
            return np.array((0,-1,0))
        elif direction == 'up' or direction == 'above':
            return np.array((0,0,1))
        elif direction == 'down' or direction == 'downward':
            return np.array((0,0,-1))
        else:
            return np.array((1,0,0))
    else:
        if REAL_ROBOT:
            vec = queried_obj.vec.squeeze()
        else:
            vec = np.array((1.,0.,0.))
        vec = vec/np.linalg.norm(vec)
        if direction == 'forward' or direction == 'along' or direction == 'towards':
            vec = vec if vec[0] >=0 else -vec
            return vec
        elif direction == 'backward' or direction == 'back' or direction == 'away':
            vec = -vec if vec[0] >=0 else vec
            return vec
        elif direction == 'left':
            normal_vec = np.array((-vec[1], vec[0], vec[2])) if vec[0] >=0 else np.array((vec[1], -vec[0], -vec[2]))
            return normal_vec
        elif direction == 'right':
            normal_vec = np.array((-vec[1], vec[0], vec[2])) if vec[0] <=0 else np.array((vec[1], -vec[0], -vec[2]))
            return normal_vec
        elif direction == 'up' or direction == 'above':
            vec = np.array((0,0,1))
            return vec
        elif direction == 'down' or direction == 'downward':
            vec = np.array((0,0,-1))
            return vec
        else:
            return np.array((1.,0.,0.))

def get_current_pos(*kwargs):
    if REAL_ROBOT:
        pose = policy.robot_env.robot.get_ee_pose()
        ee_pos = pose[0].numpy()
        return ee_pos
    else:
        return np.array((1.,0.,0.))

def get_current_ori(*kwargs):
    if REAL_ROBOT:
        pose = policy.robot_env.robot.get_ee_pose()
        ee_ori = pose[1].numpy()
        return ee_ori
    else:
        return np.array((1.,0.,0.,0.))

def get_horizontal_ori(get_vec=False):
    if get_vec == False:
        return policy.get_horizontal_ori()
    else:
        quat = policy.get_horizontal_ori()
        vec = extract_z_axis(quat)
        vec[2] = 0.
        vec = -vec if vec[0]<0 else vec
        vec /= np.linalg.norm(vec)
        return vec

def get_vertical_ori(get_vec=False):
    if get_vec==False:
        return policy.get_vertical_ori()
    else:
        return np.array((0.,0.,-1.))

def _correct_past_detection(obj_name, obj):
    global load_detected_objs, popped_detected_objs
    obj_name = parse_obj_name(obj_name)
    for i in range(len(load_detected_objs[obj_name])):
        if (obj.pcd.shape == load_detected_objs[obj_name][i][0].shape):
            popped_elem = load_detected_objs[obj_name].pop(i)
            popped_detected_objs.append(popped_elem)
            break

# ---------------------------------- Parsing ----------------------------------

def get_obj_name_from_task(task_name):
    prompt = prompt_get_obj_name_from_task + '\n' + f'Input: {task_name}' + '\n' + 'Output: '
    response = query_LLM(prompt, [], "cache/llm_parse_task_name.pkl")
    obj_name = response.text
    return obj_name

def parse_obj_name(obj_name):
    considered_classes, _ = get_considered_classes()
    with open(prompt_parse_name_str, 'r') as file:
        template_content = file.read()
    values = {"object_class": considered_classes, "object_name": obj_name}
    prompt = template_content.format(**values)
    response = query_LLM(prompt, [], "cache/llm_parse_obj_name.pkl")
    obj_class_idx = int(response.text)
    parsed_obj_name = considered_classes[obj_class_idx][0]
    return parsed_obj_name

def _parse_pos(pos_description, r_frame='object'):
    global reference_frame
    reference_frame = r_frame
    mapping_dict, little_left, little_up, little_forward, more_back = calculate_pos_scale()
    with open(prompt_parse_pos, "r") as template_file:
        template_content = template_file.read()
    values = {"mapping_dict": mapping_dict, "little_left": little_left, "little_up": little_up, "little_forward": little_forward, "more_back": more_back}
    prompt = template_content.format(**values)
    whole_prompt = prompt + '\n' + '\n' + f"""# Query: {pos_description}.""" + "\n"
    response = query_LLM(whole_prompt, ["# Query:"], "cache/llm_response_parse_pos.pkl")
    code_as_policies = response.text
    add_to_log('-'*80)
    add_to_log('*at parse_pos*')
    add_to_log(code_as_policies)
    add_to_log('-'*80)
    locals = {}
    exec(code_as_policies, globals(), locals)
    return locals['ret_val'], code_as_policies

def calculate_pos_scale():
    global queried_obj
    if REAL_ROBOT:
        whole_prompt = prompt_get_pos_scale + '\n\n' + 'Bounding box: ' + str(queried_obj.bounding_box)
    else:
        whole_prompt = prompt_get_pos_scale + '\n\n' + 'Bounding box: (1, 8, 3)'
    response = query_LLM(whole_prompt, ["Bounding box:"], "cache/llm_get_pos_scale.pkl")
    response_json = response.text
    mapping_dict = json.loads(response_json)
    add_to_log(mapping_dict)
    little_forward = mapping_dict["'A little bit' for 'forward' or 'backward'"]/100
    little_left = mapping_dict["'A little bit' for 'left' or 'right'"]/100
    little_up = mapping_dict["'A little bit' for 'up' or 'down'"]/100
    more_back = mapping_dict["'More' for 'forward' or 'back'"]/100
    return mapping_dict, little_left, little_up, little_forward, more_back

# ---------------------------------- Detection ----------------------------------

def _get_task_detection(task_name):
    """Given task name, find corresponding feature"""
    global saved_detected_obj

    obj_name = get_obj_name_from_task(task_name)

    if obj_name in saved_detected_obj.keys():
        compare_feature(saved_detected_obj[obj_name], obj_name, threshold=1)
        return saved_detected_obj[obj_name]

    whole_prompt = prompt_get_pose_from_str + '\n' + "# Query: " + f"""the centroid of "{obj_name}"."""
    response = query_LLM(whole_prompt, ["INPUT:"], "cache/llm_response_get_pose_from_str.pkl")
    code_as_policies = response.text
    add_to_log('-'*80 + '*at get_pose_from_str*' + code_as_policies + '-'*80)
    while True:
        locals = {}
        exec(code_as_policies, globals(), locals)
        obj = locals['ret_val']
        if REAL_ROBOT:
            pcd_merged = o3d.io.read_point_cloud(f"log/{TASK}/pcd_merged.pcd")
            mesh_frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=0.1, origin=[0, 0, 0])
            visualized_geometries = [pcd_merged, mesh_frame]
            sphere = o3d.geometry.TriangleMesh.create_sphere(radius=RADIUS); sphere.paint_uniform_color([0.0, 1.0, 0.0])
            sphere.translate(calculate_centroid(obj.pcd))
            visualized_geometries.append(sphere)
            o3d.visualization.draw_geometries(visualized_geometries)
        correct_detection = input("Is this detection correct? (y/n)")
        if 'y' in correct_detection.lower():
            saved_detected_obj[obj_name] = obj.dino_feature
            compare_feature(obj.dino_feature, obj_name, threshold=1)
            return obj.dino_feature
        else:
            _correct_past_detection(obj_name, obj)

def detect(obj_name, visualize=False):
    global load_detected_objs, saved_detected_obj

    parsed_obj_name = parse_obj_name(obj_name)
    obj_name = parsed_obj_name

    if load_detected_objs[obj_name] == [] or load_detected_objs[obj_name] is None:
        while True:
            a = input("It seems that the correct object was not detected, please make sure to place the drawer in a correct position. Reply 'y' once you finish it.")
            if "y" in a:
                break
        if REAL_ROBOT:
            detect_objs(load_from_cache=False, save_to_cache=False, visualize=visualize)
            load_detected_objs = pickle.load(open(f"log/{TASK}/detected_objs.pkl", "rb"))
        else:
            load_detected_objs = create_loaded_objs()

    all_pcds = []
    probs = []
    clip_features = []
    dino_features = []

    for pcd, prob, clip_feature, dino_feature in load_detected_objs[obj_name]:
        all_pcds.append(pcd)
        probs.append(prob)
        clip_features.append(clip_feature)
        dino_features.append(dino_feature)
    if visualize:
        pcd_merged = o3d.io.read_point_cloud(f"log/{TASK}/pcd_merged.pcd")
        mesh_frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=0.1, origin=[0, 0, 0])
        visualized_geometries = [pcd_merged, mesh_frame]
        for pcd in all_pcds:
            sphere = o3d.geometry.TriangleMesh.create_sphere(radius=RADIUS)
            sphere.paint_uniform_color([0.0, 1.0, 0.0])
            sphere.translate(np.mean(pcd,axis=0))
            visualized_geometries.append(sphere)
        o3d.visualization.draw_geometries(visualized_geometries)

    all_obj = detected_object()
    all_obj.pcd = all_pcds
    all_obj.prob = probs
    all_obj.clip_feature = clip_features
    all_obj.dino_feature = dino_features
    return all_obj

def sort_from_high_to_low(all_obj, key):
    all_obj_pos = [np.mean(i, axis=0) for i in all_obj.pcd]
    probs = all_obj.prob
    clip_features = all_obj.clip_feature
    dino_features = all_obj.dino_feature
    if key == 'probability':
        order = np.argsort(np.array(probs))[::-1]
    elif key == 'x':
        x_pos = [x[0] for x in all_obj_pos]
        order = np.argsort(np.array(x_pos))[::-1]
    elif key == 'y':
        y_pos = [x[1] for x in all_obj_pos]
        order = np.argsort(np.array(y_pos))[::-1]
    elif key == 'z':
        z_pos = [x[2] for x in all_obj_pos]
        order = np.argsort(np.array(z_pos))[::-1]
    obj_list = []
    for ind in order:
        obj = detected_object()
        obj.pcd = all_obj.pcd[ind]
        obj.prob = probs[ind]
        obj.clip_feature = clip_features[ind]
        obj.dino_feature = dino_features[ind]
        obj_list.append(obj)
    return obj_list

def detect_objs(load_from_cache=False, run_on_server=False, save_to_cache=True, visualize=False):
    global clip_model, clip_preprocess
    parent_folder = 'cache'
    items = os.listdir(parent_folder)
    for item in items:
        item_path = os.path.join(parent_folder, item)
        if os.path.isdir(item_path) and (item[0].isdigit() or item == TASK):
            try:
                shutil.rmtree(item_path)
                print(f"Folder '{item_path}' and its contents deleted successfully.")
            except OSError as e:
                print(f"Error deleting '{item_path}': {e}")
    try:
        shutil.rmtree(f'log/{TASK}')
        print(f"Folder 'log/{TASK}' and its contents deleted successfully.")
    except OSError as e:
        print(f"Error deleting '{item_path}': {e}")

    print("Loading CLIP model...")
    init_time = time.time()
    if clip_model is None:
        clip_model, _, clip_preprocess = open_clip.create_model_and_transforms('ViT-g-14', pretrained='laion2b_s34b_b88k')
    print("time taken:", time.time() - init_time)

    print("Loading DINOv2 model...")
    init_time = time.time()
    dinov2_model = torch.hub.load("facebookresearch/dinov2", "dinov2_vits14")
    dinov2_model.train(False)
    for param in dinov2_model.parameters():
        param.requires_grad = False
    dinov2_preprocess = DinoImageTransform()
    print("time take: ", time.time() - init_time)

    clip_candidates = setup_clip_words()
    text_features = get_text_features(clip_candidates, clip_model)
    considered_classes, other_classes = get_considered_classes()
    detected_objs = {words[0]: [] for words in considered_classes}
    if load_from_cache:
        print('Loading data...')
        init_time = time.time()
        bgr_images = pickle.load(open(f"cache/{TASK}/bgr_images.pkl", "rb"))
        pcd_merged = o3d.io.read_point_cloud(f"cache/{TASK}/pcd_merged.pcd")
        detected_objs = pickle.load(open(f"cache/{TASK}/detected_objs.pkl", "rb"))
        print("time taken:", time.time() - init_time)
    else:
        print("Computing point cloud...")
        init_time = time.time()
        used_camera_idx = 0
        # TODO: use this after you implement utils/perception/camera.py
        from utils.perception.camera import multi_cam
        bgr_images, pcd_merged, raw_points, _ = multi_cam.take_bgrd(visualize=visualize)
        image = bgr_images[realsense_serial_numbers[used_camera_idx]]
        coord2point_3d = raw_points[used_camera_idx]
        del raw_points
        print("time taken:", time.time() - init_time)

        print("Loading SAM model...")
        init_time = time.time()
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        sam = sam_model_registry["vit_h"](checkpoint="cache/sam_vit_h_4b8939.pth")
        mask_generator = SamAutomaticMaskGenerator(
            model=sam,
            points_per_side=32,
            pred_iou_thresh=0.95,
            stability_score_thresh=0.95,
            box_nms_thresh=0.2,
            crop_n_layers=1,
            crop_overlap_ratio=0.6,
            min_mask_region_area=100
        )
        print("time taken:", time.time() - init_time)

        print("Cropping picture...")
        init_time = time.time()
        height, width = image.shape[:2]
        print(width, height)
        coord2point_3d = np.reshape(coord2point_3d, (height, width, -1))
        print("time taken:", time.time() - init_time)

        print("Generating masks...")
        init_time = time.time()
        sam.to(device=device)
        masks = mask_generator.generate(image)
        mask_points = []
        mask_images = []
        mask_preprocessed_images = []
        extra_classes, num = get_extra_classes()
        if extra_classes is not None:
            extra_masks, boxes = clip_with_owl(image, num, sam, obj_name=extra_classes[0], visualize=False)
            for mask, box in zip(extra_masks, boxes):
                tmp_mask = {'segmentation': mask, 'bbox': box, 'area': mask.sum()}
                masks.append(tmp_mask)
            if extra_classes not in considered_classes:
                considered_classes.append(extra_classes)
        masks = sorted(masks, key=(lambda x: x['area']), reverse=True)
        del sam
        del mask_generator
        print("time taken:", time.time() - init_time)

        print("Processing masks...")
        init_time = time.time()
        for i, mask in enumerate(masks):
            # if mask['area'] < 600:
            #     break
            duplicate = False
            for prev_mask in masks[:i]:
                intersection = np.sum(mask['segmentation'] & prev_mask['segmentation'])
                union = np.sum(mask['segmentation'] | prev_mask['segmentation'])
                if intersection / union > 0.1:
                    duplicate = True
                    break
            if duplicate:
                continue

            x, y, w, h = map(int, mask['bbox'])
            print(x, y, w, h)
            # bbox = (y,x,y+h,x+w)
            # project(used_camera_idx,bbox,max_num=300)
            background = 255
            im = np.ones((h, w, 3), dtype=np.uint8) * background
            mask_points.append(coord2point_3d[mask['segmentation']])
            new_mask = mask['segmentation'][y:min(y+h, height), x:min(x+w, width)]
            print(new_mask.shape)
            im[new_mask] = image[y:min(y+h, height), x:min(x+w, width), :][new_mask]
            square_size = max(h, w) + 6
            im = np.pad(im, (((square_size - h) // 2, (square_size - h) // 2), ((square_size - w) // 2, (square_size - w) // 2), (0, 0)), 'constant', constant_values=background)
            # im = image[y:y+h, x:x+w, :]
            mask_images.append(im)
            mask_preprocessed_images.append(clip_preprocess(Image.fromarray(im)).to(device))
        mask_preprocessed_images = torch.stack(mask_preprocessed_images)
        print("time taken:", time.time() - init_time)

        print("Getting labels with CLIP...")
        init_time = time.time()
        clip_model.to(device=device)
        with torch.no_grad(), torch.cuda.amp.autocast():
            image_features = clip_model.encode_image(mask_preprocessed_images)
            image_features /= image_features.norm(dim=-1, keepdim=True)
            raw_probs = (100.0 * image_features @ text_features.T).softmax(dim=-1).cpu()
            print(image_features.shape, text_features.shape)
        clip_model.to(device="cpu")
        image_features = image_features.cpu().numpy()
        print("time taken:", time.time() - init_time)

        print("Compute DINOv2 features...")
        init_time = time.time()
        for i in range(len(mask_images)):
            mask_images[i] = cv2.resize(mask_images[i], (224, 224))
        mask_dino_preprocessed_images = torch.from_numpy(np.transpose(np.stack(mask_images), (0, 3, 1, 2))).to(device)
        dinov2_model.to(device=device)
        dino_image_features = dinov2_model.forward_features(dinov2_preprocess(mask_dino_preprocessed_images))["x_norm_clstoken"]
        dino_image_features = F.normalize(dino_image_features, dim=-1)
        dinov2_model.to(device="cpu")
        dino_image_features = dino_image_features.cpu().numpy()
        print("time taken:", time.time() - init_time)

        print("Postprocessing raw CLIP probabilities...")
        init_time = time.time()
        len_list = [len(words) for words in considered_classes + other_classes]
        probs = torch.zeros(*raw_probs.shape[:1], len(considered_classes) + len(other_classes), dtype=raw_probs.dtype, device=raw_probs.device)
        for i in range(len(considered_classes)):
            probs[..., i] = (raw_probs[..., sum(len_list[:i]) : sum(len_list[:i + 1])].max(dim=-1).values)
        for i in range(len(other_classes)):
            probs[..., i + len(considered_classes)] = (raw_probs[..., sum(len_list[:len(considered_classes) + i]) : sum(len_list[:len(considered_classes) + i + 1])].max(dim=-1).values)
        del raw_probs
        del clip_model
        del clip_preprocess
        print("time taken:", time.time() - init_time)

        print("Filtering detection results...")
        init_time = time.time()
        probs = probs.numpy()
        for i, p in enumerate(probs):
            predicted_class_idx = np.argmax(p)
            predicted_class_name = considered_classes[predicted_class_idx][0] if predicted_class_idx < len(considered_classes) else other_classes[predicted_class_idx - len(considered_classes)][0]
            if detected_objs.get(predicted_class_name) is not None:
                # detected_objs[predicted_class_name].append((mask_points[i], p[predicted_class_idx], norm_vec[i]))
                detected_objs[predicted_class_name].append((mask_points[i], p[predicted_class_idx], image_features[i], dino_image_features[i]))
            if not os.path.exists(f"cache/{predicted_class_idx}_{predicted_class_name}"):
                os.makedirs(f"cache/{predicted_class_idx}_{predicted_class_name}")
            cv2.imwrite(f"cache/{predicted_class_idx}_{predicted_class_name}/{i}_{p[predicted_class_idx]:.4f}.png", mask_images[i][:, :, ::-1])
        print("time taken:", time.time() - init_time)

        if save_to_cache:
            pickle.dump(bgr_images, open_file(f"cache/{TASK}/bgr_images.pkl", "wb"))
            o3d.io.write_point_cloud(f"cache/{TASK}/pcd_merged.pcd", pcd_merged)
            pickle.dump(detected_objs, open_file(f"cache/{TASK}/detected_objs.pkl", "wb"))

    delete_file(f"log/{TASK}/bgr_images.pkl")
    delete_file(f"log/{TASK}/pcd_merged.pcd")
    delete_file(f"log/{TASK}/detected_objs.pkl")
    pickle.dump(bgr_images, open_file(f"log/{TASK}/bgr_images.pkl", "wb"))
    o3d.io.write_point_cloud(f"log/{TASK}/pcd_merged.pcd", pcd_merged)
    pickle.dump(detected_objs, open_file(f"log/{TASK}/detected_objs.pkl", "wb"))

    if visualize:
        print("Plotting point cloud...")
        init_time = time.time()
        mesh_frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=0.1, origin=[0, 0, 0])
        detected_objects_geomertries = o3d.geometry.TriangleMesh()
        for key, detected_points in detected_objs.items():
            for pcd, _, _ in detected_points:
                sphere = o3d.geometry.TriangleMesh.create_sphere(radius=RADIUS)
                sphere.translate(np.mean(pcd, axis=0))
                detected_objects_geomertries += sphere
        detected_objects_geomertries.paint_uniform_color([0.0, 1.0, 0.0])
        o3d.visualization.draw_geometries([pcd_merged, mesh_frame, detected_objects_geomertries])
        print("time taken:", time.time() - init_time)

def get_text_features(clip_candidates, clip_model):
    global tokenizer
    try:
        text_features = pickle.load(open(f"cache/{TASK}/text_features.pkl", "rb"))
        return text_features
    except:
        if tokenizer is None:
            tokenizer = open_clip.get_tokenizer('ViT-g-14')
        texts = tokenizer([f"A photo of {c}." for c in clip_candidates]).to(device)
        clip_model.to(device=device)
        with torch.no_grad(), torch.cuda.amp.autocast():
            text_features = clip_model.encode_text(texts)
            text_features /= text_features.norm(dim=-1, keepdim=True)
        clip_model.to(device="cpu")
        pickle.dump(text_features, open_file(f"cache/{TASK}/text_features.pkl", "wb"))
        return text_features

# ---------------------------------- Get task pose ----------------------------------

def _get_task_pose(task_name, visualize=True):
    task_name = task_name[:-1] if task_name[-1] == '.' else task_name
    whole_prompt = prompt_get_task_pose + '\n' + '\n' + f"INPUT: the robot's task is to {task_name.lower()}." + "\n"
    add_to_log(whole_prompt)
    response = query_LLM(whole_prompt, ["INPUT:"], "cache/llm_response_get_task_pose.pkl")
    if "response" in response.text.lower():
        response_json = response.text.split('\n', 1)[1]
        add_to_log(response_json)
        pose_dict = json.loads(response_json)
    else:
        pose_dict = json.loads(response.text)
    add_to_log('-'*80)
    add_to_log('*at get_task_pose*')
    add_to_log('get_task_pose:', pose_dict)
    add_to_log('-'*80)
    pos_str = pose_dict['gripper position']
    ori_str = pose_dict['gripper orientation']
    _get_task_detection(task_name)
    pos, ori = get_pose_from_str(pos_str, ori_str, visualize=visualize)
    return pos, ori

def get_pose_from_str(pos_str, ori_str, visualize=False):
    global queried_obj
    pos_feature = pos_str.split('.')[1]
    obj_name = pos_str.split('.')[0]
    if 'absolute_vertical' in ori_str or 'absolute_horizontal' in ori_str:
        ori_feature = ori_str
        whole_prompt = prompt_get_pose_from_str + '\n' + "# Query: " + f"""the {pos_feature} of "{obj_name}"."""
    else:
        ori_feature = ori_str.split('.')[1]
        whole_prompt = prompt_get_pose_from_str + '\n' + "# Query: " + f"""the {pos_feature} and {ori_feature} of "{obj_name}"."""
    response = query_LLM(whole_prompt, ["INPUT:"], "cache/llm_response_get_pose_from_str.pkl")
    code_as_policies = response.text
    add_to_log('-'*80)
    add_to_log('*at get_pose_from_str*')
    add_to_log(code_as_policies)
    add_to_log('-'*80)
    locals = {}
    exec(code_as_policies, globals(), locals)
    obj = locals['ret_val']
    if ori_feature == 'absolute_vertical':
        if 'drawer' in obj_name:
            vertical_quat = get_vertical_ori()
            minor_axis = get_minor_axis_from_pcd(obj.pcd)
            rotate_angle = math.degrees(-math.atan2(minor_axis[1], minor_axis[0]))
            tar_quat = policy.rotate_around_gripper_z_axis(rotate_angle, vertical_quat)
            obj.rotation = tar_quat
            obj.vec = minor_axis
        elif 'shelf' in obj_name:
            obj.rotation = get_current_ori()
            minor_axis = get_minor_axis_from_pcd(obj.pcd)
            obj.vec = minor_axis
        else:
            vertical_quat = get_vertical_ori()
            main_axis = get_x_axis_from_pcd(obj.pcd)
            rotate_angle = math.degrees(-math.atan2(main_axis[1], main_axis[0]))
            tar_quat = policy.rotate_around_gripper_z_axis(rotate_angle, vertical_quat)
            obj.rotation = tar_quat
            obj.vec = main_axis
    elif ori_feature == 'absolute_horizontal':
        obj.rotation = get_horizontal_ori(get_vec=False)
        obj.vec = get_horizontal_ori(get_vec=True)
    else:
        obj.vec = obj.rotation.copy()
        obj.rotation = policy.align_z_axis_with_vector(obj.vec)
    if REAL_ROBOT:
        pcd_merged = o3d.io.read_point_cloud(f"log/{TASK}/pcd_merged.pcd")
        mesh_frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=0.1, origin=[0, 0, 0])
        visualized_geometries = [pcd_merged, mesh_frame]
        for i in range(50):
            sphere = o3d.geometry.TriangleMesh.create_sphere(radius=RADIUS/3)
            sphere.paint_uniform_color([0.0, 1.0, 0.0])
            sphere.translate(calculate_centroid(obj.pcd)+i/100*obj.vec)
            visualized_geometries.append(sphere)
        o3d.visualization.draw_geometries(visualized_geometries)
    obj.position = policy.robot_fingertip_pos_to_ee(obj.position, obj.rotation)
    obj.bounding_box = get_bounding_box_from_pcd(obj.pcd, obj.vec)
    queried_obj.bounding_box = obj.bounding_box
    queried_obj.vec = obj.vec
    if queried_obj.detected_pose is None:
        queried_obj.clip_feature = obj.clip_feature.copy()
        queried_obj.dino_feature = obj.dino_feature.copy()
        queried_obj.detected_pose = (obj.position, obj.rotation)
    return obj.position, obj.rotation

def calculate_centroid(point_cloud):
    centroid = np.mean(point_cloud, axis=0)
    return centroid

def calculate_normal_vector(point_cloud):
    cov_matrix = np.cov(point_cloud, rowvar=False)
    eigenvalues, eigenvectors = np.linalg.eigh(cov_matrix)
    min_eigenvalue_index = np.argmin(eigenvalues)
    normal_vector = eigenvectors[:, min_eigenvalue_index]
    if normal_vector[0]<0:
        normal_vector = -normal_vector
    normal_vector[2] = 0.
    normal_vector /= np.linalg.norm(normal_vector)
    return normal_vector

def calculate_major_axis(point_cloud):
    covariance_matrix = np.cov(point_cloud, rowvar=False)
    eigenvalues, eigenvectors = np.linalg.eigh(covariance_matrix)
    main_axis = eigenvectors[:, np.argmax(eigenvalues)]
    return main_axis
    # return calculate_normal_vector(point_cloud)

def calculate_minor_axis(point_cloud):
    return calculate_normal_vector(point_cloud)

def calculate_tail_point(point_cloud):
    x_axis = get_x_axis_from_pcd(point_cloud)
    centroid = calculate_centroid(point_cloud)
    tail_point = centroid - 0.5 * x_axis
    return tail_point

def get_x_axis_from_pcd(point_cloud):
    if REAL_ROBOT:
        o3d_point_cloud = o3d.geometry.PointCloud()
        o3d_point_cloud.points = o3d.utility.Vector3dVector(point_cloud)
        bbox = o3d_point_cloud.get_oriented_bounding_box()
        eight_vertices = np.asarray(bbox.get_box_points())
        sorted_indices = np.argsort(eight_vertices[:,2])[::-1]
        sorted_array_z = eight_vertices[sorted_indices]
        sorted_array_z = sorted_array_z[:4,:]
        len_xy = np.linalg.norm(sorted_array_z[0]-sorted_array_z[1])
        axis_xy = sorted_array_z[0]-sorted_array_z[1]
        len_yx = np.linalg.norm(sorted_array_z[0]-sorted_array_z[2])
        axis_yx = sorted_array_z[0]-sorted_array_z[2]
        if len_xy > len_yx:
            main_axis = axis_xy
        else:
            main_axis = axis_yx
        main_axis[2] = 0.
        main_axis /= np.linalg.norm(main_axis)
        main_axis = main_axis if main_axis[0]>0 else -main_axis
        return main_axis
    else:
        return np.array((1.,0.,0.))

def get_minor_axis_from_pcd(point_cloud):
    if REAL_ROBOT:
        o3d_point_cloud = o3d.geometry.PointCloud()
        o3d_point_cloud.points = o3d.utility.Vector3dVector(point_cloud)
        bbox = o3d_point_cloud.get_oriented_bounding_box()
        eight_vertices = np.asarray(bbox.get_box_points())
        sorted_indices = np.argsort(eight_vertices[:,2])[::-1]
        sorted_array_z = eight_vertices[sorted_indices]
        sorted_array_z = sorted_array_z[:4,:]
        len_xy = np.linalg.norm(sorted_array_z[0]-sorted_array_z[1])
        axis_xy = sorted_array_z[0]-sorted_array_z[1]
        len_yx = np.linalg.norm(sorted_array_z[0]-sorted_array_z[2])
        axis_yx = sorted_array_z[0]-sorted_array_z[2]
        if len_xy > len_yx:
            minor_axis = axis_yx
        else:
            minor_axis = axis_xy
        minor_axis[2] = 0.
        minor_axis /= np.linalg.norm(minor_axis)
        minor_axis = minor_axis if minor_axis[0]>0 else -minor_axis
        return minor_axis
    else:
        return (1., 0., 0.)

def get_bounding_box_from_pcd(point_cloud, forward_vec):
    if REAL_ROBOT:
        o3d_point_cloud = o3d.geometry.PointCloud()
        o3d_point_cloud.points = o3d.utility.Vector3dVector(point_cloud)
        bbox = o3d_point_cloud.get_oriented_bounding_box()
        eight_vertices = np.asarray(bbox.get_box_points())
        sorted_indices = np.argsort(eight_vertices[:,2])[::-1]
        sorted_array_z = eight_vertices[sorted_indices]
        len_z = sorted_array_z[:4,2].mean() - sorted_array_z[4:,2].mean()
        sorted_array_z = sorted_array_z[:4,:]
        len_xy = np.linalg.norm(sorted_array_z[0]-sorted_array_z[1])
        axis_xy = sorted_array_z[0]-sorted_array_z[1]
        len_yx = np.linalg.norm(sorted_array_z[0]-sorted_array_z[2])
        axis_yx = sorted_array_z[0]-sorted_array_z[2]
        if len_xy > len_yx:
            main_axis = axis_xy
            len_x = len_xy
            len_y = len_yx
        else:
            main_axis = axis_yx
            len_x = len_yx
            len_y = len_xy
        main_axis[2] = 0.
        main_axis /= np.linalg.norm(main_axis)
        minor_axis = np.array((main_axis[1], -main_axis[0], 0.))
        if np.abs(forward_vec.dot(main_axis)) > np.abs(forward_vec.dot(minor_axis)):
            return (round(len_x*100), round(len_y*100), round(len_z*100))
        else:
            return (round(len_y*100), round(len_x*100), round(len_z*100))
    else:
        return (1., 8., 3.)

# ---------------------------------- Image feature ----------------------------------

def compare_text_image_sim(text, vis_feature):
    global clip_model, tokenizer
    vis_feature = torch.tensor(vis_feature, device=device)
    if tokenizer is None:
        tokenizer = open_clip.get_tokenizer('ViT-g-14')
    if clip_model is None:
        clip_model, _, _ = open_clip.create_model_and_transforms('ViT-g-14', pretrained='laion2b_s34b_b88k')
    clip_model = clip_model.to(device=device)
    text = tokenizer(f"A photo of {text}.").to(device=device)
    with torch.no_grad(), torch.cuda.amp.autocast():
        text_feature = clip_model.encode_text(text).squeeze()
        text_feature /= text_feature.norm(dim=-1, keepdim=True)
        vis_feature /= vis_feature.norm(dim=-1, keepdim=True)
        raw_probs = vis_feature.dot(text_feature).item()
    del clip_model
    return raw_probs

def get_detected_feature():
    global queried_obj
    if REAL_ROBOT:
        image_feature = queried_obj.clip_feature.copy()
        dino_image_feature = queried_obj.dino_feature.copy()
        detected_pose = queried_obj.detected_pose
        return (image_feature, dino_image_feature, detected_pose)
    else:
        return (np.random.randn(1024), np.random.randn(512), (np.array((1.,0.,0.)), np.array((1.,0.,0.,0.))))

def compare_feature(img_feature, task_name, visualize=False, threshold=10, metric='L2'):
    """ Delete all irrelevant features """
    global load_detected_objs
    obj_name = get_obj_name_from_task(task_name)
    obj_name = parse_obj_name(obj_name)
    objs_specified = load_detected_objs[obj_name]
    similarity = []
    for _, _, _, new_dino_feature in objs_specified:
        if metric == 'L2':
            similarity.append(((new_dino_feature-img_feature)**2).sum())
        elif metric == 'cosine':
            similarity.append(1-new_dino_feature.dot(img_feature)/(np.linalg.norm(new_dino_feature)*np.linalg.norm(img_feature)))
        else:
            raise NotImplementedError
    add_to_log(similarity)
    order = np.argsort(np.array(similarity))
    picked_obj_idx = list(order[:threshold])
    picked_obj = [objs_specified[i] for i in picked_obj_idx]
    load_detected_objs[obj_name] = picked_obj
    if visualize:
        pcd_merged = o3d.io.read_point_cloud(f'log/{TASK}/pcd_merged.pcd')
        mesh_frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=0.1, origin=[0, 0, 0])
        detected_objects_geomertries = o3d.geometry.TriangleMesh()
        for obj in picked_obj:
            pcd, _, _, _ = obj
            p = calculate_centroid(pcd)
            sphere = o3d.geometry.TriangleMesh.create_sphere(radius=RADIUS)
            sphere.translate(p)
            detected_objects_geomertries += sphere
            detected_objects_geomertries.paint_uniform_color([0.0, 1.0, 0.0])
        o3d.visualization.draw_geometries([pcd_merged, mesh_frame, detected_objects_geomertries])

# ---------------------------------- Others ----------------------------------

def _update_object_state(ini_state, task_name):
    with open(prompt_update_obj_state_file, "r") as template_file:
        template_content = template_file.read()
    values = {"initial_state": ini_state, "task_name": task_name}
    prompt = template_content.format(**values)
    response = query_LLM(prompt, [], "cache/llm_update_obj_name.pkl")
    new_obj_state = response.text
    return new_obj_state

def segment_image(image):
    from segment_anything import SamAutomaticMaskGenerator, sam_model_registry
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    print("Loading SAM model...")
    init_time = time.time()
    sam = sam_model_registry["vit_h"](checkpoint="cache/sam_vit_h_4b8939.pth")
    mask_generator = SamAutomaticMaskGenerator(
        model=sam,
        points_per_side=32,
        pred_iou_thresh=0.95,
        stability_score_thresh=0.95,
        box_nms_thresh=0.2,
        crop_n_layers=1,
        crop_overlap_ratio=0.6,
        min_mask_region_area=100
    )
    print("time taken:", time.time() - init_time)
    print("Generating masks...")
    sam.to(device="cuda")
    masks = mask_generator.generate(image)
    masks = sorted(masks, key=(lambda x: x['area']), reverse=False)
    print("Processing masks...")
    init_time = time.time()
    mask_images = []
    for i, mask in enumerate(masks):
        duplicate = False
        for prev_mask in masks[:i]:
            intersection = np.sum(mask['segmentation'] & prev_mask['segmentation'])
            union = np.sum(mask['segmentation'] | prev_mask['segmentation'])
            if intersection / union > 0.1:
                duplicate = True
                break
        if duplicate:
            continue
        x, y, w, h = map(int, mask['bbox'])
        background = 255
        im = np.ones((h, w, 3), dtype=np.uint8) * background
        new_mask = mask['segmentation'][y:y+h, x:x+w]
        im[new_mask] = image[y:y+h, x:x+w, :][new_mask]
        im = np.pad(im, ((3, 3), (3, 3), (0, 0)), 'constant', constant_values=background)
        # im = image[y:y+h, x:x+w, :]
        mask_images.append(im)

def save_current_image(directory_path, visualize=False):
    if REAL_ROBOT:
        pass

def test_image(text):
    clip_model, _, clip_preprocess = open_clip.create_model_and_transforms('ViT-g-14', pretrained='laion2b_s34b_b88k')
    files = ['cache/image_for_retrieval/cup1.png', 'cache/image_for_retrieval/cup2.png', 'cache/image_for_retrieval/cup3.png', 'cache/image_for_retrieval/cup4.png', 'cache/image_for_retrieval/cup5.png']
    images = [clip_preprocess(Image.open(file)).to(device).to(torch.float) for file in files]
    preprocessed_images = torch.stack(images)
    clip_model = clip_model.to(device)
    with torch.no_grad(), torch.cuda.amp.autocast():
        image_features = clip_model.encode_image(preprocessed_images)
    tokenizer = open_clip.get_tokenizer('ViT-g-14')
    text = tokenizer(f"A photo of {text}.").to(device)
    with torch.no_grad(), torch.cuda.amp.autocast():
        text_feature = clip_model.encode_text(text).squeeze()
        text_feature /= text_feature.norm(dim=-1, keepdim=True)
    for i in image_features:
        i /= i.norm(dim=-1, keepdim=True)
        raw_probs = i.dot(text_feature).item()
        print(raw_probs)
    return raw_probs


import heapq
import random
import time
from typing import Any, Callable, Iterator, List, NamedTuple, Set, Tuple

from ..core.models import Groups, UserKey

Genome = List[List[UserKey]]
Rating = float
RatingFunc = Callable[[Genome, Any], Rating]

class RatedGenome(NamedTuple):
    """A rating together with its genome."""

    rating: Rating
    genome: Genome

Population = List[RatedGenome]

def build_genome(users: List[UserKey], sizes: List[int]) -> Genome:
    """Build new genome from list of users."""
    genome = []
    pos = 0
    for size in sizes:
        genome.append(users[pos:pos+size])
        pos += size
    return genome

def iter_genome_count(genome: Genome, count: int) -> Iterator[UserKey]:
    """Iterate over all members until number of member is reached."""
    for group in genome:
        for member in group:
            if count < 1:
                return
            yield member
            count -= 1

def iter_genome_set(genome: Genome, ignore_members: Set[UserKey]) -> Iterator[UserKey]:
    """Iterate over all members that will not be ignored."""
    for group in genome:
        for member in group:
            if member not in ignore_members:
                yield member

def crossover_genomes(
        genome_1: Genome, genome_2: Genome, user_count: int) -> Tuple[bool, Genome]:
    """Combine two genomes to a new one."""
    if genome_1 == genome_2 or len(genome_1) < 2:
        return (False, genome_1)
    split_pos = random.randint(1, user_count - 1)  # nosec
    users = list(iter_genome_count(genome_1, split_pos))
    users.extend(iter_genome_set(genome_2, set(users)))
    genome = build_genome(users, [len(group) for group in genome_1])
    return (True, genome)

def mutate_genome(genome: Genome) -> Genome:
    """Return a new genome value that is a permutation of the given one."""
    group_1 = random.randint(0, len(genome) - 1)  # nosec
    group_2 = random.randint(0, len(genome) - 1)  # nosec
    if group_1 == group_2:
        group_2 = (group_1 + 1) % len(genome)

    pos_1 = random.randint(0, len(genome[group_1]) - 1)  # nosec
    user_1 = genome[group_1][pos_1]
    pos_2 = random.randint(0, len(genome[group_2]) - 1)  # nosec
    user_2 = genome[group_2][pos_2]
    new_genome = []
    for number, group in enumerate(genome):
        if number == group_1:
            new_group = list(group)
            new_group[pos_1] = user_2
            new_genome.append(new_group)
        elif number == group_2:
            new_group = list(group)
            new_group[pos_2] = user_1
            new_genome.append(new_group)
        else:
            new_genome.append(group)
    return new_genome

def build_initial_population_genomes(
        population_size: int, users: List[UserKey], sizes: List[int]) -> List[Genome]:
    """Build an initial population of genomes."""
    user_list = list(users)
    result = []
    for _ in range(population_size):
        random.shuffle(user_list)
        result.append(build_genome(user_list, sizes))
    return result

def build_initial_population(
        population_size: int, users: List[UserKey], sizes: List[int],
        rating_func: RatingFunc, rating_data: Any) -> Population:
    """Build an initial population, ordered by rating."""
    genome_list = build_initial_population_genomes(population_size, users, sizes)
    population = [
        RatedGenome(rating_func(genome, rating_data), genome)
        for genome in genome_list]
    heapq.heapify(population)
    return population

def mutate_population(
        population: Population, num_mutation: int) -> List[Tuple[Genome, Rating]]:
    """Perform a number of mutations on given population."""
    result = []
    for _ in range(num_mutation):
        rated_genome = random.choice(population)  # nosec
        mutated_genome = mutate_genome(rated_genome.genome)
        result.append((mutated_genome, rated_genome.rating))
    return result

def crossover_population(
        population: Population, num_crossover: int, user_count: int) -> List[Genome]:
    """Perform a number of crossover operations on given population."""
    result = []
    for _ in range(num_crossover):
        parent_genomes = random.sample(population, 2)
        gave_birth, child_genome = crossover_genomes(
            parent_genomes[0].genome, parent_genomes[1].genome, user_count)
        if gave_birth:
            result.append(child_genome)
    return result

def reduce_population(population: Population, population_size: int) -> None:
    """Reduce the population to be of given size."""
    while len(population) > population_size:
        genome_1_pos = random.randint(0, len(population) - 1)  # nosec
        genome_2_pos = random.randint(0, len(population) - 1)  # nosec
        if genome_1_pos == genome_2_pos:
            continue

        if population[genome_1_pos].rating < population[genome_2_pos].rating:
            del population[genome_2_pos]
        else:
            del population[genome_1_pos]

def trivial_groups(users: List[UserKey], num_groups) -> Groups:
    """There is only one groups, so the result is the trivial one."""
    result = [frozenset(users)]
    while len(result) < num_groups:
        result.append(frozenset())
    return tuple(result)

def best_groups(population: Population, num_groups: int) -> Groups:
    """Select the best genome and produces resulting groups."""
    result = [frozenset(group) for group in population[0].genome]
    while len(result) < num_groups:
        result.append(frozenset())
    return tuple(result)

class StopStrategy:  # pylint: disable=too-many-instance-attributes
    """Encapsulate the descision when to stop the algorithm."""

    def __init__(
            self, initial_rating: float,
            max_rounds: int, max_best_rounds: int, max_seconds: float):
        """Initialize the object."""
        self.initial_rating: float = initial_rating
        self.max_rounds: int = max_rounds
        self.max_best_rounds: int = max_best_rounds
        self.max_seconds: float = max_seconds
        self.start()

    def start(self) -> None:
        """Signal that algorithm starts."""
        self.best_rating: float = self.initial_rating
        self.best_rounds: int = 0
        self.rounds: int = 0
        self.time_start: float = time.monotonic()

    def should_continue(self, rating: float) -> bool:
        """Return True if algorithm should continue its work."""
        self.rounds += 1
        if self.max_rounds <= self.rounds:
            return False
        if rating < self.best_rating:
            self.best_rating = rating  # pylint: disable=attribute-defined-outside-init
            self.best_rounds = 0  # pylint: disable=attribute-defined-outside-init
        else:
            self.best_rounds += 1
            if self.max_best_rounds < self.best_rounds:
                return False
        return (time.monotonic() - self.time_start) < self.max_seconds

def generic_genetic_policy(
        users: List[UserKey], sizes: List[int],
        rating_func: RatingFunc, rating_data: Any,
        stop_strategy: StopStrategy) -> Groups:
    """Calculate a group that suits the rating at best."""
    num_groups = len(sizes)
    while sizes[-1] < 1:
        sizes = sizes[:-1]
    if len(sizes) < 2:
        return trivial_groups(users, num_groups)

    population_size = min(100, len(sizes) * len(users) * 2)

    population = build_initial_population(
        population_size, users, sizes, rating_func, rating_data)

    num_crossover = len(users)
    num_mutation = len(users) + population_size

    stop_strategy.start()
    while stop_strategy.should_continue(population[0].rating):
        for genome in crossover_population(population, num_crossover, len(users)):
            population.append(RatedGenome(rating_func(genome, rating_data), genome))

        successful_mutations = 0
        for genome, old_rating in mutate_population(population, num_mutation):
            new_rating = rating_func(genome, rating_data)
            if new_rating < old_rating:
                successful_mutations += 1
            population.append(RatedGenome(new_rating, genome))

        reduce_population(population, population_size)
        heapq.heapify(population)

    return best_groups(population, num_groups)


from itertools import islice
from math import ceil
from instaloader import Instaloader, Profile
import time
import json

#Retrieves percentage amount of photos for a given a username. Max percentage is such that max number of
#photos to be downloaded is 60 (in main())
def retrieve_posts(username, num):
    PROFILE = username        # profile to download from

    L = Instaloader(download_videos = False, download_video_thumbnails = False, download_comments = True)

    profile = Profile.from_username(L.context, PROFILE)

    posts = profile.get_posts()

    posts_sorted_by_comments = sorted(posts,
                                   key=lambda p: p.comments,
                                   reverse=True)

    post_shortcodes = [x.shortcode for x in posts_sorted_by_comments]

    return post_shortcodes[:num]

#Cleans the numberings from the num_posts.txt (Removes commas and ks)
def clean_num(num_posts_all):
    num_posts_cleaned = []
    for num in num_posts_all:
        if "," in num:
            num_posts_cleaned.append(int(num.replace(",", "")))

        elif "k" in num:
            num  = num.replace("k", "")
            num = int(float(num)*1000)
            num_posts_cleaned.append(num)

        else:
            num_posts_cleaned.append(int(num))

    return num_posts_cleaned

#Main function
def main():
    #Starting Timer
    t0 = time.time()

    with open('../Data/usernames.txt') as f:
        usernames = f.readlines()

    usernames = [x.strip() for x in usernames]

    with open('../Data/num_posts.txt') as f:
        num_posts_all = f.readlines()

    num_posts_all = [x.strip() for x in num_posts_all]
    num_posts_all = clean_num(num_posts_all)

    #Sorting posts by comments and saving the list of posts
    for i in range(0, len(usernames)):
        username = usernames[i]
        if num_posts_all[i] <= 100:
            posts = retrieve_posts(username, num_posts_all[i])
        else:
            posts = retrieve_posts(username, 100)

        with open('../Data/Posts_list/' + username + ".json" , 'w') as f:
            json.dump(posts,f)

        print("Posts for %s added. | " %(username), end = "")
        print("Time elapsed: ", end = "")
        print(time.time()-t0)

    t1 = time.time()
    total_time = t1-t0

    #Ending the timer
    print("Total Time Elapsed", end = ' ')
    print(total_time)

#Calling Main function
if __name__== "__main__":
    main()


from app.auth import get_user
from app.controllers.games import GameController
from app.database import get_session
from app.exceptions import AlreadyExists
from app.models.users import User
from app.schemas.games import Game, GameCreate, GameFirstStageCreate, GameUpdate, GameGuestCreate, GameFirstStage
from app.schemas.names import Name
from app.schemas.users import User as UserSchema
from fastapi import APIRouter, Depends, HTTPException, status
from sqlalchemy.orm import Session
from sqlalchemy.orm.exc import NoResultFound
from typing import List
from uuid import UUID

router = APIRouter()

@router.get('', response_model=List[Game])
async def get_games(
    db: Session = Depends(get_session),
    user: User = Depends(get_user)
) -> List[Game]:
    """
    Get games list
    """
    return GameController.get_games(db, user)

@router.post('', status_code=status.HTTP_201_CREATED, response_model=Game)
async def create_game(
    payload: GameCreate,
    db: Session = Depends(get_session),
    user: User = Depends(get_user)
) -> Game:
    """
    Create a game
    """
    return GameController.create_game(db, payload, user)

@router.put('/{game_id}', status_code=status.HTTP_200_OK, response_model=Game)
async def update_game(
    game_id: UUID,
    payload: GameUpdate,
    db: Session = Depends(get_session),
    user: User = Depends(get_user)
) -> Game:
    """
    Update a game
    """
    try:
        return GameController.update_game(db, game_id, payload, user)
    except NoResultFound:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail='game not found')

@router.delete('/{game_id}', status_code=status.HTTP_204_NO_CONTENT)
async def delete_game(
    game_id: UUID,
    db: Session = Depends(get_session),
    user: User = Depends(get_user)
) -> None:
    """
    Delete a game
    """
    try:
        return GameController.delete_game(db, game_id, user)
    except NoResultFound:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail='game not found')

@router.get('/{game_id}/guests', status_code=status.HTTP_200_OK, response_model=List[UserSchema])
async def get_game_guests(
    game_id: UUID,
    db: Session = Depends(get_session),
    user: User = Depends(get_user)
) -> Game:
    """
    Get game guests list
    """
    try:
        return GameController.get_game_guests(db, game_id, user)
    except NoResultFound:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail='game not found')

@router.post('/{game_id}/guests', status_code=status.HTTP_200_OK, response_model=List[UserSchema])
async def add_game_guests(
    game_id: UUID,
    payload: GameGuestCreate,
    db: Session = Depends(get_session),
    user: User = Depends(get_user)
) -> Game:
    """
    Add a guest to game
    """
    try:
        return GameController.create_game_guest(db, game_id, payload.user_id, user)
    except NoResultFound:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail='game not found')
    except AlreadyExists as e:
        raise HTTPException(status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, detail=[
            {
                'loc': ['body', e.fields[0]],
                'msg': "this guest already exists",
                'type': 'type_error.already_exists',
            }
        ])

@router.delete('/{game_id}/guests/{user_id}', status_code=status.HTTP_200_OK, response_model=List[UserSchema])
async def remove_game_guests(
    game_id: UUID,
    user_id: UUID,
    db: Session = Depends(get_session),
    user: User = Depends(get_user)
) -> Game:
    """
    Remove a guest from game
    """
    try:
        return GameController.delete_game_guest(db, game_id, user_id, user)
    except NoResultFound:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail='game not found')

@router.post('/{game_id}/stage-1', status_code=status.HTTP_201_CREATED, response_model=GameFirstStage)
async def create_first_stage(
    game_id: UUID,
    payload: GameFirstStageCreate,
    db: Session = Depends(get_session),
    user: User = Depends(get_user)
) -> Game:
    """
    Post a choice for the first stage
    """
    try:
        return GameController.create_first_stage(db, game_id, payload, user)
    except NoResultFound:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail='game not found')
    except AlreadyExists as e:
        raise HTTPException(status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, detail=[
            {
                'loc': ['body', e.fields[0]],
                'msg': "this choice already exists",
                'type': 'type_error.already_exists',
            }
        ])

@router.get('/{game_id}/stage-1/next', status_code=status.HTTP_200_OK, response_model=Name)
async def get_first_stage_next(
    game_id: UUID,
    db: Session = Depends(get_session),
    user: User = Depends(get_user)
) -> Game:
    """
    Get the next name for the first stage
    """
    try:
        return GameController.get_first_stage_next(db, game_id, user)
    except NoResultFound:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail='game not found')

@router.get('/{game_id}/stage-1/result', status_code=status.HTTP_200_OK, response_model=List[Name])
async def get_first_stage_result(
    game_id: UUID,
    db: Session = Depends(get_session),
    user: User = Depends(get_user)
) -> Game:
    """
    Get the result of the first stage of a game
    """
    try:
        return GameController.get_first_stage_result(db, game_id, user)
    except NoResultFound:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail='game not found')


import torch
from torch.nn import Module

class ContrastLoss(Module):
    def __init__(self, temperature=0.1):
        super(ContrastLoss, self).__init__()
        self.temperature = temperature

    def forward(self, feat):
        anchors, positives, negatives = feat['anchor'], feat['positive'], feat['negative']
        anchor = anchors[:, 0]
        loss_1 = self.loss(anchor, positives, negatives)
        positive = positives[:, 0]
        loss_2 = self.loss(positive, anchors, negatives)
        return loss_1 + loss_2

    def loss(self, anchor, positives, negatives):
        pos = torch.exp(torch.bmm(positives, anchor[..., None]) / self.temperature).squeeze(-1)
        neg = torch.exp(torch.bmm(negatives, anchor[..., None]) / self.temperature).squeeze(-1)
        dominator = pos.sum(dim=1) + neg.sum(dim=1)
        loss = - torch.log(pos).mean(dim=1) + torch.log(dominator)
        return loss.mean()


import numpy as np
from PIL import Image

def read_pixel_data(png_path):
    im = Image.open(png_path)
    pixels = np.array(im.getdata())
    width, height = im.size
    pixels = pixels.reshape(height, width, 4)
    return pixels


"""empty message

Revision ID: 258421aeb463
Revises:
Create Date: 2023-06-19 10:14:42.127886

"""
import sqlalchemy as sa
from alembic import op

# revision identifiers, used by Alembic.
revision = "258421aeb463"
down_revision = None
branch_labels = None
depends_on = None

def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table(
        "nonebot_plugin_follow_withdraw_originmessage",
        sa.Column("message_id", sa.String(), nullable=False),
        sa.Column("adapter_name", sa.String(), nullable=False),
        sa.Column("channel_id", sa.String(), nullable=True),
        sa.PrimaryKeyConstraint("message_id"),
    )
    op.create_table(
        "nonebot_plugin_follow_withdraw_followmessage",
        sa.Column("message_id", sa.String(), nullable=False),
        sa.Column("adapter_name", sa.String(), nullable=False),
        sa.Column("channel_id", sa.String(), nullable=True),
        sa.Column("origin_message_id", sa.String(), nullable=False),
        sa.ForeignKeyConstraint(
            ["origin_message_id"],
            ["nonebot_plugin_follow_withdraw_originmessage.message_id"],
        ),
        sa.PrimaryKeyConstraint("message_id"),
    )
    # ### end Alembic commands ###

def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_table("nonebot_plugin_follow_withdraw_followmessage")
    op.drop_table("nonebot_plugin_follow_withdraw_originmessage")
    # ### end Alembic commands ###


from rest_auth.registration.serializers import RegisterSerializer
from rest_framework import serializers, exceptions
from rest_framework.exceptions import ValidationError
from django.contrib.auth import get_user_model, authenticate
from rest_auth.serializers import TokenSerializer, LoginSerializer
from rest_auth.models import TokenModel
from .models import DealerDetail, ClientDetail, CustomUser

class CustomRegisterSerializer(RegisterSerializer):
    username = serializers.CharField(required=True)
    email = serializers.EmailField(required=True)
    password1 = serializers.CharField(write_only=True)
    number = serializers.CharField(required=True)

    user_type = serializers.ChoiceField(
        choices=(
            ('Client', 'Client'),
            ('Dealer', 'Dealer'),

        ),
        style={'base_template': 'radio.html'},
        required=True, write_only=True)

    def get_cleaned_data(self):
        super(CustomRegisterSerializer, self).get_cleaned_data()
        print(self.validated_data.get('number', ''))
        return {
            'password1': self.validated_data.get('password1', ''),
            'email': self.validated_data.get('email', ''),
            'username': self.validated_data.get('username', ''),
            'user_type': self.validated_data.get('user_type', ''),
            'number': self.validated_data.get('number', ''),
        }

class DealerProfileSerializer(serializers.ModelSerializer):
    dealer_name=serializers.CharField(read_only=True)

    class Meta:
        model = DealerDetail
        fields = ('id', 'dealer_name','extra_info', 'thumbnail',
                  'image', 'latitude', 'longitude', 'on_hour','off_hour','has_bike'
                  ,'add_1','add_2','area','city')
        read_only_fields = ('thumbnail', )
        #depth=1

class ClientProfileSerializer(serializers.ModelSerializer):

    class Meta:
        model = ClientDetail
        fields = '__all__'
        read_only_fields = ('thumbnail', )
        # depth=1

class CheckSerializer(serializers.ModelSerializer):
    class Meta:
        model = CustomUser
        fields = ("user_type",)

class TokenSerializer(serializers.ModelSerializer):
    """
    Serializer for Token model.
    """
    user = CheckSerializer(
        many=False, read_only=True)  # this is add by myself.

    class Meta:
        model = TokenModel
        fields = ('key', 'user')


import torch.nn as nn
from collections import namedtuple
import torch
from torchvision import models

class Vgg16(torch.nn.Module):
    def __init__(self, requires_grad=False):
        super(Vgg16, self).__init__()
        vgg_pretrained_features = models.vgg16(weights=models.VGG16_Weights.DEFAULT).features
        self.slice1 = torch.nn.Sequential()
        self.slice2 = torch.nn.Sequential()
        self.slice3 = torch.nn.Sequential()
        self.slice4 = torch.nn.Sequential()
        for x in range(4):
            self.slice1.add_module(str(x), vgg_pretrained_features[x])
        for x in range(4, 9):
            self.slice2.add_module(str(x), vgg_pretrained_features[x])
        for x in range(9, 16):
            self.slice3.add_module(str(x), vgg_pretrained_features[x])
        for x in range(16, 23):
            self.slice4.add_module(str(x), vgg_pretrained_features[x])
        if not requires_grad:
            for param in self.parameters():
                param.requires_grad = False

    def forward(self, X):
        h = self.slice1(X)
        h_relu1_2 = h
        h = self.slice2(h)
        h_relu2_2 = h
        h = self.slice3(h)
        h_relu3_3 = h
        h = self.slice4(h)
        h_relu4_3 = h
        vgg_outputs = namedtuple("VggOutputs", ['relu1_2', 'relu2_2', 'relu3_3', 'relu4_3'])
        out = vgg_outputs(h_relu1_2, h_relu2_2, h_relu3_3, h_relu4_3)
        return out

class Alexnet(torch.nn.Module):
    def __init__(self, requires_grad=False):
        super(Alexnet, self).__init__()
        alex_pretrained_features = models.alexnet(weights=models.AlexNet_Weights.DEFAULT).features
        self.slice1 = torch.nn.Sequential()
        self.slice2 = torch.nn.Sequential()
        self.slice3 = torch.nn.Sequential()
        self.slice4 = torch.nn.Sequential()
        self.slice5 = torch.nn.Sequential()
        for x in range(2):
            self.slice1.add_module(str(x), alex_pretrained_features[x])
        for x in range(2, 5):
            self.slice2.add_module(str(x), alex_pretrained_features[x])
        for x in range(5, 8):
            self.slice3.add_module(str(x), alex_pretrained_features[x])
        for x in range(8, 10):
            self.slice4.add_module(str(x), alex_pretrained_features[x])
        for x in range(10, 12):
            self.slice5.add_module(str(x), alex_pretrained_features[x])
        if not requires_grad:
            for param in self.parameters():
                param.requires_grad = False

    def forward(self, X):
        h = self.slice1(X)
        h_relu1 = h
        h = self.slice2(h)
        h_relu2 = h
        h = self.slice3(h)
        h_relu3 = h
        h = self.slice4(h)
        h_relu4 = h
        h = self.slice5(h)
        h_relu5 = h
        alex_outputs = namedtuple("AlexnetOutputs", ['relu1', 'relu2', 'relu3', 'relu4', 'relu5'])
        out = alex_outputs(h_relu1, h_relu2, h_relu3, h_relu4, h_relu5)
        return out

class Lpips(nn.Module):
    def __init__(self, net, device):
        super(Lpips, self).__init__()
        self.net = net
        self.vgg = Vgg16(requires_grad=False).to(device)
        self.alex = Alexnet(requires_grad=False).to(device)
        self.criterion = nn.MSELoss()

    def forward(self, rec, label):
        if self.net == 'vgg16':
            features_rec = self.vgg(rec)
            features_label = self.vgg(label)
            lpips_loss = self.criterion(features_rec.relu2_2, features_label.relu2_2) + \
                         self.criterion(features_rec.relu4_3, features_label.relu4_3)

        if self.net == 'alex':
            features_rec = self.alex(rec)
            features_label = self.alex(label)
            lpips_loss = self.criterion(features_rec.relu2, features_label.relu2) + \
                         self.criterion(features_rec.relu4, features_label.relu4)

        return lpips_loss


import requests
import json
import pandas as pd
import streamlit as st
from os import getcwd, listdir
from json_normalize import json_transformation
from db_connection import cursor
from db_load import load_json_data
from duckdb_con import cursor

# Create a title
st.title("Pandascore API Data Retrieval")

# Encrypted Text Input
api_token = st.text_input("Paste in your Pandascore API Token: ", type='password')

# # Customization Text Inputs
# game_options = ['Counter Strike: Global Offensive', 'Call of Duty: Modern Warfare', 'Dota 2', 'FIFA', 'League of Legends', 'Overwatch', 'Public Underground Battle Ground (PUBG)', 'Rainbow Six Siege', 'Rocket League', 'Valorant', 'King of Glory', 'League of Legends: Wild Rift', 'Starcraft 2', 'Starcraft: Brood War']
# game_short_forms = ['csgo', 'codmw', 'dota2', 'fifa', 'lol', 'ow', 'pubg', 'r6siege', 'rl', 'valorant', 'kog', 'lol-wild-rift', 'starcraft-2', 'starcraft-brood-war']
game_df = pd.read_sql('games', con=cursor)

game_options = game_df['name'].tolist()
game_short_forms = game_df['short_name'].tolist()
game_api_map = {game: game_short_forms[i] for i, game in enumerate(game_options)}

game_name = st.selectbox("Choose a game: ", options=game_options)

endpoint_df = pd.read_sql('endpoints', con=cursor)
endpoint_options = endpoint_df['endpoint'].tolist()
# endpoint_options = ['leagues', 'matches', 'players', 'series', 'teams', 'tournaments']
endpoint_name = st.selectbox("Choose an endpoint: ", options=endpoint_options)

st.markdown("---")
st.markdown("<h3 style='text-align:center'> Query Parameters </h3>", unsafe_allow_html=True)

# Query Parameters
page_number = st.number_input("Choose a page number: ", min_value=1, max_value=10, value=1)
per_page = st.number_input("Choose number of entries per page: ", min_value=10, max_value=100, step=10)

# Create the Logic
data_path = getcwd() + "/data"

# Get JSON File Name
json_file_name = f"{game_name}_{endpoint_name}_{page_number}_{per_page}.json"

# Customized URL
url = f"https://api.pandascore.co/{game_api_map[game_name]}/{endpoint_name}?sort=&page={page_number}&per_page={per_page}&token={api_token}"

# Function to execute API call
def api_call(url):

    # Retrieve a response
    response = requests.get(url)

    # JSON Data Variable
    json_data = response.json()

    # Create context manager
    with open(f"{data_path}/{json_file_name}", "w") as my_file:
        json.dump(json_data, my_file)

    # Save to database
    load_json_data(game_name.lower(), endpoint_name, json_data, cursor)

    return url

# Create the button
if st.button("Retrieve Data"):
    if json_file_name in listdir(data_path):
        st.info(f"Your data is already saved in {json_file_name}")
    else:
        api_call(url)
        st.success(f"URL {url} data received successfully.")

# Create button for viewing the transformation
if st.button("View Transformation"):
    # Create context manager
    with open(f"{data_path}/{json_file_name}", "r") as my_file:
        result = json.load(my_file)

    result_df = json_transformation(result)

    st.dataframe(result_df)


"""
Routine for getting and setting the appropriate charges for atoms in the transition state. Can be called by isEE as
needed, or used individually from the command line by providing an appropriate settings.pkl file.
"""

import argparse
import mdtraj
import pickle
import sys
import time
import parmed
import re
import copy
from isee import main as iseemain
from isee.process import process

def main(settings):
    """
    Perform a QM/MM simulation to determine appropriate charge distribution for atoms in the specified QM region and
    write a new topology file with those charges.

    Parameters
    ----------
    settings: argparse.Namespace
        The settings namespace object.

    Returns
    -------
    new_top: str
        Path to the newly written topology file with the appropriate charge distribution

    """
    # Write new input file based on isEE input file, with added QM settings
    # todo: this is necessarily Amber-specific. There's no way to do this in OpenMM (that I know of). If I decide to
    # todo: publish isEE as a standalone package with extensibility to other packages I'm going to need to make large
    # todo: portions of this script methods of the MDEngine abc, or else take an entirely different approach.
    isee_input_file = settings.path_to_input_files + '/isee_amber.in'
    input_lines = open(isee_input_file, 'r').readlines()

    # Check to be sure this doesn't already contain qmmm stuff
    if any([('ifqnt' in line or '&qmmm' in line) for line in input_lines]):
        raise RuntimeError('The input file ' + isee_input_file + ' contains an ifqnt setting and/or a &qmmm namelist. '
                           'These should not be present (even if it\'s ifqnt=0).')

    # Add ifqnt and &qmmm namelist # todo: this is sloppy, not really suited for publication
    new_input_file = 'initialize_charges_amber.in'
    with open(new_input_file, 'w') as f:
        added = False
        for line in input_lines:
            if '&cntrl' in line:
                line = ' &cntrl\n  ifqnt=1,\n'
            elif 'nstlim=' in line:
                line = '  nstlim=1,\n'      # only a single step
            elif 'dt=' in line:
                line = '  dt=0.0001,\n'     # extremely small timestep
            elif 'ntpr=' in line:
                line = '  ntpr=1,\n'  # write to output every step
            elif 'noshakemask=' in line:
                continue    # todo: should in fact only remove QM atoms from noshakemask...
            f.write(line)

        f.write('\n &qmmm')
        f.write('\n  qmmask=\'' + str(settings.ic_qm_mask) + '\',')
        f.write('\n  qm_theory=\'' + str(settings.ic_qm_theory) + '\',')
        f.write('\n  qmcut=' + str(settings.ic_qm_cut) + ',')
        f.write('\n  qmcharge=' + str(settings.ic_qm_charge) + ',')
        if settings.ic_dftb_telec:
            f.write('\n  dftb_telec=' + str(settings.ic_dftb_telec) + ',')
        f.write('\n  printcharges=1,')
        f.write('\n &end\n')

    # Make a thread to pass to process
    thread = iseemain.Thread()
    thread.name = 'ic'
    thread.skip_update = True
    thread.idle = False
    thread.current_name = 'ic'
    thread.terminated = False
    thread.history = argparse.Namespace()
    thread.history.inpcrd = settings.initial_coordinates
    thread.history.tops = [settings.init_topology]
    thread.history.trajs = []
    thread.history.muts = []
    thread.history.score = []
    thread.history.timestamps = []

    # Have to specify appropriate settings for a short sander run
    # todo: have to do better than this if I want to publish this code.
    temp_settings = copy.copy(settings)
    temp_settings.solver = 'sander'
    temp_settings.walltime = '01:00:00'
    temp_settings.ppn = 1
    temp_settings.nodes = 1
    temp_settings.md_engine = 'amber_init_charges'

    # Submit job with process
    running = process(thread, [], [], temp_settings, inp_override=new_input_file)

    # Use the same loop strategy as in main.main to wait for this job to finish
    while True:
        if thread.gatekeeper(running, settings):
            break
        else:
            time.sleep(30)  # to prevent too-frequent calls to batch system by thread.gatekeeper

    # Read atomic charges from output file, store them in a python-interpretable format
    # This is a little ugly because the information is stored across two separate tables: one that links a QM index to
    # the MM (standard) index, and one that actually gives the partial charge data by QM index.
    output_name = thread.name + '_' + thread.current_name + '.out'
    output = open(output_name)     # open output file as an iterator

    # First, extract lines containing first QM-index/MM-index table:
    first_table_flag = False
    first_table = []
    second_table_flag = False
    second_table = []
    try:
        while True:
            line = next(output)
            if first_table_flag:
                if '*' in line:     # indicates link atoms, these come last and we want to ignore them
                    first_table_flag = False
                    continue
                first_table.append(line)
            elif second_table_flag:
                if 'Total Mulliken Charge' in line:
                    second_table_flag = False
                    break
                second_table.append(line)
            if '  QMMM: QM_NO.   MM_NO.  ATOM         X         Y         Z' in line:
                first_table_flag = True
            elif '  Atom    Element       Mulliken Charge' in line:
                second_table_flag = True
    except StopIteration:
        raise RuntimeError('Search through initialize_charges output file ' + settings.working_directory + '/' +
                           output_name + ' terminated before the end of the charge table. It must be formatted '
                           'incorrectly or in an unexpected way.')

    # Next, extract relevant data from tables using regex
    pattern = re.compile('[0-9.-]+')
    first_table_frmt = []
    second_table_frmt = []
    for line in first_table:
        regex = pattern.findall(line)
        first_table_frmt.append([regex[0], regex[1]])   # QM number, MM number
    for line in second_table:
        regex = pattern.findall(line)
        second_table_frmt.append([regex[0], regex[1]])  # QM number, charge

    # Quick sanity check
    for i in range(len(first_table_frmt)):
        if not int(first_table_frmt[i][0]) == i + 1:
            raise RuntimeError('Unexpected QM indices in QM/MM mapping:\n' + str(first_table_frmt))

    # Combine data into third table by replacing QM number in second_table with corresponding MM number
    third_table = []
    for pair in second_table_frmt:
        try:
            new_pair = [first_table_frmt[int(pair[0]) - 1][1], pair[1]]   # MM number, charge
            third_table.append(new_pair)
        except IndexError:
            if not len(third_table) == len(first_table_frmt):
               raise RuntimeError('Encountered unexpected IndexError in trying to build MM index/charge mapping')

    # Now we want to get each index in third_table as a function of the index of one atom per residue and save that
    # information as a .pkl file for set_charges to interpret later.
    mtop = mdtraj.load_prmtop(settings.init_topology)
    residues = []           # list of residues accounted for
    first_index = []        # index for first atom come to in each residue
    get_index_strs = []     # :RES@NAME for first encountered index in residue in corresponding position in residues
    relative_indices = []   # strings to evaluate to obtain each atom index, in same order as third_table
    pattern = re.compile('[0-9]+')  # regex for getting resid integer out of mtop residue
    for atom_index in [int(item[0]) for item in third_table]:
        res = pattern.findall(str(mtop.atom(atom_index - 1).residue))[-1]   # - 1 converts from Amber to mdtraj
        if not res in residues:
            residues.append(res)
            first_index.append(atom_index)
            get_index_str = ':' + res + '@' + mtop.atom(atom_index - 1).name
            get_index_strs.append(get_index_str)
            relative_indices.append(get_index_str)
        else:
            index_in_residues = residues.index(res)
            relative_distance = atom_index - first_index[index_in_residues]
            relative_indices.append(get_index_strs[index_in_residues] + ' + ' + str(relative_distance))
    relative_charge_table = [[relative_indices[i], third_table[i][1]] for i in range(len(relative_indices))]
    pickle.dump(relative_charge_table, open('relative_charge_table.pkl', 'wb'))

    # Finally, call set_charges to actually implement modifying the topology file
    new_top = set_charges(settings.init_topology)

    ## Deprecated additional minimization/equilibration step following charge initialization
    # # Setup thread and temp_settings again for equilibration job with new topology file
    # thread.current_name = 'equil'
    # thread.skip_update = True
    # thread.history.tops = [new_top]
    # temp_settings.solver = 'pmemd.cuda'
    # temp_settings.walltime = '02:00:00'
    # temp_settings.md_engine = 'amber'
    #
    # new_input_file = 'ic_equil_amber.in'
    # with open(new_input_file, 'w') as f:
    #     added = False
    #     for line in input_lines:
    #         if 'dt=' in line:
    #             line = '  dt=0.002,\n'          # 2 fs steps
    #         elif 'nstlim=' in line:
    #             line = '  nstlim=500000,\n'     # 500000 steps -> 1 ns run
    #         f.write(line)
    #
    # # Submit new equilibration job with process
    # running = process(thread, [], temp_settings, inp_override=new_input_file)
    #
    # # Use the same loop strategy as in main.main to wait for this job to finish
    # while True:
    #     if thread.gatekeeper(running, settings):
    #         break
    #     else:
    #         time.sleep(30)  # to prevent too-frequent calls to batch system by thread.gatekeeper

    return new_top #, settings.working_directory + '/' + thread.name + '_' + thread.current_name + '.rst7'

def set_charges(top):
    """
    Implements setting the charges in a topology file 'top' to match the contents of relative_charge_table.pkl

    Parameters
    ----------
    settings: argparse.Namespace
        The settings namespace object.

    Returns
    -------
    new_top: str
        Path to the newly written topology file with the appropriate charge distribution

    """
    # First step: load relative_charge_table.pkl file and then calculate indices to produce absolute_charge_table
    relative_charge_table = pickle.load(open('relative_charge_table.pkl', 'rb'))
    evaluated = []
    pattern = re.compile('\:[0-9]+\@[A-Za-z0-9+-]+')    # supports atom names consisting of numbers, letters, and +/-
    mtop = mdtraj.load_prmtop(top)
    for item in relative_charge_table:
        relative_to = pattern.findall(item[0])[0]
        if not relative_to in [item[0] for item in evaluated]:
            parsed = relative_to.replace(':', 'resid ').replace('@', ' and name "') + '"'   # quotes for literal name
            relative_to_index = str(int(mtop.select(parsed)[0]) + 1)     # + 1 converts from mdtraj to Amber
        else:
            internal_index = [item[0] for item in evaluated].index(relative_to)
            relative_to_index = str(evaluated[internal_index][1])
        to_eval = item[0].replace(relative_to, relative_to_index)
        evaluated.append([relative_to, eval(to_eval)])

    absolute_charge_table = [[evaluated[i][1], relative_charge_table[i][1]] for i in range(len(relative_charge_table))]

    # If top is a path, dissect it into filename and path to that filename
    if '/' in top:
        path_to_top = top[:top.rindex('/') + 1]
        top = top[top.rindex('/') + 1:]
    else:
        path_to_top = ''

    parmed_top = parmed.load_file(path_to_top + top)

    new_top = 'ic_' + top
    for atom_charge_pair in absolute_charge_table:
        arg = [str(item) for item in atom_charge_pair]
        try:
            # change <property> <mask> <new_value>
            setcharge = parmed.tools.actions.change(parmed_top, 'CHARGE', '\'@' + arg[0] + '\'', arg[1])
            setcharge.execute()
        except parmed.tools.exceptions.SetParamError as e:
            raise RuntimeError('encountered parmed.tools.exceptions.SetParamError: ' + e + '\n'
                               'The offending atom/charge pair is: ' + str(arg))
    parmed_top.write_parm(path_to_top + new_top)

    return path_to_top + new_top

if __name__ == "__main__":
    # Load settings.pkl file given as command line argument
    settings = pickle.load(open(sys.argv[1], 'rb'))
    main(settings)


'''
启动websocket，发送和接收数据
redmorningcn 2020.10.19
'''
from   websocket_server import WebsocketServer
import threading
import json
from   datetime import datetime
import threading

def startServer():
    #t = threading.Timer(60, Server.run)
    t = threading.Thread(target = Server.run)
    t.start()

def server_send(text):
    #print("Server.send_message",Server.serverflg,text)
    if Server.serverflg == 2:
        Server.send_message(text)

def server_recv():
    if Server.messagelen > 0:
        #print(Server.messagelen)
        Server.messagelen = 0
        return Server.messagetxt
    return 0

class Server:
    serverflg  = 1
    messagetxt = ''
    messagelen = 0

    @classmethod
    def run(cls):
        PORT = 1234
        server = WebsocketServer(PORT, host='0.0.0.0')
        cls.myserver = server

        server.set_fn_new_client(cls.new_client)
        server.set_fn_client_left(cls.client_left)
        server.set_fn_message_received(cls.message_received)
        print("Web socket server has started at %s." % datetime.now())
        #showTimer()
        server.run_forever()

    @classmethod
    def getEvent(cls, data):
        try:
            jsonDic = json.loads(data)
            dic = dict(jsonDic)
            event = dic["event"]
        except Exception as error:
            print("Get event error:%s" % error)
            return None, {}
        return event, dic

    @classmethod
    def new_client(cls, client, server):
        print("A new client connected and was given id %d at %s." % (client['id'], datetime.now()))
        server.send_message_to_all("Hey all, a new client has joined us at %s." % datetime.now())
        Server.serverflg = 2    #连接正常
    @classmethod
    def send_message(cls,text):
        #print("A new client connected and was given id %d at %s." % (client['id'], datetime.now()))
        cls.myserver.send_message_to_all(text)

    @classmethod
    def client_left(cls, client, server):
        Server.serverflg   = 1  #连接断开
        print("Client(%d) disconnected at %s." % (client['id'], datetime.now()))

    @classmethod
    def message_received(cls, client, server, message):
        event, dic = cls.getEvent(message)
        #print(event, dic)
        #if event == 1:
        #    print(dic["tree"])

        Server.messagetxt = message
        Server.messagelen = len(message)
        #print(Server.messagelen,Server.messagetxt)
        #print("Client(%d) said: %s at %s" % (client['id'], message, datetime.now()))

import   time
if __name__=='__main__':
    startServer()      #启动服务

    while True:
        time.sleep(1)

        txt =  server_recv()
        if txt:
            print(txt)


# -*- coding: utf-8 -*-

from biocluster.workflow import Workflow
import glob
import os
from bson import ObjectId
import re
import json
from biocluster.core.exceptions import OptionError
from mbio.packages.metagbin.common_function import link_dir

class FaprotaxWorkflow(Workflow):
    def __init__(self, wsheet_object):
        self._sheet = wsheet_object
        super(FaprotaxWorkflow, self).__init__(wsheet_object)
        options = [
            {"name": "otu_table", "type": "infile","format": "meta.otu.otu_table, meta.otu.tax_summary_dir"},
            {"name": "group_table", "type": "infile","format": "meta.otu.group_table"},
            {"name": "main_id", "type": "string"},
            {"name": "update_info", "type": "string"},
            {"name": "group_id", "type": "string"},
            {"name": "asv_id", "type": 'string'},
            {"name": "group_detail", "type": "string"},
            ]
        self.add_option(options)
        self.set_options(self._sheet.options())
        self.faprotax = self.add_tool('meta.faprotax_predict')
        self.IMPORT_REPORT_DATA = True
        self.IMPORT_REPORT_AFTER_END = False

    def run_faprotax_predict(self):
        with open(self.option("group_table").prop['path'], "r") as f, open(self.work_dir + "/group.txt", "w") as t:
            data = f.readlines()
            t.write("#sample\t#group\n")
            group_detail = json.loads(self.option("group_detail"))
            if len(group_detail) == 1 and "All" in group_detail:
                for i in data[1:]:
                    t.write(i.split("\t")[0] + "\t" + "All" + "\n")
            else:
                for i in data[1:]:
                    t.write(i)
        mapping_file = self.work_dir + "/mapping_file"
        self.faprotax.set_options({
            "otu_table": self.option("otu_table")
            })
        self.faprotax.on("end", self.set_db)
        self.faprotax.run()

    def run(self):
        self.run_faprotax_predict()
        super(FaprotaxWorkflow, self).run()

    def set_db(self):
        """
        保存结果表到mongo数据库中
        """
        link_dir(self.faprotax.output_dir,self.output_dir)
        self.logger.info("正在写入mongo数据库")
        self.api_faprotax = self.api.api("metaasv.faprotax")
        self.api_faprotax.add_Faprotax(prediction_file=self.output_dir + "/function_prediction.txt",group_file=self.work_dir +"/group.txt",main_id=self.option("main_id"))
        self.end()

    def end(self):
        result_dir = self.add_upload_dir(self.output_dir)
        result_dir.add_relpath_rules([
            [".", "", "Faprotax功能预测结果目录",0],
            ["./function_prediction.txt", "txt", "Faprotax功能预测结果表",0],
            ["./report.txt", "txt", "Faprotax功能预测报告文件",0],
        ])
        super(FaprotaxWorkflow, self).end()


from parameter_learner import *
import tensorflow as tf
import numpy as np

slim = tf.contrib.slim

def get_random_dynamic_param(dynamic_parameters_shape_map, batch_size):
  dynamic_parameters_map = dict()
  for name, shape in dynamic_parameters_shape_map.iteritems():
    dynamic_parameters_map[name] = tf.random_normal([batch_size] + shape)
  return dynamic_parameters_map

def weight_hashing_parameter_learner_net():
  num_classes = 1
  k_shot = 5

  meta_batch_size = 1
  height = 1
  width = 1
  depth_in = 5

  params = dict(is_training = True,
    after_hashing_scale_hyperparameters={},
    parameter_prediction_convline=None,
    fc_hyperparameters ={},
    decompression_factor=num_classes*k_shot)

  inputs = tf.random_normal([
      meta_batch_size, num_classes*k_shot,
      height, width, depth_in])
  dynamic_params_shape_map = dict(w=[3,3,5,50], b=[50])
  param_learner = WeightHashingParameterLearner(**params)
  net = param_learner.learn(inputs, dynamic_params_shape_map)
  return net

from IPython import embed
if __name__ == '__main__':
  net = weight_hashing_parameter_learner_net()
  init_op = tf.global_variables_initializer()
  with tf.Session() as sess:
    sess.run(init_op)
    out = sess.run(net)
    embed()