import numpy
import openravepy
import logging
from .exceptions import UnsupportedTypeSerializationException

TYPE_KEY = '__type__'

serialization_logger = logging.getLogger('prpy.serialization')
deserialization_logger = logging.getLogger('prpy.deserialization')

# Serialization.
def serialize(obj):
    from numpy import ndarray
    from openravepy import Environment, KinBody, Robot, Trajectory
    from prpy.tsr import TSR, TSRChain

    NoneType = type(None)

    if isinstance(obj, (int, float, basestring, NoneType)):
        return obj
    elif isinstance(obj, (list, tuple)):
        return [ serialize(x) for x in obj ]
    elif isinstance(obj, dict):
        obj = { serialize(k): serialize(v) for k, v in obj.iteritems() }
        obj[TYPE_KEY] = dict.__name__
        return obj
    elif isinstance(obj, ndarray):
        return {
            TYPE_KEY: ndarray.__name__,
            'data': serialize(obj.tolist())
        }
    elif isinstance(obj, Environment):
        return {
            TYPE_KEY: Environment.__name__,
            'data': serialize_environment(obj)
        }
    elif isinstance(obj, KinBody):
        return {
            TYPE_KEY: KinBody.__name__,
            'name': obj.GetName()
        }
    elif isinstance(obj, Robot):
        return {
            TYPE_KEY: Robot.__name__,
            'name': obj.GetName()
        }
    elif isinstance(obj, KinBody.Link):
        return {
            TYPE_KEY: KinBody.Link.__name__,
            'name': obj.GetName(),
            'parent_name': obj.GetParent().GetName()
        }
    elif isinstance(obj, KinBody.Joint):
        return {
            TYPE_KEY: KinBody.Joint.__name__,
            'name': obj.GetName(),
            'parent_name': obj.GetParent().GetName()
        }
    elif isinstance(obj, Robot.Manipulator):
        return {
            TYPE_KEY: KinBody.Manipulator.__name__,
            'name': obj.GetName(),
            'parent_name': obj.GetParent().GetName()
        }
    elif isinstance(obj, Trajectory):
        return {
            TYPE_KEY: Trajectory.__name__,
            'data': obj.serialize(0)
        }
    elif isinstance(obj, TSR):
        return {
            TYPE_KEY: TSR.__name__,
            'data': obj.to_dict()
        }
    elif isinstance(obj, TSRChain):
        return {
            TYPE_KEY: TSRChain.__name__,
            'data': obj.to_dict()
        }
    else:
        raise UnsupportedTypeSerializationException(obj)

def serialize_environment(env):
    return {
        'bodies': [ serialize_kinbody(body) for body in env.GetBodies() ],
    }

def serialize_environment_file(env, path, writer=None):
    if writer is None:
        import json
        writer = json.dump

    data = serialize_environment(env)

    if path is not None:
        with open(path, 'wb') as output_file:
            writer(data, output_file)
            serialization_logger.debug('Wrote environment to "%s".', path)

    return data

def serialize_kinbody(body):
    all_joints = []
    all_joints.extend(body.GetJoints())
    all_joints.extend(body.GetPassiveJoints())

    data = {
        'is_robot': body.IsRobot(),
        'name': body.GetName(),
        'uri': body.GetXMLFilename(),
        'links': map(serialize_link, body.GetLinks()),
        'joints': map(serialize_joint, all_joints),
    }
    data['kinbody_state'] = serialize_kinbody_state(body)

    if body.IsRobot():
        data.update(serialize_robot(body))

    return data

def serialize_robot(robot):
    return {
        'manipulators': map(serialize_manipulator, robot.GetManipulators()),
        'robot_state': serialize_robot_state(robot),
    }

def serialize_kinbody_state(body):
    data = {
        name: get_fn(body)
        for name, (get_fn, _) in KINBODY_STATE_MAP.iteritems()
    }

    link_transforms, dof_branches = body.GetLinkTransformations(True)
    data.update({
        'link_transforms': map(serialize_transform, link_transforms),
        'dof_branches': dof_branches.tolist(),
        'dof_values': body.GetDOFValues().tolist(),
    })

    return data

def serialize_robot_state(body):
    data = {
        name: get_fn(body)
        for name, (get_fn, _) in ROBOT_STATE_MAP.iteritems()
    }
    data['grabbed_bodies'] = map(serialize_grabbed_info, body.GetGrabbedInfo())
    return data

def serialize_link(link):
    data = { 'info': serialize_link_info(link.GetInfo()) }

    # Bodies loaded from ".kinbody.xml" do not have GeometryInfo's listed in
    # their LinkInfo class. We manually read them from GetGeometries().
    # TODO: This may not correctly preserve non-active geometry groups.
    data['info']['_vgeometryinfos'] = [
        serialize_geometry_info(geometry.GetInfo()) \
        for geometry in link.GetGeometries()
    ]
    return data

def serialize_joint(joint):
    return { 'info': serialize_joint_info(joint.GetInfo()) }

def serialize_manipulator(manipulator):
    return { 'info': serialize_manipulator_info(manipulator.GetInfo()) }

def serialize_with_map(obj, attribute_map):
    return {
        key: serialize_fn(getattr(obj, key))
        for key, (serialize_fn, _) in attribute_map.iteritems()
    }

def serialize_link_info(link_info):
    return serialize_with_map(link_info, LINK_INFO_MAP)
    
def serialize_joint_info(joint_info):
    return serialize_with_map(joint_info, JOINT_INFO_MAP)

def serialize_manipulator_info(manip_info):
    return serialize_with_map(manip_info, MANIPULATOR_INFO_MAP)

def serialize_geometry_info(geom_info):
    return serialize_with_map(geom_info, GEOMETRY_INFO_MAP)

def serialize_grabbed_info(grabbed_info):
    return serialize_with_map(grabbed_info, GRABBED_INFO_MAP)

def serialize_transform(t):
    from openravepy import quatFromRotationMatrix

    return {
        'position': list(map(float,t[0:3, 3])),
        'orientation': list(map(float,quatFromRotationMatrix(t[0:3, 0:3]))),
    }

# Deserialization.
def _deserialize_internal(env, data, data_type):
    from numpy import array, ndarray
    from openravepy import (Environment, KinBody, Robot, Trajectory,
                            RaveCreateTrajectory)
    from prpy.tsr import TSR, TSRChain
    from .exceptions import UnsupportedTypeDeserializationException

    if data_type == dict.__name__:
        return {
            deserialize(env, k): deserialize(env, v)
            for k, v in data.iteritems()
            if k != TYPE_KEY
        }
    elif data_type == ndarray.__name__:
        return array(data['data'])
    elif data_type in [ KinBody.__name__, Robot.__name__ ]:
        body = env.GetKinBody(data['name'])
        if body is None:
            raise ValueError('There is no body with name "{:s}".'.format(
                data['name']))

        return body
    elif data_type == KinBody.Link.__name__:
        body = env.GetKinBody(data['parent_name'])
        if body is None: