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


class Solution:
    def swimInWater(self, grid: List[List[int]]) -> int:
        import heapq
        n = len(grid)
        queue = []  # python 小顶堆
        heapq.heappush(queue , (grid[0][0], 0*n + 0)) # x*n+y
        visited = [0]*(n**2)
        dirs = [-1,0,1,0,-1]
        visited[0] = 1
        while queue:
            pos = heapq.heappop(queue)
            t = pos[0]
            x = pos[1]//n
            y = pos[1]%n
            if x == n-1 and y == n-1:
                return t
            for i in range(4):
                tx = x + dirs[i]
                ty = y + dirs[i+1]
                if tx < 0 or tx >=n or ty < 0 or ty >= n:
                    continue
                if visited[tx*n+ty] == 1:
                    continue
                visited[tx*n+ty] = 1
                heapq.heappush(queue , (max(t,grid[tx][ty]) , tx*n+ty) )
        return -1


#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import re
import sys
import json
import uuid
import base64
import socket
import string
import binascii
import datetime

import dpkt, pcap

import argparse
import urllib.parse

def common_tcp_data( ip ) :
    tcp = ip.data
    src = ip.src
    dst = ip.dst
    src_a = socket.inet_ntoa(src)
    dst_a = socket.inet_ntoa(dst)

    # TCP情報収集
    tcp_data = {}
    tcp_data["seq"] = tcp.seq
    tcp_data["seq_raw"] = tcp.seq
    tcp_data["window_size_value"] = tcp.win
    tcp_data["window_size"] = tcp.win
    tcp_data["window_size_scalefactor"] = tcp.win

    return { "uuid": str(uuid.uuid4()), "destination_ip": dst_a, "destination_port": tcp.dport,
            "source_ip": src_a, "source_port": tcp.sport, "tcp":tcp_data }

def parse_tcp_tls( ip, tcp_buffer ) :
    TLS_HANDSHAKE = 22

    records, bytes_used = dpkt.ssl.tls_multi_factory(tcp_buffer)
    results = []
    for record in records:
        if record.type != TLS_HANDSHAKE:
            continue
        if len(record.data) == 0:
            continue
        client_hello = bytearray(record.data)
        if client_hello[0] != 1:
            # We only want client HELLO
            continue
    return ""

def parse_tcp_http( ip, tcp_buffer ) :
    cmd_list = ["wget","curl","rm","chmod","busybox","mv","ps","kill","xargs","iptables","grep","pkill","netstat","pgrep",
            "chattr","service","systemctl","crontab","apt","apt-get","unhide","ufw","userdel","adduser","useradd"]
    pattern = "https?://[\w/:%#\$&\?\(\)~\.=\+\-]+"

    http = dpkt.http.Request(tcp_buffer)

    # HTTP情報収集
    http_data = {}
    http_data["http_http_request_method"] = http.method
    http_data["http_http_request_uri"] = http.uri
    http_data["http_http_request_version"] = "HTTP/"+http.version
    http_data["http_http_request"] = True
    http_data["http_http_request_number"] = "1"
    for key in http.headers :
        http_data["http_http_"+key] = http.headers[key]
    #http_data["http_http_request_full_uri"] = "http://"+http_data["http_http_host"]+"/"+http_data["http_http_request_uri"]
    #http_data["http_http_request_full_uri"] = "http:/"+os.path.join(http_data["http_http_host"],http_data["http_http_request_uri"])
    http_data["http_http_request_full_uri"] = "http://"+http_data["http_http_host"]+"/"+http_data["http_http_request_uri"]

    http_data["http_http_params"] = urllib.parse.urlparse(http_data["http_http_request_full_uri"]).query
    http_data["http_http_body"] = http.body.decode('utf-8')

    # URIの分割
    dirname, basename = os.path.split(http_data["http_http_request_uri"])
    #http_data["http_http_request_uri_detail"] = {"urlpath": dirname, "resourcename": basename}
    http_data["http_http_request_uri_urlpath"]      = dirname
    http_data["http_http_request_uri_resourcename"] = basename

    # 攻撃で利用されるコマンド文字列の検索
    http_data["http_http_cmd"] = []

    check_target = http_data["http_http_params"]+" "+http_data["http_http_body"]
    for cmd in cmd_list :
        if (cmd in check_target) :
            http_data["http_http_cmd"].append( cmd )

    # ダウンロードURLの抽出
    http_data["http_http_download"] = re.findall(pattern, check_target )

    # ダウンロードURLの分解
    http_data["http_http_download_detail"] = []
    for url in http_data["http_http_download"] :
        u_parse = urllib.parse.urlparse(url)
        result = {
                "method"      : u_parse.scheme,
                "domain"      : u_parse.netloc,
                "path"        : u_parse.path,
                "params"      : u_parse.params,
                "query"       : u_parse.query,
                "query_detail": [] }
        result["query_detail"] = urllib.parse.parse_qs(u_parse.query)
        http_data["http_http_download_detail"].append(result["query_detail"])

    return str(http), http_data

def parse_tcp_telnet( ip, tcp_buffer ) :
    datas = dpkt.telnet.strip_options(tcp_buffer)
    tmp = []
    for item in datas[0] :
        tmp.append( item.decode() )
    return "\\n".join(tmp)

def parse_tcp_smb( ip, tcp_buffer ) :
    smb = dpkt.smb.SMB(tcp_buffer)
    return str(repr(smb))

def parse_udp_stun( ip, udp_buffer ) :
    datas = dpkt.stun.parse_attrs(udp_buffer)
    print(datas)

def detect_tcp( ip, tcp_info ) :
    result = common_tcp_data( ip )
    result["timestamp"] = str(datetime.datetime.utcfromtimestamp(tcp_info["start_time"]))
    result["connect_time"] = str(datetime.datetime.utcfromtimestamp(tcp_info["end_time"]))
    result["duration"] = tcp_info["end_time"]-tcp_info["start_time"]
    try :
        result["payload"] = base64.b64encode(tcp_info["buffer"]).decode()
    except :
        result["payload"] = "base64 except"

    try :
        ret_data, http_data = parse_tcp_http( ip, tcp_info["buffer"] )
        result["app_proto"] = "http"
        result["payload_printable"] = ret_data
        result["http"] = http_data
        return result
    except :
        import traceback
        traceback.print_exc()

    try :
        parse_tcp_tls( ip, tcp_info["buffer"] )
        result["app_proto"] = "ssl"
        return result
    except :
        import traceback
        traceback.print_exc()

    try :
        ret_data = parse_tcp_smb( ip, tcp_info["buffer"] )
        result["app_proto"] = "smb"
        result["payload_printable"] = ret_data
        return None
    except :
        import traceback
        traceback.print_exc()

    try :
        ret_data = parse_tcp_telnet( ip, tcp_info["buffer"] )
        result["app_proto"] = "telnet"
        result["payload_printable"] = ret_data
        return result
    except :
        import traceback
        traceback.print_exc()
#        pass

def main( device, pcap_option, pcap_file, output_file ) :
    packet_buffer = {}

    if pcap_file != None :
        print( pcap_file )
        pc = dpkt.pcap.Reader(open(pcap_file,'rb'))
    else :
        print( device )
        #pc = pcap.pcap( )
        pc = pcap.pcap( device, promisc=True, immediate=True )
        if pcap_option != None :
            pc.setfilter( 'host %s'%(pcap_option) )
        print( "Device Name:", pc.name, "|Filter:",pc.filter )

    for t,buf in pc :
        # Packetの取得
        eth = dpkt.ethernet.Ethernet(buf)

        # IPプロトコル確認
        if type(eth.data) != dpkt.ip.IP:
            continue

        ip = eth.data
        #TCPデータ
        if type(ip.data) == dpkt.tcp.TCP:
            tcp = ip.data
            fin_flag = int(( tcp.flags & dpkt.tcp.TH_FIN ) != 0)
            syn_flag = int(( tcp.flags & dpkt.tcp.TH_SYN ) != 0)
            rst_flag = int(( tcp.flags & dpkt.tcp.TH_RST ) != 0)
            psh_flag = int(( tcp.flags & dpkt.tcp.TH_PUSH) != 0)
            ack_flag = int(( tcp.flags & dpkt.tcp.TH_ACK ) != 0)
            urg_flag = int(( tcp.flags & dpkt.tcp.TH_URG ) != 0)

            src = ip.src
            dst = ip.dst
            src_a = socket.inet_ntoa(src)
            dst_a = socket.inet_ntoa(dst)

            source_info = "%s:%d"%(src_a, tcp.sport)

            if syn_flag == 1 and ack_flag == 0 :
                if source_info not in packet_buffer :
                    packet_buffer[source_info] = {"start_time": t, "end_time": None, "buffer": None}
                continue

            if source_info not in packet_buffer :
                continue

            if fin_flag == 0 :
                if packet_buffer[source_info]["buffer"] == None :
                    packet_buffer[source_info]["buffer"] = tcp.data
                else :
                    packet_buffer[source_info]["buffer"] += tcp.data
                continue
            else :
                packet_buffer[source_info]["end_time"] = t
                tcp_info = packet_buffer[source_info]
                del packet_buffer[source_info]

                result = detect_tcp( ip, tcp_info )
                if result == None :
                    continue
                print(result, flush=True)

                dt_now = datetime.datetime.now()
                if output_file != None :
                    filename = dt_now.strftime( output_file )
                    f = open(filename, 'a')
                    json.dump(result, f)
                    f.write( '\n' )
                    f.close()
                else :
                    print( json.dumps(result) )

        elif type(ip.data) == dpkt.udp.UDP:
            udp = ip.data
            if len(udp.data) != 0:
                continue

            print(udp)
            try :
                parse_udp_stun( ip, tcp )
                print("parse_udp_stun")
                continue
            except :
                pass
        elif type(ip.data) == dpkt.icmp.ICMP:
            icmp = ip.data
            if len(icmp.data) != 0:
                continue

            print(icmp)
            print("icmp")

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='sniffer')
    parser.add_argument('-i', '--interface', help='target device')
    parser.add_argument('-t', '--target', help='target host 1.1.1.1 or 2.2.2.2')
    parser.add_argument('-r', '--read', help='pcapファイル名')
    parser.add_argument('-o', '--output', help='出力ファイル名 ', default="/iot_honey/logs/rfpf_%Y%m%d.json")

    args = parser.parse_args()

    if args.read != None :
        main( None, None, args.read, args.output )
    else :
        main( args.interface, args.target, None, args.output )
#        main( args.interface, args.target, None, None )


# -*- coding: utf-8 -*-
# __author__ = 'fwy'
# last_modify:20200609

from biocluster.core.exceptions import OptionError
from biocluster.module import Module
import gevent
import re
import os
import unittest
import gevent.subprocess as subprocess
import json
import glob
from collections import OrderedDict
from biocluster.config import Config
import pandas as pd
import numpy  as np

class DiffGenesetAllModule(Module):
    """
    该Module用于基因融合分析，默认使用方法
    """
    def __init__(self, work_id):
        super(DiffGenesetAllModule, self).__init__(work_id)
        options = [
            # {"name": "geneset_names", "type": "infile", "format": "ref_rna_v2.common"},
            {"name": "diff_path", "type": "string", "default": None},
            {"name": "annot_result", "type": "string", "default": None},
            {"name": "diff_method", "type": "string"},
            {"name": "transcript_exp_file", "type": "string"},
            {"name": "level", "type": "string", "default": "G"},
            {"name": "kegg_version", "type": "string", "default": None},
            {"name": "go_version", "type": "string", "default": '20200628'},
            {"name": "reactome_version", "type": "string", "default": None},
            {"name": "species", "type": "string", "default": "Homo_sapiens"}
        ]
        self.add_option(options)
        self.geneset_prepare_tools=[]
        self.genest_infos = OrderedDict()
        self.total_dict = OrderedDict()
        self.run_genesets_analysis = True
        self.diff_geneset_prepare = ""
        self.selecte_geneset = ""

    def check_options(self):
        return True

    def set_step(self, event):
        if 'start' in event['data'].keys():
            event['data']['start'].start()
        if 'end' in event['data'].keys():
            event['data']['end'].finish()
        self.step.update()

    def get_genesets_info(self):
        diff_dir = self.option("diff_path")
        all_detail = pd.read_table(os.path.join(diff_dir, 'all_detail.txt'), header=0, sep='\t')
        all_detail['type'] = np.where(all_detail['seq_id'].str.contains("sRNA"), 'sRNA', 'mRNA')
        all_detail.fillna('', inplace=True)
        diff_dict_list = all_detail.to_dict('records')
        all_diff_list = list()
        for i in all_detail.groupby('compare'):
            compare = i[0]
            diff_pd = i[1]
            ctrl, test = compare.split('|')
            name = ctrl + '_vs_' + test + '_mRNA'
            sig_seqs = diff_pd[(~diff_pd['seq_id'].str.contains("sRNA")) & (diff_pd['significant'] == 'yes')][
                'seq_id'].tolist()
            sig_regulate = list(diff_pd['regulate'][diff_pd['significant'] == 'yes'])
            all_diff_list += sig_seqs
        all_diff_list = list(set(all_diff_list))
        geneset_name = 'All_Diff_mRNA'
        length = len(all_diff_list)
        self.genest_infos[geneset_name] = {
                "geneset_name": geneset_name,
                "compare": "All",
                "gene_num" :length,
                "regulate": "all",
                "compare_path": os.path.join(diff_dir, 'all_detail.txt'),
                'all_diff_list':all_diff_list
        }
        with open(os.path.join(self.work_dir, "sele_geneset_info"), "w") as f:
            result_infos = self.genest_infos[geneset_name]
            json.dump(result_infos, f)
        self.prepare_common_file()

    def prepare_common_file(self):
        self.logger.info("开始准备差异数据挖掘公共文件")
        self.common = self.add_tool("prok_rna.diffgt4work.annot_prepare")
        opts = {
            "annot_result": self.option("annot_result"),
            "level": self.option("level"),
            "kegg_version": self.option("kegg_version"),
            'go_version': self.option("go_version"),
            "transcript_exp_file": self.option("transcript_exp_file"),
            "species": self.option("species"),
        }
        self.common.set_options(opts)
        self.common.on("end", self.prepare_geneset_infos)
        # self.common.on("end", self.prepare_geneset_infos)
        self.common.run()

    def get_selected_genset_info(self):
        # self.get_genesets_info()
        geneset_name = self.genest_infos.keys()[0]
        if self.genest_infos[geneset_name]["gene_num"] < 50:
            return False
        else:
            return geneset_name

    def prepare_geneset_infos(self):
        self.selecte_geneset= self.get_selected_genset_info()
        if not self.selecte_geneset:
            self.set_output()
        else:
            self.diff_geneset_prepare = self.add_tool("prok_rna.diffgt4work.diff_geneset_prepare")
            opts = {
                "annot_result": self.option("annot_result"),
                "compare": self.genest_infos[self.selecte_geneset]["compare"],
                "regulate": self.genest_infos[self.selecte_geneset]["regulate"],
                "geneset_name": self.genest_infos[self.selecte_geneset]["geneset_name"],
                "compare_path": self.genest_infos[self.selecte_geneset]["compare_path"],
                "level": self.option("level"),
                "species": self.option("species"),
                "geneset_infos" :os.path.join(self.work_dir, "sele_geneset_info"),
            }
            self.diff_geneset_prepare.set_options(opts)
            self.diff_geneset_prepare.on('end', self.set_output)
            self.diff_geneset_prepare.run()

    def set_output(self):
        if not self.run_genesets_analysis :
            self.total_dict["has_geneset_pipline_results"] = False
        else:
            common_file_path = self.common.option("common_file_json").prop["path"]
            common_file_dict = json.load(open(self.common.option("common_file_json").prop["path"]))
            self.total_dict["common_file"] = common_file_dict
            self.total_dict["genesets"] =OrderedDict()
            geneset_json = self.diff_geneset_prepare.option("geneset_json").prop["path"]
            geneset_dict = json.load(open(geneset_json))
            geneset_name = geneset_dict["geneset_name"]
            self.total_dict["genesets"][geneset_name] = geneset_dict
        with open(os.path.join(self.output_dir,"prepare_json"),"w") as f:
            json.dump(self.total_dict,f,indent=2)
        self.end()

    def run(self):
        super(DiffGenesetAllModule, self).run()
        self.logger.info("开始运行差异基因集数据挖掘数据准备")
        self.logger.info("首先解析基因集列表文件")
        self.get_genesets_info()

    def end(self):
        super(DiffGenesetAllModule, self).end()

class TestFunction(unittest.TestCase):
    """
    This is test for the tool. Just run this script to do test.
    """
    def test(self):
        import random
        from mbio.workflows.single import SingleWorkflow
        from biocluster.wsheet import Sheet
        test_dir='/mnt/ilustre/users/sanger-dev/workspace/20190322/Snp_tsg_33538_3123_8568/SnpRna'
        data = {
            "id": "geneset_prepare" + str(random.randint(1, 10000)),
            "type": "module",
            "name": "medical_transcriptome.workflow_diffgt.diff_geneset_all",
            "instant": False,
            "options": dict(
                diff_path= "/mnt/ilustre/users/sanger-dev/workspace/20200923/MedicalTranscriptome_medical_transcriptome_workflow_20200923_153802009/DiffexpBatch/output",
                annot_result ="/mnt/ilustre/users/sanger-dev/workspace/20200923/MedicalTranscriptome_medical_transcriptome_workflow_20200923_153802009/AnnotMerge__1/output",
                diff_method="DESeq2",
                kegg_version ="202007",
                reactome_version="202007",
                level = "G",
                gene_count_file ="/mnt/ilustre/users/sanger-dev/workspace/20200923/MedicalTranscriptome_medical_transcriptome_workflow_20200923_153802009/Quant/output/gene.count.matrix",
            )
           }
        wsheet = Sheet(data=data)
        wf = SingleWorkflow(wsheet)
        wf.run()

if __name__ == '__main__':
    unittest.main()


from django.shortcuts import render,redirect
from .forms import ProductoForm,PedidoForm,Burger_BossForm,ExtraForm,IngredienteForm,HamburguesaForm
from django.views.generic import CreateView,DeleteView,ListView,UpdateView,DetailView
from .models import Producto,Pedido,Burger_Boss,Extra,Ingrediente,Hamburguesa
from django.urls import reverse_lazy
#Crea clases
class HomePage(ListView):
    model = Producto
    template_name = 'homepage3.html'

class Ubicanos(ListView):
    model = Producto
    template_name = 'ubicanos.html'
class Pagos(ListView):
    model = Producto
    template_name = 'pagos.html'
class ProductoList(ListView):
    model = Producto
    template_name = 'pedido.html'

    def get_context_data(self, **kwargs):
        context= super().get_context_data()
        context['extra_list'] = Extra.objects.all()
        return context

class PedidoCreate(CreateView):
    model = Pedido
    form_class = PedidoForm
    template_name = 'crear_pedido.html'
    success_url = reverse_lazy('pedido_completado')

class PedidoList(ListView):
    model = Pedido
    template_name = 'pedido_completado.html'

class ProductoListCambiarEstado(ListView):
    model = Pedido
    template_name = 'cambiar_estado.html'

class OtroEstado(UpdateView):
    model = Pedido
    form_class = PedidoForm
    template_name = 'otro_estado.html'
    success_url = reverse_lazy('cambiar_estado')

class ProductoCreate(CreateView):
    model = Producto
    form_class = ProductoForm
    template_name = 'crear_producto.html'
    success_url = reverse_lazy('pedido_admin')

class ExtraCreate(CreateView):
    model = Extra
    form_class = ExtraForm
    template_name = 'crear_extra.html'
    success_url = reverse_lazy('pedido_admin')

class IngredienteCreate(CreateView):
    model = Ingrediente
    form_class = IngredienteForm
    template_name = 'crear_ingrediente.html'
    success_url = reverse_lazy('pedido_admin')

class BurgerBossCreate(CreateView):
    model = Burger_Boss
    form_class = Burger_BossForm
    template_name = 'registro.html'
    success_url = reverse_lazy('crear_hamburguesa2')

class HamburguesaCreate(CreateView):
    model = Hamburguesa
    form_class = HamburguesaForm
    template_name = 'crear_hamburguesa.html'
    success_url = reverse_lazy('extra')

#class VerHamburguesas2(ListView):
 #   model = Hamburguesa
  #  template_name = 'ver_hamburugesa.html'

class Pagos(ListView):
    model = Producto
    template_name = 'pagos.html'

class ModificarProductoCarne(UpdateView):
    model = Producto
    form_class = ProductoForm
    template_name = 'modificar_producto_carne.html'
    success_url = reverse_lazy('pedido_admin')

class ModificarProductoQueso(UpdateView):
    model = Producto
    form_class = ProductoForm
    template_name = 'modificar_producto_queso.html'
    success_url = reverse_lazy('pedido_admin')

class ModificarIngrediente(UpdateView):
    model = Ingrediente
    form_class = IngredienteForm
    template_name = 'modificar_ingrediente.html'
    success_url = reverse_lazy('pedido_admin')

class ModificarExtra(UpdateView):
    model = Extra
    form_class = ExtraForm
    template_name = 'modificar_extra.html'
    success_url = reverse_lazy('pedido_admin')

class EliminarCarne(DeleteView):
    model = Producto
    template_name = 'eliminar_carne.html'
    success_url = reverse_lazy('pedido_admin')

class EliminarQueso(DeleteView):
    model = Producto
    template_name = 'eliminar_queso.html'
    success_url = reverse_lazy('pedido_admin')

class EliminarIngrediente(DeleteView):
    model = Ingrediente
    template_name = 'eliminar_ingrediente.html'
    success_url = reverse_lazy('pedido_admin')

class EliminarExtra(DeleteView):
    model = Extra
    template_name = 'eliminar_extra.html'
    success_url = reverse_lazy('pedido_admin')


from django.http import HttpResponse
from django.shortcuts import render, redirect
from .models import Item

def home_page(request):
    # if request.method == 'POST':
    #     return HttpResponse(request.POST['item_text'])
    # return render(request, 'lists/home.html')
    # item = Item()
    # item.text = request.POST.get('item_text', '')
    # item.save()

    if request.method == 'POST':
        # new_item_text = request.POST['item_text']
        Item.objects.create(text=request.POST['item_text'])
        return redirect('/')
    # else:
    #     new_item_text = ''

    items = Item.objects.all()
    context = {
        'items': items,
        }
    return render(request, 'lists/home.html', context)


import numpy as np

import mne

from plotly.graph_objs import Figure, Scattergl

from helperfunctions.annotation_helperfunctions import get_annotations

import constants as c
import globals

def _get_scaling(EEG_scale):
    """Calculates scaling factor to multiply data with for given scale.

    Args:
        EEG_scale (float): Desired scale of data. Defaults to 1e-6.

    Returns:
        float: Scaling factor to multiply data with.
    """
    if EEG_scale:
        temp = EEG_scale
        temp_string = str(EEG_scale)

        if 'e' in temp_string:
            e_index = temp_string.index('e')
            zeros = int(temp_string[e_index + 2:])
            temp *= 10 ** zeros
            temp *= 10 ** zeros
        else:
            zeros = 0
            while temp < 1:
                temp *= 10
                zeros += 1
            temp *= 10 ** zeros

        scaling_factor = temp

        # print('Apply custom scaling {}'.format(scaling_factor))
    else:
        scaling_factor = c.CONVERSION_VALUE_VOLTS_TO_MICROVOLTS

    return scaling_factor

def _get_offset(channel_offset):
    """Returns offset factor to multiply data with.

    Args:
        channel_offset (int): Offset between channels. Defaults to 40 (μV).

    Returns:
        int: Offset factor to multiply data with.
    """
    if channel_offset or channel_offset == 0:
        offset_factor = channel_offset
    else:
        offset_factor = c.DEFAULT_Y_AXIS_OFFSET

    return offset_factor

def _get_time(EEG_data, sample_frequency):
    """Calculates timescale and recording length of given EEG data.

    Args:
        EEG_data (array): EEG data.
        sample_frequency (float): Sample frequency of data.

    Returns:
        tuple(array, float): Array of all timepoints, recording length.
    """
    # Create time-scale for x-axis of visualization
    timestep = 1 / sample_frequency

    timescale = [timestep * i for i in range(EEG_data.shape[0])]

    recording_length = float(timescale[-1] + timestep)
    print('Length of recording in seconds: {}'.format(recording_length))

    return timescale, recording_length

def _get_plotting_data(raw, file_name, selected_channel_names, EEG_scale, channel_offset, model_output=[], model_channels=[]):
    """Generates dict holding all relevant data from raw object and model outputs for plotting.

    Args:
        raw (mne.io.Raw): Raw object to get data from.
        file_name (string): File-name.
        selected_channel_names (list): List of strings of selected channel names to plot.
        EEG_scale (float): Desired scaling.
        channel_offset (int): Desired channel offset.
        model_output (list, optional): List of arrays of model outputs. Defaults to [].
        model_channels (list, optional): List of lists of strings of channel names model output is based on. Defaults to [].

    Returns:
        dict: Dict holding all relevant data from raw object and model outputs for plotting.
    """
    plotting_data = {'EEG': {}, 'model': [], 'plot': {}}

    plotting_data['EEG']['file_name'] = file_name

    plotting_data['EEG']['scaling_factor'] = _get_scaling(EEG_scale)
    plotting_data['plot']['offset_factor'] = _get_offset(channel_offset)

    if selected_channel_names:
        check = all(channel in raw.ch_names for channel in selected_channel_names)
        if check:
            plotting_data['EEG']['EEG_data'] = np.transpose(raw.get_data(selected_channel_names))
            plotting_data['EEG']['channel_names'] = selected_channel_names
    else:
        print('Displaying all channels')
        plotting_data['EEG']['EEG_data'] = np.transpose(raw.get_data())
        plotting_data['EEG']['channel_names'] = raw.ch_names

    plotting_data['EEG']['timescale'], plotting_data['EEG']['recording_length'] = _get_time(plotting_data['EEG']['EEG_data'], raw.info['sfreq'])

    # Calculate offset for y-axis
    offset_EEG = plotting_data['EEG']['EEG_data'].copy()
    offset_EEG = offset_EEG * plotting_data['EEG']['scaling_factor']

    eog_channels_indices = mne.pick_types(raw.info, eog=True)
    eog_channels = []
    for channel_index in eog_channels_indices:
        eog_channels.append(raw.ch_names[channel_index])
    # print(eog_channels)

    default_channel_colors = []
    highlighted_channel_colors = []
    channel_visibility = []

    for channel_index in range(offset_EEG.shape[1]):
        # Calculate offset for y-axis
        offset_EEG[:, channel_index] = offset_EEG[:, channel_index] + ((plotting_data['plot']['offset_factor']) * (len(plotting_data['EEG']['channel_names']) - 1 - channel_index))  # First channel goes to top of the plot

        # Channel colors
        if plotting_data['EEG']['channel_names'][channel_index] in raw.info['bads']:
            default_channel_colors.append(c.BAD_CHANNEL_COLOR)
            channel_visibility.append(False)
        else:
            if plotting_data['EEG']['channel_names'][channel_index] in eog_channels:
                default_channel_colors.append('blue')
            else:
                default_channel_colors.append('black')
            channel_visibility.append(True)

        model_channel = False
        for model in range(len(model_channels)):
            if model_channels[model]:
                if plotting_data['EEG']['channel_names'][channel_index] in model_channels[model]:
                    model_channel = True

        if model_channel:
            highlighted_channel_colors.append('green')
        elif plotting_data['EEG']['channel_names'][channel_index] in raw.info['bads']:
            highlighted_channel_colors.append(c.BAD_CHANNEL_COLOR)
        else:
            highlighted_channel_colors.append('black')

    plotting_data['EEG']['offset_EEG_data'] = offset_EEG

    for model_index, model_array in enumerate(model_output):
        default_channel_colors.append(model_array)
        highlighted_channel_colors.append(model_array)
        channel_visibility.append(True)

        plotting_data['model'].append({})
        plotting_data['model'][model_index]['model_data'] = model_array
        plotting_data['model'][model_index]['model_channels'] = model_channels[model_index]

        plotting_data['model'][model_index]['model_timescale'] = np.linspace(0, plotting_data['EEG']['recording_length'], num=model_array.shape[0])

        plotting_data['model'][model_index]['offset_model_data'] = [-((2 + model_index) * (plotting_data['plot']['offset_factor'])) for i in range(len(plotting_data['model'][model_index]['model_timescale']))]

    plotting_data['EEG']['default_channel_colors'] = default_channel_colors
    plotting_data['EEG']['highlighted_channel_colors'] = highlighted_channel_colors
    plotting_data['EEG']['channel_visibility'] = channel_visibility

    y_ticks_model_output = np.arange((-len(plotting_data['model']) - 1), -1)
    y_ticks_channels = np.arange(0, len(plotting_data['EEG']['channel_names']))
    y_ticks = np.concatenate((y_ticks_model_output, y_ticks_channels))
    y_ticks = y_ticks * (plotting_data['plot']['offset_factor'])

    plotting_data['plot']['y_ticks'] = y_ticks

    y_tick_labels = [channel_name for channel_name in plotting_data['EEG']['channel_names']]
    for model_id in range(len(plotting_data['model'])):
        y_tick_labels.append('M{}'.format(model_id))
    y_tick_labels.reverse()

    plotting_data['plot']['y_tick_labels'] = y_tick_labels

    return plotting_data

def get_EEG_figure(file_name, raw, selected_channel_names, annotation_label, EEG_scale=None, channel_offset=None, model_output=None, model_channels=[], use_slider=False, show_annotations_only=False):
    """Generates initial EEG figure.

    Args:
        file_name (string): File name.
        raw (mne.io.Raw): Raw object to plot data from.
        selected_channel_names (list): List of strings of selected channel names to plot.
        annotation_label (string); Label for new annotations.
        EEG_scale (float): Desired scaling.
        channel_offset (int): Desired channel offset.
        model_output (list, optional): List of arrays of model outputs. Defaults to [].
        model_channels (list, optional): List of lists of strings of channel names model output is based on. Defaults to [].
        use_slider (bool, optional): Whether or not to activate view-slider. Defaults to False.

    Returns:
        plotly.graph_objs.Figure: Plot of EEG data.
    """
    fig = Figure()

    globals.plotting_data = _get_plotting_data(raw, file_name, selected_channel_names, EEG_scale, channel_offset, model_output, model_channels)
    # globals.plotting_data = plotting_data.copy()

    fig = get_EEG_plot(globals.plotting_data, globals.x0, globals.x1, annotation_label, use_slider, show_annotations_only)

    return fig

def get_EEG_plot(plotting_data, x0, x1, annotation_label, use_slider=False, show_annotations_only=False, skip_hoverinfo=False):
    """Generates EEG plots.

    Args:
        plotting_data (dict): Dict holding all relevant data from raw object and model outputs for plotting.
        x0 (float): X-coordinate (in seconds) to start plot.
        x1 (float): X-coordinate (in seconds) to end plot.
        annotation_label (string); Label for new annotations.
        use_slider (bool, optional): Whether or not to activate view-slider. Defaults to False.

    Returns:
        plotly.graph_objs.Figure: Plot of EEG data.
    """
    fig = Figure()

    # plotting_data = data_to_plot.copy()

    index_0 = None
    index_1 = None

    for index, timepoint in enumerate(plotting_data['EEG']['timescale']):
        if timepoint > x0:
            index_0 = index
            break

    for index, timepoint in enumerate(plotting_data['EEG']['timescale']):
        if timepoint > x1:
            index_1 = index
            break

    for channel_index in range(plotting_data['EEG']['offset_EEG_data'].shape[1]):
        fig.add_trace(
            Scattergl(
                x=plotting_data['EEG']['timescale'][index_0:index_1],
                y=plotting_data['EEG']['offset_EEG_data'][index_0:index_1, channel_index],
                marker=dict(color=plotting_data['EEG']['default_channel_colors'][channel_index], size=0.1),
                name=plotting_data['EEG']['channel_names'][channel_index],
                customdata=plotting_data['EEG']['EEG_data'][index_0:index_1, channel_index] * plotting_data['EEG']['scaling_factor'],  # y-data without offset
                hoverinfo='none' if skip_hoverinfo else 'all',
                hovertemplate='' if skip_hoverinfo else '%{fullData.name} | Time (in seconds)=%{x:.2f}, Amplitude (in μV)=%{customdata:.3f}' + '<extra></extra>' if plotting_data['EEG']['scaling_factor'] == c.CONVERSION_VALUE_VOLTS_TO_MICROVOLTS else '%{fullData.name} | Time (in seconds)=%{x:.2f}, Amplitude (scaled)=%{customdata:.3f}' + '<extra></extra>',
                mode='lines+markers'  # 'lines+markers'
            )
        )

    default_channel_colors = plotting_data['EEG']['default_channel_colors'].copy()
    highlighted_channel_colors = plotting_data['EEG']['highlighted_channel_colors'].copy()

    model_index_0 = None
    model_index_1 = None

    # Model predictions
    for model_index in range(len(plotting_data['model'])):
        for index, timepoint in enumerate(plotting_data['model'][model_index]['model_timescale']):
            if timepoint > x0:
                model_index_0 = index
                break

        for index, timepoint in enumerate(plotting_data['model'][model_index]['model_timescale']):
            if timepoint > x1:
                model_index_1 = index
                break

        trace_number = plotting_data['EEG']['offset_EEG_data'].shape[1] + model_index

        default_channel_colors[trace_number] = plotting_data['EEG']['default_channel_colors'][trace_number][model_index_0:model_index_1]
        highlighted_channel_colors[trace_number] = plotting_data['EEG']['highlighted_channel_colors'][trace_number][model_index_0:model_index_1]

        fig.add_trace(
            Scattergl(
                x=plotting_data['model'][model_index]['model_timescale'][model_index_0:model_index_1],
                y=plotting_data['model'][model_index]['offset_model_data'][model_index_0:model_index_1],
                marker=dict(
                    size=10,
                    cmax=1,
                    cmin=0,
                    color=plotting_data['model'][model_index]['model_data'][model_index_0:model_index_1],
                    # colorbar=dict(
                    #     title='Colorbar',
                    #     yanchor="top",
                    #     y=1,
                    #     x=1
                    # ),
                    colorscale='RdBu_r'
                ),
                name='M{}'.format(model_index),
                mode='markers',
                customdata=plotting_data['model'][model_index]['model_data'][model_index_0:model_index_1],
                hoverinfo='none' if skip_hoverinfo else 'all',
                hovertemplate='' if skip_hoverinfo else 'Time=%{x:.2f}, Prediction=%{customdata:.2f}<extra>%{fullData.name}</extra>'
            )
        )

    longest_channel_name_length = len(max(plotting_data['EEG']['channel_names'], key=len))

    fig.update_layout(
        plot_bgcolor='#fafafa',
        paper_bgcolor='#dfdfdf',
        title=dict(
            text=plotting_data['EEG']['file_name'],
            y=0.98,
            x=0.5 if not plotting_data['model'] else 0.6,
            xanchor='center',
            yanchor='top'
        ),
        xaxis=dict(
            # title_text='Time (in seconds)',
            rangeslider=dict(
                visible=True if use_slider else False,
                thickness=0.04
            ),
            type="linear",
        ),

        legend=dict(
            # itemwidth=30
            # orientation="h",
            # yanchor="bottom",
            # y=1,
            # xanchor="left",
            x=1.01
        ),

        autosize=False,
        margin=dict(
            autoexpand=False,
            l=longest_channel_name_length * 6 + 15,  #30
            r=longest_channel_name_length * 6 + 80,  #115
            # b=0,
            t=50,
            pad=5,
        ),

        dragmode='drawrect',
        newshape=dict(
            fillcolor=globals.annotation_label_colors[annotation_label],
            opacity=0.6,
            drawdirection='vertical',
            layer='below',
            line_width=0
        ),
    )

    # y_ticks_model_output = np.arange((-len(plotting_data['model']) - 1), -1)
    # y_ticks_channels = np.arange(0, len(plotting_data['EEG']['channel_names']))
    # y_ticks = np.concatenate((y_ticks_model_output, y_ticks_channels))
    # y_ticks = y_ticks * (plotting_data['plot']['offset_factor'])

    # plotting_data['plot']['y_ticks'] = y_ticks

    # y_tick_labels = [channel_name for channel_name in plotting_data['EEG']['channel_names']]
    # for model_id in range(len(plotting_data['model'])):
    #     y_tick_labels.append('M{}'.format(model_id))
    # y_tick_labels.reverse()

    # plotting_data['plot']['y_tick_labels'] = y_tick_labels

    fig.update_yaxes(
        tickmode='array',
        tickvals=plotting_data['plot']['y_ticks'],
        ticktext=plotting_data['plot']['y_tick_labels'],
        showgrid=False,
        zeroline=False,
        fixedrange=False,
        range=((-(2 + len(plotting_data['model'])) * (c.DEFAULT_Y_AXIS_OFFSET)), ((len(plotting_data['EEG']['channel_names']) + 1) * (c.DEFAULT_Y_AXIS_OFFSET)))  # Start y-axis range to cut off peaks
    )
    fig.update_xaxes(
        # title_text='Time (in seconds)'
        showgrid=True,
        zeroline=False,
        constrain='domain',
        range=(x0, x1) if (not use_slider or show_annotations_only) else (x0, x0 + 11),
    )

    # Add annotations
    marked_annotations = get_annotations(globals.raw)
    # print(marked_annotations)

    for annotation in marked_annotations:
        # if not ((annotation[0] < globals.x0 and annotation[1] < globals.x0) or (annotation[0] > globals.x1 and annotation[1] > globals.x1)):
        fig.add_vrect(
            editable=True,
            x0=annotation[0],
            x1=annotation[1],
            # annotation_text=annotation[2],
            fillcolor=globals.annotation_label_colors[annotation[2]] if annotation[2] in globals.annotation_label_colors.keys() else 'red',
            opacity=0.6,
            layer='below',
            line_width=0
        )

    fig.update_layout(
        updatemenus=list([
            dict(
                buttons=list([
                    dict(label="Reset time-axis",
                        method="relayout",
                        args=[{
                            "xaxis.range[0]": x0,
                            "xaxis.range[1]": x1 if (not use_slider) or show_annotations_only else x0 + 11,
                        }]
                    ),
                    dict(label="Reset channel-axis",
                        method="relayout",
                        args=[{
                            "yaxis.range[0]": (-(2 + len(plotting_data['model'])) * (c.DEFAULT_Y_AXIS_OFFSET)),
                            "yaxis.range[1]": ((len(plotting_data['EEG']['channel_names']) + 1) * (c.DEFAULT_Y_AXIS_OFFSET))
                        }]
                    ),
                    dict(label='Hide/show bad channels',
                        method='restyle',
                        args2=[{'visible': True}],
                        args=[{'visible': plotting_data['EEG']['channel_visibility']}]
                    ),
                    dict(label='Highlight model-channels',
                        method='restyle',
                        args2=[{'marker.color': default_channel_colors}],
                        args=[{'marker.color': highlighted_channel_colors}],
                        visible=True if plotting_data['model'] else False
                    ),
                ]),
                direction = 'left',
                # active=4,
                # pad = {'l': 100, 't': 10},
                showactive = False,
                type = 'buttons',
                xanchor = 'left',
                yanchor = 'top',
                x = 0,
                y = 1.08,
                bgcolor = '#fafafa',
                bordercolor = '#c7c7c7'
            )
        ])
    )

    return fig


#!/usr/bin/python3
"""
Simple API module
"""
from datetime import datetime, timezone
from flask import Flask, jsonify, request
from os import getenv

app = Flask(__name__)

@app.route('/api', strict_slashes=False)
def index():
    """
    Reads URL parameters and returns JSON content
    """
    days = ["Monday", "Tuesday", "Wednesday",
            "Thursday", "Friday", "Saturday", "Sunday"]

    slack_name = request.args.get('slack_name')
    track = request.args.get('track')

    print(request.query_string)

    repo_url = 'https://github.com/Michael-Maina/hng-projects'
    file_url = 'https://github.com/Michael-Maina/hng-projects/blob/master/stage_1/api.py'

    utc_time = datetime.now(timezone.utc)
    today = days[utc_time.weekday()]
    formatted_time = utc_time.strftime("%Y-%m-%dT%H:%M:%SZ")

    output = {
        'slack_name': slack_name,
        'current_day': today,
        'utc_time': formatted_time,
        'track': track,
        'github_file_url': file_url,
        'github_repo_url': repo_url,
        'status_code': 200
    }
    return jsonify(output), 200

if __name__ == "__main__":
    API_HOST = getenv('API_HOST')
    API_PORT = getenv('API_PORT')

    host = API_HOST if API_HOST else '0.0.0.0'
    port = API_PORT if API_PORT else 5000

    app.run(port=port, host=host)


# A Function to get the Integers for the Main List i.e. nums
def Input_List():
    list = []
    while True:
        user_input = input('Please enter the numbers for the list or "Done" to finish the list: ')
        if user_input == "done":
            break
        else:
            try:
                number = int(user_input)
                list.append(number)
            except ValueError:
                print(f'Invalid Input. Please enter an integer or "Done"')
    print(f'The Final List is: {list}')

    return list

# A Function to pause the program and get the user permisson to proceed
def press_enter_to_continue():
    while True:
        user_input = input("Are you Ready to find the Pivot Index of the Given List?(Press Enter to Proceed)")
        if user_input == "":
            break
        else:
            print("Invalid input. Please press Enter without typing anything.")

# Program to find the Pivot Index
nums = []
obtain_list = Input_List() # Get the return value of the predefined Function
nums.extend(obtain_list)  # Get the values to the nums variable
press_enter_to_continue() # Pause for the process

Pivot_Index = -1 # Initialize the value to -1 as if there is no Pivot Index

Total_Sum = sum(nums) # Plan 1
R_Sum = Total_Sum # Plan 1
L_Sum = 0 # Plan 1

for i in range(len(nums)): # Plan 2
    Current_Value = nums[i] # Plan 2
    R_Sum -= Current_Value # Plan 2

    if L_Sum == R_Sum: # Plan 3
        Pivot_Index = i # Plan 3

    L_Sum += Current_Value # Plan 4

print(f'The PIVOT INDEX of the given List is: {Pivot_Index}')

# ! this is a test comment
# ? what is this???


basket = ['a', 'x', 'b', 'c', 'd', 'e', 'd']

for item in basket:
    print(item)

for i, char in enumerate('Helloooooo'):
    print(i, char)

def loop_fifty_times(index):

    '''
    This is a docstring that just provides documentaion for methd
    '''

    i = index
    while i < 50:
        print(i)
        i += 1

# fund duplicates


#-*- coding:utf-8 -*-
# 栈特性:
# 队列最大的特性就是先进后出，后进先出，最形象的比喻就是摞在一起的盘子
# 用数组可以实现栈，链表也可以实现栈
# 栈主要有两个动作：压栈push和出栈pop
#
# Author: AllenGFLiu
# #####################################################

class ArrayStack():
    """使用Python内部容器list来模拟数组，实现栈的数据结构。
    基于数组的栈，一般都有容量限制。

    push:压入栈，若是已满则返回False,压入成功则返回True.

    pop:出栈，正常出栈返回出栈的值，若无数据可出，则返回None.
    """
    def __init__(self, capacity):
        self.array = ['*']*capacity  # 对self.array初始化，只是为了方便展示结构体
        self.capacity = capacity
        self.count = 0

    def push(self, value):
        if self.count == self.capacity:
            return False

        self.array[self.count]=value
        self.count += 1
        return True

    def dynamic_push(self, value):
        if self.count == self.capacity:
            new_list = ['*'] * self.capacity
            self.capacity *= 2
            self.array.extend(new_list)

        self.array[self.count] = value
        self.count += 1
        return True

    def pop(self):
        if self.count == 0:
            return None

        pop_value = self.array[self.count-1]
        self.array[self.count-1]='*'
        self.count -= 1
        return pop_value

    def __repr__(self):
        return '|'.join(self.array)

if __name__ == '__main__':
    my_stack = ArrayStack(3)
    my_stack.push('a')
    my_stack.push('b')
    my_stack.push('c')
    my_stack.dynamic_push('d')
    print(my_stack)
    value = my_stack.pop()
    print(value)
    print(my_stack)


import pymortar
import time
import pytz
import datetime
from rfc3339 import rfc3339
import xbos_services_getter as xbos

pymortar_client = pymortar.Client()

def get_outside_temp_data(building):

    interval = 3600

    outside_temperature_query = """SELECT ?temp WHERE {
        ?temp rdf:type brick:Weather_Temperature_Sensor .
    };"""

    weather_stations_view = pymortar.View(
        name="weather_stations_view",
        sites=[building],
        definition=outside_temperature_query,
    )

    weather_stations_stream = pymortar.DataFrame(
        name="weather_stations",
        aggregation=pymortar.MEAN,
        window=str(int(interval)) + 's',
        timeseries=[
            pymortar.Timeseries(
                view="weather_stations_view",
                dataVars=["?temp"],
            )
        ]
    )

    weather_stations_time_params = pymortar.TimeParams(
        start=rfc3339(start),
        end=rfc3339(end),
    )

    request = pymortar.FetchRequest(
        sites=[building],
        views=[
            weather_stations_view
        ],
        dataFrames=[
            weather_stations_stream
        ],
        time=weather_stations_time_params
    )

    df = outside_temperature_data = pymortar_client.fetch(request)['weather_stations']

    return df

#This query doesn't return any data for "hayward-station-1", but returns data for all other buildings
end = datetime.datetime.now().replace(tzinfo=pytz.utc) - datetime.timedelta(weeks=52)
start = end - datetime.timedelta(days=10)

for building in xbos.get_buildings(xbos.get_building_zone_names_stub()):
    #Looping through all buildings and getting response
    df = get_outside_temp_data(building)
    if df is None:
        print(building) #The only building that doesn't return data is "hayward-station-1"