import os
import sys
import math
import random
import pickle
import skimage
import numpy as np
import tensorflow as tf
import skimage.io as imgio
from datetime import datetime
import skimage.transform as imgtf

class dataset_dsrg():
    def __init__(self,config={}):
        self.config = config
        self.w,self.h = self.config.get("input_size",(321,321))
        self.categorys = self.config.get("categorys",["train"])
        self.category_num = self.config.get("category_num",21)
        self.main_path = self.config.get("main_path",os.path.join("data","VOCdevkit","VOC2012"))
        self.ignore_label = self.config.get("ignore_label",255)
        self.default_category = self.config.get("default_category",self.categorys[0])
        self.img_mean = np.ones((self.w,self.h,3))
        self.img_mean[:,:,0] *= 104.00698793
        self.img_mean[:,:,1] *= 116.66876762
        self.img_mean[:,:,2] *= 122.67891434

        self.data_f,self.data_len = self.get_data_f()

    def get_data_f(self):
        self.cues_data = pickle.load(open("data/localization_cues-sal.pickle","rb"),encoding="iso-8859-1")
        data_f = {}
        data_len = {}
        for category in self.categorys:
            data_f[category] = {"img":[],"gt":[],"label":[],"id":[],"id_for_slice":[]}
            data_len[category] = 0
        for one in self.categorys:
           assert one == "train", "extra category found!"
           with open(os.path.join("data","input_list.txt"),"r") as f:
               for line in f.readlines():
                   line = line.rstrip("\n")
                   id_name,id_identy = line.split(" ")
                   id_name = id_name[:-4] # then id_name is like '2007_007028'
                   data_f[one]["id"].append(id_name)
                   data_f[one]["id_for_slice"].append(id_identy)
                   data_f[one]["img"].append(os.path.join(self.main_path,"JPEGImages","%s.jpg" % id_name))
                   data_f[one]["gt"].append(os.path.join(self.main_path,"SegmentationClassAug","%s.png" % id_name))

               if "length" in self.config:
                   length = self.config["length"]
                   data_f[one]["id"] = data_f[one]["id"][:length]
                   data_f[one]["id_for_slice"] = data_f[one]["id_for_slice"][:length]
                   data_f[one]["img"] = data_f[one]["img"][:length]
                   data_f[one]["gt"] = data_f[one]["gt"][:length]
                   print("id:%s" % str(data_f[one]["id"]))
                   print("img:%s" % str(data_f[one]["img"]))
                   print("id_for_slice:%s" % str(data_f[one]["id_for_slice"]))

           data_len[one] = len(data_f[one]["id"])

        print("len:%s" % str(data_len))
        return data_f,data_len

    def next_batch(self,category=None,batch_size=None,epoches=-1):
        if category is None: category = self.default_category
        if batch_size is None:
            batch_size = self.config.get("batch_size",1)
        dataset = tf.data.Dataset.from_tensor_slices({
            "id":self.data_f[category]["id"],
            "id_for_slice":self.data_f[category]["id_for_slice"],
            "img_f":self.data_f[category]["img"],
            "gt_f":self.data_f[category]["gt"],
            })
        def m(x):
            id_ = x["id"]

            img_f = x["img_f"]
            img_raw = tf.read_file(img_f)
            img = tf.image.decode_image(img_raw)
            gt_f = x["gt_f"]
            gt_raw = tf.read_file(gt_f)
            gt = tf.image.decode_image(gt_raw)

            id_for_slice = x["id_for_slice"]
            def get_data(identy):
                identy = identy.decode()
                tag = np.zeros([self.category_num])
                tag[self.cues_data["%s_labels" % identy]] = 1.0
                tag[0] = 1.0
                cues = np.zeros([41,41,21])
                cues_i = self.cues_data["%s_cues" % identy]
                cues[cues_i[1],cues_i[2],cues_i[0]] = 1.0
                return tag.astype(np.float32),cues.astype(np.float32)

            tag,cues = tf.py_func(get_data,[id_for_slice],[tf.float32,tf.float32])
            img,gt,cues = self.image_preprocess(img,gt,cues,flip=True)

            img = tf.reshape(img,[self.h,self.w,3])
            gt = tf.reshape(gt,[self.h,self.w,1])
            tag.set_shape([21])
            cues.set_shape([41,41,21])

            return img,gt,tag,cues,id_

        dataset = dataset.repeat(epoches)
        dataset = dataset.shuffle(self.data_len[category])
        dataset = dataset.map(m)
        dataset = dataset.batch(batch_size)
        iterator = dataset.make_initializable_iterator()
        img,gt,tag,cues,id_ = iterator.get_next()

        return img,gt,tag,cues,id_,iterator

    def image_preprocess(self,img,gt,cues,flip=False):
        #img = tf.image.resize_image_with_crop_or_pad(img, self.h, self.w)
        #gt = tf.image.resize_image_with_crop_or_pad(gt, self.h, self.w)
        img = tf.expand_dims(img,axis=0)
        img = tf.image.resize_bilinear(img,(self.h,self.w))
        img = tf.squeeze(img,axis=0)
        gt = tf.expand_dims(gt,axis=0)
        gt = tf.image.resize_nearest_neighbor(gt,(self.h,self.w))
        gt = tf.squeeze(gt,axis=0)

        r,g,b = tf.split(axis=2,num_or_size_splits=3,value=img)
        img = tf.cast(tf.concat([b,g,r],2),dtype=tf.float32)
        img -= self.img_mean

        if flip is True:
            r = tf.random_uniform([1])
            r = tf.reduce_sum(r)
            img = tf.cond(r < 0.5, lambda:tf.image.flip_left_right(img),lambda:img)
            gt = tf.cond(r < 0.5, lambda:tf.image.flip_left_right(gt),lambda:gt)
            cues = tf.cond(r < 0.5, lambda:tf.image.flip_left_right(cues),lambda:cues)

        return img,gt,cues


# import os
# import numpy as np
#
# import torch
# import torch.nn as nn
# from torch.autograd import Variable
#
# from gensim.models.word2vec import Word2Vec
#
# from pymongo import MongoClient
#
# from tempfile import gettempdir
#
# from obonet import read_obo
#
# from src.python.preprocess import *
#
# import argparse
#
#
# # CNN Model (2 conv layer)
# class CNN_AC(nn.Module):
#     def __init__(self, n_classes):
#         super(CNN_AC, self).__init__()
#         self.n_classes = n_classes
#         self.layer1 = nn.Sequential(
#             nn.Conv2d(1, 16, kernel_size=5, padding=2),
#             nn.BatchNorm2d(16),
#             nn.ReLU(),
#             nn.MaxPool2d(2))
#         self.layer2 = nn.Sequential(
#             nn.Conv2d(16, 32, kernel_size=5, padding=2),
#             nn.BatchNorm2d(32),
#             nn.ReLU(),
#             nn.MaxPool2d(2))
#         self.fc = nn.Linear(8 * 5 * 32, n_classes)
#         self.sf = nn.Softmax()
#
#     def forward(self, x):
#         out = self.layer1(x)
#         out = self.layer2(out)
#         out = out.view(out.size(0), -1)
#         out = self.fc(out)
#         out = self.sf(out)
#         return out
#
#
# def train(model, train_dataset, test_dataset, output_dir):
#
#     print("Training Model: #TRN=%d , #VLD=%d, #CLS=%d" %
#           (len(train_dataset), len(test_dataset), model.n_classes))
#
#     # Hyper Parameters
#
#     batch_size = 32
#     num_epochs = 5
#     learning_rate = 0.003
#
#     # Loss and Optimizer
#     criterion = nn.MultiLabelMarginLoss()
#     optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
#
#     for epoch in range(num_epochs):
#
#         # pbar = tqdm(range(len(train_dataset)), "training ... ")
#
#         loader = DataLoader(train_dataset, batch_size)
#         for i, (seqs, lbls) in enumerate(loader):
#
#             model.train()
#             sequences = Variable(torch.from_numpy(seqs).float())
#             labels = Variable(torch.from_numpy(lbls).long())
#
#             # Forward + Backward + Optimize
#             optimizer.zero_grad()
#             outputs = model(sequences)
#
#             train_loss = criterion(outputs, labels)
#             train_loss.backward()
#             optimizer.step()
#
#             if (i + 1) % 100 == 0:
#
#                 test_loss = eval(model, criterion, test_dataset)
#                 print('Epoch [%d/%d], Step [%d/%d], Train Loss: %.4f, Test Loss: %.4f'
#                                      % (epoch + 1, num_epochs, i + 1, len(train_dataset) // batch_size,
#                                         train_loss.data[0], test_loss))
#             # pbar.update(batch_size)
#
#         # pbar.close()
#
#         # Save the Trained Model
#         torch.save(model.state_dict(), '%s/cnn.pkl' % output_dir)
#
#
# def eval(model, criterion, dataset):
#     model.eval()
#     batch_size = 32
#     loss = 0.0
#     loader = DataLoader(dataset, batch_size)
#     for i, (seqs, lbls) in enumerate(loader):
#         sequences = Variable(torch.from_numpy(seqs).float())
#         labels = Variable(torch.from_numpy(lbls).long())
#         outputs = model(sequences)
#         loss += criterion(outputs, labels).data[0]
#     return loss
#
#
# def main():
#
#     output_dir = args.out_dir
#     if not os.path.exists(output_dir):
#         os.mkdir(output_dir)
#
#     client = MongoClient("mongodb://127.0.0.1:27017")
#     db = client['prot2vec']
#
#     exp_codes = ["EXP", "IDA", "IPI", "IMP", "IGI", "IEP"]
#     queries = [{'Evidence': {'$in': exp_codes}, 'Aspect': asp, 'DB': 'UniProtKB'} for asp in ['F', 'C', 'P']]
#
#     num_mfo, num_cco, num_bpo = [db.goa_uniprot.count(q) for q in queries]
#     src_mfo, src_cco, src_bpo = [db.goa_uniprot.find(q).limit(10000) for q in queries]
#     # src_mfo, src_cco, src_bpo = [db.goa_uniprot.find(q) for q in queries]
#
#     print("#ANNO MF=%s CC=%s BP=%s" % (num_mfo, num_cco, num_bpo))
#     seqid2goid, goid2seqid = GoAnnotationCollectionLoader(src_mfo, 10000, 'F').load()
#     # seqid2goid, goid2seqid = GoAnnotationCollectionLoader(src_mfo, num_mfo, 'F').load()
#
#     query = {"_id": {"$in": unique(list(seqid2goid.keys())).tolist()}}
#     num_seq = db.uniprot.count(query)
#     src_seq = db.uniprot.find(query)
#     print("#SEQ=%s" % num_seq)
#     seqid2seq = UniprotCollectionLoader(src_seq, num_seq).load()
#
#     dataset = Dataset(seqid2seq, seqid2goid)
#     print(dataset)
#     p99 = np.percentile(list(map(lambda r: len(r.seq), dataset.records)), 99)
#     print("99 percentile:\t%d" % p99)
#
#     print("Reading GeneOntology")
#     GO = read_obo("http://purl.obolibrary.org/obo/go/go-basic.obo")
#
#     dataset.augment(GO)
#     print(dataset)
#
#     train_set, test_set = dataset.split(0.2)
#
#
#
# if __name__ == "__main__":
#
#     parser = argparse.ArgumentParser()
#
#     parser.add_argument("-o", "--out_dir", type=str, required=False,
#                         default=gettempdir(), help="Specify the output directory.")
#     parser.add_argument("--mongo_url", type=str, default='mongodb://localhost:27017/',
#                         help="Supply the URL of MongoDB")
#     parser.add_argument("--device", type=str, default='cpu',
#                         help="Specify what device you'd like to use e.g. 'cpu', 'gpu0' etc.")
#
#     args = parser.parse_args()
#
#     main()


# @author Simon Stepputtis <sstepput@asu.edu>, Interactive Robotics Lab, Arizona State University

from pyrep import PyRep
from pyrep.objects.vision_sensor import VisionSensor
from ros1compat import kdl_urdf_parser as kdl_parser_py
import PyKDL as kdl
import numpy as np
import random
import math
import scipy
from scipy.interpolate import interp1d
import hashids
import time
import json
from voice import Voice
from joblib import Parallel, delayed
import os

# How many processes shold collect data in parallel?
# A good measure is to put the number of CPU cores you have (Note that each process needs ~1GB of RAM also)
PROCESSES           = 4
# How many demonstrations (picking and pouring) should each process collect?
SAMPLES_PER_PROCESS = 10
# Ever n demonstrations, VRep will be restarted entirely, not just the simulation. You don't need to change this
RESET_EACH          = 20
# If you run more than 1 process, you should run VRep headless
VREP_HEADLESS       = True
# Default position of the UR5 robot. You do not need to change this
DEFAULT_UR5_JOINTS  = [105.0, -30.0, 120.0, 90.0, 60.0, 90.0]
# Path to the UR5 URDF file
ROBOT_URDF          = "../GDrive/ur5_robot.urdf"
# General speed of the robot. Lower values will increase the robot's movement speed
TGEN_SPEED_FACTOR   = 150
# Height at which to grasp the cups. You do not need to change this.
GRASP_HEIGHT        = 0.115
# Output directory of the collected data
DATA_PATH           = "../GDrive/dain/collected/"
# Where to find the VRep scene file. This has to be an absolute path.
VREP_SCENE          = "../GDrive/testscene.ttt"
VREP_SCENE          = os.getcwd() + "/" + VREP_SCENE

class SimulatorState(object):
    def __init__(self, array):
        self.array = array
        self.fromArray(array)

    def fromArray(self, array):
        if len(array) != 32:
            printError("Expected state of length 32, but got state of length " + str(len(array)))
        self.array = array
        self.data = {}
        self.data["joint_robot_position"]   = array[0:6]
        self.data["joint_robot_velocity"]   = array[6:12]
        self.data["tcp_position"]           = array[12:15]
        self.data["tcp_orientation"]        = array[15:18]
        self.data["tcp_linear_velocity"]    = array[18:21]
        self.data["tcp_angular_veloctiy"]   = array[21:24]
        self.data["tcp_target_position"]    = array[24:27]
        self.data["tcp_target_orientation"] = array[27:30]
        self.data["joint_gripper"]          = [array[30]]
        self.data["joint_gripper_velocity"] = [array[31]]

    def toArray(self):
        return self.array

class Robot(object):
    def __init__(self):
        kdl_tree           = kdl_parser_py.treeFromFile(ROBOT_URDF)[1]
        self.offsets       = np.deg2rad([90.0, -90.0, 0.0, -90.0, 0.0, 0.0])
        self.base_joints   = np.deg2rad(DEFAULT_UR5_JOINTS)
        transform          = kdl_tree.getChain("world", "tcp")

        self.kdl_fk        = kdl.ChainFkSolverPos_recursive(transform)
        self.kdl_ik        = kdl.ChainIkSolverPos_LMA(transform)
        self.kdl_input     = kdl.JntArray(transform.getNrOfJoints())
        self.kdl_output    = kdl.JntArray(transform.getNrOfJoints())
        self.num_joints    = transform.getNrOfJoints()
        if self.num_joints != len(self.base_joints):
            printError("Extracted chain has " + str(self.num_joints) + " joints, but " + len(self.base_joints) + " were expected.")

    def getJointAnglesFromCurrent(self, loc, rot, current):
        self._reset(current)
        return self._getJointAngles(loc, rot)

    def getJointAngles(self, loc, rot):
        self._reset()
        return self._getJointAngles(loc, rot)

    def getTcpFromAngles(self, angles):
        self._reset(angles)
        goal = kdl.Frame()
        self.kdl_fk.JntToCart(self.kdl_input, goal)
        return goal

    def _reset(self, target=None):
        if target is None:
            target = self.base_joints

        for joint_idx in range(self.num_joints):
            self.kdl_input[joint_idx] = target[joint_idx] + self.offsets[joint_idx]

    def _getJointAngles(self, loc, rot):
        goal   = kdl.Frame()
        goal.p = kdl.Vector(*loc)
        goal.M = kdl.Rotation.EulerZYX(rot[0], rot[1], rot[2],)

        self.kdl_ik.CartToJnt(self.kdl_input, goal, self.kdl_output)
        return [v - self.offsets[i] for i, v in enumerate(self.kdl_output)]

def genPosition(prev):
    px = 0
    py = 0
    done = False
    while not done:
        done = True
        px = random.uniform(-0.9, 0.35)
        py = random.uniform(-0.9, 0.35)
        dist = np.sqrt(px**2 + py**2)
        if dist < 0.5 or dist > 0.9:
            done = False
        for o in prev:
            if np.sqrt((px - o[0])**2 + (py - o[1])**2) < 0.25:
                done = False
        if px > 0 and py > 0:
            done = False
        angle = -45
        r_px  = px * np.cos(np.deg2rad(angle)) + py * np.sin(np.deg2rad(angle))
        r_py  = py * np.cos(np.deg2rad(angle)) - px * np.sin(np.deg2rad(angle))
        if r_py > 0.075:
            done = False
    return [px, py]

def _getCameraImage(camera):
    rgb_obs = camera.capture_rgb()
    rgb_obs = (np.asarray(rgb_obs) * 255).astype(dtype=np.uint8)
    rgb_obs = np.flip(rgb_obs, (2))
    return rgb_obs

def _getSimulatorState(pyrep):
    _, s, _, _ = pyrep.script_call(function_name_at_script_name="getState@control_script",
                                    script_handle_or_type=1,
                                    ints=(), floats=(), strings=(), bytes="")
    state = SimulatorState(s)
    return state

def _setRobotJoints(pyrep, joints):
        result = pyrep.script_call(function_name_at_script_name="setRobotJoints@control_script",
                                        script_handle_or_type=1,
                                        ints=(), floats=joints, strings=(), bytes="")

def _setJointVelocityFromTarget(pyrep, joints):
    _, s, _, _ = pyrep.script_call(function_name_at_script_name="setJointVelocityFromTarget@control_script",
                                    script_handle_or_type=1,
                                    ints=(), floats=joints, strings=(), bytes="")

def _setJointVelocityFromTarget_Direct(pyrep, joints):
        _, s, _, _ = pyrep.script_call(function_name_at_script_name="setJointVelocityFromTarget_Direct@control_script",
                                       script_handle_or_type=1,
                                       ints=(), floats=joints, strings=(), bytes="")

def _releaseObject(self):
    _, _, _, _ = self.pyrep.script_call(function_name_at_script_name="releaseObject@control_script",
                                    script_handle_or_type=1,
                                    ints=(), floats=(), strings=(), bytes="")

def _moveL(robot, current, goal, pad_one=False):
    g_pos, g_rot   = goal
    if type(g_pos) == np.ndarray:
        g_pos = g_pos.tolist()
    if type(g_rot) == np.ndarray:
        g_rot = g_rot.tolist()
    if type(current) == np.ndarray:
        current = current.tolist()
    padding        = 1.0 if pad_one else 0.0
    kdl_frame      = robot.getTcpFromAngles(current)
    s_pos          = [kdl_frame.p.x(), kdl_frame.p.y(), kdl_frame.p.z()]
    s_rot          = kdlFrameToRot(kdl_frame).tolist()
    distance       = np.sqrt( np.power(s_pos[0] - g_pos[0], 2) + np.power(s_pos[1] - g_pos[1], 2) + np.power(s_pos[2] - g_pos[2], 2) )
    steps          = max(2, int(np.ceil(distance * TGEN_SPEED_FACTOR))) # at least 2 steps
    tcp_trj        = np.linspace(s_pos + s_rot, g_pos + g_rot, num=steps, endpoint=True, axis=0)
    joint_trj      = np.zeros((steps, 7), dtype=np.float32)
    joint_trj[0,:] = current + [padding]
    for i in range(1,steps,1):
        joint_trj[i] = robot.getJointAnglesFromCurrent(loc=tcp_trj[i,:3], rot=tcp_trj[i,3:], current=joint_trj[i-1,:6]) + [padding]
    return joint_trj

def _moveJ(robot, current, goal, pad_one=False):
    g_pos, g_rot   = goal
    if type(g_pos) == np.ndarray:
        g_pos = g_pos.tolist()
    if type(g_rot) == np.ndarray:
        g_rot = g_rot.tolist()
    if type(current) == np.ndarray:
        current = current.tolist()
    kdl_frame      = robot.getTcpFromAngles(current)
    s_pos          = [kdl_frame.p.x(), kdl_frame.p.y(), kdl_frame.p.z()]
    goal_joints    = robot.getJointAnglesFromCurrent(loc=g_pos, rot=g_rot, current=current)
    distance       = np.sqrt( np.power(s_pos[0] - g_pos[0], 2) + np.power(s_pos[1] - g_pos[1], 2) + np.power(s_pos[2] - g_pos[2], 2) )
    steps          = max(2, int(np.ceil(distance * TGEN_SPEED_FACTOR))) # at least 2 steps
    joint_trj      = np.linspace(current, goal_joints, num=steps, endpoint=True, axis=0)
    padding        = np.ones((joint_trj.shape[0], 1)) if pad_one else np.zeros((joint_trj.shape[0], 1))
    joint_trj      = np.hstack((joint_trj, padding))
    return joint_trj

def createEnvironment(pyrep):
    _setRobotJoints(pyrep, np.deg2rad(DEFAULT_UR5_JOINTS))

    ncups  = np.random.randint(1,6)
    nbowls = np.random.randint(1,3)
    bowls  = np.random.choice(3, size=nbowls, replace=False) + 1
    cups   = np.random.choice(6, size=ncups, replace=False) + 1
    ints   = [nbowls, ncups] + bowls.tolist() + cups.tolist()
    floats = []

    prev   = []
    for i in range(nbowls + ncups):
        prev.append(genPosition(prev))
        floats += prev[-1]
        # if i < nbowls and bowls[i] > 10:
        #     floats += [np.random.uniform(-math.pi/4.0,  math.pi/4.0)]
        # else:
        #     floats += [0.0]
        floats += [0.0]

    result = pyrep.script_call(
        function_name_at_script_name="generateScene@control_script",
        script_handle_or_type=1,
        ints=ints,
        floats=floats,
        strings=(),
        bytes=""
    )
    return ints, floats

def kdlFrameToRot(frame):
    rotation = frame.M.GetEulerZYX()
    rotation = np.asarray(rotation, dtype=np.float32)
    # Weird fix...
    rotation[0] += np.deg2rad(360.0) if rotation[0] < 0 else 0.0
    return rotation

def _getObjectInfo(ints, floats, t_id, iscup):
    nbowls        = ints[0]
    ncups         = ints[1]
    bowl_ids      = ints[2:2+nbowls]
    cup_ids       = ints[2+nbowls:2+nbowls+ncups]

    index = None
    if iscup:
        index = np.argwhere(np.asarray(cup_ids)==t_id)[0,0] + nbowls
    else:
        index = np.argwhere(np.asarray(bowl_ids)==t_id)[0,0]

    data = floats[index * 3 : index * 3 + 3]
    return data

def _graspClosestObject():
    result = pyrep.script_call(
        function_name_at_script_name="graspObject@control_script",
        script_handle_or_type=1,
        ints=[],
        floats=[],
        strings=(),
        bytes=""
    )

def getCollisionWaypoints(robot, start, goal):
    x1 = start[0]
    y1 = start[1]
    x2 = goal[0]
    y2 = goal[1]

    distance = np.abs(x2*y1 - y2*x1) / np.sqrt(np.power(y2-y1,2) + np.power(x2-x1,2))

    if distance < 0.15:
        kdl_frame       = robot.getTcpFromAngles(np.deg2rad(DEFAULT_UR5_JOINTS))
        original_rot    = kdlFrameToRot(kdl_frame)
        original_pos    = [kdl_frame.p.x(), kdl_frame.p.y(), kdl_frame.p.z()]
        original_pos[2] = goal[2]
        return [("J", original_pos, original_rot)]
    else:
        return []

def _generatePouring(robot, current, rot):
    trj          = np.tile(np.expand_dims(current, 0),reps=[int(np.floor(75*rot)),1])
    steps        = trj.shape[0]
    target       = trj[0,5] + rot
    kdl_frame    = robot.getTcpFromAngles(current)
    current_pos  = [kdl_frame.p.x(), kdl_frame.p.y(), kdl_frame.p.z()]

    inter    = interp1d([0.0, 0.4, 0.6, 1.0], [current_pos[2], current_pos[2] - 0.075, current_pos[2] - 0.075, current_pos[2]], kind='cubic')
    zpos     = inter(np.linspace(0.0, 1.0, num=steps, endpoint=True))

    for i in range(steps):
        kdl_frame = robot.getTcpFromAngles(trj[i,:6].tolist())
        step_pos  = [kdl_frame.p.x(), kdl_frame.p.y(), kdl_frame.p.z()]
        step_rot  = kdlFrameToRot(kdl_frame)

        step_pos[2] = zpos[i]
        trj[i,:6]   = robot.getJointAnglesFromCurrent(loc=step_pos, rot=step_rot, current=trj[0,:6])

    inter    = interp1d([0.0, 0.4, 0.6, 1.0], [trj[0,5], target, target, trj[-1,5]], kind='cubic')
    trj[:,5] = inter(np.linspace(0.0, 1.0, num=steps, endpoint=True))

    return trj

def _setupTask(env, robot, current):
    # TODO edit
    task           = {}
    cur_phase      = 0
    cur_round      = 0
    # task["amount"] = np.random.choice([180, 110])
    task["amount"] = np.random.choice(["all", "one"])
    ints, floats   = env
    nbowls         = ints[0]
    ncups          = ints[1]
    bowl_ids       = ints[2:2+nbowls]
    cup_ids        = ints[2+nbowls:2+nbowls+ncups]
    task["num_rounds"] = 2 * np.random.randint(1, ncups) if task["amount"] == "all" else 2

    waypoints    = []
    kdl_frame    = robot.getTcpFromAngles(current)
    current_rot  = kdlFrameToRot(kdl_frame)
    current_pos  = [kdl_frame.p.x(), kdl_frame.p.y(), kdl_frame.p.z()]

    kdl_frame    = robot.getTcpFromAngles(np.deg2rad(DEFAULT_UR5_JOINTS))
    original_rot = kdlFrameToRot(kdl_frame)
    original_pos = [kdl_frame.p.x(), kdl_frame.p.y(), kdl_frame.p.z()]

    while cur_round < num_rounds:
        if cur_phase == 0:
            target    = np.random.choice(cup_ids)
            task["target/id"] = target
            task["target/type"] = "cup"
            target    = _getObjectInfo(ints, floats, target, iscup=True)
            target[2] = GRASP_HEIGHT
            rot       = [r for r in current_rot]
            rot[0]    += _calculateAngle(target[0], target[1])
            waypoints.append(("L", [t for t in target], rot))

            norm        = np.linalg.norm(target[:2], ord=2)
            factor      = norm / (norm * 0.85)
            target2     = [t for t in target]
            target2[0] /= factor
            target2[1] /= factor
            waypoints.insert(0, ("J", target2, rot))
            waypoints.append(("G", None, None))
            target3     = [t for t in target2]
            target3[2] += 0.10
            waypoints.append(("L", target3, rot))
            cur_phase += 1
        elif cur_phase == 1:
            target    = np.random.choice(bowl_ids)
            task["target/id"] = target
            task["target/type"] = "bowl"
            target    = _getObjectInfo(ints, floats, target, iscup=False)
            target[2] = current_pos[2]
            target    = adjustTargetForPouring(target[0], target[1], target[2])
            rot       = [r for r in original_rot]
            rot[0]    += _calculateAngle(target[0], target[1])
            waypoints += getCollisionWaypoints(robot, current_pos, target)
            waypoints.append(("J", target, rot))
            # waypoints.append(("P", None, np.deg2rad(task["amount"])))
            waypoints.append(("D", None, None))
            waypoints.append(("I", 40, None))
            cur_phase -= 1

        trajectory       = np.zeros((1,7), dtype=np.float32)
        trajectory[0,:6] = current
        grasp_active     = False if cur_phase == 0 else True
        if grasp_active:
            trajectory[0,-1] = 1.0

        for i, wp in enumerate(waypoints):
            if wp[0] == "L": # Move linear in tool space
                part = _moveL(robot, trajectory[-1,:6], wp[1:], pad_one=grasp_active)
            elif wp[0] == "J": # Move linear in joint space
                part = _moveJ(robot, trajectory[-1,:6], wp[1:], pad_one=grasp_active)
            elif wp[0] == "G": # Close gripper
                part = np.tile(trajectory[-1,:],reps=[30,1])
                part[:,-1] = 1.0
                grasp_active = True
            elif wp[0] == "P": # Do pouring motion
                part = _generatePouring(robot, trajectory[-1,:], wp[2])
            elif wp[0] == "D": # Drop object
                part = np.tile(trajectory[-1,:],reps=[30,1])
                part[:,-1] = 0.0
                grasp_active = False
            elif wp[0] == "I": # Idle a little
                part = np.tile(np.expand_dims(trajectory[-1,:], 0),reps=[wp[1],1])

            trajectory = np.vstack((trajectory, part))

        cur_round += 1

    task["trajectory"] = trajectory

    return task

def _calculateAngle(x, y):
    # Rotate:
    alpha = np.deg2rad(-135)
    point = np.dot( np.asarray([[np.cos(-alpha), np.sin(-alpha)],[-np.sin(-alpha), np.cos(-alpha)]]), np.asarray([x,y]) )

    angle = np.arctan(point[0] / point[1]) * -1.0
    return angle

def adjustTargetForPouring(x, y, z):
    for i in range(10):
        alpha    = np.deg2rad((i+1)/2.0)
        point    = np.dot( np.asarray([[np.cos(-alpha), np.sin(-alpha)],[-np.sin(-alpha), np.cos(-alpha)]]), np.asarray([x,y]) )
        distance = np.sqrt( np.power(x-point[0],2) + np.power(y-point[1],2) )
        if np.sqrt( np.power(x-point[0],2) + np.power(y-point[1],2) ) >= 0.025:
            break
    return [point[0], point[1], z]

def saveTaskToFile(path, task):
    for k, v in task.items():
        if type(v) == np.ndarray:
            task[k] = v.tolist()
        elif type(v) == np.int64:
            task[k] = int(v)
        elif type(v) in [str, list, int]:
            pass
        else:
            print("Serializing unhandled type", k, type(v))

    with open(path, "w") as fh:
        json.dump(task, fh)

def _generateVoice(voice, task):
    sentence =  voice.generateSentence(task)
    print("-> " + sentence)
    return sentence

def collectSingleSample(pyrep):
    robot = Robot()
    voice = Voice(load=False)
    rgb_camera = VisionSensor("kinect_rgb_full")
    pyrep.start()

    frame       = 0
    done        = False
    task        = None
    environment = None
    # phase       = 0
    _round      = 0
    num_rounds  = 0
    trj_step    = 0
    hid         = hashids.Hashids()
    task_name   = hid.encode(int(time.time() * 1000000))
    while not done:
        state     = _getSimulatorState(pyrep)
        if frame == 0:
            environment = createEnvironment(pyrep)
        elif task is None:
            task               = _setupTask(environment, robot, state.data["joint_robot_position"])
            task["name"]       = task_name
            # task["phase"]      = phase
            task["image"]      = _getCameraImage(rgb_camera)
            task["ints"]       = environment[0]
            task["floats"]     = environment[1]
            task["state/raw"]  = []
            task["state/dict"] = []
            task["voice"]      = _generateVoice(voice, task)
            num_rounds = task["num_rounds"]
        else:
            task["state/raw"].append(state.toArray())
            task["state/dict"].append(state.data)
            try:
                angles    = task["trajectory"][trj_step,:]
                trj_step += 1
            except IndexError:
                angles   = task["trajectory"][-1,:]
                _round  += 1
                name     = task_name + "_" + str(phase) + ".json"
                # saveTaskToFile(DATA_PATH + name, task)
                print('TASK COMPLETE')
                print(task)
                task     = None
                trj_step = 0

                if _round >= num_rounds:
                    done = True
            _setJointVelocityFromTarget(pyrep, angles)

        pyrep.step()
        frame += 1

    pyrep.stop()

def run():
    pyrep = None
    for i in range(SAMPLES_PER_PROCESS):
        if i % RESET_EACH == 0:
            if pyrep is not None:
                pyrep.shutdown()
            pyrep = PyRep()
            pyrep.launch(VREP_SCENE, headless=VREP_HEADLESS)

        collectSingleSample(pyrep)

    if pyrep is not None:
        pyrep.shutdown()

if __name__ == "__main__":
    # processes = [Process(target=run, args=()) for i in range(PROCESSES)]
    # [p.start() for p in processes]
    # [p.join() for p in processes]

    # Parallel(n_jobs=PROCESSES)(delayed(run)() for i in range(PROCESSES))
    run()


import torch
import torch.optim as optim
import numpy as np
from torch.distributions import Normal, Categorical
import torch.nn.functional as F
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
class PPO(object):
    def __init__(self, memory, acnet, clip_param_start, clip_param_now,
                 num_epi, clip_linear_jud, max_grad_norm,
                 ppo_update_time, buffer_capacity, learn_start,
                 batch_size, learning_rate, gamma):
        self.acnet = acnet
        self.gamma = gamma
        self.learning_rate = learning_rate
        self.batch_size = batch_size
        self.clip_param_now = clip_param_now
        self.memory = memory
        self.counter = 0
        self.training_step = 0
        self.optimizers = optim.Adam(self.acnet.parameters(), self.learning_rate)
        self.scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(self.optimizers, 300, eta_min=0.0001, last_epoch=-1)

    def select_action(self, state, pre_action, batch_size=1, ht=None, ct=None, test=False):
        with torch.no_grad():
            action_prob,_, hc, cc = self.acnet(state, pre_action, batch_size, 1, ht=ht, ct=ct, test=False)
        c = Categorical(action_prob)
        action = c.sample()

        return action.item(), action_prob[:, action.item()].item(), action_prob, hc, cc

    def update(self):
        loss = 0
        actorloss = 0
        criticloss = 0
        state, action, reward, old_action_log_prob , done = self.memory.getEpiBatch()

        for o in range(1):
            for m in range(3):
                s_t_cur = state[m]
                action_cur = action[m]
                reward_cur = reward[m]
                policy_old = old_action_log_prob[m]
                done_cur = done[m]
                num = len(reward_cur)
                if num < 1:
                    continue
                rnn_length = m + 1
                rnn_batch = num // rnn_length

                action_in = np.zeros([num, 12])
                for k in range(rnn_batch):
                    for p in range(rnn_length):
                        if p > 0:
                            idx = k * rnn_length + p
                            action_in[idx, action_cur[idx - 1]] = 1.
                R = 0
                Gt = []
                for i in reversed(range(int(len(reward_cur)))):
                    R = reward_cur[i] + self.gamma * R
                    Gt.insert(0, R)
                    if done_cur[i] == True and i != (len(reward_cur)-1):
                        R = 0
                Gt = torch.FloatTensor(Gt).cuda()
                s_t_cur = torch.FloatTensor(s_t_cur.transpose(0, 3, 1, 2)).cuda()
                action_in = torch.FloatTensor(action_in).cuda()
                action_cur = torch.LongTensor(action_cur).cuda()
                policy_old = torch.FloatTensor(policy_old).cuda()

                Gt_index = Gt.view(-1, 1)
                action_prob, V,_,_ = self.acnet(s_t_cur, action_in, s_t_cur.size(0), length=rnn_length, off=True)
                delta = Gt_index - V
                advantage = delta.detach()
                if advantage.size(0) > 1:
                    advantage = (advantage - advantage.mean(0).item())/(advantage.std(0).item() + 1e-10)
                action_cur = action_cur.view(s_t_cur.size(0), 1)
                action_prob = action_prob.gather(1, action_cur)  # new policy

                ratio = (action_prob / policy_old)

                surr1 = ratio * advantage
                surr2 = torch.clamp(ratio, 1 - self.clip_param_now, 1 + self.clip_param_now) * advantage

                loss_entropy = torch.mean(torch.log(action_prob) * action_prob)
                action_loss = -torch.min(surr1, surr2).mean()
                value_loss = F.mse_loss(Gt_index.detach(), V)
                loss = action_loss + 0.5 * value_loss + 0.01 * loss_entropy
                actorloss = action_loss
                criticloss = value_loss
                self.optimizers.zero_grad()
                loss.backward()
                self.optimizers.step()
                self.training_step += 1
        return loss, actorloss, criticloss


#!/bin/env python
import math
#import simtk.openmm.app as soa
#import simtk.unit as su
import sys,os
import numpy as np
from .fhet import hetnames,cofactors
import copy

MASS =    {'O':15.999,  'N':14.010,
           'C':12.010,  'H': 1.008,
           'F':19.000,
           'Na':22.99,  'NA':22.99,
           'P':30.970,  'S':32.060,
           'Cl':35.45,  'CL':35.45,
           'Br':79.90,  'BR':79.90,
           'CU':63.55,
           }
ATOM_NUM = {'O': 8,   'N': 7,
            'C': 6,   'H': 1,
            'F': 9,
            'Na':11,  'NA':11,
            'P':15,   'S':16,
            'Cl':17,  'CL':17,
            'Br':35,  'BR':35,
            'CU':29,
            'X':999
            }

MAX = 99999
TMPFILE = "FPDB.SOA.TMPPDBFILE.PDB"
kJ_to_kcal = 1/4.184
BIG_NUM = 1.e10
HBOND_DISTANCE_CUTOFF = 3.5  # angstron
HBOND_ANGLE_CUTOFF = 0.666667*math.pi # pi

PROGS = { "I-interpret":"/home/fuqy/Software/I-interpret/bin/I-interpret" ,
          "pdbconvert":"/opt/schrodinger2016-2/utilities/pdbconvert",
          "hetgrp_ffgen":"/opt/schrodinger2016-2/utilities/hetgrp_ffgen",
          "opls_to_gmx":"~/.ffallrain/fqy_scripts/opls2005_to_gmx.py",
        }

if True: ### residue names
    standard_protein_residues = ('ARG','LYS','HIS','HIP','HIE',
         'HID','ASP','GLU',
         'ASN','GLN','SER','THR','CYS','CYX','GLY','PRO','ALA',
         'VAL','LEU','ILE','MET','PHE','TYR','TRP' )

    standard_ion_redsidues = ('FE','MN','NI','CA','ZN','MG',"CL","NA")
    standard_water_residues = ('HOH','SOL','WAT')
    standard_residues = standard_protein_residues + standard_water_residues + standard_ion_redsidues

    protein_hbond_donors = (
        ('*','N','H'),
        ('*','N','HN'),
        ('ARG', 'NE','HE'),
        ('ARG','NH1','HH11'),
        ('ARG','NH1','HH12'),
        ('ARG','NH2','HH21'),
        ('ARG','NH2','HH22'),
        ('LYS', 'NZ', 'HZ1'),
        ('LYS', 'NZ', 'HZ2'),
        ('LYS', 'NZ', 'HZ3'),
        ('HIS','ND1', 'HD1'),
        ('HIS','NE2', 'HE2'),
        ('HIP','ND1', 'HD1'),
        ('HIP','NE2', 'HE2'),
        ('HID','ND1', 'HD1'),
        ('HIE','NE2', 'HE2'),
        ('ASN','ND2','HD21'),
        ('ASN','ND2','HD22'),
        ('GLN','NE2','HE22'),
        ('GLN','NE2','HE21'),
        ('SER','OG','HG'),
        ('THR','OG1','HG1'),
        ('CYS','SG','HG'),
        ('TYR','OH','HH'),
        ('TRP','NE1','HE1'),
    )

    one_letter_code = {
               'ala'.upper() : 'A' ,
               'arg'.upper() : 'R' ,
               'asn'.upper() : 'N' ,
               'asp'.upper() : 'D' ,
               'asx'.upper() : 'B' ,
               'cys'.upper() : 'C' ,
               'glu'.upper() : 'E' ,
               'gln'.upper() : 'Q' ,
               'glx'.upper() : 'Z' ,
               'gly'.upper() : 'G' ,
               'his'.upper() : 'H' ,
               'ile'.upper() : 'I' ,
               'leu'.upper() : 'L' ,
               'lys'.upper() : 'K' ,
               'met'.upper() : 'M' ,
               'phe'.upper() : 'F' ,
               'pro'.upper() : 'P' ,
               'ser'.upper() : 'S' ,
               'thr'.upper() : 'T' ,
               'trp'.upper() : 'W' ,
               'tyr'.upper() : 'Y' ,
               'val'.upper() : 'V' ,
    }
    three_letter_code = dict()
    for key,value in list(one_letter_code.items()):
        three_letter_code[value] = key

if True: ### Global varieties
    newnew = '/home/fuqiuyu/.ffallrain/newnew'
    TMPFILE = 'FIND_RING.PDB'
    TMPFILE_MOL2 = 'FIND_RING.mol2'
    babel = '/usr/bin/babel'

class fATOM():
    def __init__(self,atom_line = None):
        if atom_line == None:
            atom_line = "ATOM      1  X   DEF     1       0.000   0.000   0.000  1.00  0.00           X"

        tmpname = atom_line[12:16]
        if tmpname[0] in "0123456789" and tmpname[3]!=" ":
            tmpname = tmpname[1:] + tmpname[0]
        tmpname = tmpname.strip()
        self.name = tmpname
        try:
            if len(atom_line)>=78:
                self.element = atom_line[76:78].strip()
            else:
                self.element = tmpname[0]
            self.atom_num = ATOM_NUM[self.element]
        except:
            # print("DEBUG",tmpname)
            self.element = tmpname[0]  ### !!!! NOT FINISHED
        self.resi_name = atom_line[17:20].strip()
        try:
            self.resi_index = int(atom_line[22:26])
        except:
            self.resi_index = 0
        self.charge = None
        self.sig = None
        self.eps = None
        self.conf = atom_line[16]
        self.index = int(atom_line[6:11])
        x = float(atom_line[30:38])
        y = float(atom_line[38:46])
        z = float(atom_line[46:54])
        self.posi = [x,y,z]
        try:
            self.bf =  float(atom_line[60:66])
        except:
            self.bf = 0
        try:
            self.occ =  float(atom_line[54:60])
        except:
            self.occ = 0
        self.connect_count = 0 ## for bond
        self.bond = list()
        self.var1 = None
        self.var2 = None
        self.var3 = None

    def addparm(self,charge,sig,eps):
        self.charge = charge
        self.sig = sig
        self.eps = eps

    def _addconnect(self):
        self.connect_count += 1

    def _delconnect(self):
        self.connect_count -= 1
        if self.connect_count < 0 :
            sys.stderr.write("WARNING,BOND CONNECTION LESS THAN ZERO.\n")
    def addbond(self,bond):
        self.bond.append(bond)
        self._addconnect()
    def delbond(self,bond):
        self.bond.remove(bond)
        self._delconnect()

class fCHEMO():
    @staticmethod
    def _next_atom_line(resi_lines):
        for atom_line in resi_lines:
            if len(atom_line)>=6 and atom_line[:6] in ('HETATM','ATOM  '):
                yield atom_line

    ### Did not finished !
    def addH(self,keep_current = False,nc = 0 ):
        sys.stderr.write("##### FPDB WARNING: Currently, using fCHEMO.addH will remove the names/indexes of the exist atoms !!\n")
        parameter_text = '''
                       SP1_CUTOFF_ANGLE = 155.0
                       SP2_CUTOFF_ANGLE = 115.0
                      TORSION_RING_FIVE = 7.5
                       TORSION_RING_SIX = 15.0
                     FLAT_TORSION_ANGLE = 30.0

                       CHARGED_CARBOXYL = YES
                     ALIPHATIC_NITROGEN = YES
                    CHARGED_GUANIDINIUM = YES
                   ODD_RING_AROMATICITY = YES
                   CONVERT_DATIVE_BONDS = YES

                   LARGEST_SUBSTRUCTURE = NO
                     ADD_HYDROGEN_ATOMS = YES
                  HEAVY_ATOM_STATISTICS = NO
        '''
        ofp = open("parameter.txt",'w')
        ofp.write(parameter_text)
        ofp.close()
        ofp = open("FPDB_TMPLIG.pdb",'w')
        self.write_pdb(ofp)
        ofp.close()
        os.system("%s FPDB_TMPLIG.pdb FPDB_OUTLIG.pdb"%PROGS["I-interpret"])
        os.system("rm FPDB_TMPLIG.pdb parameter.txt")
        lines = list()
        for line in open("FPDB_OUTLIG.pdb"):
            if len(line)>=6 and line[:6] in ("HETATM","ATOM  "):
                lines.append(line)
        self._update(resi_lines = lines)
        for atom in self.atoms:
            atom.conf = ' '

    def _update(self,resi_lines):
        for line in resi_lines:
            H_index = 1
            tmp_atom = fATOM(line)
            exist_flag = False
            for atom in self.atoms:
                if dist_2(atom,tmp_atom)<0.1:
                    exist_flag = True
                    break
            if exist_flag:
                pass
            else:
                name = "H%d"%H_index
                while ( name in list(self.atoms_d.keys()) ):
                    H_index += 1
                    name = "H%d"%H_index
                tmp_atom.name = name
                self.atoms.append(tmp_atom)
                self.atoms_d[name] = tmp_atom

    def rename_atom(self):
        element_index = dict()
        for atom in self.atoms:
            if atom.element not in element_index:
                element_index[atom.element] = 0
                newname = atom.element+"%d"%element_index[atom.element]
            else:
                element_index[atom.element] += 1
                newname = atom.element+"%d"%element_index[atom.element]
            atom.name = newname
            self.atoms_d[atom.name] = atom

    def generate_OPLS_parameters(self, nc = 0, version = '2005', rename_atom = False):
        if rename_atom:
            self.rename_atom()

        self.removeH()
        self.addH( keep_current = False,nc = nc)
        resname = self.name.lower()
        self.write_pdb("%s.pdb"%resname)
        os.system("%s -ipdb %s.pdb -omae %s.mae"%(PROGS['pdbconvert'],resname,resname))
        os.system("%s %s %s.mae"%(PROGS['hetgrp_ffgen'],version,resname))
        os.system("%s -f %s.pdb -o %s.top -p %s"%(PROGS['opls_to_gmx'],resname,resname,resname))
        os.system("mkdir %s_parameter"%resname)
        os.system("mv %s %s.mae %s.pdb %s.top %s_nb.itp %s_parameter/"%(resname,resname,resname,resname,resname,resname))
        sys.stdout.write(">>>>> FPDB, parameters of compound %s generated. \n"%resname)

        # load parameters
        nb_itp = "%s_parameter/%s_nb.itp"%(resname,resname)
        top = "%s_parameter/%s.top"%(resname,resname)
        self.load_parameters_from_file(nb_itp,top)

    def load_parameters_from_file(self,nb_itp,top):
        try :
            # vdw
            flag = False
            for line in open(nb_itp):
                if "atomtypes" in line:
                    flag = True
                    continue
                if flag :
                    if "name" not in line:
                        if line.strip() != "":
                            name = line.split()[0]
                            sig = float(line.split()[5])
                            eps = float(line.split()[6])
                            self.atoms_d[name].sig = sig
                            self.atoms_d[name].eps = eps
            # charge
            flag = False
            for line in open(top):
                if "atoms" in line:
                    flag = True
                    continue
                elif "bonds" in line:
                    flag = False
                    continue
                if flag :
                    if line.strip() != "":
                        if line.strip()[0] != ";":
                            name = line.split()[1]
                            charge = float(line.split()[6])
                            self.atoms_d[name].charge = charge
        except Exception as e:
            print(("##### Load parameter Fail: %s"%self.name))
            print((e.message))

    def __init__(self,resi_lines=None):
        if resi_lines == None:
            resi_lines = list()
        try:
            self.name = resi_lines[0][17:20].strip()
            self.index = int(resi_lines[0][22:26])
            self.chain = resi_lines[0][21]
            self.insertion = resi_lines[0][26]
        except:
            self.name = "UND"
            self.index = 0
            self.chain = " "
            self.insertion = " "

        self.atoms = list()
        self.index_shift = 0

        for atom_line in fRESIDUE._next_atom_line(resi_lines):
            self.atoms.append( fATOM(atom_line) )
        if len(self.atoms)>0:
            self.index_shift = self.atoms[0].index - 1
        else:
            self.index_shift = 0
        self.atoms_d = dict()
        for atom in self.atoms:
            self.atoms_d[atom.name] = atom
        self.var1 = None
        self.var2 = None
        self.var3 = None
        self.del_duplicate_atoms()

    def getCOM(self):
        mass = np.array([MASS[i.element] for i in self.atoms])
        coors = self.getxyzs
        return np.dot(mass, coors) / sum(mass)

    def getxyzs(self):
      return np.array([i.posi for i in self.atoms])

    def updatexyzs(self, xyzs):
      for i, xyz in zip(self.atoms, xyzs):
        i.posi = xyz

    def add_atom(self,atom):
        if hasattr(atom,'name'):
            self.atoms.append(atom)
            self.atoms_d[atom.name] = atom
        else :
            try:
                a = fATOM(atom)
                self.atoms.append( a )
                self.atoms_d[a.name] = a
            except:
                sys.stderr.write("##### Error, Error loading atom %s"%atom)

    def remove_atom(self,atom):
        self.atoms.remove(atom)
        if atom.name in self.atoms_d:
            del(self.atoms_d[atom.name])

    def removeH(self):
        to_del = [ x for x in self.atoms if x.element == "H" ]
        for atom in to_del:
            self.remove_atom(atom)

    def find_atoms(self, name = None , index = None ):
        result = list()
        ignore_name = False
        ignore_index = False
        if name == None:
            ignore_name = True
        if index == None:
            ignore_index = True
        for atom in self.atoms:
            if ignore_name or atom.name in name:
                if ignore_index or atom.index in index:
                        result.append(atom)
        return result

    def write_pdb(self,ofile=None):
        pdbstr = ""
        for atom in self.atoms:
            x,y,z = atom.posi

            tmpname = atom.name
            if len(tmpname) == 4:
                tmpname = tmpname[-1]+tmpname[:-1]
            else:
                tmpname = " "+tmpname+" "*(3-len(tmpname))

            line = 'ATOM  %5d %4s%1s%-4s%1s%4d    %8.3f%8.3f%8.3f%6.2f%6.2f%12s\n'%(
                    atom.index,tmpname,atom.conf,self.name,self.chain,self.index,
                    x,y,z,atom.occ,atom.bf,atom.element)
            pdbstr += line
        if ofile is None:
            return pdbstr
        elif hasattr(ofile,'write'):
            ofile.write(pdbstr)
        else:
            ofp = open(ofile,'w')
            ofp.write(pdbstr)

    def write_pdb_plop(self,ofile=None):
        # make atomline : case if protein or metal
        if self.name in standard_protein_residues:
            atomline='ATOM  %5d  %-3s%1s%3s %1s%4d    %8.3f%8.3f%8.3f  1.00  0.00\n'
        else:
            atomline='HETATM%5d  %-3s%1s%3s %1s%4d    %8.3f%8.3f%8.3f  1.00  0.00\n' # qiuyu Fu
        if self.name in standard_ion_redsidues:
            atomline='HETATM%5d %-4s%1s%3s %1s%4d    %8.3f%8.3f%8.3f  1.00  0.00\n'
        # Case: no ofp
        if ofile is None:
            ofp = str()
            if self.name not in standard_protein_residues: ofp +='TER\n'
            for atom in self.atoms:
                x,y,z = atom.posi
                line = atomline%(atom.index,atom.name,atom.conf,self.name,self.chain,self.index,x,y,z)
                ofp += line
            if self.name not in standard_protein_residues: ofp +='TER\n'
            return ofp
        # Case: file handle
        elif hasattr(ofile,'write'):
            ofp = ofile
            if self.name not in standard_protein_residues:ofp.write('TER\n')
            for atom in self.atoms:
                pass
                if len(atom.name) != 4 :
                    x,y,z = atom.posi
                    line = atomline%(atom.index,atom.name,atom.conf,self.name,self.chain,self.index,x,y,z)
                    ofp.write(line)
                else:
                    tmpatomline='HETATM%5d %-4s%1s%3s %1s%4d    %8.3f%8.3f%8.3f  1.00  0.00\n'
                    x,y,z = atom.posi
                    line = tmpatomline%(atom.index,atom.name,atom.conf,self.name,self.chain,self.index,x,y,z)
                    ofp.write(line)
            if self.name not in standard_protein_residues:ofp.write('TER\n')
        # Case: String ( filename )
        else:
            ofp = open(ofile,'a')
            if self.name not in standard_protein_residues:ofp.write('TER\n')
            for atom in self.atoms:
                pass
                if len(atom.name) != 4 :
                    x,y,z = atom.posi
                    line = atomline%(atom.index,atom.name,atom.conf,self.name,self.chain,self.index,x,y,z)
                    ofp.write(line)
                else:
                    tmpatomline='HETATM%5d %-4s%1s%3s %1s%4d    %8.3f%8.3f%8.3f  1.00  0.00\n'
                    x,y,z = atom.posi
                    line = tmpatomline%(atom.index,atom.name,atom.conf,self.name,self.chain,self.index,x,y,z)
                    ofp.write(line)
            if self.name not in standard_protein_residues:ofp.write('TER\n')
            ofp.close()

    def debug(self):
        print(('name',self.name))
        print(('index',self.index))
        print(('insertion',self.insertion))
        print('atoms:')
        for atom in self.atoms:
            print((atom.name,))
        print()

    def find_h(self):
        h_list = list()
        for atom in self.atoms:
            if atom.element == 'H':
                h_list.append(atom)
        #print "DEBUG:",h_list
        return h_list

    def _find_bond_atoms(self,atom,cutoff):
        cutoff_2 = cutoff ** 2
        a_list = list()
        for a in self.atoms:
            if a == atom:
                continue
            if dist_2(atom,a) < cutoff_2:
                a_list.append(a)
        return a_list

    def find_polar_h(self):
        polar_h_list = list()
        for atom in self.find_h():
            for a in self._find_bond_atoms(atom, 1.2): # bond length with H is usually < 1.0 A
                if a.element in ("N","O","S"): # currently only take O,N and S into account.
                    polar_h_list.append(atom)
        return polar_h_list

    def find_hbond_acceptor(self):
        l = list()
        for atom in self.atoms:
            if  atom.element in ("N","O"):
                l.append(atom)
        return l

    def find_hbond_donar(self):
        donar = list()
        for atom in self.find_h():
            for a in self._find_bond_atoms(atom, 1.2): # bond length with H is usually < 1.0 A
                if a.element in ("N","O","S"): # currently only take O,N and S into account.
                    donar.append( (a,atom) )

        return donar

    def find_hbond_with(self,resi): ## use three
        return self.find_hbond(resi)

    def find_hbond(self,resi): ## use three
        if not hasattr(resi,'atoms'):
            print("##### Error. Currently the only usage of find_h_bond() is : self.find_h_bond( another_residue ). Will stop ")
            sys.exit()
        if self == resi:
            print("##### Warning: find hbond between identical residues. Will return empty list.")
            return list(),list()

        as_donar = list()
        for d,h in self.find_hbond_donar():
            for a in resi.find_hbond_acceptor():
                if fHbond.is_hydrogen_bond(d,h,a):
                    as_donar.append( fHbond(d,h,a) )
        as_acceptor = list()
        for a in self.find_hbond_acceptor():
            for d,h in resi.find_hbond_donar():
                if fHbond.is_hydrogen_bond(d,h,a):
                    as_acceptor.append( fHbond(d,h,a) )
        return as_donar,as_acceptor

    def del_duplicate_atoms(self):
        #
        names = set()
        tmp_atoms = list()
        for atom in self.atoms:
            if atom.name not in names:
                tmp_atoms.append(atom)
                names.add(atom.name)
            else:
                pass
        self.atom = tmp_atoms

class fSDF_MOL(fCHEMO):
    def __init__(self, sdf_frame):
        fCHEMO.__init__(self)
        n_atom = int(sdf_frame[3].split()[0])
        atomlines = sdf_frame[4:4+n_atom]
        self.name = "___"
        self.index = 0
        for atomline in atomlines:
            tmpatom = fATOM()
            items = atomline.split()
            x = float(items[0])
            y = float(items[1])
            z = float(items[2])
            name = items[3]
            tmpatom.posi = (x,y,z)
            tmpatom.index = 0
            tmpatom.name = name
            tmpatom.element = name[0]
            self.add_atom(tmpatom)

class fRESIDUE(fCHEMO):
    pass

class fRING(fCHEMO):
    def __init__(self,atom_list):
        fCHEMO.__init__(self,[])
        for atom in atom_list:
            self.atom.append(atom)

class fHbond(fCHEMO):
    def __init__(self,d,h,a):
        fCHEMO.__init__(self)
        self.add_atom(d)
        self.add_atom(h)
        self.add_atom(a)
        self.d = d
        self.h = h
        self.a = a
        self.angle = angle(d,h,a)
        self.dist = dist(d,a)
        if self.angle  < HBOND_ANGLE_CUTOFF:
            sys.stderr.write( "##### Warning ! You registed a hydrogen bond, but angle(d,h,a) is smaller than HBOND_ANGLE_CUTOFF(%f)\n"%HBOND_ANGLE_CUTOFF )
        if self.dist  > HBOND_DISTANCE_CUTOFF:
            sys.stderr.write( "##### Warning ! You registed a hydrogen bond, but dist(d,a) is larger than HBOND_DISTANCE_CUTOFF(%f)\n"%HBOND_DISTANCE_CUTOFF )

    @staticmethod
    def is_hydrogen_bond(d,h,a):
        tmp_angle = angle(d,h,a)
        tmp_dist = dist(d,a)
        if tmp_angle >=  HBOND_ANGLE_CUTOFF and tmp_dist <= HBOND_DISTANCE_CUTOFF:
            return True
        else:
            return False

class fBOND:
    def __init__(self,atoma,atomb):
        # print "DEBUG ADD BOND BETWEEN",atoma.name,'and',atomb.name
        if atoma.name<atomb.name:
            atoma.addbond(self)
            atomb.addbond(self)
            self._ = (atoma,atomb)
        elif atoma.name>atomb.name:
            atoma.addbond(self)
            atomb.addbond(self)
            self._ = (atomb,atoma)
        else:
            pass
            sys.stderr.write("WARNING,INTENDED TO CONNECT BOND BETWEEN IDENTICAL ATOMS\n")
    def del_me(self):
        self._[0].delbond(self)
        self._[1].delbond(self)

class fCOMPOUND(fCHEMO):
    BOND_CUTOFF_H = 1.5 # Angstrom
    BOND_CUTOFF_H_2 = 2.25 # Angstrom
    BOND_CUTOFF = 2.0 # Angstrom
    BOND_CUTOFF_2 = 4.0 # Angstrom
    def __init__(self,compound_lines):
        fCHEMO.__init__(self,compound_lines)
        self._atom_classification()
        self._makebond()
    def _atom_classification(self):
        self.h_list = list()
        self.heavy_list = list()
        for atom in self.atoms:
            if len(atom.name)>=1 and atom.name[0] == 'H' :
                self.h_list.append(atom)
            else:
                self.heavy_list.append(atom)
    def _makebond(self):
        self.bonds = set()
        self._make_h_bond()
        self._make_heavy_bond()
    def _make_h_bond(self):
        for h in self.h_list:
            for heavy in self.heavy_list:
                if dist_2(h,heavy) <= self.BOND_CUTOFF_H_2:
                    self.bonds.add(fBOND(h,heavy))
                    break
    def _make_heavy_bond(self):
        for heavy_1 in self.heavy_list:
            for heavy_2 in self.heavy_list:
                if heavy_1.name < heavy_2.name :
                    if dist_2(heavy_1,heavy_2) <= self.BOND_CUTOFF_2:
                        self.bonds.add(fBOND(heavy_1,heavy_2))
    def truncate_leaf(self):
        while True:
            to_be_deleted = list()
            for atom in self.atoms:
                assert atom.connect_count > 0
                if atom.connect_count == 1:
                    to_be_deleted.append(atom)
            if len(to_be_deleted) == 0 :
                break
            else:
                for atom in to_be_deleted:
                    bonds = atom.bond
                    for bond in bonds:
                        bond.del_me()
                        self.bonds.remove(bond)
                    self.atoms.remove(atom)
        self._atom_classification()
    def find_ring(self):
        pass
        self.write_mol2(TMPFILE_MOL2)
        answer = os.popen("%s %s "%(newnew,TMPFILE_MOL2)).readlines()
        os.remove(TMPFILE_MOL2)
        rings = list()
        for line in answer:
            tmpring = list()
            numbers = [ int(x)+self.index_shift for x in line.replace(',' , '').split()[2:] ]
            tmpring = self.find_atoms(index = numbers)
            rings.append(tmpring)
        self.rings = self._merge_ring(rings)
        return self.rings
        #analyze answer and tranlate to my class  ( fRING class )
    def _merge_ring(self,rings):
        new_rings = list()
        for ring in rings:
            flag_keep = True
            for i in range(len(new_rings)):
                cp_ring = new_rings[i]
                if len(ring)+len(cp_ring)-len( set( ring + cp_ring )) >=2 :
                    new_rings[i] = list( set(ring+cp_ring) )
                    flag_keep = False
                    break
            if flag_keep:
                new_rings.append(ring)
        return new_rings
    def write_mol2(self,ofile):
        self.write_pdb(TMPFILE)
        os.system("%s -ipdb %s -omol2 %s"%(babel,TMPFILE,ofile))
        os.system("echo '@<TRIPOS>ATOM' >> %s"%ofile)
        os.remove(TMPFILE)
    def list_connect_count(self):
        for atom in self.atoms:
            print((atom.name,atom.connect_count))
    def debug(self):
        fCHEMO.debug(self)
        print("Hydrogen atoms:")
        for h in self.h_list:
            print((h.name))
        print("Heavy atoms:")
        for heavy in self.heavy_list:
            print((heavy.name))
        print(("Num of Bonds:",len(self.bonds)))
        print("Atom index:")
        for atom in self.atoms:
            print((atom.index,))
        print()

class fCHAIN():
    def __init__(self,chain_index = 'A'):
        self.residues = list()
        self.chain_name = chain_index
    def add_resi(self,resi):
        self.residues.append(resi)
    def delete_resi(self,resi):
        self.residues.remove(resi)

class fTOPOLOGY():
    @staticmethod
    def _next_resi_lines(lines):
        resi_lines = list()
        oldresindex = None
        for line in lines:
            if len(line)<6 or line[:6] not in ('ATOM  ','HETATM'):
                continue
            else:
                resindex = line[22:26].strip()+line[21]
                if resindex == oldresindex or oldresindex == None:
                    resi_lines.append(line)
                else:
                    yield resi_lines
                    resi_lines = list()
                    resi_lines.append(line)
            oldresindex = resindex
        yield resi_lines

    def __init__(self,lines):
        self.residues = list()
        for resi_lines in fTOPOLOGY._next_resi_lines(lines):
            self.residues.append( fRESIDUE(resi_lines) )
        self.residues_d = dict()
        for resi in self.residues:
            self.residues_d[resi.index] = resi

        self.chains = dict()
        for resi in self.residues:
            if resi.chain in list(self.chains.keys()):
                self.chains[resi.chain].add_resi(resi)
            else:
                self.chains[resi.chain] = fCHAIN(resi.chain)
                self.chains[resi.chain].add_resi(resi)

    def add_residue(self, resi):
        self.residues.append(resi)
        self.residues_d[resi.index] = resi
        if resi.chain in list(self.chains.keys()):
            self.chains[resi.chain].add_resi(resi)
        else:
            self.chains[resi.chain] = fCHAIN(resi.chain)
            self.chains[resi.chain].add_resi(resi)
    def remove_residue(self,resi):
        self.residues.remove(resi)
        del(self.residues_d[resi.index])
        try:
            self.chains[resi.chain].delete_resi(resi)
        except Exception as e:
            print((e.message))

    def get_protein_residues(self):
        prot_residues = list()
        for resi in self.residues:
            if resi.name in standard_protein_residues:
                prot_residues.append(resi)
        return prot_residues

    def get_water_residues(self):
        water_residues = list()
        for resi in self.residues:
            if resi.name in standard_water_residues:
                water_residues.append(resi)
        return water_residues

    def find_residues(self, name = None , index = None ,atom_index = None ):
        result = list()
        ignore_name = False
        ignore_index = False
        ignore_atom_index = False
        if name == None:
            ignore_name = True
        if index == None:
            ignore_index = True
        if atom_index == None:
            ignore_atom_index = True
        for residue in self.residues:
            if ignore_name or residue.name in name:
                if ignore_index or residue.index in index:
                    if ignore_atom_index or sum( [ atom_index.count(x.index) for x in residue.atoms] ) > 0:
                        result.append(residue)
        return result

    def write_model(self,ofp):
        for residue in self.residues:
            residue.write_pdb(ofp)

class fPDB:
    def __init__(self,frame = None, fragmentation=False):
        lines = None
        if hasattr(frame,'isalpha'):
            lines = open(frame).readlines()
        elif frame == None:
            lines = list()
        else:
            lines = frame

        for line in lines:
            if len(line)>=6 and line[:6]=="CRYST1":
                x = float(line.split()[1])
                y = float(line.split()[2])
                z = float(line.split()[3])
                self.box = x,y,z
                break

        if not hasattr(self,'box'):
            self.box = 0,0,0

        self.model_n = 1
        for line in lines:
            if len(line)>=6 and line[:6]=="MODEL ":
                try:
                    if len(line.split())==2:
                        self.model_n = int(line.split()[-1])
                    else:
                        self.model_n = int(line[5:14])
                except:
                    pass
        if fragmentation:
          from .frag import fFRAGTOPO
          self.topology = fFRAGTOPO(lines)
        else:
          self.topology = fTOPOLOGY(lines)

    @staticmethod
    def load_ff_param_resi(resi,gmxtop):
        resi_atoms = set( [ x.name for x in resi.atoms ] )
        gmx_resi_set = set(gmxtop.get_resilist())
        gmx_atoms = None

        if resi.name in standard_water_residues:
            tmpname = 'HOH'
        else:
            tmpname = resi.name

        if tmpname in gmx_resi_set:
            _ = dict()
            for x in gmxtop.get_resi(tmpname) :
                _[x[0]] = x
            # _ = { x[0]:x for x in gmxtop.get_resi(tmpname) }
            if resi_atoms == set(_.keys()):
                gmx_atoms = _
            else:
                gmx_atoms = None

        if gmx_atoms == None:
            for gmx_resi in gmx_resi_set:
             #   _ = { x[0]:x for x in gmxtop.get_resi(gmx_resi)  }
                _ = dict()
                for x in gmxtop.get_resi(gmx_resi):
                    _[x[0]] = x
                if resi_atoms == set(_.keys()):
                    # print(>>>>> Loading GMX parameters : %s"%gmx_resi)
                    if tmpname != gmx_resi :
                        gmx_atoms = _
                        pass
                        if tmpname != 'HOH':
                            sys.stderr.write("===== Warning : different residue names while loading parameters %s ( parameters: %s).\n"%(tmpname,gmx_resi))

        if gmx_atoms is not None:
            for atom in resi.atoms:
                gmx_atom = gmx_atoms[atom.name]
                atom.addparm( gmx_atom[3],gmx_atom[1],gmx_atom[2] )
            return

        # for gmx_resi in gmxtop.get_resilist_amber():
        #     gmx_atoms = set( [ x[0] for x in gmxtop.get_resi_amber(gmx_resi) ] )

        #     if resi_atoms == gmx_atoms:
        #         pass
        #         print(">>>>> Loading AMBER parameters : %s"%gmx_resi)
        #         if resi.name != gmx_resi :
        #             pass
        #             print("##### Warning : Different residue names while loading parameters %s and %s"%(resi.name,gmx_resi))
        #         for atom in resi.atoms:
        #             for gmx_atom in gmxtop.get_resi_amber(gmx_resi):
        #                 if atom.name == gmx_atom[0]:
        #                     atom.addparm( gmx_atom[3],gmx_atom[1],gmx_atom[2] )
        #         return

        print(("##### Warning : Error in loading residue parameters ",resi.name,resi.index, " set all parameter within this residue to zero "))
        for atom in resi.atoms:
            atom.addparm(0,0,0)
        # print("Error loading residue parameters",resi.name,resi.index)
        return

    def load_ff_params(self, gmxtop ):
        residues = list(self.topology.residues)
        for resi in residues :
            fPDB.load_ff_param_resi(resi,gmxtop)

    def check_params(self):
        for resi in self.topology.residues:
            for atom in resi.atoms:
                print((atom.name,atom.charge,atom.sig,atom.eps))

    def find_protein_center(self):
        n = 0
        xs = 0.
        ys = 0.
        zs = 0.
        for residue in self.topology.residues:
            if residue.name in standard_protein_residues:
                for atom in residue.atoms:
                    x,y,z = atom.posi
                    xs += x
                    ys += y
                    zs += z
                    n += 1
        if n == 0 :
            return 0,0,0
        else:
            return xs/n,ys/n,zs/n

    def center(self, center_posi = (0,0,0) ):
        for resi in self.topology.residues:
            for atom in resi.atoms:
                for i in range(3):
                    if atom.posi[i] - center_posi[i] > 0:
                        atom.posi[i] = atom.posi[i] - self.box[i]*int( ( atom.posi[i]-center_posi[i] )/self.box[i] + 0.5 )
                    else :
                        atom.posi[i] = atom.posi[i] + self.box[i]*int( ( - atom.posi[i] + center_posi[i] )/self.box[i] + 0.5 )

    def write_pdb(self,ofp):
        if hasattr(ofp,'write'):
            flag_is_handle = True
            flag_is_str = False
        else:
            ofp = open(ofp,'w')
            flag_is_handle = False
            flag_is_str = True

        ofp.write("MODEL     %4d\n"%self.model_n)
        ofp.write("CRYST1%9.3f%9.3f%9.3f  90.00  90.00  90.00 P 65 2 2     12\n"%self.box)
        self.topology.write_model(ofp)
        ofp.write("ENDMDL\n")

        if flag_is_str:
            ofp.close()

    def optimize_lys_H(self):
        """ A specified method for proprocessing SPA graphic.pdb, fix three hydrogen postion on LYS NZ"""
        R_HN = 1.0
        for resi in self.topology.find_residues(name = ("LYS",) ):
            ce = resi.atoms_d['CE']
            nz = resi.atoms_d['NZ']
            h1 = resi.atoms_d['HZ1']
            h2 = resi.atoms_d['HZ2']
            h3 = resi.atoms_d['HZ3']
            r_tmp = dist(ce,nz)
            x = (nz.posi[0] - ce.posi[0])/r_tmp*R_HN
            y = (nz.posi[1] - ce.posi[1])/r_tmp*R_HN
            z = (nz.posi[2] - ce.posi[2])/r_tmp*R_HN
            from . import frotate
            matrix1 = frotate.build_matrix(h2.posi,h3.posi,-1.318)
            newcoord1 = frotate.rotate_atom( (x,y,z),(0,0,0),matrix1 )
            matrix2 = frotate.build_matrix(h3.posi,h1.posi,-1.318)
            newcoord2 = frotate.rotate_atom( (x,y,z),(0,0,0),matrix2 )
            matrix3 = frotate.build_matrix(h1.posi,h2.posi,-1.318)
            newcoord3 = frotate.rotate_atom( (x,y,z),(0,0,0),matrix3 )
            for i in (0,1,2):
                h1.posi[i] = newcoord1[i] + nz.posi[i]
                h2.posi[i] = newcoord2[i] + nz.posi[i]
                h3.posi[i] = newcoord3[i] + nz.posi[i]
    def choose_conformation(self,conf=None):
        if 'conf' == None:
            'conf' == 'A'

        tmpdb = copy.deepcopy(self)
        tmpdb.topology = fTOPOLOGY(list())
        for resi in self.topology.residues:
            tmpresi = copy.deepcopy(resi)
            tmpresi.atoms = list()
            tmpresi.atoms_d = dict()
            for atom in resi.atoms :
                if atom.conf in (" ",conf):
                    tmpresi.add_atom(atom)
            tmpdb.topology.add_residue(tmpresi)

        return tmpdb

###### Function to compute vdw
def calc_vdw(a,b,dist_2):
    sig1,eps1=a.sig,a.eps
    sig2,eps2=b.sig,b.eps
    sig = 0.5 * ( sig1 + sig2 )
    sig = sig * sig
    eps = math.sqrt(eps1*eps2)
    _dist = dist_2 / 100.
    _ = ( sig/_dist ) ** 3
    return  4 * eps * ( _ * _ - _ )

###### Function to compute charge
def calc_chg(a,b,dist_2):
    #f_charge = 138.935485
    charge1 = a.charge
    charge2 = b.charge
    return 138.935485*charge1*charge2/(math.sqrt(dist_2)/10.)

###### Function to compute vdw
def atomicEF(a,b):
    k = 138.935485
    sig1,eps1=a.sig,a.eps
    sig2,eps2=b.sig,b.eps
    sig = 0.5 * ( sig1 + sig2 )
    sig = sig * sig
    sig6 = sig ** 3
    sig12 = sig6 * sig6
    eps = np.sqrt(eps1*eps2)
    x0 = np.array(a.posi)
    x1 = np.array(b.posi)
    x0 /= 10
    x1 /= 10
    r2 = sum((x1-x0)**2)
    r = np.sqrt(r2)
    r6 = r2**3
    r8 = r6 * r2
    r12 = r6 * r6
    r14 = r12 * r2
    Evdw = 4 * eps * (sig12/r12 - sig6/r6)
    Fvdw = 24 * eps * (sig6/r8 - 2*sig12/r14) * (x1-x0)
    chg1 = a.charge
    chg2 = b.charge
    Eelec = k*chg1*chg2/r
    Felec = - k*chg1*chg2*(x1-x0)/(r*r2)
    return (Evdw + Eelec)*kJ_to_kcal, (Fvdw + Felec)*kJ_to_kcal/10 # to A

###### dist_2 atom atom
def dist_2(a,b):
    if hasattr(a,'posi'):
        x1,y1,z1 = a.posi
    else:
        x1,y1,z1 = a
    if hasattr(b,'posi'):
        x2,y2,z2 = b.posi
    else:
        x2,y2,z2 = b
    return (x1-x2)**2 + (y1-y2)**2 + (z1-z2)**2

def dist(a,b):
    return math.sqrt( dist_2(a,b) )
def dist_resi_resi_2(a,b):
    assert len(a.atoms)>0 and len(b.atoms)>0
    dist = BIG_NUM
    for i in a.atoms:
        for j in b.atoms:
            tmp = dist_2(i,j)
            if tmp < dist:
                dist = tmp
    return dist
def dist_resi_resi(a,b):
    return math.sqrt(dist_resi_resi_2(a,b))
def dist_atom_resi_2(a,b):
    dist = BIG_NUM
    if hasattr(b,'atoms'):
        assert len(b.atoms)>0
    else:
        assert len(b) > 0
    if hasattr(b,'atoms'):
        for j in b.atoms:
            tmp = dist_2(a,j)
            if tmp< dist:
                dist = tmp
    else:
        for j in b:
            tmp = dist_2(a,j)
            if tmp< dist:
                dist = tmp
    return dist

def dist_atom_resi(a,b):
    return math.sqrt( dist_atom_resi_2(a,b) )

def angle( a1,b,a2):
    dist_a1_a2_2 = dist_2(a1,a2)
    dist_a1_b = dist(a1,b)
    dist_a2_b = dist(a2,b)
    cos_angle = ( - dist_a1_a2_2 + dist_a1_b**2 + dist_a2_b**2 )/( 2*dist_a1_b*dist_a2_b )
    angle = math.acos(cos_angle)
    return angle

def dihedral(a1,b1,b2,a2):
    "Method at https://math.stackexchange.com/questions/47059/how-do-i-calculate-a-dihedral-angle-given-cartesian-coordinates"

    vector_b1 = (np.array( b1.posi ) - np.array( a1.posi ))[:3]
    vector_b2 = (np.array( b2.posi ) - np.array( b1.posi ))[:3]
    vector_b3 = (np.array( a2.posi ) - np.array( b2.posi ))[:3]

    mod_b2 = np.sqrt( vector_b2.dot(vector_b2) )
    unit_b2 = vector_b2/mod_b2

    vector_n1 = np.cross(vector_b1,vector_b2)
    vector_n2 = np.cross(vector_b2,vector_b3)

    vector_m1 = np.cross( vector_n1, unit_b2 )

    x = vector_n1.dot(vector_n2)
    y = vector_m1.dot(vector_n2)

    angle = math.atan2(y,x)
    return angle

def potential_atom_atom( atoma,atomb ):
    d_2 = dist_2(atoma,atomb)
    vdw =  calc_vdw(atoma,atomb,d_2)
    chg =  calc_chg(atoma,atomb,d_2)
    return vdw*kJ_to_kcal,chg*kJ_to_kcal

def potential_resi( resia,resib ):
    if resia == resib:
        # sys.stderr.write("WARNING, YOU ARE ATTEMPING TO COMPUTE THE POTENTIAL BETWEEN IDENTICAL RESIDUES. THE FUNCTION WILL IGNORE IT AND RETURN ZERO.\n")
        return 0,0
    vdw = 0
    chg = 0
    for atoma in resia.atoms:
        for atomb in resib.atoms:
            v,c = potential_atom_atom(atoma,atomb)
            vdw += v
            chg += c
    return vdw, chg

def resiEFT( resia,resib ):
    if resia == resib:
        # sys.stderr.write("WARNING, YOU ARE ATTEMPING TO COMPUTE THE POTENTIAL BETWEEN IDENTICAL RESIDUES. THE FUNCTION WILL IGNORE IT AND RETURN ZERO.\n")
        return 0,0,0
    com1 = resib.getCOM() # not finish !!! jc
    E = 0
    F = np.zeros(3)
    T = np.zeros(3)
    for atoma in resia.atoms:
        for atomb in resib.atoms:
            e, f = atomicEF(atoma,atomb)
            E += e
            #f = -f # get negative for grid database. jc
            F += f # force a --> b
            x = (np.array(atomb.posi) - com1) # A
            T += np.cross(x, f)
    return E, F, T

def next_frame(filename):
    just_yield = True
    frame = list()
    for line in open(filename):
        if just_yield : # or ( len(line)>=6 and line[:6] in ("MODEL ","TITLE ") ):
            frame = list()
            frame.append(line)
            just_yield = False
        elif len(line)>=3 and line[:3] == "END":
            frame.append(line)
            yield frame
            just_yield = True
            frame = list()
        else:
            frame.append(line)

def next_frame_sdf(filename):
    just_yield = True
    frame = list()
    for line in open(filename):
        if just_yield : # or ( len(line)>=6 and line[:6] in ("MODEL ","TITLE ") ):
            frame = list()
            frame.append(line)
            just_yield = False
        elif "$$$$" in line:
            frame.append(line)
            yield frame
            just_yield = True
            frame = list()
        else:
            frame.append(line)

if False:
    def dist_atom(a,b):
        return math.sqrt( (a[0]-b[0])**2 + (a[1]-b[1])**2 + (a[2]-b[2])**2 )

    def dist_atom_resi(atom, resi):
        assert len(atom)>=3
        assert len(resi)>0
        dist = MAX
        node = None
        for atomb in resi:
            tmp = dist_atom(atom,atomb)
            if tmp < dist:
                dist = tmp
                node = atomb
        assert node != None
        return dist,node

# Test
if __name__ == '__main__':
    pass


import csv
import mne
import pandas as pd
import matplotlib.pyplot as plt
import torch
import math
import numpy as np
import json
import re
import ast

def make_file_list(edf_list: str, csv_list: str, data_dir: str) -> list:
    file_list = []
    file1 = open(edf_list)
    file2 = open(csv_list)
    reader1 = csv.reader(file1) # reader for the edf file locations
    reader2 = csv.reader(file2) # reader for the rec file location

    for i in zip(reader1, reader2):
        first = data_dir + i[0][0][2:]
        second = data_dir + i[1][0][2:-3] + 'lbl'
        file_list.append([first, second])

    file1.close()
    file2.close()
    return file_list # returns list of file locations

# https://stackoverflow.com/questions/20910213/loop-over-two-generator-together

def read_and_export_files(file_list: list, montage: dict, save_loc: str):
    global nr
    for direct in file_list:
        edf_dir = direct[0]
        csv_dir = direct[1]

        data = mne.io.read_raw_edf(edf_dir, preload=True) # read edf file

        if data.__len__() < 82500: # if the file has less than 5,5 mins of
            continue               # recorded data then it's discarded

        data = data.filter(0.1, 100) # use filter on data
        data = data.notch_filter(60) # use filter on data

        sfreq = int(data.info['sfreq']) # get the sampling freqency

        df = data.to_data_frame() # make pandas dataframe

        inv_map = {v[0]: k for k, v in montage1.items()}
        # to make sure, that the correct targets are given to the right
        # channels, the index order is used.

        which_montages = set()
        target = []

        with open(csv_dir, "r") as file: # read rec file
            ls = csv.reader(file, delimiter='å')
            for rows in ls:
                l = 0
                try:
                    l = rows[0]
                except:
                    continue

                if l:
                    if l[:7] == "symbols":
                        m = re.match(r".+(\{.+\})", l)
                        string = m.group(1).replace("'", '"')
                        string = re.sub(r'(\d+)', r'"\1"', string)

                        sym_dict = json.loads(string)
                        # make dict with the keys to the artifacts

                    elif l[:5] == "label":
                        m = re.match(r".+\{(.+)(\[.+\])\}", l)
                        first = m.group(1).split(", ")
                        second = m.group(2).strip('][').split(", ")
                        ind = second.index('1.0')
                        obs = sym_dict[str(ind)] # get the encoding of the data
                        if (obs in ['eyem', 'artf', 'chew',
                                    'shiv', 'musc', 'elpp', 'bckg']):
                            target.append([int(first[4]),
                                           float(first[2]),
                                           float(first[3])])

                        which_montages.add(int(first[4]))

        sorted_index = sorted(list(which_montages)) # sort the montage index

        first = True
        for i in sorted_index: # using the montage information we make the new
            col_names = montage.get(i) # data-frame using only the channels
                                       # that has been labeled

            if (col_names[0] == "EKG"): # & first # special case that is removed
                continue
        #        df_new = df[col_names[1]]
        #        df_new = df_new.rename(col_names[0])
        #        first = False
        #    elif (col_names[0] == "EKG"): # special case for montage 2
        #        list1 = df[col_names[1]]
        #        list1 = list1.rename(col_names[0])
        #        df_new = pd.concat([df_new, diff], axis=1, join='inner')

            if first:
                list1 = df[col_names[1]] # get the first series
                list2 = df[col_names[2]] # get the second series
                df_new = list1 - list2
                df_new = pd.DataFrame(df_new.rename(col_names[0])) # Rename
                first = False
            else:
                list1 = df[col_names[1]]
                list2 = df[col_names[2]]
                diff = list1 - list2
                diff = diff.rename(col_names[0]) # Rename
                df_new = pd.concat([df_new, diff], axis=1, join='inner')

        tar = torch.zeros(df_new.shape[1], df_new.shape[0]) # make target data

        for i in target: # i = [montage_channel, start, end, type_artifact]
            index = sorted_index.index(i[0]) # Find the correct index in the target
            tar[index][sfreq * math.floor(i[1]): sfreq * math.ceil(i[2])] = 1
                # Make the artifacts = 1

        ind = torch.tensor(df_new.values.T) # data-frame to tensor
        torch.save(ind, save_loc + f'/model_input ({nr}).pt') # save input
        torch.save(tar, save_loc + f'/model_target ({nr}).pt') # save target

        with open(save_loc + "/data_encoding.csv", "a", newline='') as f:
            write = csv.writer(f) # save information that link the nr of the
                                  # .pt files with the .edf files.
            write.writerow([edf_dir, csv_dir, nr])

        nr += 1

montage1 = {
0: ["FP1-F7", "EEG FP1-REF", "EEG F7-REF"],
1: ["F7-T3", "EEG F7-REF", "EEG T3-REF"],
2: ["T3-T5", "EEG T3-REF", "EEG T5-REF"],
3: ["T5-O1", "EEG T5-REF", "EEG O1-REF"],
4: ["FP2-F8", "EEG FP2-REF", "EEG F8-REF"],
5: ["F8-T4", "EEG F8-REF", "EEG T4-REF"],
6: ["T4-T6", "EEG T4-REF", "EEG T6-REF"],
7: ["T6-O2", "EEG T6-REF", "EEG O2-REF"],
8: ["A1-T3", "EEG A1-REF", "EEG T3-REF"],
9: ["T3-C3", "EEG T3-REF", "EEG C3-REF"],
10: ["C3-CZ", "EEG C3-REF", "EEG CZ-REF"],
11: ["CZ-C4", "EEG CZ-REF", "EEG C4-REF"],
12: ["C4-T4", "EEG C4-REF", "EEG T4-REF"],
13: ["T4-A2", "EEG T4-REF", "EEG A2-REF"],
14: ["FP1-F3", "EEG FP1-REF", "EEG F3-REF"],
15: ["F3-C3", "EEG F3-REF", "EEG C3-REF"],
16: ["C3-P3", "EEG C3-REF", "EEG P3-REF"],
17: ["P3-O1", "EEG P3-REF", "EEG O1-REF"],
18: ["FP2-F4", "EEG FP2-REF", "EEG F4-REF"],
19: ["F4-C4", "EEG F4-REF", "EEG C4-REF"],
20: ["C4-P4", "EEG C4-REF", "EEG P4-REF"],
21: ["P4-O2", "EEG P4-REF", "EEG O2-REF"]}
montage2 = {
0: ["FP1-F7", "EEG FP1-LE", "EEG F7-LE"],
1: ["F7-T3", "EEG F7-LE", "EEG T3-LE"],
2: ["T3-T5", "EEG T3-LE", "EEG T5-LE"],
3: ["T5-O1", "EEG T5-LE", "EEG O1-LE"],
4: ["FP2-F8", "EEG FP2-LE", "EEG F8-LE"],
5: ["F8-T4", "EEG F8-LE", "EEG T4-LE"],
6: ["T4-T6", "EEG T4-LE", "EEG T6-LE"],
7: ["T6-O2", "EEG T6-LE", "EEG O2-LE"],
8: ["A1-T3", "EEG A1-LE", "EEG T3-LE"],
9: ["T3-C3", "EEG T3-LE", "EEG C3-LE"],
10: ["C3-CZ", "EEG C3-LE", "EEG CZ-LE"],
11: ["CZ-C4", "EEG CZ-LE", "EEG C4-LE"],
12: ["C4-T4", "EEG C4-LE", "EEG T4-LE"],
13: ["T4-A2", "EEG T4-LE", "EEG A2-LE"],
14: ["FP1-F3", "EEG FP1-LE", "EEG F3-LE"],
15: ["F3-C3", "EEG F3-LE", "EEG C3-LE"],
16: ["C3-P3", "EEG C3-LE", "EEG P3-LE"],
17: ["P3-O1", "EEG P3-LE", "EEG O1-LE"],
18: ["FP2-F4", "EEG FP2-LE", "EEG F4-LE"],
19: ["F4-C4", "EEG F4-LE", "EEG C4-LE"],
20: ["C4-P4", "EEG C4-LE", "EEG P4-LE"],
21: ["P4-O2", "EEG P4-LE", "EEG O2-LE"],
22: ["EKG", "EEG EKG-LE", "None"]
}
montage3 = {
0: ["FP1-F7", "EEG FP1-REF", "EEG F7-REF"],
1: ["F7-T3", "EEG F7-REF", "EEG T3-REF"],
2: ["T3-T5", "EEG T3-REF", "EEG T5-REF"],
3: ["T5-O1", "EEG T5-REF", "EEG O1-REF"],
4: ["FP2-F8", "EEG FP2-REF", "EEG F8-REF"],
5: ["F8-T4", "EEG F8-REF", "EEG T4-REF"],
6: ["T4-T6", "EEG T4-REF", "EEG T6-REF"],
7: ["T6-O2", "EEG T6-REF", "EEG O2-REF"],
8: ["T3-C3", "EEG T3-REF", "EEG C3-REF"],
9: ["C3-CZ", "EEG C3-REF", "EEG CZ-REF"],
10: ["CZ-C4", "EEG CZ-REF", "EEG C4-REF"],
11: ["C4-T4", "EEG C4-REF", "EEG T4-REF"],
12: ["FP1-F3", "EEG FP1-REF", "EEG F3-REF"],
13: ["F3-C3", "EEG F3-REF", "EEG C3-REF"],
14: ["C3-P3", "EEG C3-REF", "EEG P3-REF"],
15: ["P3-O1", "EEG P3-REF", "EEG O1-REF"],
16: ["FP2-F4", "EEG FP2-REF", "EEG F4-REF"],
17: ["F4-C4", "EEG F4-REF", "EEG C4-REF"],
18: ["C4-P4", "EEG C4-REF", "EEG P4-REF"],
19: ["P4-O2", "EEG P4-REF", "EEG O2-REF"]
}

dir1_edf_list = "C:/Users/Marc/Desktop/data/v1.0.0/lists/edf_01_tcp_ar.list"
dir2_edf_list = "C:/Users/Marc/Desktop/data/v1.0.0/lists/edf_02_tcp_le.list"
dir3_edf_list = "C:/Users/Marc/Desktop/data/v1.0.0/lists/edf_03_tcp_ar_a.list"

dir1_csv_list = "C:/Users/Marc/Desktop/data/v1.0.0/lists/rec_01_tcp_ar.list"
dir2_csv_list = "C:/Users/Marc/Desktop/data/v1.0.0/lists/rec_02_tcp_le.list"
dir3_csv_list = "C:/Users/Marc/Desktop/data/v1.0.0/lists/rec_03_tcp_ar_a.list"

if __name__ == "__main__":
    montage_list = [montage1, montage2, montage3]
    dir_edf_list = [dir1_edf_list, dir2_edf_list, dir3_edf_list]
    dir_csv_list = [dir1_csv_list, dir2_csv_list, dir3_csv_list]

    data_dir = "C:/Users/Marc/Desktop/data/v1.0.0"
    nr = 1

    save_loc = "C:/Users/Marc/Desktop/model_data" # location to save files
    f = open(save_loc + "/data_encoding.csv", 'w')
    f.close() # file for the encoding

    for i in range(0, 3):
        dir_edf = dir_edf_list[i]
        dir_rec = dir_csv_list[i]

        file_list = make_file_list(dir_edf, dir_rec, data_dir)

        montage = montage_list[i] # find the correct montage
        read_and_export_files(file_list, montage, save_loc)

Older functions & test files/Data_work/Transform_data_v1.py


from simpleland.utils import TickPerSecCounter
from typing import Dict, Any

import pygame
from .object import GObject
from .common import Vector, Line, Circle, Polygon
from .object_manager import GObjectManager
from PIL import Image
import numpy as np
import os

from .config import RendererConfig
from .asset_bundle import AssetBundle
from .player import Player
from . import gamectx
import math
import time
import pkg_resources
import logging

def to_pygame(p, surface):
    """Convenience method to convert pymunk coordinates to pygame surface
    local coordinates.

    Note that in case positive_y_is_up is False, this function wont actually do
    anything except converting the point to integers.
    """

    return int(p[0]), surface.get_height() - int(p[1])

class Renderer:

    def __init__(self, config: RendererConfig, asset_bundle: AssetBundle):
        if config.sdl_audio_driver:
            os.environ['SDL_AUDIODRIVER'] = config.sdl_audio_driver
        if config.sdl_video_driver:
            os.environ["SDL_VIDEODRIVER"] = config.sdl_video_driver
        self.asset_bundle = asset_bundle
        self.config:RendererConfig = config
        self.format = config.format
        self._display_surf = None
        self.resolution = self.width, self.height = config.resolution
        self.aspect_ratio = (self.width * 1.0) / self.height

        self.initialized = False
        self.frame_cache = None
        self.debug=config.debug_render_bodies

        # These will be source object properties eventually
        self.view_height = 60000.0
        self.view_width = self.view_height * self.aspect_ratio
        logging.info(f"{self.view_width} by {self.view_width}")
        self.center = Vector.zero()

        self.images = {}
        self.sounds = {}

        # FPS Counter
        self.fps_counter = TickPerSecCounter(2)
        self.log_info = None

    def set_log_info(self,log_info):
        self.log_info = log_info

    def load_sounds(self):
        if self.config.sound_enabled:
            for k,path in self.asset_bundle.sound_assets.items():
                sound = pygame.mixer.Sound(pkg_resources.resource_filename(__name__,path))
                sound.set_volume(0.06)
                self.sounds[k] = sound

    def load_images(self):
        self.images = {}
        for k, path in self.asset_bundle.image_assets.items():
            image = pygame.image.load(pkg_resources.resource_filename(__name__,path)).convert_alpha()
            # image = pygame.transform.scale(image,(200,200))
            self.images[k] = image

    def play_sounds(self, sound_ids):
        if self.config.sound_enabled:
            for sid in sound_ids:
                self.sounds[sid].play()

    def get_image_by_id(self, image_id):
        return self.images.get(image_id)

    def update_view(self, view_height):
        self.view_height = max(view_height, 0.001)
        self.view_width = self.view_height * self.aspect_ratio

    def initialize(self):
        if self.config.sound_enabled:
            pygame.mixer.pre_init(  44100, -16, 2, 1024)
        pygame.init()

        self._display_surf = pygame.display.set_mode(self.resolution)  # , pygame.HWSURFACE | pygame.DOUBLEBUF)
        self.load_sounds()
        self.load_images()
        self.initialized = True

    def render_to_console(self, lines, x, y, fsize=18, spacing=12, color=(180, 180, 180)):
        font = pygame.font.SysFont(None, fsize)
        for i, l in enumerate(lines):
            font.render(l, False, (0, 0, 0))
            self._display_surf.blit(font.render(l, True, color), (x, y + spacing * i))

    def _draw_line(self,
                  obj: GObject,
                  line: Line,
                  color,
                  screen_factor,
                  screen_view_center,
                  angle,
                  center):
        obj_pos = obj.get_body().position
        obj_location1 = (obj_pos + line.a.rotated(obj.get_body().angle) - center)
        obj_location1 = obj_location1.rotated(-angle)
        p1 = screen_factor * (obj_location1 + screen_view_center)
        p1 = to_pygame(p1, self._display_surf)

        obj_location2 = (obj_pos + line.b.rotated(obj.get_body().angle) - center)
        obj_location2 = obj_location2.rotated(-angle)
        p2 = screen_factor * (obj_location2 + screen_view_center)
        p2 = to_pygame(p2, self._display_surf)

        pygame.draw.line(self._display_surf,
                         color,
                         p1,
                         p2,
                         1)

    def _draw_circle(self,
                    obj: GObject,
                    circle: Circle,
                    color,
                    screen_factor,
                    screen_view_center,
                    angle,
                    center):
        obj_pos = obj.get_body().position + circle.offset # TODO: Doesn't work with offset
        obj_location = (obj_pos - center)
        obj_location = obj_location.rotated(-angle)
        p = screen_factor * (obj_location + screen_view_center)
        p = to_pygame(p, self._display_surf)
        size = int(circle.radius * screen_factor[0])
        pygame.draw.circle(self._display_surf,
                           color,
                           p,
                           size,
                           1)

    def _draw_polygon(self,
                     obj: GObject,
                     shape: Polygon,
                     color,
                     screen_factor,
                     screen_view_center,
                     angle,
                     center):
        obj_pos = obj.get_body().position

        body_angle =obj.get_body().angle
        verts = shape.get_vertices()
        new_verts = []
        for v in verts:
            obj_location = (obj_pos + v.rotated(body_angle) - center)
            obj_location = obj_location.rotated(-angle)
            p = screen_factor * (obj_location + screen_view_center)
            p = to_pygame(p, self._display_surf)
            new_verts.append(p)
        pygame.draw.polygon(self._display_surf,
                color,
                new_verts,
                1)

    def _draw_object(self, center, obj, angle, screen_factor, screen_view_center, color=None):

        color = (255, 50, 50)

        image_id = obj.get_data_value("image")
        image_used = image_id is not None and not self.config.disable_textures
        if image_used:
            body_angle = obj.get_body().angle
            obj_pos = obj.get_body().position
            if len(obj.get_shapes()) == 0:
                return
            main_shape = list(obj.get_shapes())[0]
            body_w, body_h = main_shape.radius *2 , main_shape.radius *2
            image = self.get_image_by_id(image_id)
            image_size = int(body_w*screen_factor[0]),int(body_h*screen_factor[1])

            if image_size[0]> 5000:
                print("zoom out/ to close {}".format(image_size))

            image = pygame.transform.scale(image,image_size)
            image = pygame.transform.rotate(image,((body_angle-angle) * 57.2957795)%360)
            rect = image.get_rect()

            image_loc = screen_factor * ((obj_pos- center).rotated(-angle)  + screen_view_center)
            image_loc = to_pygame(image_loc, self._display_surf)
            rect.center = image_loc
            self._display_surf.blit(image,rect)

        if not image_used or self.config.render_shapes:
            for i, shape in enumerate(obj.get_shapes()):
                if i == 1:
                    color = (100, 200, 200)
                else:
                    color = (255, 50, 50)

                if type(shape) == Line:
                    self._draw_line(obj,
                                shape,
                                color,
                                screen_factor,
                                screen_view_center,
                                angle = angle,
                                center = center)
                elif type(shape) == Circle:
                    self._draw_circle(obj,
                                    shape,
                                    color,
                                    screen_factor,
                                    screen_view_center,
                                    angle = angle,
                                    center = center)
                elif type(shape) == Polygon:
                    self._draw_polygon(obj,
                                    shape,
                                    color,
                                    screen_factor,
                                    screen_view_center,
                                    angle = angle,
                                    center=center)

    def _draw_grid(self, center, angle, screen_factor, screen_view_center, color = (100, 80, 80),size = 20):

        for line in range(0,size*2):

            obj_pos1 = Vector(-size,line-size)
            obj_location1 = (obj_pos1 - center)
            obj_location1 = obj_location1.rotated(-angle)
            p1 = screen_factor * (obj_location1 + screen_view_center)
            p1 = to_pygame(p1, self._display_surf)

            obj_pos2 = Vector(size,line-size)
            obj_location2 = (obj_pos2 - center)
            obj_location2 = obj_location2.rotated(-angle)
            p2 = screen_factor * (obj_location2 + screen_view_center)
            p2 = to_pygame(p2, self._display_surf)

            pygame.draw.line(self._display_surf,
                            color,
                            p1,
                            p2,
                            1)

        for line in range(0,size*2):

            obj_pos1 = Vector(line-size,-size)
            obj_location1 = (obj_pos1 - center)
            obj_location1 = obj_location1.rotated(-angle)
            p1 = screen_factor * (obj_location1 + screen_view_center)
            p1 = to_pygame(p1, self._display_surf)

            obj_pos2 = Vector(line-size,size)
            obj_location2 = (obj_pos2 - center)
            obj_location2 = obj_location2.rotated(-angle)
            p2 = screen_factor * (obj_location2 + screen_view_center)
            p2 = to_pygame(p2, self._display_surf)

            pygame.draw.line(self._display_surf,
                            color,
                            p1,
                            p2,
                        1)

    # TODO: Clean this up
    def process_frame(self,
                    render_time,
                    player:Player,
                    additional_data: Dict[str, Any]={}):

        if not self.initialized:
            self.initialize()

        if player is None:
            return
        object_manager = gamectx.object_manager
        # import pdb;pdb.set_trace()
        self._display_surf.fill((0, 0, 0))

        if self.debug:
            import pymunk.pygame_util
            draw_options = pymunk.pygame_util.DrawOptions(self._display_surf)
            gamectx.physics_engine.space.debug_draw(draw_options)

        camera = player.get_camera()
        # TODO: make max customizeable
        self.update_view(max(camera.get_distance(),1))

        view_obj = None
        center = Vector(self.view_width/2,self.view_height/2)
        angle=0
        if self.config.view_type == 0:
            view_obj = object_manager.get_by_id(player.get_object_id(), render_time)
            if view_obj is not None:
                center = view_obj.get_body().position
                angle = view_obj.get_body().angle

        center = center - camera.position_offset
        screen_factor = Vector(self.width / self.view_width, self.height / self.view_height)
        screen_view_center = Vector(self.view_width, self.view_height) / 2.0

        obj_list_sorted_by_depth= object_manager.get_objects_for_timestamp_by_depth(render_time)
        if self.config.draw_grid:
            self._draw_grid(center, angle, screen_factor, screen_view_center, self.config.grid_size)
        for depth, render_obj_dict in enumerate(obj_list_sorted_by_depth):
            for k, obj in render_obj_dict.items():
                if view_obj is not None and k == view_obj.get_id():
                    continue
                # elif abs((center - obj.get_body().position).length) > camera.get_distance() and obj.get_data_value('type') != 'static':
                #     continue
                self._draw_object(center, obj, angle, screen_factor, screen_view_center)
            if view_obj is not None and depth == view_obj.depth:
                self._draw_object(center, view_obj, angle, screen_factor, screen_view_center)

        if self.config.show_console:
            console_info = ["FPS:{}".format(self.fps_counter.avg())]
            if self.log_info is not None:
                console_info.append(self.log_info)
            self.render_to_console(console_info, self. config.resolution[0]-100, 5)

    def render_frame(self):
        self.fps_counter.tick()
        if self.config.save_observation:
            self.get_last_frame()
        frame = self.frame_cache
        self.frame_cache = None
        if self.config.render_to_screen:
            pygame.display.flip()
        return frame

    def get_last_frame(self):

        img_st = pygame.image.tostring(self._display_surf, self.format)
        data = Image.frombytes(self.format, self.config.resolution, img_st)
        np_data = np.array(data)

        # cache
        self.frame_cache = np_data
        return np_data