#!/usr/bin/env python
"""
This file is part of the package FUNtoFEM for coupled aeroelastic simulation
and design optimization.

Copyright (C) 2015 Georgia Tech Research Corporation.
Additional copyright (C) 2015 Kevin Jacobson, Jan Kiviaho and Graeme Kennedy.
All rights reserved.

FUNtoFEM is licensed under the Apache License, Version 2.0 (the "License");
you may not use this software except in compliance with the License.
You may obtain a copy of the License at

   http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

import numpy as np

class PyOptOptimization(object):
    def __init__(self,comm,eval_forward,eval_gradient,number_of_steps=1e9,read_history=True,unscale_design_variables=None):
        """
        A class to save PyOpt optimization histories. It can also be used for problems with FUN3D that only can take one
        a few passes through the solvers at a time to avoid memory leaks stopping the run. The history is saved to the
        disk. If history files exist, They will be read in for all previous steps so that the optimizer has the full
        history even if it only takes one step this time.

        Parameters
        ----------
        comm: MPI communicator
            MPI communicator for the optimization problem
        eval_forward: function
            function that evaluates the forward analysis. The format of eval_forward matches what PyOpt expects:
            obj, con, fail = eval_forward(x)
        eval_gradient:function
            function that evaluates the gradients. The format of eval_forward matches what PyOpt expects:
            g, A, fail = eval_gradient(x,obj,con)
        number_of_steps: int
            This class will feed the optimizer any steps from the history then will allow the optimizer to take a
            maximum number of number_of_steps new adjoint steps. This class then quit(). The default value of 1e9 will
            essentially let the optimization run without stopping it.
        read_history: bool
            If false, existing history files will be ignored and overwritten, i.e., the optimization problem will
            restart.
        unscale_design_variables: function
            Function that takes the scaled design variables from the optimizer and returns the dimensional dv's.
            If this function is provided, the unscaled variables will be saved to the history.
            If no unscale_design_variable function is provided, the scaled design variables will be saved.
            x_unscaled = unscale_design_variables(x)
        """

        self.comm = comm

        self.number_of_steps = number_of_steps
        self.forward_step = 0
        self.adjoint_step = 0

        # Get the histories if they exist
        if read_history and self.comm.Get_rank()==0:
            try:
                obj_hist          = np.load('obj_hist.npy')
                con_hist          = np.load('con_hist.npy')
                dv_hist           = np.load('dv_hist.npy')
                fail_hist         = np.load('fail_hist.npy')
                obj_grad_hist     = np.load('obj_grad_hist.npy')
                con_grad_hist     = np.load('con_grad_hist.npy')
                forward_hist_step = obj_hist.shape[1]
                adjoint_hist_step = obj_grad_hist.shape[2]
            except:
                obj_hist          = None
                con_hist          = None
                dv_hist           = None
                fail_hist         = None
                obj_grad_hist     = None
                con_grad_hist     = None
                forward_hist_step = 0
                adjoint_hist_step = 0
        else:
            obj_hist          = None
            con_hist          = None
            dv_hist           = None
            fail_hist         = None
            obj_grad_hist     = None
            con_grad_hist     = None
            forward_hist_step = 0
            adjoint_hist_step = 0

        self.dv_hist           = comm.bcast(dv_hist,root=0)
        self.obj_hist          = comm.bcast(obj_hist,root=0)
        self.con_hist          = comm.bcast(con_hist,root=0)
        self.fail_hist         = comm.bcast(fail_hist,root=0)
        self.obj_grad_hist     = comm.bcast(obj_grad_hist,root=0)
        self.con_grad_hist     = comm.bcast(con_grad_hist,root=0)
        self.forward_hist_step = comm.bcast(forward_hist_step,root=0)
        self.adjoint_hist_step = comm.bcast(adjoint_hist_step,root=0)
        if self.comm.Get_rank()==0 and self.obj_hist is None:
            print("PyOpt Optimization: Couldn't/didn't read history files. Starting from the beginning")

        # hold onto the forward and gradient evaluation functions
        self.eval_forward = eval_forward
        self.eval_gradient = eval_gradient

        if unscale_design_variables is None:
            self.unscale_design_variables = self._noscaling_of_design_variables
        else:
            self.unscale_design_variables = unscale_design_variables

    def eval_obj_con(self,x):
        """
        Wrapper for the function evaluation that saves/reads history data.
        If taking a new step, evaluate the functions and save to the disk.
        If not taking a new step, read the functions' values from the disk.
        This is the function to give the PyOpt optimizer as the function evaluation.

        Parameters
        ----------
        x: numpy array or list
            design variable list

        Returns
        -------
        obj: numpy array
            objective value(s)
        con: numpy array
            constraint value(s)
        fail: int
            whether or not the simulation failed
        """

        fail = 0
        if self.adjoint_step>=self.adjoint_hist_step+self.number_of_steps:
            if self.comm.Get_rank()==0:
                print("PyOpt Optimization: Reached the requested number of adjoint steps... Stopping")
            quit()

        elif self.forward_step >= self.forward_hist_step:
            obj, con, fail = self.eval_forward(x)

            # Write the history of the objective and constraint evaluations
            if np.size(obj)==1:
                obj_array = np.ones(1) * obj
            else:
                obj_array = np.array(obj)
            obj_array=np.expand_dims(obj_array,axis=1)

            if self.obj_hist is None:
                self.obj_hist = obj_array
            else:
                self.obj_hist = np.concatenate((self.obj_hist,obj_array),axis=1)

            con_array = np.expand_dims(np.array(con),axis=1)
            if self.con_hist is None:
                self.con_hist = con_array
            else:
                self.con_hist = np.concatenate((self.con_hist,con_array),axis=1)

            fail_array = np.ones(1,dtype=int) * fail
            if self.fail_hist is None:
                self.fail_hist = fail_array
            else:
                self.fail_hist = np.concatenate((self.fail_hist,fail_array),axis=0)

            dv_array = np.expand_dims(np.array(self.unscale_design_variables(x)),axis=1)

            if self.dv_hist is None:
                self.dv_hist = dv_array
            else:
                self.dv_hist = np.concatenate((self.dv_hist,dv_array),axis=1)

            if self.comm.Get_rank()==0:
                np.save('obj_hist.npy',self.obj_hist)
                np.save('con_hist.npy',self.con_hist)
                np.save('fail_hist.npy',self.fail_hist)
                np.save('dv_hist.npy'  ,self.dv_hist)

        else:
            if self.comm.Get_rank()==0:
                print("PyOpt Optimization: Reading from history, forward step",self.forward_step)
            obj = self.obj_hist[:,self.forward_step]
            con = self.con_hist[:,self.forward_step]
            fail= self.fail_hist[self.forward_step]

        self.forward_step += 1

        return obj, con, fail

    def eval_obj_con_grad(self,x,obj,con):
        """
        Wrapper for the function gradient evaluation that saves/reads history data.
        If taking a new step, evaluate the function gradient and save to the disk.
        If not taking a new step, read the function gradient values from the disk.
        This is the function to give the PyOpt optimizer as the gradient evaluation.

        Parameters
        ----------
        x: numpy array or list
            design variable list
        obj: numpy array
            objective value(s)
        con: numpy array
            constraint value(s)

        Returns
        -------
        g: numpy array
            objective gradient(s)
        a: numpy array
            constraint gradient(s)
        fail: int
            whether or not the simulation failed

        """

        fail = 0
        if self.adjoint_step >= self.adjoint_hist_step:
            g, a, fail = self.eval_gradient(x,obj,con)

            # Write the history of the objective and constraint gradients
            obj_grad_array = np.expand_dims(g,axis=2)
            if self.obj_grad_hist is None:
                self.obj_grad_hist = obj_grad_array
            else:
                self.obj_grad_hist = np.concatenate((self.obj_grad_hist,obj_grad_array),axis=2)

            con_grad_array = np.expand_dims(a,axis=2)
            if self.con_grad_hist is None:
                self.con_grad_hist = con_grad_array
            else:
                self.con_grad_hist = np.concatenate((self.con_grad_hist,con_grad_array),axis=2)

            if self.comm.Get_rank()==0:
                np.save('obj_grad_hist.npy',self.obj_grad_hist)
                np.save('con_grad_hist.npy',self.con_grad_hist)

        else:
            if self.comm.Get_rank()==0:
                print("PyOpt Optimization: Reading from history, adjoint step",self.adjoint_step)
            g = self.obj_grad_hist[:,:,self.adjoint_step]
            a = self.con_grad_hist[:,:,self.adjoint_step]

        self.adjoint_step += 1

        return g, a, fail

    def _noscaling_of_design_variables(self,x):
        """
        Dummy function for design variable unscaling for the case where an unscaling function isn't provided

        Parameters
        ----------
        x: list or np array
            design variables

        Returns
        -------
        x: list or np array
            design variables

        """
        return x


import io
import sys
import uuid
import zipfile

import collections
import numpy as np

from .. import util
from .. import graph

from ..constants import log

try:
    import networkx as nx
except BaseException as E:
    # create a dummy module which will raise the ImportError
    # or other exception only when someone tries to use networkx
    from ..exceptions import ExceptionModule
    nx = ExceptionModule(E)

def load_3MF(file_obj,
             postprocess=True,
             **kwargs):
    """
    Load a 3MF formatted file into a Trimesh scene.

    Parameters
    ------------
    file_obj : file-like
      Contains 3MF formatted data

    Returns
    ------------
    kwargs : dict
      Constructor arguments for `trimesh.Scene`
    """
    # dict, {name in archive: BytesIo}
    archive = util.decompress(file_obj, file_type='zip')
    # get model file
    model = archive['3D/3dmodel.model']

    # read root attributes only from XML first
    event, root = next(etree.iterparse(model, tag=('{*}model'), events=('start',)))
    # collect unit information from the tree
    if 'unit' in root.attrib:
        metadata = {'units': root.attrib['unit']}
    else:
        # the default units, defined by the specification
        metadata = {'units': 'millimeters'}

    # { mesh id : mesh name}
    id_name = {}
    # { mesh id: (n,3) float vertices}
    v_seq = {}
    # { mesh id: (n,3) int faces}
    f_seq = {}
    # components are objects that contain other objects
    # {id : [other ids]}
    components = collections.defaultdict(list)
    # load information about the scene graph
    # each instance is a single geometry
    build_items = []

    consumed_names = set()
    # iterate the XML object and build elements with an LXML iterator
    # loaded elements are cleared to avoid ballooning memory
    model.seek(0)
    for event, obj in etree.iterparse(model, tag=('{*}object', '{*}build')):
        # parse objects
        if 'object' in obj.tag:
            # id is mandatory
            index = obj.attrib['id']

            # start with stored name
            name = obj.attrib.get('name', str(index))
            # apparently some exporters name multiple meshes
            # the same thing so check to see if it's been used
            if name in consumed_names:
                name = name + str(index)
            consumed_names.add(name)
            # store name reference on the index
            id_name[index] = name

            # if the object has actual geometry data parse here
            for mesh in obj.iter('{*}mesh'):
                vertices = mesh.find('{*}vertices')
                v_seq[index] = np.array([[i.attrib['x'],
                                          i.attrib['y'],
                                          i.attrib['z']] for
                                         i in vertices.iter('{*}vertex')],
                                        dtype=np.float64)
                vertices.clear()
                vertices.getparent().remove(vertices)

                faces = mesh.find('{*}triangles')
                f_seq[index] = np.array([[i.attrib['v1'],
                                          i.attrib['v2'],
                                          i.attrib['v3']] for
                                         i in faces.iter('{*}triangle')],
                                        dtype=np.int64)
                faces.clear()
                faces.getparent().remove(faces)

            # components are references to other geometries
            for c in obj.iter('{*}component'):
                mesh_index = c.attrib['objectid']
                transform = _attrib_to_transform(c.attrib)
                components[index].append((mesh_index, transform))

        # parse build
        if 'build' in obj.tag:
            # scene graph information stored here, aka "build" the scene
            for item in obj.iter('{*}item'):
                # get a transform from the item's attributes
                transform = _attrib_to_transform(item.attrib)
                # the index of the geometry this item instantiates
                build_items.append((item.attrib['objectid'], transform))

        # free resources
        obj.clear()
        obj.getparent().remove(obj)
        del obj

    # have one mesh per 3MF object
    # one mesh per geometry ID, store as kwargs for the object
    meshes = {}
    for gid in v_seq.keys():
        name = id_name[gid]
        meshes[name] = {'vertices': v_seq[gid],
                        'faces': f_seq[gid],
                        'metadata': metadata.copy()}
        meshes[name].update(kwargs)

    # turn the item / component representation into
    # a MultiDiGraph to compound our pain
    g = nx.MultiDiGraph()
    # build items are the only things that exist according to 3MF
    # so we accomplish that by linking them to the base frame
    for gid, tf in build_items:
        g.add_edge('world', gid, matrix=tf)
    # components are instances which need to be linked to base
    # frame by a build_item
    for start, group in components.items():
        for i, (gid, tf) in enumerate(group):
            g.add_edge(start, gid, matrix=tf)

    # turn the graph into kwargs for a scene graph
    # flatten the scene structure and simplify to
    # a single unique node per instance
    graph_args = []
    parents = collections.defaultdict(set)
    for path in graph.multigraph_paths(G=g,
                                       source='world'):
        # collect all the transform on the path
        transforms = graph.multigraph_collect(G=g,
                                              traversal=path,
                                              attrib='matrix')
        # combine them into a single transform
        if len(transforms) == 1:
            transform = transforms[0]
        else:
            transform = util.multi_dot(transforms)

        # the last element of the path should be the geometry
        last = path[-1][0]
        # if someone included an undefined component, skip it
        if last not in id_name:
            log.debug('id {} included but not defined!'.format(last))
            continue
        # frame names unique
        name = id_name[last] + util.unique_id()
        # index in meshes
        geom = id_name[last]

        # collect parents if we want to combine later
        if len(path) > 2:
            parent = path[-2][0]
            parents[parent].add(last)

        graph_args.append({'frame_from': 'world',
                           'frame_to': name,
                           'matrix': transform,
                           'geometry': geom})

    # solidworks will export each body as its own mesh with the part
    # name as the parent so optionally rename and combine these bodies
    if postprocess and all('body' in i.lower() for i in meshes.keys()):
        # don't rename by default
        rename = {k: k for k in meshes.keys()}
        for parent, mesh_name in parents.items():
            # only handle the case where a parent has a single child
            # if there are multiple children we would do a combine op
            if len(mesh_name) != 1:
                continue
            # rename the part
            rename[id_name[next(iter(mesh_name))]] = id_name[parent].split(
                '(')[0]

        # apply the rename operation meshes
        meshes = {rename[k]: m for k, m in meshes.items()}
        # rename geometry references in the scene graph
        for arg in graph_args:
            if 'geometry' in arg:
                arg['geometry'] = rename[arg['geometry']]

    # construct the kwargs to load the scene
    kwargs = {'base_frame': 'world',
              'graph': graph_args,
              'geometry': meshes,
              'metadata': metadata}

    return kwargs

def export_3MF(mesh,
               batch_size=4096,
               compression=zipfile.ZIP_DEFLATED,
               compresslevel=5):
    """
    Converts a Trimesh object into a 3MF file.

    Parameters
    ---------
    mesh trimesh.trimesh
      Mesh or Scene to export.
    batch_size : int
      Number of nodes to write per batch.
    compression : zipfile.ZIP_*
      Type of zip compression to use in this export.
    compresslevel : int
      For Python > 3.7 specify the 0-9 compression level.

    Returns
    ---------
    export : bytes
      Represents geometry as a 3MF file.
    """

    if sys.version_info < (3, 6):
        # Python only added 'w' mode to `zipfile` in Python 3.6
        # and it is not worth the effort to work around
        raise NotImplementedError(
            "3MF export requires Python >= 3.6")
    from ..scene.scene import Scene

    if not isinstance(mesh, Scene):
        mesh = Scene(mesh)

    geometry = mesh.geometry
    graph = mesh.graph.to_networkx()
    base_frame = mesh.graph.base_frame

    # xml namespaces
    model_nsmap = {
        None: "http://schemas.microsoft.com/3dmanufacturing/core/2015/02",
        "m": "http://schemas.microsoft.com/3dmanufacturing/material/2015/02",
        "p": "http://schemas.microsoft.com/3dmanufacturing/production/2015/06",
        "b": "http://schemas.microsoft.com/3dmanufacturing/beamlattice/2017/02",
        "s": "http://schemas.microsoft.com/3dmanufacturing/slice/2015/07",
        "sc": "http://schemas.microsoft.com/3dmanufacturing/securecontent/2019/04",
    }

    rels_nsmap = {
        None: "http://schemas.openxmlformats.org/package/2006/relationships"
    }

    # model ids
    def model_id(x, models=[]):
        if x not in models:
            models.append(x)
        return str(models.index(x) + 1)

    # 3mf archive dict {path: BytesIO}
    file_obj = io.BytesIO()

    # specify the parameters for the zip container
    zip_kwargs = {'compression': compression}
    # compresslevel was added in Python 3.7
    if sys.version_info >= (3, 7):
        zip_kwargs['compresslevel'] = compresslevel

    with zipfile.ZipFile(file_obj, mode='w', **zip_kwargs) as z:
        # 3dmodel.model
        with z.open("3D/3dmodel.model", mode='w') as f, etree.xmlfile(
                f, encoding="utf-8"
        ) as xf:
            xf.write_declaration()

            # stream elements
            with xf.element("model", {"unit": "millimeter"}, nsmap=model_nsmap):
                # objects with mesh data and/or references to other objects
                with xf.element("resources"):
                    # stream objects with actual mesh data
                    for i, (name, m) in enumerate(geometry.items()):
                        # attributes for object
                        attribs = {
                            "id": model_id(name),
                            "name": name,
                            "type": "model",
                            "p:UUID": str(uuid.uuid4())
                        }
                        with xf.element("object", **attribs):
                            with xf.element("mesh"):
                                with xf.element("vertices"):
                                    # vertex nodes are writed directly to the file
                                    # so make sure lxml's buffer is flushed
                                    xf.flush()
                                    for i in range(0, len(m.vertices), batch_size):
                                        batch = m.vertices[i: i + batch_size]
                                        fragment = (
                                            '<vertex x="{}" y="{}" z="{}" />'
                                            * len(batch)
                                        )
                                        f.write(
                                            fragment.format(*batch.flatten()).encode(
                                                "utf-8"
                                            )
                                        )
                                with xf.element("triangles"):
                                    xf.flush()
                                    for i in range(0, len(m.faces), batch_size):
                                        batch = m.faces[i: i + batch_size]
                                        fragment = (
                                            '<triangle v1="{}" v2="{}" v3="{}" />'
                                            * len(batch)
                                        )
                                        f.write(
                                            fragment.format(*batch.flatten()).encode(
                                                "utf-8"
                                            )
                                        )

                    # stream components
                    for node in graph.nodes:
                        if node == base_frame or node.startswith("camera"):
                            continue
                        if len(graph[node]) == 0:
                            continue

                        attribs = {
                            "id": model_id(node),
                            "name": node,
                            "type": "model",
                            "p:UUID": str(uuid.uuid4())
                        }
                        with xf.element("object", **attribs):
                            with xf.element("components"):
                                for next, data in graph[node].items():
                                    transform = " ".join(
                                        str(i)
                                        for i in np.array(data["matrix"])[
                                            :3, :4
                                        ].T.flatten()
                                    )
                                    xf.write(
                                        etree.Element(
                                            "component",
                                            {
                                                "objectid": model_id(data["geometry"])
                                                if "geometry" in data
                                                else model_id(next),
                                                "transform": transform,
                                            },
                                        )
                                    )

                # stream build (objects on base_frame)
                with xf.element("build", {"p:UUID": str(uuid.uuid4())}):
                    for node, data in graph[base_frame].items():
                        if node.startswith("camera"):
                            continue
                        transform = " ".join(
                            str(i) for i in np.array(data["matrix"])[:3, :4].T.flatten()
                        )
                        uuid_tag = "{{{}}}UUID".format(model_nsmap['p'])
                        xf.write(
                            etree.Element(
                                "item",
                                {
                                    "objectid": model_id(node),
                                    "transform": transform,
                                    uuid_tag: str(uuid.uuid4())
                                },
                                nsmap=model_nsmap
                            )
                        )

        # .rels
        with z.open("_rels/.rels", "w") as f, etree.xmlfile(f, encoding="utf-8") as xf:
            xf.write_declaration()
            # stream elements
            with xf.element("Relationships", nsmap=rels_nsmap):
                rt = "http://schemas.microsoft.com/3dmanufacturing/2013/01/3dmodel"
                xf.write(
                    etree.Element(
                        "Relationship",
                        Type=rt,
                        Target="/3D/3dmodel.model",
                        Id="rel0",
                    )
                )

        # [Content_Types].xml
        with z.open("[Content_Types].xml", "w") as f, etree.xmlfile(
            f, encoding="utf-8"
        ) as xf:
            xf.write_declaration()
            # xml namespaces
            nsmap = {
                None: "http://schemas.openxmlformats.org/package/2006/content-types"
            }

            # stream elements
            types = [
                ("jpeg", "image/jpeg"),
                ("jpg", "image/jpeg"),
                ("model", "application/vnd.ms-package.3dmanufacturing-3dmodel+xml"),
                ("png", "image/png"),
                ("rels", "application/vnd.openxmlformats-package.relationships+xml"),
                (
                    "texture",
                    "application/vnd.ms-package.3dmanufacturing-3dmodeltexture",
                ),
            ]
            with xf.element("Types", nsmap=nsmap):
                for ext, ctype in types:
                    xf.write(etree.Element("Default", Extension=ext, ContentType=ctype))

    return file_obj.getvalue()

def _attrib_to_transform(attrib):
    """
    Extract a homogeneous transform from a dictionary.

    Parameters
    ------------
    attrib: dict, optionally containing 'transform'

    Returns
    ------------
    transform: (4, 4) float, homogeonous transformation
    """

    transform = np.eye(4, dtype=np.float64)
    if 'transform' in attrib:
        # wangle their transform format
        values = np.array(
            attrib['transform'].split(),
            dtype=np.float64).reshape((4, 3)).T
        transform[:3, :4] = values
    return transform

# do import here to keep lxml a soft dependency
try:
    from lxml import etree
    _three_loaders = {'3mf': load_3MF}
except ImportError:
    _three_loaders = {}


import pybullet as pb
import numpy as np
import time,os,math,inspect,re,errno
import random,glob,math
from shutil import copyfile

import sensenet
from sensenet import spaces
from sensenet.error import Error
from sensenet.envs.handroid.hand_env import HandEnv

class IndexFingerHandEnv(HandEnv):

    def __init__(self,options={}):
        self.options = options
        #print("options",options)
        self.steps = 0

        currentdir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))
        parentdir = os.path.dirname(os.path.dirname(currentdir))
        os.sys.path.insert(0,parentdir)
        #TODO check if options is a string, so we know which environment to load
        if 'render' in self.options and self.options['render'] == True:
            pb.connect(pb.GUI)
        else:
            pb.connect(pb.DIRECT)
        self.timeStep = 0.6 # 0.04 by default
        pb.setPhysicsEngineParameter(fixedTimeStep=self.timeStep)
        pb.setGravity(0,0,0)
        pb.setRealTimeSimulation(0)
        self.move = 0.1 #0.05 # 0.01
        self.pi = 3.1415926535
        self.tinyForce = 5*10e-6
        self.load_object()
        self.load_agent()

        self.pinkId = 0
        self.middleId = 1
        self.indexId = 2
        self.thumbId = 3
        self.ring_id = 4
        self.indexEndID = 21 # Need get position and orientation from index finger parts
        self.offset = 0.02 # Offset from basic position
        self.downCameraOn = False
        self.prev_distance = 10000000
        self.action_space = spaces.Discrete(8)
        self.touch_width = 200
        self.touch_height = 200

    def load_agent(self):
        dir_path = os.path.dirname(os.path.realpath(__file__))
        agent_path = dir_path + "/data/MPL/MPL.xml"
        objects = pb.loadMJCF(agent_path,flags=0)
        self.agent=objects[0]  #1 total
        #if self.obj_to_classify:
        obj_po = pb.getBasePositionAndOrientation(self.obj_to_classify)
        self.hand_cid = pb.createConstraint(self.agent,-1,-1,-1,pb.JOINT_FIXED,[0,0,0],[0,0,0],[0,0,0])
        #hand_po = pb.getBasePositionAndOrientation(self.agent)
        #distance = math.sqrt(sum([(xi-yi)**2 for xi,yi in zip(obj_po[0],hand_po[0])])) #TODO faster euclidean
        pb.resetBasePositionAndOrientation(self.agent,(obj_po[0][0],obj_po[0][1]+0.5,obj_po[0][2]),obj_po[1])

        hand_po = pb.getBasePositionAndOrientation(self.agent)

        ho = pb.getQuaternionFromEuler([0.0, 0.0, 0.0])
        pb.changeConstraint(self.hand_cid,(hand_po[0][0],hand_po[0][1],hand_po[0][2]),ho, maxForce=200)

        joint7Pos = pb.getJointState(self.agent, 7)[0]
        joint24Pos = pb.getJointState(self.agent, 24)[0]
        joint32Pos = pb.getJointState(self.agent, 32)[0]
        joint40Pos = pb.getJointState(self.agent, 40)[0]

        # these are pretty close
        joint7tgt = self.pi/4 # thumb
        joint9tgt = self.pi/2
        joint11tgt = self.pi/2
        joint13tgt = self.pi/2
        # pretty close
        joint24tgt = self.pi/2.25 # middle finger
        joint26tgt = self.pi/2
        joint28tgt = self.pi/2
        # working
        joint32tgt = self.pi/2.5 # ring finger
        joint34tgt = self.pi/2
        joint36tgt = self.pi/2
        # pretty close
        joint40tgt = self.pi/3 # pinky
        joint42tgt = self.pi/2
        joint44tgt = self.pi/2

        while (joint7Pos < joint7tgt) and (joint24Pos < joint24tgt) and (joint32Pos < joint32tgt) and (joint40Pos < joint40tgt):

            pb.setJointMotorControl2(self.agent,7,pb.POSITION_CONTROL,joint7tgt)
            pb.setJointMotorControl2(self.agent,9,pb.POSITION_CONTROL,joint9tgt)
            pb.setJointMotorControl2(self.agent,11,pb.POSITION_CONTROL,joint11tgt)
            pb.setJointMotorControl2(self.agent,13,pb.POSITION_CONTROL,joint13tgt)

            pb.setJointMotorControl2(self.agent,24,pb.POSITION_CONTROL,joint24tgt)
            pb.setJointMotorControl2(self.agent,26,pb.POSITION_CONTROL,joint26tgt)
            pb.setJointMotorControl2(self.agent,28,pb.POSITION_CONTROL,joint28tgt)

            pb.setJointMotorControl2(self.agent,32,pb.POSITION_CONTROL,joint32tgt)
            pb.setJointMotorControl2(self.agent,34,pb.POSITION_CONTROL,joint34tgt)
            pb.setJointMotorControl2(self.agent,36,pb.POSITION_CONTROL,joint36tgt)

            pb.setJointMotorControl2(self.agent,40,pb.POSITION_CONTROL,joint40tgt)
            pb.setJointMotorControl2(self.agent,42,pb.POSITION_CONTROL,joint42tgt)
            pb.setJointMotorControl2(self.agent,44,pb.POSITION_CONTROL,joint44tgt)

            pb.stepSimulation()

            joint7Pos = pb.getJointState(self.agent, 7)[0]
            joint24Pos = pb.getJointState(self.agent, 24)[0]
            joint32Pos = pb.getJointState(self.agent, 32)[0]
            joint40Pos = pb.getJointState(self.agent, 40)[0]
            #return random.random()

    def _step(self,action):
        done = False
        #reward (float): amount of reward achieved by the previous action. The scale varies between environments, but the goal is always to increase your total reward.
        #done (boolean): whether it's time to reset the environment again. Most (but not all) tasks are divided up into well-defined episodes, and done being True indicates the episode has terminated. (For example, perhaps the pole tipped too far, or you lost your last life.)
        #observation (object): an environment-specific object representing your observation of the environment. For example, pixel data from a camera, joint angles and joint velocities of a robot, or the board state in a board game.
        #info (dict): diagnostic information useful for debugging. It can sometimes be useful for learning (for example, it might contain the raw probabilities behind the environment's last state change). However, official evaluations of your agent are not allowed to use this for learning.
        def convertSensor(finger_index):
            if finger_index == self.indexId:
                return random.uniform(-1,1)
                #return 0
            else:
                #return random.uniform(-1,1)
                return 0
        def convertAction(action):
            #converted = (action-30)/10
            #converted = (action-16)/10
            if action == 6:
                converted = -1
            elif action == 25:
                converted = 1
            #print("action ",action)
            #print("converted ",converted)
            return converted

        aspect = 1
        camTargetPos = [0,0,0]
        yaw = 40
        pitch = 10.0
        roll=0
        upAxisIndex = 2
        camDistance = 4
        pixelWidth = 320
        pixelHeight = 240
        nearPlane = 0.0001
        farPlane = 0.022
        lightDirection = [0,1,0]
        lightColor = [1,1,1]#optional
        fov = 50  #10 or 50
        hand_po = pb.getBasePositionAndOrientation(self.agent)
        ho = pb.getQuaternionFromEuler([0.0, 0.0, 0.0])

        # So when trying to modify the physics of the engine, we run into some problems. If we leave
        # angular damping at default (0.04) then the hand rotates when moving up and dow, due to torque.
        # If we set angularDamping to 100.0 then the hand will bounce off into the background due to
        # all the stored energy, when it makes contact with the object. The below set of parameters seem
        # to have a reasonably consistent performance in keeping the hand level and not inducing unwanted
        # behavior during contact.

        pb.changeDynamics(self.agent, linkIndex=-1, angularDamping=0.9999)

        if action == 65298 or action == 0: #down
            pb.changeConstraint(self.hand_cid,(hand_po[0][0]+self.move,hand_po[0][1],hand_po[0][2]),ho, maxForce=self.tinyForce)
        elif action == 65297 or action == 1: #up
            pb.changeConstraint(self.hand_cid,(hand_po[0][0]-self.move,hand_po[0][1],hand_po[0][2]),ho, maxForce=self.tinyForce)
        elif action == 65295 or action == 2: #left
            pb.changeConstraint(self.hand_cid,(hand_po[0][0],hand_po[0][1]+self.move,hand_po[0][2]),ho, maxForce=self.tinyForce)
        elif action== 65296 or action == 3: #right
            pb.changeConstraint(self.hand_cid,(hand_po[0][0],hand_po[0][1]-self.move,hand_po[0][2]),ho, maxForce=self.tinyForce)
        elif action == 44 or action == 4: #<
            pb.changeConstraint(self.hand_cid,(hand_po[0][0],hand_po[0][1],hand_po[0][2]+self.move),ho, maxForce=self.tinyForce)
        elif action == 46 or action == 5: #>
            pb.changeConstraint(self.hand_cid,(hand_po[0][0],hand_po[0][1],hand_po[0][2]-self.move),ho, maxForce=self.tinyForce)
        elif action >= 6 and action <= 7:
            # keeps the hand from moving towards origin
            pb.changeConstraint(self.hand_cid,(hand_po[0][0],hand_po[0][1],hand_po[0][2]),ho, maxForce=200)
            if action == 7:
                action = 25 #bad kludge redo all this code

            index = convertAction(action) # action = 6 or 25 due to kludge -> return -1 or 1

            # Getting positions of the index joints to use for moving to a relative position
            joint17Pos = pb.getJointState(self.agent, 17)[0]
            joint19Pos = pb.getJointState(self.agent, 19)[0]
            joint21Pos = pb.getJointState(self.agent, 21)[0]
            # need to keep the multiplier relatively small otherwise the joint will continue to move
            # when you take other actions
            finger_jump = 0.1
            newJoint17Pos = joint17Pos + index*finger_jump
            newJoint19Pos = joint19Pos + index*finger_jump
            newJoint21Pos = joint21Pos + index*finger_jump

            # following values found by experimentation
            if newJoint17Pos <= -0.7:
                newJoint17Pos = -0.7
            elif newJoint17Pos >= 0.57:
                newJoint17Pos = 0.57
            if newJoint19Pos <= 0.13:
                newJoint19Pos = 0.13
            elif newJoint19Pos >= 0.42:
                newJoint19Pos = 0.42
            if newJoint21Pos <= -0.8:
                newJoint21Pos = -0.8
            elif newJoint21Pos >= 0.58:
                newJoint21Pos = 0.58

            pb.setJointMotorControl2(self.agent,17,pb.POSITION_CONTROL,newJoint17Pos)
            pb.setJointMotorControl2(self.agent,19,pb.POSITION_CONTROL,newJoint19Pos)
            pb.setJointMotorControl2(self.agent,21,pb.POSITION_CONTROL,newJoint21Pos)

        if self.downCameraOn: viewMatrix = down_view()
        else: viewMatrix = self.ahead_view()
        projectionMatrix = pb.computeProjectionMatrixFOV(fov,aspect,nearPlane,farPlane)
        w,h,img_arr,depths,mask = pb.getCameraImage(200,200, viewMatrix,projectionMatrix, lightDirection,lightColor,renderer=pb.ER_TINY_RENDERER)
        #w,h,img_arr,depths,mask = pb.getCameraImage(200,200, viewMatrix,projectionMatrix, lightDirection,lightColor,renderer=pb.ER_BULLET_HARDWARE_OPENGL)
        new_obs = np.absolute(depths-1.0)
        new_obs[new_obs > 0] =1
        self.current_observation = new_obs.flatten()
        info = [42] #answer to life,TODO use real values
        pb.stepSimulation()

        self.steps +=1
        #reward if moving towards the object or touching the object
        reward = 0
        max_steps = 1000
        if self.is_touching():
            touch_reward = 10
            if 'debug' in self.options and self.options['debug'] == True:
                print("TOUCHING!!!!")
        else:
            touch_reward = 0
        obj_po = pb.getBasePositionAndOrientation(self.obj_to_classify)
        distance = math.sqrt(sum([(xi-yi)**2 for xi,yi in zip(obj_po[0],hand_po[0])])) #TODO faster euclidean
        #distance =  np.linalg.norm(obj_po[0],hand_po[0])
        #print("distance:",distance)
        if distance < self.prev_distance:
            reward += 1 * (max_steps - self.steps)
        elif distance > self.prev_distance:
            reward -= 10
        reward -= distance
        reward += touch_reward
        self.prev_distance = distance
        if self.steps >= max_steps or self.is_touching():
            done = True
        return self.current_observation,reward,done,info


from envs.rover_lander_1 import rover_lander_1
from envs.rover_lander_2 import rover_lander_2
import os
import random
import datetime
import requests
from collections import deque
import numpy as np
from tqdm.auto import tqdm

import tensorflow as tf
from tensorflow.keras import Sequential, layers, optimizers, activations

import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--master-endpoint", help="Endpoint for train_master")
parser.add_argument("--worker-name", help="Worker name")
parser.add_argument("--env", help="Env name")
parser.add_argument("--show-preview", help="Show preview", action='store_true', default=False)
parser.add_argument("--mem-size", help="Replay memory size", type=int, default=1000000)
parser.add_argument("--max-ep", help="Maximum episodes", type=int, default=5000)
parser.add_argument("--lr", help="Learning rate", type=float, default=0.001)
parser.add_argument("--loss", help="Loss function", default="mse")
args = parser.parse_args()

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

MASTER_ENDPOINT = args.master_endpoint
WORKER_NAME = args.worker_name
MAX_EPISODES = args.max_ep
SHOW_PREVIEW = args.show_preview
REPLAY_MEMORY_SIZE = args.mem_size

RENDER_EVERY = 10
SAVE_MODEL_EVERY = 500
SAVE_MODEL_LOCAL_EVERY = 20

TRAIN_PARAMS = {'learning_rate':  args.lr,
                'loss':  args.loss,
                'batch_size': 64
                }

Q_PARAMS = {'epsilon': 1.0,
            'gamma': 0.99,
            'epsilon_min': 0.01,
            'epsilon_decay': 0.996}

id = None

def connect():
    global id
    id = int(requests.get(MASTER_ENDPOINT + f"/master/connect?worker_name={WORKER_NAME}&env={args.env}&max_episodes={MAX_EPISODES}&current_episode=0").text)

def update(properties=[]):
    global id
    requests.get(MASTER_ENDPOINT + f"/master/update?id={str(id)}&" + "&".join([str(p[0]) + "=" + str(p[1]) for p in properties]))

def send_model(model_path):
    global id
    files = {'model': open(model_path,'rb')}
    r = requests.post(MASTER_ENDPOINT + f"/master/send_model?id={str(id)}", files=files)

class customCallback(tf.keras.callbacks.Callback):
    def on_epoch_end(self, epoch, logs={}):
        update([("acc", logs.get('acc')),
                ("loss", logs.get('loss')),
                ("mse", logs.get('mse'))])

    # def on_train_end(self, logs=None):
    #     update([("last_trained", "curr_time")])

class Agent:
    def __init__(self):
        self.replay_memory = deque(maxlen=REPLAY_MEMORY_SIZE)
        self.model = self.create_model()
        self.model_v = 0
    def create_model(self):
        model = Sequential([layers.Conv2D(64, kernel_size=3, input_shape=(rover_lander_1.observation_space[0], rover_lander_1.observation_space[1], 3), activation=activations.relu),
                            layers.MaxPooling2D(pool_size=(2, 2)),
                            layers.Conv2D(32, kernel_size=3, activation=activations.relu),
                            layers.MaxPooling2D(pool_size=(2, 2)),
                            layers.Flatten(),
                            layers.Dense(120, activation=activations.relu),
                            layers.Dense(rover_lander_1.action_space, activation=activations.linear)])
        model.compile(loss=TRAIN_PARAMS['loss'], optimizer=optimizers.Adam(lr=TRAIN_PARAMS['learning_rate']), metrics=['mse', 'acc'])
        return model

    def load_model(self, model_path):
        self.model = tf.keras.models.load_model(model_path)

    def save_model(self, local_only=True):
        self.model.save(f"models/{WORKER_NAME}_v{self.model_v}.h5")
        self.model_v = self.model_v + 1
        if not local_only:
            send_model(f"models/{WORKER_NAME}.h5")

    def update_replay_memory(self, state, action, reward, next_state, done):
        self.replay_memory.append((state, action, reward, next_state, done))

    def replay(self):
        if len(self.replay_memory) < TRAIN_PARAMS['batch_size']:
            return

        minibatch = random.sample(self.replay_memory, TRAIN_PARAMS['batch_size'])
        states = np.array([i[0] for i in minibatch])
        actions = np.array([i[1] for i in minibatch])
        rewards = np.array([i[2] for i in minibatch])
        next_states = np.array([i[3] for i in minibatch])
        dones = np.array([i[4] for i in minibatch])

        # states = states.reshape(states.shape[0], np.prod(states.shape[1:]))/255
        # next_states = next_states.reshape(next_states.shape[0], np.prod(next_states.shape[1:]))/255

        states = states/255
        next_states = next_states/255

        targets = rewards + Q_PARAMS['gamma']*(np.amax(self.model.predict_on_batch(next_states), axis=1))*(1-dones)
        targets_full = self.model.predict_on_batch(states)
        ind = np.array([i for i in range(TRAIN_PARAMS['batch_size'])])
        targets_full[[ind], [actions]] = targets

        self.model.fit(states, targets_full, epochs=1, verbose=0, callbacks=[customCallback()])
        if Q_PARAMS['epsilon'] > Q_PARAMS['epsilon_min']:
            Q_PARAMS['epsilon'] *= Q_PARAMS['epsilon_decay']

    def qs(self, state):
        state = state.reshape(1, *state.shape)/255
        return np.argmax(self.model.predict(state))

if __name__ == '__main__':
    if not SHOW_PREVIEW:
        os.environ["SDL_VIDEODRIVER"] = "dummy"

    if args.env == 'rover_lander_1':
        env = rover_lander_1()
    elif args.env == 'rover_lander_2':
        env = rover_lander_2()

    agent = Agent()
    connect()
    episode_rewards = []
    for episode in tqdm(range(0, MAX_EPISODES), ascii=True, unit='episodes'):
        episode_reward = 0
        step = 1
        current_state = env.reset()
        done = False
        while not done:
            if np.random.random() > Q_PARAMS['epsilon']:
                action = agent.qs(current_state)
            else:
                action = env.random_action_sample()

            if SHOW_PREVIEW and not episode % RENDER_EVERY:
                    env.render()

            new_state, reward, done = env.step(action)
            episode_reward += reward

            agent.update_replay_memory(current_state, action, reward, new_state, done)
            current_state = new_state
            agent.replay()
            step = step + 1

        episode_rewards.append(episode_reward)
        update([("last_ep_score", episode_reward),
            ("avg_ep_score", sum(episode_rewards) / len(episode_rewards)),
            ("episodes", episode)])

        if not episode % SAVE_MODEL_EVERY:
            pass
            # agent.save_model(local_only=False)

        if not episode % SAVE_MODEL_LOCAL_EVERY:
            agent.save_model(local_only=True)


#!/usr/bin/env python
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.

import os
import sys
import subprocess

import torch
import numpy as np

from model.model import Model

def get_a3m_feat(path):
    with open(path) as fp:
        line = fp.readline()
        if line.startswith(">"):
            line = fp.readline()
        L = len(line.strip())
    program = [
        os.path.join(os.path.dirname(__file__), "bin/a3m_to_feat"),
        "--input",
        path,
        "--max_gap",
        "7",
        "--max_keep",
        "5000",
        "--sample_ratio",
        "1.0",
    ]
    process = subprocess.run(program, capture_output=True)
    assert process.returncode == 0, "Invalid A3M file"
    x = np.copy(np.frombuffer(process.stdout, dtype=np.int8))
    x = x.reshape((-1, L, 7 * 2 + 3)).transpose((0, 2, 1))
    assert (x < 23).all(), "Internal error"
    seq = x[0][0]
    return {
        "seq": torch.tensor(seq).long()[None],
        "msa": torch.tensor(x).long()[None],
        "index": torch.arange(seq.shape[0]).long()[None],
    }

def load_model():
    model = Model().eval()
    weight_path = os.path.join(os.path.dirname(__file__), "weights.chk")
    state_dict = torch.load(weight_path, map_location=torch.device("cpu"))
    model.load_state_dict(state_dict)
    return model

def main():
    if len(sys.argv) != 3:
        print("Usage: python run_inference.py <a3m> <output_file>")
        sys.exit(1)
    feat = get_a3m_feat(sys.argv[1])
    model = load_model()
    with torch.no_grad():
        output = model(feat)
    contact = torch.sum(output[..., :12], dim=-1)[0]
    np.savetxt(sys.argv[2], contact)

if __name__ == "__main__":
    main()