"""Load transformations from URDF files.

See :doc:`transform_manager` for more information.
"""
import os
import numpy as np
from bs4 import BeautifulSoup
from .transform_manager import TransformManager
from .transformations import transform_from, concat
from .rotations import (
    active_matrix_from_extrinsic_roll_pitch_yaw, matrix_from_axis_angle,
    norm_vector)

class UrdfTransformManager(TransformManager):
    """Transformation manager that can load URDF files.

    URDF is the `Unified Robot Description Format <http://wiki.ros.org/urdf>`_.
    URDF allows to define joints between links that can be rotated about one
    axis. This transformation manager allows to set the joint angles after
    joints have been added or loaded from an URDF.

    .. warning::

        Note that this module requires the Python package beautifulsoup4.

    .. note::

        Joint angles must be given in radians.

    Parameters
    ----------
    strict_check : bool, optional (default: True)
        Raise a ValueError if the transformation matrix is not numerically
        close enough to a real transformation matrix. Otherwise we print a
        warning.

    check : bool, optional (default: True)
        Check if transformation matrices are valid and requested nodes exist,
        which might significantly slow down some operations.
    """
    def __init__(self, strict_check=True, check=True):
        super(UrdfTransformManager, self).__init__(strict_check, check)
        self._joints = {}
        self.collision_objects = []
        self.visuals = []

    def add_joint(self, joint_name, from_frame, to_frame, child2parent, axis,
                  limits=(float("-inf"), float("inf")), joint_type="revolute"):
        """Add joint.

        Parameters
        ----------
        joint_name : str
            Name of the joint

        from_frame : Hashable
            Child link of the joint

        to_frame : Hashable
            Parent link of the joint

        child2parent : array-like, shape (4, 4)
            Transformation from child to parent

        axis : array-like, shape (3,)
            Rotation axis of the joint (defined in the child frame)

        limits : pair of float, optional (default: (-inf, inf))
            Lower and upper joint angle limit

        joint_type : str, optional (default: 'revolute')
            Joint type: revolute or prismatic (continuous is the same as
            revolute)
        """
        self.add_transform(from_frame, to_frame, child2parent)
        self._joints[joint_name] = (
            from_frame, to_frame, child2parent, norm_vector(axis), limits,
            joint_type)

    def set_joint(self, joint_name, value):
        """Set joint position.

        Note that joint values are clipped to their limits.

        Parameters
        ----------
        joint_name : str
            Name of the joint

        value : float
            Joint angle in radians in case of revolute joints or position
            in case of prismatic joint.

        Raises
        ------
        KeyError
            If joint_name is unknown
        """
        if joint_name not in self._joints:
            raise KeyError("Joint '%s' is not known" % joint_name)
        from_frame, to_frame, child2parent, axis, limits, joint_type = \
            self._joints[joint_name]
        # this is way faster than np.clip:
        value = min(max(value, limits[0]), limits[1])
        if joint_type == "revolute":
            joint_rotation = matrix_from_axis_angle(
                np.hstack((axis, (value,))))
            joint2A = transform_from(
                joint_rotation, np.zeros(3), strict_check=self.strict_check)
        else:
            assert joint_type == "prismatic"
            joint_offset = value * axis
            joint2A = transform_from(
                np.eye(3), joint_offset, strict_check=self.strict_check)
        self.add_transform(from_frame, to_frame, concat(
            joint2A, child2parent, strict_check=self.strict_check,
            check=self.check))

    def get_joint_limits(self, joint_name):
        """Get limits of a joint.

        Parameters
        ----------
        joint_name : str
            Name of the joint

        Returns
        -------
        limits : pair of float
            Lower and upper joint angle limit

        Raises
        ------
        KeyError
            If joint_name is unknown
        """
        if joint_name not in self._joints:
            raise KeyError("Joint '%s' is not known" % joint_name)
        return self._joints[joint_name][4]

    def load_urdf(self, urdf_xml, mesh_path=None, package_dir=None):
        """Load URDF file into transformation manager.

        Parameters
        ----------
        urdf_xml : str
            Robot definition in URDF

        mesh_path : str, optional (default: None)
            Path in which we search for meshes that are defined in the URDF.
            Meshes will be ignored if it is set to None and no 'package_dir'
            is given.

        package_dir : str, optional (default: None)
            Some URDFs start file names with 'package://' to refer to the ROS
            package in which these files (textures, meshes) are located. This
            variable defines to which path this prefix will be resolved.
        """
        robot_name, links, joints = parse_urdf(
            urdf_xml, mesh_path, package_dir, self.strict_check)
        initialize_urdf_transform_manager(self, robot_name, links, joints)

    def plot_visuals(self, frame, ax=None, ax_s=1, wireframe=False,
                     convex_hull_of_mesh=True, alpha=0.3):  # pragma: no cover
        """Plot all visuals in a given reference frame.

        Visuals can be boxes, spheres, cylinders, or meshes. Note that visuals
        that cannot be connected to the reference frame are omitted.

        Parameters
        ----------
        frame : Hashable
            Reference frame

        ax : Matplotlib 3d axis, optional (default: None)
            If the axis is None, a new 3d axis will be created

        ax_s : float, optional (default: 1)
            Scaling of the new matplotlib 3d axis

        wireframe : bool, optional (default: False)
            Plot wireframe (surface otherwise)

        convex_hull_of_mesh : bool, optional (default: True)
            Displays convex hull of meshes instead of the original mesh. This
            makes plotting a lot faster with complex meshes.

        alpha : float, optional (default: 0.3)
            Alpha value of the surface / wireframe that will be plotted

        Returns
        -------
        ax : Matplotlib 3d axis
            New or old axis
        """
        return self._plot_objects(
            self.visuals, frame, ax, ax_s, wireframe, convex_hull_of_mesh,
            alpha)

    def plot_collision_objects(
            self, frame, ax=None, ax_s=1, wireframe=True,
            convex_hull_of_mesh=True, alpha=1.0):  # pragma: no cover
        """Plot all collision objects in a given reference frame.

        Collision objects can be boxes, spheres, cylinders, or meshes. Note
        that collision objects that cannot be connected to the reference frame
        are omitted.

        Parameters
        ----------
        frame : Hashable
            Reference frame

        ax : Matplotlib 3d axis, optional (default: None)
            If the axis is None, a new 3d axis will be created

        ax_s : float, optional (default: 1)
            Scaling of the new matplotlib 3d axis

        wireframe : bool, optional (default: True)
            Plot wireframe (surface otherwise)

        convex_hull_of_mesh : bool, optional (default: True)
            Displays convex hull of meshes instead of the original mesh. This
            makes plotting a lot faster with complex meshes.

        alpha : float, optional (default: 1)
            Alpha value of the surface / wireframe that will be plotted

        Returns
        -------
        ax : Matplotlib 3d axis
            New or old axis
        """
        return self._plot_objects(
            self.collision_objects, frame, ax, ax_s, wireframe,
            convex_hull_of_mesh, alpha)

    def _plot_objects(self, objects, frame, ax=None, ax_s=1, wireframe=True,
                      convex_hull_of_mesh=True, alpha=1.0):  # pragma: no cover
        """Plot all objects in a given reference frame.

        Objects can be boxes, spheres, cylinders, or meshes. Note that objects
        that cannot be connected to the reference frame are omitted.

        Parameters
        ----------
        objects : list
            Objects that will be plotted

        frame : Hashable
            Reference frame

        ax : Matplotlib 3d axis, optional (default: None)
            If the axis is None, a new 3d axis will be created

        ax_s : float, optional (default: 1)
            Scaling of the new matplotlib 3d axis

        wireframe : bool, optional (default: True)
            Plot wireframe (surface otherwise)

        convex_hull_of_mesh : bool, optional (default: True)
            Displays convex hull of meshes instead of the original mesh. This
            makes plotting a lot faster with complex meshes.

        alpha : float, optional (default: 1)
            Alpha value of the surface / wireframe that will be plotted

        Returns
        -------
        ax : Matplotlib 3d axis
            New or old axis
        """
        if ax is None:
            from .plot_utils import make_3d_axis
            ax = make_3d_axis(ax_s)
        for obj in objects:
            ax = obj.plot(
                self, frame, ax, wireframe=wireframe,
                convex_hull=convex_hull_of_mesh, alpha=alpha)
        return ax

def parse_urdf(urdf_xml, mesh_path=None, package_dir=None, strict_check=True):
    """Parse information from URDF file.

    Parameters
    ----------
    urdf_xml : str
        Robot definition in URDF

    mesh_path : str, optional (default: None)
        Path in which we search for meshes that are defined in the URDF.
        Meshes will be ignored if it is set to None and no 'package_dir'
        is given.

    package_dir : str, optional (default: None)
        Some URDFs start file names with 'package://' to refer to the ROS
        package in which these files (textures, meshes) are located. This
        variable defines to which path this prefix will be resolved.

    strict_check : bool, optional (default: True)
        Raise a ValueError if the transformation matrix is not numerically
        close enough to a real transformation matrix. Otherwise we print a
        warning.

    Returns
    -------
    robot_name : str
        Name of the robot

    links : list of Link
        Links of the robot

    joints : list of Joint
        Joints of the robot

    Raises
    ------
    UrdfException
        If URDF is not valid
    """
    urdf = BeautifulSoup(urdf_xml, "xml")

    # URDF XML schema:
    # https://github.com/ros/urdfdom/blob/master/xsd/urdf.xsd

    robot = urdf.find("robot")
    if robot is None:
        raise UrdfException("Robot tag is missing.")

    if not robot.has_attr("name"):
        raise UrdfException("Attribute 'name' is missing in robot tag.")

    robot_name = robot["name"]

    materials = dict([
        _parse_material(material)
        for material in robot.findAll("material", recursive=False)])

    links = [_parse_link(link, materials, mesh_path, package_dir, strict_check)
             for link in robot.findAll("link", recursive=False)]

    link_names = [link.name for link in links]
    joints = [_parse_joint(joint, link_names, strict_check)
              for joint in robot.findAll("joint", recursive=False)]

    return robot_name, links, joints

def initialize_urdf_transform_manager(tm, robot_name, links, joints):
    """Initializes transform manager from previously parsed URDF data.

    Parameters
    ----------
    tm : UrdfTransformManager
        Transform manager

    robot_name : str
        Name of the robot

    links : list of Link
        Links of the robot

    joints : list of Joint
        Joints of the robot
    """
    tm.add_transform(links[0].name, robot_name, np.eye(4))
    _add_links(tm, links)
    _add_joints(tm, joints)

def _parse_material(material):
    """Parse material."""
    if not material.has_attr("name"):
        raise UrdfException("Material name is missing.")
    colors = material.findAll("color")
    if len(colors) not in [0, 1]:
        raise UrdfException("More than one color is not allowed.")
    if len(colors) == 1:
        color = _parse_color(colors[0])
    else:
        color = None
    # TODO texture is currently ignored
    return material["name"], color

def _parse_color(color):
    """Parse color."""
    if not color.has_attr("rgba"):
        raise UrdfException("Attribute 'rgba' of color tag is missing.")
    return np.fromstring(color["rgba"], sep=" ")

def _parse_link(link, materials, mesh_path, package_dir, strict_check):
    """Create link."""
    if not link.has_attr("name"):
        raise UrdfException("Link name is missing.")

    result = Link()
    result.name = link["name"]

    visuals, visual_transforms = _parse_link_children(
        link, "visual", materials, mesh_path, package_dir, strict_check)
    result.visuals = visuals
    result.transforms.extend(visual_transforms)

    collision_objects, collision_object_transforms = _parse_link_children(
        link, "collision", dict(), mesh_path, package_dir, strict_check)
    result.collision_objects = collision_objects
    result.transforms.extend(collision_object_transforms)

    inertial = link.find("inertial")
    if inertial is not None:
        result.inertial_frame[:, :] = _parse_origin(inertial, strict_check)
        result.mass = _parse_mass(inertial)
        result.inertia[:, :] = _parse_inertia(inertial)
        result.transforms.append(
            ("inertial_frame:%s" % result.name, result.name,
             result.inertial_frame))

    return result

def _parse_link_children(link, child_type, materials, mesh_path, package_dir,
                         strict_check):
    """Parse collision objects or visuals."""
    children = link.findAll(child_type)
    shape_objects = []
    transforms = []
    for i, child in enumerate(children):
        if child.has_attr("name"):
            name = "%s:%s/%s" % (child_type, link["name"], child["name"])
        else:
            name = "%s:%s/%s" % (child_type, link["name"], i)

        color = None
        if child_type == "visual":
            material = child.find("material")
            if material is not None:
                material_name, color = _parse_material(material)
                if color is None and material_name in materials:
                    color = materials[material_name]

        child2link = _parse_origin(child, strict_check)
        transforms.append((name, link["name"], child2link))

        shape_objects.extend(_parse_geometry(
            child, name, color, mesh_path, package_dir))
    return shape_objects, transforms

def _parse_geometry(child, name, color, mesh_path, package_dir):
    """Parse geometric primitives (box, cylinder, sphere) or meshes."""
    geometry = child.find("geometry")
    if geometry is None:
        raise UrdfException("Missing geometry tag in link '%s'" % name)
    result = []
    for shape_type in ["box", "cylinder", "sphere", "mesh"]:
        shapes = geometry.findAll(shape_type)
        Cls = shape_classes[shape_type]
        for shape in shapes:
            shape_object = Cls(
                name, mesh_path=mesh_path, package_dir=package_dir,
                color=color)
            shape_object.parse(shape)
            result.append(shape_object)
    return result

def _parse_origin(entry, strict_check):
    """Parse transformation."""
    origin = entry.find("origin")
    translation = np.zeros(3)
    rotation = np.eye(3)
    if origin is not None:
        if origin.has_attr("xyz"):
            translation = np.fromstring(origin["xyz"], sep=" ")
        if origin.has_attr("rpy"):
            roll_pitch_yaw = np.fromstring(origin["rpy"], sep=" ")
            # URDF and KDL use the active convention for rotation matrices.
            # For more details on how the URDF parser handles the
            # conversion from Euler angles, see this blog post:
            # https://orbitalstation.wordpress.com/tag/quaternion/
            rotation = active_matrix_from_extrinsic_roll_pitch_yaw(
                roll_pitch_yaw)
    return transform_from(
        rotation, translation, strict_check=strict_check)

def _parse_mass(inertial):
    """Parse link mass."""
    mass = inertial.find("mass")
    if mass is not None and mass.has_attr("value"):
        result = float(mass["value"])
    else:
        result = 0.0
    return result

def _parse_inertia(inertial):
    """Parse inertia matrix."""
    inertia = inertial.find("inertia")

    result = np.zeros((3, 3))
    if inertia is None:
        return result

    if inertia.has_attr("ixx"):
        result[0, 0] = float(inertia["ixx"])
    if inertia.has_attr("ixy"):
        ixy = float(inertia["ixy"])
        result[0, 1] = ixy
        result[1, 0] = ixy
    if inertia.has_attr("ixz"):
        ixz = float(inertia["ixz"])
        result[0, 2] = ixz
        result[2, 0] = ixz
    if inertia.has_attr("iyy"):
        result[1, 1] = float(inertia["iyy"])
    if inertia.has_attr("iyz"):
        iyz = float(inertia["iyz"])
        result[1, 2] = iyz
        result[2, 1] = iyz
    if inertia.has_attr("izz"):
        result[2, 2] = float(inertia["izz"])
    return result

def _parse_joint(joint, link_names, strict_check):
    """Create joint object."""
    j = Joint()

    if not joint.has_attr("name"):
        raise UrdfException("Joint name is missing.")
    j.joint_name = joint["name"]

    if not joint.has_attr("type"):
        raise UrdfException("Joint type is missing in joint '%s'."
                            % j.joint_name)

    parent = joint.find("parent")
    if parent is None:
        raise UrdfException("No parent specified in joint '%s'"
                            % j.joint_name)
    if not parent.has_attr("link"):
        raise UrdfException("No parent link name given in joint '%s'."
                            % j.joint_name)
    j.parent = parent["link"]
    if j.parent not in link_names:
        raise UrdfException("Parent link '%s' of joint '%s' is not "
                            "defined." % (j.parent, j.joint_name))

    child = joint.find("child")
    if child is None:
        raise UrdfException("No child specified in joint '%s'"
                            % j.joint_name)
    if not child.has_attr("link"):
        raise UrdfException("No child link name given in joint '%s'."
                            % j.joint_name)
    j.child = child["link"]
    if j.child not in link_names:
        raise UrdfException("Child link '%s' of joint '%s' is not "
                            "defined." % (j.child, j.joint_name))

    j.joint_type = joint["type"]

    if j.joint_type in ["planar", "floating"]:
        raise UrdfException("Unsupported joint type '%s'" % j.joint_type)
    if j.joint_type not in ["revolute", "continuous", "prismatic",
                            "fixed"]:
        raise UrdfException("Joint type '%s' is not allowed in a URDF "
                            "document." % j.joint_type)

    j.child2parent = _parse_origin(joint, strict_check)

    j.joint_axis = np.array([1, 0, 0])
    if j.joint_type in ["revolute", "continuous", "prismatic"]:
        axis = joint.find("axis")
        if axis is not None and axis.has_attr("xyz"):
            j.joint_axis = np.fromstring(axis["xyz"], sep=" ")

    j.limits = _parse_limits(joint)
    return j

def _parse_limits(joint):
    """Parse joint limits."""
    limit = joint.find("limit")
    lower, upper = float("-inf"), float("inf")
    if limit is not None:
        if limit.has_attr("lower"):
            lower = float(limit["lower"])
        if limit.has_attr("upper"):
            upper = float(limit["upper"])
    return lower, upper

def _add_links(tm, links):
    """Add previously parsed links.

    Parameters
    ----------
    tm : UrdfTransformManager
        Transform manager

    links : list of Link
        Joint information from URDF
    """
    for link in links:
        tm.visuals.extend(link.visuals)
        tm.collision_objects.extend(link.collision_objects)

        for from_frame, to_frame, transform in link.transforms:
            tm.add_transform(from_frame, to_frame, transform)

def _add_joints(tm, joints):
    """Add previously parsed joints.

    Parameters
    ----------
    tm : UrdfTransformManager
        Transform manager

    joints : list of Joint
        Joint information from URDF
    """
    for joint in joints:
        if joint.joint_type in ["revolute", "continuous"]:
            tm.add_joint(
                joint.joint_name, joint.child, joint.parent,
                joint.child2parent, joint.joint_axis, joint.limits,
                "revolute")
        elif joint.joint_type == "prismatic":
            tm.add_joint(
                joint.joint_name, joint.child, joint.parent,
                joint.child2parent, joint.joint_axis, joint.limits,
                "prismatic")
        else:
            assert joint.joint_type == "fixed"
            tm.add_transform(joint.child, joint.parent, joint.child2parent)

class Link(object):
    """Link from URDF file.

    This class is only required temporarily while we parse the URDF.

    Attributes
    ----------
    name : str
        Link name

    visuals : list of Geometry
        Visual geometries

    collision_objects : list of Geometry
        Geometries for collision calculation

    transforms : list
        Transformations given as tuples: name of frame A, name of frame B,
        transform A2B

    inertial_frame : array, shape (4, 4)
        Pose of inertial frame with respect to the link

    mass : float
        Mass of the link

    inertia : array, shape (3, 3)
        Inertia matrix
    """
    def __init__(self):
        self.name = None
        self.visuals = []
        self.collision_objects = []
        self.transforms = []
        self.inertial_frame = np.eye(4)
        self.mass = 0.0
        self.inertia = np.zeros((3, 3))

class Joint(object):
    """Joint from URDF file.

    This class is only required temporarily while we parse the URDF.

    Attributes
    ----------
    child : str
        Name of the child

    parent : str
        Name of the parent frame

    child2parent : array-like, shape (4, 4)
        Transformation from child to parent

    joint_name : str
        Name of the joint that defines the transformation

    joint_axis : array-like, shape (3,)
        Rotation axis of the joint (defined in the child frame)

    joint_type : str
        Either 'fixed' or 'revolute'

    limits : pair of float
        Lower and upper joint angle limit
    """
    def __init__(self):
        self.child = None
        self.parent = None
        self.child2parent = np.eye(4)
        self.joint_name = None
        self.joint_axis = None
        self.joint_type = "fixed"
        self.limits = float("-inf"), float("inf")

class Geometry(object):
    """Geometrical object."""
    def __init__(self, frame, mesh_path, package_dir, color):
        self.frame = frame
        self.mesh_path = mesh_path
        self.package_dir = package_dir
        self.color = color

    def parse(self, xml):
        """Parse parameters of geometry."""

    def plot(self, tm, frame, ax=None, alpha=0.3, wireframe=True,
             convex_hull=True):
        """Plot geometry."""

class Box(Geometry):
    """Geometrical object: box."""
    def __init__(self, frame, mesh_path, package_dir, color):
        super(Box, self).__init__(frame, mesh_path, package_dir, color)
        self.size = np.zeros(3)

    def parse(self, xml):
        """Parse box size."""
        if xml.has_attr("size"):
            self.size[:] = np.fromstring(xml["size"], sep=" ")

    def plot(self, tm, frame, ax=None, alpha=0.3, wireframe=True,
             convex_hull=True):  # pragma: no cover
        """Plot box."""
        A2B = tm.get_transform(self.frame, frame)
        color = self.color if self.color is not None else "k"
        from .plot_utils import plot_box
        return plot_box(
            ax, self.size, A2B, wireframe=wireframe, alpha=alpha, color=color)

class Sphere(Geometry):
    """Geometrical object: sphere."""
    def __init__(self, frame, mesh_path, package_dir, color):
        super(Sphere, self).__init__(frame, mesh_path, package_dir, color)
        self.radius = 0.0

    def parse(self, xml):
        """Parse sphere radius."""
        if not xml.has_attr("radius"):
            raise UrdfException("Sphere has no radius.")
        self.radius = float(xml["radius"])

    def plot(self, tm, frame, ax=None, alpha=0.3, wireframe=True,
             convex_hull=True):  # pragma: no cover
        """Plot sphere."""
        center = tm.get_transform(self.frame, frame)[:3, 3]
        color = self.color if self.color is not None else "k"
        from .plot_utils import plot_sphere
        return plot_sphere(
            ax, self.radius, center, wireframe=wireframe, alpha=alpha,
            color=color)

class Cylinder(Geometry):
    """Geometrical object: cylinder."""
    def __init__(self, frame, mesh_path, package_dir, color):
        super(Cylinder, self).__init__(frame, mesh_path, package_dir, color)
        self.radius = 0.0
        self.length = 0.0

    def parse(self, xml):
        """Parse cylinder radius and length."""
        if not xml.has_attr("radius"):
            raise UrdfException("Cylinder has no radius.")
        self.radius = float(xml["radius"])
        if not xml.has_attr("length"):
            raise UrdfException("Cylinder has no length.")
        self.length = float(xml["length"])

    def plot(self, tm, frame, ax=None, alpha=0.3, wireframe=True,
             convex_hull=True):  # pragma: no cover
        """Plot cylinder."""
        A2B = tm.get_transform(self.frame, frame)
        color = self.color if self.color is not None else "k"
        from .plot_utils import plot_cylinder
        return plot_cylinder(
            ax, self.length, self.radius, 0.0, A2B, wireframe=wireframe,
            alpha=alpha, color=color)

class Mesh(Geometry):
    """Geometrical object: mesh."""
    def __init__(self, frame, mesh_path, package_dir, color):
        super(Mesh, self).__init__(frame, mesh_path, package_dir, color)
        self.filename = None
        self.scale = np.ones(3)

    def parse(self, xml):
        """Parse mesh filename and scale."""
        if self.mesh_path is None and self.package_dir is None:
            self.filename = None
        else:
            if not xml.has_attr("filename"):
                raise UrdfException("Mesh has no filename.")
            if self.mesh_path is not None:
                self.filename = os.path.join(self.mesh_path, xml["filename"])
            else:
                assert self.package_dir is not None
                self.filename = xml["filename"].replace(
                    "package://", self.package_dir)
            if xml.has_attr("scale"):
                self.scale = np.fromstring(xml["scale"], sep=" ")

    def plot(self, tm, frame, ax=None, alpha=0.3, wireframe=True,
             convex_hull=True):  # pragma: no cover
        """Plot mesh."""
        from .plot_utils import plot_mesh
        A2B = tm.get_transform(self.frame, frame)
        color = self.color if self.color is not None else "k"
        return plot_mesh(
            ax, self.filename, A2B, self.scale, wireframe=wireframe,
            convex_hull=convex_hull, alpha=alpha, color=color)

shape_classes = {"box": Box,
                 "sphere": Sphere,
                 "cylinder": Cylinder,
                 "mesh": Mesh}

class UrdfException(Exception):
    """Exception while parsing URDF files."""


# -*- coding: utf-8 -*-
"""
Created on Sun Feb 26 16:41:01 2017

@author: andi
"""

# exec(open("/home/andi/dev/visr/src/python/scripts/instantiate_multi_renderer.py").read())

import visr
import panning
import rcl
import pml

import numpy as np

# from multi_renderer import OutputSwitch
from multi_renderer import MultiRenderer

class RealTimeMultiRenderer( visr.CompositeComponent ):
    def __init__( self, context, name, parent,
                 loudspeakerConfigFiles,
                 numberOfInputs,
                 numberOfOutputs,
                 interpolationPeriod,
                 diffusionFilterFile,
                 udpReceivePort=8888,
                 controlReceivePort=8889,
                 trackingConfiguration='' ):
        super(RealTimeMultiRenderer,self).__init__( context, name, parent)
        self.input = visr.AudioInputFloat( "in", self, numberOfInputs )
        self.output = visr.AudioOutputFloat( "out", self, numberOfOutputs )

        rendererConfigs = []
        for cfgFile in loudspeakerConfigFiles:
            rendererConfigs.append( panning.LoudspeakerArray(cfgFile) )

        diffFilters = np.array(pml.MatrixParameterFloat.fromAudioFile( diffusionFilterFile ))

        self.multiRenderer = MultiRenderer(context, name, self,
                                           loudspeakerConfigs=rendererConfigs,
                                           numberOfInputs=numberOfInputs,
                                           numberOfOutputs=numberOfOutputs,
                                           interpolationPeriod=interpolationPeriod,
                                           diffusionFilters=diffFilters,
                                           trackingConfiguration='',
                                           controlDataType=pml.Float )
        self.audioConnection( self.input, self.multiRenderer.audioPort("in" ) )
        self.audioConnection( self.multiRenderer.audioPort("out" ), self.output )

        self.sceneReceiver = rcl.UdpReceiver( context, "SceneReceiver", self,
                                             port=udpReceivePort,
                                             mode=rcl.UdpReceiver.Mode.Asynchronous )
        self.sceneDecoder = rcl.SceneDecoder( context, "SceneDecoder", self )
        self.parameterConnection( self.sceneReceiver.parameterPort("messageOutput"),
                                 self.sceneDecoder.parameterPort("datagramInput") )
        self.parameterConnection( self.sceneDecoder.parameterPort( "objectVectorOutput"),
                                 self.multiRenderer.parameterPort( "objectIn" ) )

        self.controlReceiver = rcl.UdpReceiver( context, "ControlReceiver", self,
                                               port=controlReceivePort,
                                               mode=rcl.UdpReceiver.Mode.Asynchronous)
        self.controlDecoder = rcl.ScalarOscDecoder( context, "ControlDecoder", dataType='float')
        self.parameterConnection( self.controlReceiver.parameterPort("messageOutput"),
                                 self.controlDecoder.parameterPort("datagramInput") )
        self.parameterConnection( self.controlDecoder.parameterPort( "dataOut"),
                                 self.multiRenderer.parameterPort( "controlIn" ) )


import os
import io
import logging
import subprocess
import time
import math
import numpy as np
import torch
import torch.nn as nn
from third_party import wavfile
from third_party import kaldi_io as kio

TENSORBOARD_LOGGING = 1

def cleanup_ckpt(expdir, num_last_ckpt_keep):
    ckptlist = [t for t in os.listdir(expdir) if t.endswith('.pt') and t != 'last-ckpt.pt']
    ckptlist = sorted(ckptlist)
    ckptlist_rm = ckptlist[:-num_last_ckpt_keep]
    logging.info("Clean up checkpoints. Remain the last {} checkpoints.".format(num_last_ckpt_keep))
    for name in ckptlist_rm:
       os.remove(os.path.join(expdir, name))

def get_command_stdout(command, require_zero_status=True):
    """ Executes a command and returns its stdout output as a string.  The
        command is executed with shell=True, so it may contain pipes and
        other shell constructs.

        If require_zero_stats is True, this function will raise an exception if
        the command has nonzero exit status.  If False, it just prints a warning
        if the exit status is nonzero.

        See also: execute_command, background_command
    """
    p = subprocess.Popen(command, shell=True,
                         stdout=subprocess.PIPE)

    stdout = p.communicate()[0]
    if p.returncode is not 0:
        output = "Command exited with status {0}: {1}".format(
            p.returncode, command)
        if require_zero_status:
            raise Exception(output)
        else:
            logger.warning(output)
    return stdout

def load_wave(path):
    """
    path can be wav filename or pipeline
    """

    # parse path
    items = path.strip().split(":", 1)
    if len(items) != 2:
        raise ValueError("Unknown path format.")
    tag = items[0]
    path = items[1]
    if tag == "file":
        sample_rate, data = wavfile.read(path)
    elif tag == "pipe":
        path = path[:-1]
        out = get_command_stdout(path, require_zero_status=True)
        sample_rate, data = wavfile.read(io.BytesIO(out))
    elif tag == "ark":
        fn, offset = path.split(":", 1)
        offset = int(offset)
        with open(fn, 'rb') as f:
            f.seek(offset)
            sample_rate, data = wavfile.read(f, offset=offset)
    else:
        raise ValueError("Unknown file tag.")
    data = data.astype(np.float32)
    return sample_rate, data

def load_feat(path):
    items = path.strip().split(":", 1)
    if len(items) != 2:
        raise ValueError("Unknown path format.")
    tag = items[0]
    path = items[1]
    if tag == "ark":
        return kio.read_mat(path)
    else:
        raise ValueError("Unknown file tag.")

def parse_scp(fn):
    dic = {}
    with open(fn, 'r') as f:
        cnt = 0
        for line in f:
            cnt += 1
            items = line.strip().split(' ', 1)
            if len(items) != 2:
                logging.warning('Wrong formated line {} in scp {}, skip it.'.format(cnt, fn))
                continue
            dic[items[0]] = items[1]
    return dic

def str2bool(v):
    if v.lower() in ('yes', 'true', 't', 'y', '1'):
        return True
    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
        return False
    else:
        raise argparse.ArgumentTypeError('Unsupported value encountered.')

class Timer(object):
    def __init__(self):
        self.start = 0.

    def tic(self):
        self.start = time.time()

    def toc(self):
        return time.time() - self.start

# ==========================================
# auxilary functions for sequence
# ==========================================

def get_paddings(src, lengths):
    paddings = torch.zeros_like(src).to(src.device)
    for b in range(lengths.shape[0]):
        paddings[b, lengths[b]:, :] = 1
    return paddings

def get_paddings_by_shape(shape, lengths, device="cpu"):
    paddings = torch.zeros(shape).to(device)
    if shape[0] != lengths.shape[0]:
        raise ValueError("shape[0] does not match lengths.shape[0]:"
            " {} vs. {}".format(shape[0], lengths.shape[0]))
    T = shape[1]
    for b in range(shape[0]):
        if lengths[b] < T:
            l = lengths[b]
            paddings[b, l:] = 1
    return paddings

def get_transformer_padding_byte_masks(B, T, lengths):
    masks = get_paddings_by_shape([B, T], lengths).byte()
    return masks

def get_transformer_casual_masks(T):
    masks = -torch.triu(
            torch.ones(T, T), diagonal=1)*9e20
    return masks

# ==========================================
# visualization
# ==========================================
if TENSORBOARD_LOGGING == 1:
    import logging
    mpl_logger = logging.getLogger("matplotlib")
    mpl_logger.setLevel(logging.WARNING)

    import matplotlib as mpl
    mpl.use('Agg')
    import matplotlib.pyplot as plt
    from tensorboardX import SummaryWriter

    class Visualizer(object):
        def __init__(self):
            self.writer = None
            self.fig_step = 0

        def set_writer(self, log_dir):
            if self.writer is not None:
                raise ValueError("Dont set writer twice.")
            self.writer = SummaryWriter(log_dir)

        def add_scalar(self, tag, value, step):
            self.writer.add_scalar(tag=tag,
                scalar_value=value, global_step=step)

        def add_graph(self, model):
            self.writer.add_graph(model)

        def add_image(self, tag, img, data_formats):
            self.writer.add_image(tag,
                img, 0, dataformats=data_formats)

        def add_img_figure(self, tag, img, step=None):
            fig, axes = plt.subplots(1,1)
            axes.imshow(img)
            self.writer.add_figure(tag, fig, global_step=step)

        def close(self):
            self.writer.close()

    visualizer = Visualizer()


import os,  inspect
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)

import math
import gym
from gym import spaces
from gym.utils import seeding
import numpy as np
import time
import pybullet
from . import bullet_client
from . import racecar
import random
import pybullet_data
from pkg_resources import parse_version

RENDER_HEIGHT = 720
RENDER_WIDTH = 960

class RacecarZEDGymEnv(gym.Env):
  metadata = {
      'render.modes': ['human', 'rgb_array'],
      'video.frames_per_second' : 50
  }

  def __init__(self,
               urdfRoot=pybullet_data.getDataPath(),
               actionRepeat=10,
               isEnableSelfCollision=True,
               isDiscrete=False,
               renders=True):
    print("init")
    self._timeStep = 0.01
    self._urdfRoot = urdfRoot
    self._actionRepeat = actionRepeat
    self._isEnableSelfCollision = isEnableSelfCollision
    self._ballUniqueId = -1
    self._envStepCounter = 0
    self._renders = renders
    self._width = 100
    self._height = 10

    self._isDiscrete = isDiscrete
    if self._renders:
      self._p = bullet_client.BulletClient(
          connection_mode=pybullet.GUI)
    else:
      self._p = bullet_client.BulletClient()

    self.seed()
    self.reset()
    observationDim = len(self.getExtendedObservation())
    #print("observationDim")
    #print(observationDim)

    observation_high = np.array([np.finfo(np.float32).max] * observationDim)
    if (isDiscrete):
      self.action_space = spaces.Discrete(9)
    else:
       action_dim = 2
       self._action_bound = 1
       action_high = np.array([self._action_bound] * action_dim)
       self.action_space = spaces.Box(-action_high, action_high, dtype=np.float32)
    self.observation_space = spaces.Box(low=0, high=255, shape=(self._height, self._width, 4), dtype=np.uint8)

    self.viewer = None

  def reset(self):
    self._p.resetSimulation()
    #p.setPhysicsEngineParameter(numSolverIterations=300)
    self._p.setTimeStep(self._timeStep)
    #self._p.loadURDF(os.path.join(os.path.dirname(__file__),"../data","plane.urdf"))
    stadiumobjects = self._p.loadSDF(os.path.join(self._urdfRoot,"stadium.sdf"))
    #move the stadium objects slightly above 0
    for i in stadiumobjects:
      pos,orn = self._p.getBasePositionAndOrientation(i)
      newpos = [pos[0],pos[1],pos[2]+0.1]
      self._p.resetBasePositionAndOrientation(i,newpos,orn)

    dist = 5 +2.*random.random()
    ang = 2.*3.1415925438*random.random()

    ballx = dist * math.sin(ang)
    bally = dist * math.cos(ang)
    ballz = 1

    self._ballUniqueId = self._p.loadURDF(os.path.join(self._urdfRoot,"sphere2red.urdf"),[ballx,bally,ballz])
    self._p.setGravity(0,0,-10)
    self._racecar = racecar.Racecar(self._p,urdfRootPath=self._urdfRoot, timeStep=self._timeStep)
    self._envStepCounter = 0
    for i in range(100):
      self._p.stepSimulation()
    self._observation = self.getExtendedObservation()
    return np.array(self._observation)

  def __del__(self):
    self._p = 0

  def seed(self, seed=None):
    self.np_random, seed = seeding.np_random(seed)
    return [seed]

  def getExtendedObservation(self):
     carpos,carorn = self._p.getBasePositionAndOrientation(self._racecar.racecarUniqueId)
     carmat = self._p.getMatrixFromQuaternion(carorn)
     ballpos,ballorn = self._p.getBasePositionAndOrientation(self._ballUniqueId)
     invCarPos,invCarOrn = self._p.invertTransform(carpos,carorn)
     ballPosInCar,ballOrnInCar = self._p.multiplyTransforms(invCarPos,invCarOrn,ballpos,ballorn)
     dist0 = 0.3
     dist1 = 1.
     eyePos = [carpos[0]+dist0*carmat[0],carpos[1]+dist0*carmat[3],carpos[2]+dist0*carmat[6]+0.3]
     targetPos = [carpos[0]+dist1*carmat[0],carpos[1]+dist1*carmat[3],carpos[2]+dist1*carmat[6]+0.3]
     up = [carmat[2],carmat[5],carmat[8]]
     viewMat = self._p.computeViewMatrix(eyePos,targetPos,up)
     #viewMat = self._p.computeViewMatrixFromYawPitchRoll(carpos,1,0,0,0,2)
     #print("projectionMatrix:")
     #print(self._p.getDebugVisualizerCamera()[3])
     projMatrix = [0.7499999403953552, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, -1.0000200271606445, -1.0, 0.0, 0.0, -0.02000020071864128, 0.0]
     img_arr = self._p.getCameraImage(width=self._width,height=self._height,viewMatrix=viewMat,projectionMatrix=projMatrix)
     rgb=img_arr[2]
     np_img_arr = np.reshape(rgb, (self._height, self._width, 4))
     self._observation = np_img_arr
     return self._observation

  def step(self, action):
    if (self._renders):
      basePos,orn = self._p.getBasePositionAndOrientation(self._racecar.racecarUniqueId)
      #self._p.resetDebugVisualizerCamera(1, 30, -40, basePos)

    if (self._isDiscrete):
            fwd = [-1,-1,-1,0,0,0,1,1,1]
            steerings = [-0.6,0,0.6,-0.6,0,0.6,-0.6,0,0.6]
            forward = fwd[action]
            steer = steerings[action]
            realaction = [forward,steer]
    else:
      realaction = action

    self._racecar.applyAction(realaction)
    for i in range(self._actionRepeat):
      self._p.stepSimulation()
      if self._renders:
        time.sleep(self._timeStep)
      self._observation = self.getExtendedObservation()

      if self._termination():
        break
      self._envStepCounter += 1
    reward = self._reward()
    done = self._termination()
    #print("len=%r" % len(self._observation))

    return np.array(self._observation), reward, done, {}

  def render(self, mode='human', close=False):
    if mode != "rgb_array":
      return np.array([])
    base_pos,orn = self._p.getBasePositionAndOrientation(self._racecar.racecarUniqueId)
    view_matrix = self._p.computeViewMatrixFromYawPitchRoll(
        cameraTargetPosition=base_pos,
        distance=self._cam_dist,
        yaw=self._cam_yaw,
        pitch=self._cam_pitch,
        roll=0,
        upAxisIndex=2)
    proj_matrix = self._p.computeProjectionMatrixFOV(
        fov=60, aspect=float(RENDER_WIDTH)/RENDER_HEIGHT,
        nearVal=0.1, farVal=100.0)
    (_, _, px, _, _) = self._p.getCameraImage(
        width=RENDER_WIDTH, height=RENDER_HEIGHT, viewMatrix=view_matrix,
        projectionMatrix=proj_matrix, renderer=pybullet.ER_BULLET_HARDWARE_OPENGL)
    rgb_array = np.array(px)
    rgb_array = rgb_array[:, :, :3]
    return rgb_array

  def _termination(self):
    return self._envStepCounter>1000

  def _reward(self):
    closestPoints = self._p.getClosestPoints(self._racecar.racecarUniqueId,self._ballUniqueId,10000)

    numPt = len(closestPoints)
    reward=-1000
    #print(numPt)
    if (numPt>0):
      #print("reward:")
      reward = -closestPoints[0][8]
      #print(reward)
    return reward

  if parse_version(gym.__version__) < parse_version('0.9.6'):
    _render = render
    _reset = reset
    _seed = seed
    _step = step


"""
A TESS TPF contains the following data:

ColDefs(
    name = 'TIME'; format = 'D'; unit = 'BJD - 2457000, days'; disp = 'D14.7'
    name = 'TIMECORR'; format = 'E'; unit = 'd'; disp = 'E14.7'
    name = 'CADENCENO'; format = 'J'; disp = 'I10'
    name = 'RAW_CNTS'; format = '143J'; unit = 'count'; null = -1; disp = 'I8'; dim = '(11,13)'
    name = 'FLUX'; format = '143E'; unit = 'e-/s'; disp = 'E14.7'; dim = '(11,13)'
    name = 'FLUX_ERR'; format = '143E'; unit = 'e-/s'; disp = 'E14.7'; dim = '(11,13)'
    name = 'FLUX_BKG'; format = '143E'; unit = 'e-/s'; disp = 'E14.7'; dim = '(11,13)'
    name = 'FLUX_BKG_ERR'; format = '143E'; unit = 'e-/s'; disp = 'E14.7'; dim = '(11,13)'
    name = 'QUALITY'; format = 'J'; disp = 'B16.16'
    name = 'POS_CORR1'; format = 'E'; unit = 'pixel'; disp = 'E14.7'
    name = 'POS_CORR2'; format = 'E'; unit = 'pixel'; disp = 'E14.7'
)
"""
from datetime import datetime
import numpy as np
from numpy import array, ndarray

from astropy.wcs import WCS
from astropy.io import fits

class TargetPixelFile:
    def __init__(
        self,
        time: array,
        cadenceno: array,
        flux: ndarray,
        flux_err: ndarray,
        quality: array,
        wcs: WCS = None,
        meta: dict = None,
    ):
        self.time = time
        self.cadenceno = cadenceno
        self.flux = flux
        self.flux_err = flux_err
        self.quality = quality
        self.wcs = wcs
        self.meta = meta

    @property
    def n_cadences(self):
        return self.time.shape[0]

    @property
    def n_columns(self):
        return self.flux.shape[1]

    @property
    def n_rows(self):
        return self.flux.shape[2]

    @property
    def timecorr(self):
        if not hasattr(self, "_timecorr"):
            self._timecorr = np.zeros(self.n_cadences, dtype="float32")
        return self._timecorr

    @property
    def raw_cnts(self):
        if not hasattr(self, "_raw_cnts"):
            self._raw_cnts = np.empty(
                (self.n_cadences, self.n_rows, self.n_columns), dtype="int"
            )
            self._raw_cnts[:, :, :] = -1
        return self._raw_cnts

    @property
    def pos_corr1(self):
        if not hasattr(self, "_pos_corr1"):
            self._pos_corr1 = np.zeros(self.n_cadences, dtype="float32")
        return self._pos_corr1

    @property
    def pos_corr2(self):
        if not hasattr(self, "_pos_corr2"):
            self._pos_corr2 = np.zeros(self.n_cadences, dtype="float32")
        return self._pos_corr2

    @staticmethod
    def from_cutouts(images: list) -> "TargetPixelFile":
        shape = (len(images), images[0].flux.shape[0], images[0].flux.shape[1])
        flux = np.empty(shape)
        flux_err = np.empty(shape)
        for idx, img in enumerate(images):
            flux[idx] = img.flux
            flux_err[idx] = np.nan

        time = np.array([img.time for img in images])
        cadenceno = np.array([img.cadenceno for img in images])
        quality = np.array([img.quality for img in images])
        return TargetPixelFile(
            time=time,
            cadenceno=cadenceno,
            flux=flux,
            flux_err=flux_err,
            quality=quality,
        )

    @staticmethod
    def read(path) -> "TargetPixelFile":
        f = fits.open(path)
        tpf = TargetPixelFile(
            time=f[1].data["TIME"],
            cadenceno=f[1].data["CADENCENO"],
            flux=f[1].data["FLUX"],
            flux_err=f[1].data["FLUX_ERR"],
            quality=f[1].data["QUALITY"],
        )
        return tpf

    def write(self, *args, **kwargs):
        hdulist = self._create_hdulist()
        return hdulist.writeto(*args, **kwargs)

    def _create_hdulist(self):
        """Returns an astropy.io.fits.HDUList object."""
        return fits.HDUList(
            [self._create_primary_hdu(), self._create_table_extension()]
        )

    def _create_primary_hdu(self):
        """Returns the primary extension (#0)."""
        hdu = fits.PrimaryHDU()
        hdu.header["ORIGIN"] = "Unofficial tess-bite product"
        hdu.header["DATE"] = datetime.now().strftime("%Y-%m-%d")
        hdu.header["CREATOR"] = "tess-bite"
        return hdu

    def _create_table_extension(self):
        """Create the 'TARGETTABLES' extension (i.e. extension #1)."""
        # Turn the data arrays into fits columns and initialize the HDU
        coldim = "({},{})".format(self.n_columns, self.n_rows)
        eformat = "{}E".format(self.n_rows * self.n_columns)
        jformat = "{}J".format(self.n_rows * self.n_columns)
        cols = []
        cols.append(
            fits.Column(name="TIME", format="D", unit="BJD - 2454833", array=self.time)
        )
        cols.append(
            fits.Column(name="TIMECORR", format="E", unit="D", array=self.timecorr)
        )
        cols.append(fits.Column(name="CADENCENO", format="J", array=self.cadenceno))
        cols.append(
            fits.Column(
                name="RAW_CNTS",
                format=jformat,
                unit="count",
                dim=coldim,
                array=self.raw_cnts,
            )
        )
        cols.append(
            fits.Column(
                name="FLUX", format=eformat, unit="e-/s", dim=coldim, array=self.flux
            )
        )
        cols.append(
            fits.Column(
                name="FLUX_ERR",
                format=eformat,
                unit="e-/s",
                dim=coldim,
                array=self.flux_err,
            )
        )
        cols.append(fits.Column(name="QUALITY", format="J", array=self.quality))
        cols.append(
            fits.Column(
                name="POS_CORR1", format="E", unit="pixels", array=self.pos_corr1
            )
        )
        cols.append(
            fits.Column(
                name="POS_CORR2", format="E", unit="pixels", array=self.pos_corr2
            )
        )
        coldefs = fits.ColDefs(cols)
        hdu = fits.BinTableHDU.from_columns(coldefs)

        return hdu


#!/usr/bin/env python
'''Bridge IPython widgets with PortAudio controls and DSP callbacks'''

import logging
import numpy as np
import pyaudio
import IPython.html.widgets

LOG = logging.getLogger(__name__)

class AudioConnector(object):
    '''A class to connect a PyAudio stream to a callback function.'''

    def __init__(self, callback, sr=22050, window=1024, dtype=np.float32,
                 channels=1, width=2, output=False, **kwargs):
        '''
        :parameters:
            - callback : function
              Must support the positional arguments:
                - y : np.ndarray
                - sr : int

              If it returns a value, it must be a buffer of the
              same shape as y.

            - sr : int > 0
              The sampling rate from PyAudio

            - window : int > 0
              The number of samples to buffer

            - dtype : type
              The data type for the np.ndarray buffer

            - channels : int
              The number of channels to read

            - width : int
              The number of bytes per sample

            - output : bool
              Enable audio pass-through

            - **kwargs :
              Additional keyword arguments to pass through to `callback`
        '''
        self.channels = channels
        self.width = width

        self.callback = callback

        self.sr = sr
        self.window = window
        self.dtype = dtype

        self.output = output

        self.kwargs = kwargs

        # Helpers for interpreting the byte array as floats
        self.fmt = '<i{:d}'.format(self.width)
        self.scale = 1./float(1 << ((8 * self.width) - 1))

        # Build the portaudio interface
        self.port_ = pyaudio.PyAudio()

        my_fmt = self.port_.get_format_from_width(self.width)
        # Instantiate the stream
        self.stream_ = self.port_.open(start=False,
                                       format=my_fmt,
                                       frames_per_buffer=self.window,
                                       channels=self.channels,
                                       rate=self.sr,
                                       input=True,
                                       output=self.output,
                                       stream_callback=self.__audio_callback)

    def start(self):
        '''Start the audio stream.'''

        LOG.info('Stream started')
        self.stream_.start_stream()

    def stop(self):
        '''Stop (pause) the audio stream.'''

        LOG.info('Stream stopped')
        self.stream_.stop_stream()

    def set_state(self, active):
        '''Set the state:

        :parameters:
          - active : bool
            If true, start streaming
            If false, stop streaming
        '''
        if active:
            self.start()
        else:
            self.stop()

    def __audio_callback(self, in_data, frame_count, t_info, status):
        '''Callback function for PyAudio.

        See PyAudio.Stream documentation for details.
        '''

        # Interpret the buffer as a numpy array of floats
        y = self.scale * np.frombuffer(in_data, self.fmt).astype(self.dtype)

        # Pass data to the callback
        try:
            y_out = self.callback(y, self.sr, **self.kwargs)
        except Exception as e_callback:
            LOG.error('Exception in callback: {0}'.format(e_callback))

        if y_out is None:
            # callback was silent, passthru
            out_data = in_data
        else:
            # reinterpret the result
            out_data = (y_out / self.scale).astype(self.fmt)

        # Let pyaudio continue
        if self.output:
            return (out_data, pyaudio.paContinue)
        else:
            return (None, pyaudio.paContinue)

    def test_callback(self, y=None):
        '''Test the stream callback.
        This can be useful for debugging your callback function without
        depending on PortAudio.

        :parameters:
          - y : ndarray or None
            Buffer to pass back as data to the callback.
            If none, white noise is used.
        '''

        if y is None:
            y = np.random.randn(self.window)

        self.callback(y, self.sr, **self.kwargs)

    def __del__(self):
        '''Class destructor.

        This stops the stream and terminates the PortAudio connection.
        '''

        # Close the stream
        self.stop()

        # Close the portaudio connector
        self.port_.terminate()

def playback_widget(audio_connector):
    '''Construct a toggle widget to control an AudioConnector object.

    :usage:
        >>> audio_connector = ipyaudio.AudioConnector(my_callback)
        >>> interact_widget = ipyaudio.playback_widget(audio_connector)
        >>> IPython.display.display(interact_widget)

    :parameters:
        - audio_connector : ipyaudio.AudioConnector
          An AudioConnector object

    :returns:
        - interact : IPython container widget
          An IPython interact widget
    '''

    play_widget = IPython.html.widgets.ToggleButtonWidget(value=False)

    def wrapper_f(**kwargs):
        '''A wrapper function to handle instancemethods in interact'''
        audio_connector.set_state(**kwargs)

    return IPython.html.widgets.interactive(wrapper_f, active=play_widget)


import argparse
from email.policy import default
import gym
from matplotlib.pyplot import flag
import numpy as np
import os
import torch
import shutil
import time
import d4rl
import json

import continuous_bcq.BCQ
import continuous_bcq.utils as utils
import os.path as osp
from absl import logging, app, flags

# Trains BCQ offline
def train_BCQ(env, state_dim, action_dim, max_action, device, output_dir, args):
    path, split = osp.split(args.exp_dir)
    path, game = osp.split(path)
    # For saving files

    # Initialize policy
    policy = continuous_bcq.BCQ.BCQ(state_dim, action_dim, max_action, device, args.discount, args.tau, args.lmbda, args.phi)

    # Load buffer
    replay_buffer = utils.ReplayBuffer(state_dim, action_dim, device)
    data_folder_name = "_".join(["training_logs", "BCQ", args.method])
    data_path = osp.join(args.exp_dir, data_folder_name)
    h5path = osp.join(data_path, game + ".hdf5")
    dataset = env.get_dataset(h5path=h5path)
    N = dataset['rewards'].shape[0]
    for i in range(N-1):
        obs = dataset['observations'][i]
        new_obs = dataset['observations'][i+1]
        action = dataset['actions'][i]
        reward = dataset['rewards'][i]
        done_bool = bool(dataset['terminals'][i])
        replay_buffer.add(obs, action, new_obs, reward, done_bool)
    logging.info(f'Loaded buffer from {h5path}')

    evaluations = []
    training_iters = 0
    while training_iters < args.max_timesteps:
        logging.info(f"Training step: {training_iters}")
        tstart = time.time()
        pol_vals = policy.train(replay_buffer, iterations=int(args.eval_freq), batch_size=args.batch_size)
        tend = time.time()
        evaluations.append(eval_policy(policy, game, args.seed))
        np.save(os.path.join(output_dir, f"eval_returns.npy"), evaluations)
        training_iters += args.eval_freq
        policy.save(output_dir)

# Runs policy for X episodes and returns average reward
# A fixed seed is used for the eval environment
def eval_policy(policy, env_name, seed, eval_episodes=10):
    eval_env = gym.make(env_name)
    eval_env.seed(seed + 100)

    avg_reward = 0.
    for _ in range(eval_episodes):
        state, done = eval_env.reset(), False
        while not done:
            action = policy.select_action(np.array(state))
            state, reward, done, _ = eval_env.step(action)
            avg_reward += reward
    avg_reward /= eval_episodes
    logging.info(f"Evaluation over {eval_episodes} episodes: {avg_reward:.3f}")
    return avg_reward
flags.DEFINE_integer('seed', 0, "Sets Gym, PyTorch and Numpy seeds")
flags.DEFINE_string('buffer_name', "Robust", "Prepends name to filename")
flags.DEFINE_integer('eval_freq', 5000, "How often (time steps) we evaluate")
flags.DEFINE_integer('max_timesteps', 1000000, "Max time steps to run environment or train for (this defines buffer size)")

flags.DEFINE_integer("start_timesteps", 25000, "Time steps initial random policy is used before training behavioral")
flags.DEFINE_float("rand_action_p", 0.3, "Probability of selecting random action during batch generation")
flags.DEFINE_float("gaussian_std", 0.3, "Std of Gaussian exploration noise (Set to 0.1 if DDPG trains poorly)")
flags.DEFINE_integer("batch_size", 100, "Mini batch size")
flags.DEFINE_float("discount", 0.99, "Discount factor")
flags.DEFINE_float("tau", 0.005, "Target network update rate")
flags.DEFINE_float("lmbda", 0.75, "Weighting for clipped double Q-learning in BCQ")
flags.DEFINE_float("phi", 0.05, "Max perturbation hyper-parameter for BCQ")
flags.DEFINE_string("exp_dir", None, "path to experiments")
flags.DEFINE_string("method", None, "the reward learning method")

args = flags.FLAGS

def main(argv):
    logging.set_verbosity(logging.INFO)
    path, split = osp.split(args.exp_dir)
    path, game = osp.split(path)
    exp_base, exp_id = osp.split(path)
    data_folder_name = "_".join(["training_logs", "BCQ", args.method])
    data_path = osp.join(args.exp_dir, data_folder_name)
    if osp.exists(data_path):
        shutil.rmtree(data_path)
    os.makedirs(data_path, exist_ok=True)
    data_file = "_".join([exp_id, game, split, args.method]) + ".zip"
    shutil.copy2(osp.join(exp_base, "preference_rewards", data_file), data_path)
    shutil.unpack_archive(osp.join(data_path, data_file), extract_dir=data_path, format="zip")
    os.remove(osp.join(data_path, data_file))
    results_dir = os.path.join(args.exp_dir, f'BCQ_{args.method}')
    os.makedirs(results_dir, exist_ok=True)
    with open(os.path.join(results_dir, 'params.json'), 'w') as params_file:
        json.dump({'game': game, 'seed': args.seed}, params_file)
    env = gym.make(game)
    env.seed(args.seed)
    torch.manual_seed(args.seed)
    np.random.seed(args.seed)
    state_dim = env.observation_space.shape[0]
    action_dim = env.action_space.shape[0]
    max_action = float(env.action_space.high[0])
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    train_BCQ(env, state_dim, action_dim, max_action, device, results_dir, args)
    shutil.rmtree(data_path)

    archive_name = osp.join(exp_base, "agents", "_".join([exp_id, game, split, "BCQ", args.method]))
    shutil.make_archive(base_name=archive_name,
        root_dir=results_dir, base_dir=None, format="zip")

if __name__ == '__main__':
    flags.mark_flag_as_required('method')
    flags.mark_flag_as_required('exp_dir')
    app.run(main)


import numpy as np
from pytest import fixture

from nlisim.coordinates import Point
from nlisim.grid import RectangularGrid
from nlisim.oldmodules.epithelium import EpitheliumCellData, EpitheliumCellList
from nlisim.oldmodules.fungus import FungusCellData, FungusCellList

@fixture
def iron():
    # a 10 x 10 x 10 grid with 10 iron
    i = np.empty((10, 10, 10))
    i.fill(10)
    yield i

@fixture
def n_cyto():
    # a 10 x 10 x 10 grid with 10 iron
    i = np.empty((10, 10, 10))
    i.fill(0)
    yield i

@fixture
def m_cyto():
    # a 10 x 10 x 10 grid with 10 iron
    i = np.empty((10, 10, 10))
    i.fill(0)
    yield i

@fixture
def tissue():
    # a 10 x 10 x 10 grid of blood
    t = np.empty((10, 10, 10))
    t.fill(3)
    yield t

@fixture
def epithelium_list(grid: RectangularGrid):
    epithelium = EpitheliumCellList(grid=grid)
    yield epithelium

@fixture
def populated_epithelium(epithelium_list: EpitheliumCellList, grid: RectangularGrid):
    epithelium_list.append(
        EpitheliumCellData.create_cell(
            point=Point(x=grid.x[3], y=grid.y[3], z=grid.z[3]),
        )
    )

    yield epithelium_list

@fixture
def fungus_list(grid: RectangularGrid):
    fungus = FungusCellList(grid=grid)
    yield fungus

@fixture
def populated_fungus(fungus_list: FungusCellList, grid: RectangularGrid):
    points = []
    for i in range(int(grid.x[1]), int(grid.x[6]), 10):
        points.append(Point(x=i, y=i, z=i))

    for point in points:
        fungus_list.append(
            FungusCellData.create_cell(
                point=point,
                status=FungusCellData.Status.GROWABLE,
                form=FungusCellData.Form.HYPHAE,
                iron=0,
                mobile=False,
            )
        )
    yield fungus_list

# tests

# internalize_conidia
def test_internalize_conidia_none(
    populated_epithelium: EpitheliumCellList,
    grid: RectangularGrid,
    fungus_list: FungusCellList,
):
    cell = populated_epithelium[0]
    vox = grid.get_voxel(cell['point'])
    assert len(fungus_list.get_cells_in_voxel(vox)) == 0

    populated_epithelium.internalize_conidia(0, 10, 1, grid, fungus_list)

    assert populated_epithelium.len_phagosome(0) == 0
    for v in cell['phagosome']:
        assert v == -1

def test_internalize_conidia_1(
    epithelium_list: EpitheliumCellList, grid: RectangularGrid, fungus_list: FungusCellList
):
    point = Point(x=35, y=35, z=35)
    epithelium_list.append(EpitheliumCellData.create_cell(point=point))
    fungus_list.append(
        FungusCellData.create_cell(point=point, status=FungusCellData.Status.RESTING)
    )

    vox = grid.get_voxel(epithelium_list[0]['point'])

    epithelium_list.internalize_conidia(0, 10, 1, grid, fungus_list)

    assert grid.get_voxel(fungus_list[0]['point']) == vox
    assert epithelium_list.len_phagosome(0) == 1
    assert 0 in epithelium_list[0]['phagosome']

def test_internalize_conidia_2(
    epithelium_list: EpitheliumCellList, grid: RectangularGrid, fungus_list: FungusCellList
):
    point = Point(x=35, y=35, z=35)
    epithelium_list.append(EpitheliumCellData.create_cell(point=point))
    fungus_list.append(
        FungusCellData.create_cell(point=point, status=FungusCellData.Status.RESTING)
    )
    fungus_list.append(
        FungusCellData.create_cell(point=point, status=FungusCellData.Status.RESTING)
    )

    vox = grid.get_voxel(epithelium_list[0]['point'])

    epithelium_list.internalize_conidia(0, 10, 1, grid, fungus_list)

    assert grid.get_voxel(fungus_list[0]['point']) == vox
    assert epithelium_list.len_phagosome(0) == 2
    assert 0 in epithelium_list[0]['phagosome']
    assert 1 in epithelium_list[0]['phagosome']

def test_internalize_conidia_2b(
    epithelium_list: EpitheliumCellList, grid: RectangularGrid, fungus_list: FungusCellList
):
    point = Point(x=35, y=35, z=35)
    epithelium_list.append(EpitheliumCellData.create_cell(point=point))
    fungus_list.append(
        FungusCellData.create_cell(point=point, status=FungusCellData.Status.RESTING)
    )
    fungus_list.append(
        FungusCellData.create_cell(point=point, status=FungusCellData.Status.RESTING)
    )

    vox = grid.get_voxel(epithelium_list[0]['point'])
    fungus_list[0]['internalized'] = True  # say by macrophage

    epithelium_list.internalize_conidia(0, 10, 1, grid, fungus_list)

    assert grid.get_voxel(fungus_list[0]['point']) == vox
    assert epithelium_list.len_phagosome(0) == 1
    assert 0 not in epithelium_list[0]['phagosome']
    assert 1 in epithelium_list[0]['phagosome']

def test_internalize_conidia_max(
    epithelium_list: EpitheliumCellList, grid: RectangularGrid, fungus_list: FungusCellList
):
    point = Point(x=35, y=35, z=35)
    epithelium_list.append(EpitheliumCellData.create_cell(point=point))
    fungus_list.append(
        FungusCellData.create_cell(point=point, status=FungusCellData.Status.RESTING)
    )

    max_spores = 10
    epithelium_list[0]['phagosome'][:max_spores] = 99  # artificially fill

    epithelium_list.internalize_conidia(0, max_spores, 1, grid, fungus_list)

    assert epithelium_list.len_phagosome(0) == max_spores
    assert 0 not in epithelium_list[0]['phagosome']
    assert not fungus_list[0]['internalized']

def test_dead_conidia_1(
    epithelium_list: EpitheliumCellList, grid: RectangularGrid, fungus_list: FungusCellList
):
    point = Point(x=35, y=35, z=35)
    epithelium_list.append(EpitheliumCellData.create_cell(point=point))
    fungus_list.append(
        FungusCellData.create_cell(point=point, status=FungusCellData.Status.RESTING)
    )

    epithelium_list.internalize_conidia(0, 10, 1, grid, fungus_list)
    fungus_list[0]['dead'] = True  # simulate killing

    epithelium_list.remove_dead_fungus(fungus_list)

    assert epithelium_list.len_phagosome(0) == 0
    assert 0 not in epithelium_list[0]['phagosome']

def test_produce_cytokines_0(
    epithelium_list: EpitheliumCellList,
    grid: RectangularGrid,
    fungus_list: FungusCellList,
    m_cyto,
    n_cyto,
):
    s_det = 0
    h_det = 0
    cyto_rate = 10

    assert m_cyto[3, 3, 3] == 0
    assert n_cyto[3, 3, 3] == 0

    point = Point(x=35, y=35, z=35)
    epithelium_list.append(EpitheliumCellData.create_cell(point=point))

    fungus_list.append(
        FungusCellData.create_cell(
            point=point,
            status=FungusCellData.Status.RESTING,
            form=FungusCellData.Form.CONIDIA,
            iron=0,
            mobile=False,
        )
    )

    # fungus is not swollen or germinated
    epithelium_list.cytokine_update(s_det, h_det, cyto_rate, m_cyto, n_cyto, fungus_list, grid)

    assert m_cyto[3, 3, 3] == 0
    assert n_cyto[3, 3, 3] == 0

    fungus_list[0]['status'] = FungusCellData.Status.SWOLLEN

    epithelium_list.cytokine_update(s_det, h_det, cyto_rate, m_cyto, n_cyto, fungus_list, grid)

    assert m_cyto[3, 3, 3] == 10
    assert n_cyto[3, 3, 3] == 10

def test_produce_cytokines_0b(
    epithelium_list: EpitheliumCellList,
    grid: RectangularGrid,
    fungus_list: FungusCellList,
    m_cyto,
    n_cyto,
):
    s_det = 0
    h_det = 0
    cyto_rate = 10

    assert m_cyto[3, 3, 3] == 0
    assert n_cyto[3, 3, 3] == 0

    point = Point(x=35, y=35, z=35)
    epithelium_list.append(EpitheliumCellData.create_cell(point=point))

    fungus_list.append(
        FungusCellData.create_cell(
            point=point,
            status=FungusCellData.Status.SWOLLEN,
            form=FungusCellData.Form.CONIDIA,
            iron=0,
            mobile=False,
        )
    )

    epithelium_list.cytokine_update(s_det, h_det, cyto_rate, m_cyto, n_cyto, fungus_list, grid)

    assert m_cyto[3, 3, 3] == 10
    assert n_cyto[3, 3, 3] == 10

    fungus_list.append(
        FungusCellData.create_cell(
            point=point,
            status=FungusCellData.Status.GROWABLE,
            form=FungusCellData.Form.HYPHAE,
            iron=0,
            mobile=False,
        )
    )

    epithelium_list.cytokine_update(s_det, h_det, cyto_rate, m_cyto, n_cyto, fungus_list, grid)

    assert m_cyto[3, 3, 3] == 20
    assert n_cyto[3, 3, 3] == 30

def test_produce_cytokines_2(
    epithelium_list: EpitheliumCellList,
    grid: RectangularGrid,
    fungus_list: FungusCellList,
    m_cyto,
    n_cyto,
):
    s_det = 1
    h_det = 2
    cyto_rate = 10

    assert m_cyto[3, 3, 3] == 0
    assert n_cyto[3, 3, 3] == 0

    point = Point(x=35, y=35, z=35)
    epithelium_list.append(EpitheliumCellData.create_cell(point=point))

    spoint = Point(x=15, y=35, z=35)
    fungus_list.append(
        FungusCellData.create_cell(
            point=spoint,
            status=FungusCellData.Status.SWOLLEN,
            form=FungusCellData.Form.CONIDIA,
            iron=0,
            mobile=False,
        )
    )

    epithelium_list.cytokine_update(s_det, h_det, cyto_rate, m_cyto, n_cyto, fungus_list, grid)

    assert m_cyto[3, 3, 3] == 0
    assert n_cyto[3, 3, 3] == 0

    fungus_list.append(
        FungusCellData.create_cell(
            point=spoint,
            status=FungusCellData.Status.GROWABLE,
            form=FungusCellData.Form.HYPHAE,
            iron=0,
            mobile=False,
        )
    )

    epithelium_list.cytokine_update(s_det, h_det, cyto_rate, m_cyto, n_cyto, fungus_list, grid)

    assert m_cyto[3, 3, 3] == 0
    assert n_cyto[3, 3, 3] == 10

def test_damage_conidia(
    epithelium_list: EpitheliumCellList, grid: RectangularGrid, fungus_list: FungusCellList
):
    kill = 2
    t = 1
    health = 100

    point = Point(x=35, y=35, z=35)
    epithelium_list.append(EpitheliumCellData.create_cell(point=point))
    fungus_list.append(
        FungusCellData.create_cell(point=point, status=FungusCellData.Status.RESTING)
    )

    epithelium_list.internalize_conidia(0, 10, 1, grid, fungus_list)

    epithelium_list.damage(kill, t, health, fungus_list)

    assert fungus_list.cell_data['health'][0] == 50

    epithelium_list.damage(kill, t, health, fungus_list)

    assert fungus_list.cell_data['health'][0] == 0

def test_kill_epithelium(
    epithelium_list: EpitheliumCellList, grid: RectangularGrid, fungus_list: FungusCellList
):
    # should release all conidia

    point = Point(x=35, y=35, z=35)
    epithelium_list.append(EpitheliumCellData.create_cell(point=point))
    fungus_list.append(
        FungusCellData.create_cell(point=point, status=FungusCellData.Status.RESTING)
    )

    fungus_list.cell_data['internalized'][0] = True
    epithelium_list[0]['phagosome'][0] = 0  # internalized

    epithelium_list.die_by_germination(fungus_list)
    assert fungus_list.cell_data['internalized'][0]
    assert epithelium_list.len_phagosome(0) == 1

    fungus_list.cell_data['status'][0] = FungusCellData.Status.GERMINATED

    epithelium_list.die_by_germination(fungus_list)

    assert not fungus_list.cell_data['internalized'][0]
    assert epithelium_list.len_phagosome(0) == 0

def test_kill_epithelium_n(
    epithelium_list: EpitheliumCellList, grid: RectangularGrid, fungus_list: FungusCellList
):
    # should release all conidia

    point = Point(x=35, y=35, z=35)
    epithelium_list.append(EpitheliumCellData.create_cell(point=point))
    for _ in range(0, 10):
        fungus_list.append(
            FungusCellData.create_cell(point=point, status=FungusCellData.Status.RESTING)
        )

    epithelium_list.internalize_conidia(0, 10, 1, grid, fungus_list)

    epithelium_list.die_by_germination(fungus_list)
    for i in range(0, 10):
        assert fungus_list.cell_data['internalized'][i]
    assert epithelium_list.len_phagosome(0) == 10

    fungus_list.cell_data['status'][6] = FungusCellData.Status.GERMINATED

    epithelium_list.die_by_germination(fungus_list)

    for i in range(0, 10):
        assert not fungus_list.cell_data['internalized'][i]
    assert epithelium_list.len_phagosome(0) == 0


# pylint: disable=no-member
"""
Simulated carracing environment.
"""
import argparse
from os.path import join, exists
import torch
from torch.distributions.categorical import Categorical
from torchvision import transforms
import gym
from gym import spaces
import matplotlib.pyplot as plt
import numpy as np
from pyglet.window import key
import json
import copy
from torchvision.utils import save_image
import time
import pickle
#from .. import control
from control import Agent
from utils import reshape_to_img
from envs import get_env_params
from models import RSSModel

class SimulatedCarracing(gym.Env): # pylint: disable=too-many-instance-attributes
    """
    Simulated Car Racing.

    Gym environment using learnt RSSM to simulate the
    CarRacing-v0 environment.
    """
    def __init__(self, directory, num_action_repeats=5):
        #self.num_action_repeats = num_action_repeats

        rssm_file = join(directory, 'joint', 'rssm_best.tar')
        assert exists(rssm_file), "No RSSM model in the directory..."

        env_params = get_env_params(args.gamename)
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.IMAGE_RESIZE_DIM = env_params['IMAGE_RESIZE_DIM']

        decoder_reward_condition = False
        decoder_make_sigmas = False

        # init models
        self.rssm = RSSModel(
            env_params['ACTION_SIZE'],
            env_params['LATENT_RECURRENT_SIZE'],
            env_params['LATENT_SIZE'],
            env_params['EMBEDDING_SIZE'],
            env_params['NODE_SIZE'],
            decoder_reward_condition,
            decoder_make_sigmas,
            device=device,
        )
        state = torch.load(rssm_file, map_location={'cuda:0': str(device)})
        self.rssm.load_state_dict(state['state_dict'])

        # spaces
        self.action_space = spaces.Box(np.array([-1, 0, 0]), np.array([1, 1, 1]))
        self.observation_space = spaces.Box(low=0, high=255, shape=(64, 64, 3),
                                            dtype=np.uint8)

        self.hidden, self.state = self.rssm.init_hidden_state(1)
        #self.hidden, self.state =self.hidden.unsqueeze(0), self.state.unsqueeze(0)

        # obs
        self._obs = None
        self._visual_obs = None

        # rendering
        self.monitor = None
        self.figure = None

        self.reset()

    def reset(self):
        """ Resetting """
        import matplotlib.pyplot as plt

        '''l_and_h_starter = pickle.load(open('latent_and_hidden_starters.pkl' ,'rb'))
        self._lstate = l_and_h_starter[0] #torch.randn(1, LATENT_SIZE)
        self._hstate = l_and_h_starter[1] #2 * [torch.zeros(1, LATENT_RECURRENT_SIZE)]
        '''

        # also reset monitor
        if not self.monitor:
            self.figure = plt.figure()
            self.monitor = plt.imshow(
                np.zeros((self.IMAGE_RESIZE_DIM, self.IMAGE_RESIZE_DIM, 3),
                         dtype=np.uint8))

    def step(self, action, r):
        """ One step forward in the dream like environment!!! """
        with torch.no_grad():

            action = torch.Tensor(action).unsqueeze(0).unsqueeze(0)
            single_dream_step_dict = self.rssm.perform_rollout(action, hidden=self.hidden, state=self.state )

            #predict next reward
            pred_rewards = self.rssm.decode_sequence_reward(single_dream_step_dict['hiddens'], single_dream_step_dict['prior_states'])
            mus, logsigmas = self.rssm.decode_sequence_obs(single_dream_step_dict['hiddens'], single_dream_step_dict['prior_states'], pred_rewards )
            self._obs = reshape_to_img(mus, self.IMAGE_RESIZE_DIM)
            np_obs = self._obs.cpu().numpy()

            np_obs = np.clip(np_obs, 0, 1) * 255
            np_obs = np.transpose(np_obs, (0, 2, 3, 1))
            np_obs = np_obs.squeeze()
            np_obs = np_obs.astype(np.uint8)
            self._visual_obs = np_obs

            self.hidden = single_dream_step_dict['hiddens'].squeeze(0)
            self.state = single_dream_step_dict['prior_states'].squeeze(0)

            return np_obs, pred_rewards.item(), False #d.item() > 0

    def render(self): # pylint: disable=arguments-differ
        """ Rendering """
        import matplotlib.pyplot as plt
        if not self.monitor:
            self.figure = plt.figure()
            self.monitor = plt.imshow(
                np.zeros((self.IMAGE_RESIZE_DIM, self.IMAGE_RESIZE_DIM, 3),
                         dtype=np.uint8))
        self.monitor.set_data(self._visual_obs)
        plt.pause(.01)

        '''#print('updated cross entropies')
        # choose the best next action out of all of them.
        best_actions_ind = torch.argmax(all_particles_cum_rewards)
        worst_actions_ind = torch.argmin(all_particles_cum_rewards)

        best_action = all_particles_sequential_actions[best_actions_ind,0, :]

        best_rewards_seq = sequential_rewards[best_actions_ind]
        worst_rewards_seq = sequential_rewards[worst_actions_ind]

        print('best sequential actions', all_particles_sequential_actions[best_actions_ind] )
        print('worst sequential actions', all_particles_sequential_actions[worst_actions_ind] )

        print('best rewards seq', best_rewards_seq)
        print('worst rewards seq', worst_rewards_seq)

        sequence_best_pred_obs = all_particles_all_predicted_obs[best_actions_ind]
        sequence_worst_pred_obs = all_particles_all_predicted_obs[worst_actions_ind]

        # saving this out.
        sequence_best_pred_obs = sequence_best_pred_obs.view(sequence_best_pred_obs.shape[0], 3, 64, 64)
        sequence_worst_pred_obs = sequence_worst_pred_obs.view(sequence_worst_pred_obs.shape[0], 3, 64, 64)
        save_image(sequence_best_pred_obs,
                        'exp_dir/simulation/samples/best_sample_' + str(time_step) + '.png')
        save_image(sequence_worst_pred_obs,
                        'exp_dir/simulation/samples/worst_sample_' + str(time_step) + '.png')

        print('predicted reward of best particle', torch.max(all_particles_cum_rewards))
        print('predicted reward of worst particle', torch.min(all_particles_cum_rewards))
        #print('best action is:', best_action)
        return best_action.unsqueeze(0)'''

if __name__ == '__main__':
    # argument parsing
    parser = argparse.ArgumentParser()
    parser.add_argument('--logdir', type=str, help='Directory from which MDRNN and VAE are '
                        'retrieved.')
    parser.add_argument('--gamename', type=str, default='carracing', help='The game being played.')
    parser.add_argument('--num_action_repeats', type=int, default=5, help='Number of times action is repeated.')
    parser.add_argument('--real_obs', action='store_true',
                    help="Use the real observations not dream ones")
    parser.add_argument('--test_agent', action='store_true',
                    help="Put the trained controller to the test!")
    parser.add_argument('--render', type=int, default=1,
                    help="Turn off if want to run lots of experiments")
    args = parser.parse_args()
    args.render = bool(args.render)

    if args.gamename == 'carracing':
        game_env = 'CarRacing-v0'
    if args.gamename == 'pendulum':
        game_env = 'Pendulum-v0'

    def on_key_press(event):
        """ Defines key pressed behavior """
        if event.key == 'up':
            action[1] = 1
        if event.key == 'down':
            action[2] = .8
        if event.key == 'left':
            action[0] = -1
        if event.key == 'right':
            action[0] = 1

    def on_key_release(event):
        """ Defines key pressed behavior """
        if event.key == 'up':
            action[1] = 0
        if event.key == 'down':
            action[2] = 0
        if event.key == 'left' and action[0] == -1:
            action[0] = 0
        if event.key == 'right' and action[0] == 1:
            action[0] = 0

    # this is for the real observation based driving!
    def key_press(k, mod):
        global restart
        if k == 0xff0d: restart = True
        if k == key.LEFT:  action[0] = -1.0
        if k == key.RIGHT: action[0] = +1.0
        if k == key.UP:    action[1] = +1.0
        if k == key.DOWN:  action[2] = +0.8   # set 1.0 for wheels to block to zero rotation

    def key_release(k, mod):
        if k == key.LEFT  and action[0] == -1.0: action[0] = 0
        if k == key.RIGHT and action[0] == +1.0: action[0] = 0
        if k == key.UP:    action[1] = 0
        if k == key.DOWN:  action[2] = 0

    decoder_reward_condition = False

    if args.real_obs:
        if args.test_agent:
            # loads in the planner and rssm
            agent = Agent(args.gamename, 'exp_dir/joint/', decoder_reward_condition,
                planner_n_particles=700, cem_iters=7, model_version='best', return_plan_images=True)
            env = agent.env
        else:
            env = gym.make(game_env)
            _ = env.reset()

        if args.render:
            figure = plt.figure()
            monitor = plt.imshow(
                np.zeros((96, 96, 3),
                    dtype=np.uint8))
            env.viewer.window.on_key_press = key_press
            env.viewer.window.on_key_release = key_release

    else:
        env = SimulatedCarracing(args.logdir, num_action_repeats=args.num_action_repeats)
        env.figure.canvas.mpl_connect('key_press_event', on_key_press)
        env.figure.canvas.mpl_connect('key_release_event', on_key_release)

    obs = env.reset()
    action = np.array([0., 0., 0.])
    reward = torch.Tensor([0]).unsqueeze(0)

    IMAGE_RESIZE_DIM = 64

    cum_reward = 0
    t = 0

    transform = transforms.Compose([
                transforms.ToPILImage(),
                transforms.Resize(IMAGE_RESIZE_DIM, IMAGE_RESIZE_DIM),
                transforms.ToTensor()
            ])

    sim_rewards = []

    if args.real_obs:
        # need to generate actions!
        if args.test_agent:
            with torch.no_grad():
                if args.render:
                    agent.rollout(10, render=True, display_monitor=monitor, discount_factor=0.9 )
                else:
                    nsims = 20
                    print('Running ', nsims, 'simulations!')
                    cum_rews, terminals = agent.simulate(num_episodes=nsims)
                    cum_rews = np.array(cum_rews)
                    print('cum rewards:', cum_rews)
                    print('mean reward:', cum_rews.mean(), 'max:', cum_rews.max(), 'min:', cum_rews.min() )
        else:
            # allow person to play.
            while True:
                for _ in range(args.num_action_repeats):
                    print(action)
                    # this is the real environment.
                    obs, reward, done, _ = env.step(action)
                    cum_reward += reward
                    t+=1

                    if done:
                        print(cum_reward)
                        break

                monitor.set_data(obs)
                env.render()

    else:
        # dream like state using environment made by the simulated agent!
        while True:
            #for _ in range(args.num_action_repeats):
            reward = torch.Tensor([reward]).unsqueeze(0)
            print('action is:', action)
            # simulated environment.
            _, reward, done = env.step(action, reward)

            cum_reward += reward
            t+=1

            if done:
                print(cum_reward)
                break

            env.render()
            #time.sleep(0.2)

            print('time step of simulation', t)
            print('=============================')


import tensorflow as tf
import tqdm
from few_shot_network import FewshotsNet
import os
import glob
import numpy as np
from collections import defaultdict

path = os.getcwd()

class ExperimentBuilder:

    def __init__(self, data):
        """
        Initializes an ExperimentBuilder object. The ExperimentBuilder object takes care of setting up our experiment
        and provides helper functions such as run_training_epoch and run_validation_epoch to simplify out training
        and evaluation procedures.
        :param data: A data provider class
        """
        self.data = data

    # def whydidyounotjoinusformoredrinks(self):
    #
    #     print "we went for drinks with Ryutaor"

    def build_experiment(self, batch_size, ways, shots, query_size, image_shape, data_format):

        """

        :param batch_size: The experiment batch size
        :param ways: An integer indicating the number of classes per support set
        :param shots: An integer indicating the number of samples per class
        :param channels: The image channels
        :param fce: Whether to use full context embeddings or not
        :return: a few_shot_learning object, along with the losses, the training ops and the init op
        """
        height, width, channels = image_shape
        self.support_set_x = tf.placeholder(tf.float32, shape = [batch_size, ways, shots, height, width,
                                            channels], name = 'support_set_images')
        self.support_set_y = tf.placeholder(tf.uint8, shape = [batch_size, ways, shots], name = 'support_set_labels')
        self.query_x = tf.placeholder(tf.float32, shape = [batch_size, query_size*ways, height, width, channels], name = 'query_images')
        self.query_y = tf.placeholder(tf.uint8, shape = [batch_size, query_size*ways], name = 'query_labels')
        self.is_training = tf.placeholder(tf.bool, name='training_flag')
        self.dropout_prob = tf.placeholder(tf.float32, name='dropout_prob')
        # self.support_set_x_cand = tf.placeholder(tf.float32, shape=[20,ways,shots,height,width,channels], name="support_set_images_cand")
        # self.support_set_y_cand = tf.placeholder(tf.uint8,shape=[20,ways,shots], name="support_set_labels_cand")
        # self.query_x_cand = tf.placeholder(tf.float32,shape=[20,1,height,width,channels],name="query_images_cand")
        # self.query_y_cand =tf.placeholder(tf.uint8,shape=[20,1], name="query_labels_cand")
        self.current_learning_rate = 0.0001
        self.learning_rate = tf.placeholder(tf.float32, name='learning_rate')
        self.few_shot_miniImagenet = FewshotsNet(batch_size=batch_size, support_set_images=self.support_set_x,
                                                 support_set_labels=self.support_set_y,
                                                 query_images=self.query_x, query_labels=self.query_y, data_format = data_format,
                                                 dropout_prob=self.dropout_prob,
                                                 is_training=self.is_training, learning_rate=self.learning_rate)

        self.losses, self.ada_opts = self.few_shot_miniImagenet.init_train()
        init = tf.global_variables_initializer()
        self.total_train_iter = 0
        self.total_test_iter = 0
        self.total_val_iter = 0
        return self.few_shot_miniImagenet, self.losses, self.ada_opts, init

    def run_training_epoch(self, total_training_episodes, writer, sess):
        """
        Runs one training epoch
        :param total_train_batches: Number of batches to train on
        :param sess: Session object
        :return: mean_training_categorical_crossentropy_loss and mean_training_accuracy
        """
        losses = []
        accuracies = []
        with tqdm.tqdm(total=total_training_episodes) as pbar:

            for i in range(total_training_episodes):  # train epoch
                support_set_x, support_set_y, query_x, query_y = self.data.get_tra_batch()
                _, c_loss_value, acc = sess.run(
                    [self.ada_opts, self.losses['loss'], self.losses['accuracy']],
                    feed_dict={self.dropout_prob: 0.5, self.support_set_x: support_set_x,
                               self.support_set_y: support_set_y, self.query_x: query_x, self.query_y: query_y,
                               self.is_training: True, self.learning_rate: self.current_learning_rate})

                tf.logging.info('train_loss:{}, accuracy:{}'.format(c_loss_value, acc))
                self.total_train_iter += 1
                pbar.update(1)
                losses.append(c_loss_value)
                accuracies.append(acc)
                if self.total_train_iter % 100 == 0:
                    train_summary = tf.Summary()
                    loss_last_100 = np.mean(losses[-100:])

                    accuracy_last_100 = np.mean(accuracies[-100:])

                    train_summary.value.add(tag='loss', simple_value=loss_last_100)
                    train_summary.value.add(tag='accuracy', simple_value = accuracy_last_100)
                    writer.add_summary(train_summary, self.total_train_iter)
                if self.total_train_iter % 2000 == 0:
                    self.current_learning_rate /= 2
                    print("change learning rate", self.current_learning_rate)

        average_accuracy = np.mean(accuracies)
        average_loss = np.mean(losses)
        return average_loss, average_accuracy

    def run_validation_epoch(self, total_val_episodes, writer, sess):
        """
        Runs one validation epoch
        :param total_val_batches: Number of batches to train on
        :param sess: Session object
        :return: mean_validation_categorical_crossentropy_loss and mean_validation_accuracy
        """
        losses = []
        accuracies = []

        with tqdm.tqdm(total=total_val_episodes) as pbar:
            for i in range(total_val_episodes):  # validation epoch
                support_set_x, support_set_y, query_x, query_y = self.data.get_val_test_batch("val")

                c_loss_value, acc = sess.run(
                    [self.losses['loss'], self.losses['accuracy']],
                    feed_dict={self.dropout_prob: 1.0, self.support_set_x: support_set_x,
                               self.support_set_y: support_set_y, self.query_x: query_x, self.query_y: query_y,
                               self.is_training: False})

                tf.logging.info('val_loss:{}, accuracy:{}'.format(c_loss_value, acc))
                pbar.update(1)

                losses.append(c_loss_value)
                accuracies.append(acc)
                self.total_val_iter += 1

            mean_loss = np.mean(losses)
            mean_accuracy = np.mean(accuracies)
            val_summary = tf.Summary()
            val_summary.value.add(tag='loss', simple_value=mean_loss)
            val_summary.value.add(tag='accuracy', simple_value=mean_accuracy)
            writer.add_summary(val_summary, self.total_val_iter)

        return mean_loss, mean_accuracy

    def run_testing_epoch(self, total_test_episodes, sess, writer):
        """
        Runs one testing epoch
        :param total_test_batches: Number of batches to train on
        :param sess: Session object
        :return: mean_testing_categorical_crossentropy_loss and mean_testing_accuracy
        """
        losses = []
        accuracies = []
        with tqdm.tqdm(total=total_test_episodes) as pbar:
            for i in range(total_test_episodes):
                support_set_x, support_set_y, query_x, query_y = self.data.get_val_test_batch("test")
                c_loss_value, acc = sess.run(
                    [self.losses['loss'], self.losses['accuracy']],
                    feed_dict={self.dropout_prob: 1.0, self.support_set_x: support_set_x,
                               self.support_set_y: support_set_y, self.query_x: query_x,
                               self.query_y: query_y,
                               self.is_training: False})

                tf.logging.info('test_loss:{}, accuracy:{}'.format(c_loss_value, acc))
                losses.append(c_loss_value)
                accuracies.append(acc)
                pbar.update(1)
                self.total_test_iter += 1

            test_summary = tf.Summary()
            mean_loss = np.mean(losses)
            mean_accuracy = np.mean(accuracies)
            test_summary.value.add(tag='loss', simple_value=mean_loss)
            test_summary.value.add(tag='accuracy', simple_value=mean_accuracy)
            writer.add_summary(test_summary, self.total_train_iter)

        return mean_loss, mean_accuracy

    def run_find_cand_epoch(self,sess):
        """
        validate the current classifier on candidate training data
        to find the most uncertain several images for each class
        """
        cand_imgs_dir = path + '/miniImagenet/tra_cand_data_mini/'
        # class_ = glob.glob(cand_imgs_dir + 'n*')
        total_episode = 0
        selected_imgs = defaultdict(list)
        for i in range(5):
            imgs_dir = cand_imgs_dir + str(i)
            imgs = glob.glob(imgs_dir + '/*.png')
            uncertain_dic = defaultdict(list)
            top_50_unceratin_imgs = []
            for j, img in enumerate(imgs):

                support_set_x, support_set_y, query_x, query_y = self.data.get_cand_data(i, img)
                uncertainty = sess.run([self.losses['loss']],
                                       feed_dict={self.dropout_prob: 1.0,
                                                  self.support_set_x:support_set_x,
                                                  self.support_set_y:support_set_y,
                                                  self.query_x:query_x,
                                                  self.query_y:query_y,
                                                  self.is_training:False})

                uncertain_dic[img].append(uncertainty)
                total_episode += 1
                print(uncertainty, img)
            uncertain_dic = sorted(uncertain_dic.items(), key= lambda x: x[1], reverse=True)
            for img_ in uncertain_dic[0:20]:
                img_name = img_[0]
                top_50_unceratin_imgs.append(img_name)

            selected_imgs[i].append(top_50_unceratin_imgs)
        return selected_imgs