import numpy as np
from numpy.random import uniform
import json
import math
from utils.quaternion import Quaternion
from utils.quaternion import computeAngVel, computeAngVelRel

from Kinematic.py import KinematicCore

class JointType():
    BASE     = 0
    FIXED    = 1
    REVOLUTE = 2
    SPHERE   = 3

# pose dof for joints
Dof = {
    JointType.BASE:       4,
    JointType.FIXED:      0,
    JointType.REVOLUTE:   1,
    JointType.SPHERE:     4,
}

# exp dof for joints
ExpDof = {
    JointType.BASE:       0,
    JointType.FIXED:      0,
    JointType.REVOLUTE:   1,
    JointType.SPHERE:     3,
}

NAME2TYPE={
    "none":       JointType.BASE,
    "fixed":      JointType.FIXED,
    "revolute":   JointType.REVOLUTE,
    "spherical": JointType.SPHERE,
}

POSE2QUAT={
    JointType.BASE:       lambda x: Quaternion(x[0], x[1], x[2], x[3]),
    JointType.FIXED:      lambda x: Quaternion(1, 0, 0, 0),
    JointType.REVOLUTE:   lambda x: Quaternion(math.cos(x/2), 0.0, 0.0, math.sin(x/2)),
    JointType.SPHERE:     lambda x: Quaternion(x[0], x[1], x[2], x[3]),
}

VEL2OMG={
    JointType.BASE:       lambda x: np.array([x[0], x[1], x[2]]),
    JointType.FIXED:      lambda x: np.array([0, 0, 0]),
    JointType.REVOLUTE:   lambda x: np.array([0, 0, x[0]]),
    JointType.SPHERE:     lambda x: np.array([x[0], x[1], x[2]]),
}

class Joint(object):
    def __init__(self, info):
        """
            Inputs:
                info  json item load from character file
        """
        self._pos       = np.array([info["AttachX"], info["AttachY"], info["AttachZ"]])
        self._name      = info["Name"]
        self._parent_id = info["Parent"]
        self._parent    = None
        self._type_name = info['Type']
        self._type      = NAME2TYPE[info['Type']]
        self._dof       = Dof[self._type]
        self._expdof    = ExpDof[self._type]
        self._w         = info["DiffWeight"]
        self._is_end_effector = info["IsEndEffector"]

        if self._type is JointType.REVOLUTE:
            self._torque_lim = [info["TorqueLim"]]
        elif self._type is JointType.SPHERE:
            if "TorqueLimX" not in info:
                self._torque_lim = [info["TorqueLim"]]*4
            else:
                self._torque_lim = [info["TorqueLimX"], info["TorqueLimY"], info["TorqueLimZ"], 0]
        else:
            self._torque_lim = [0]*4

        if self._type is JointType.REVOLUTE:
            self.limlow = np.array([info["LimLow0"]])
            self.limhigh = np.array([info["LimHigh0"]])
        elif self._type is JointType.SPHERE:
            self.limlow = -3.14 * np.ones(3)
            self.limhigh = 3.14 * np.ones(3)
        else:
            self.limlow = np.array([])
            self.limhigh = np.array([])

        self.pose = np.zeros(self._dof)
        self.pose2quat = POSE2QUAT[self._type]
        self.quat = Quaternion(1, 0, 0, 0)

        self.vel = np.zeros(self._dof)
        self.vel2omg = VEL2OMG[self._type]
        self.omg = np.zeros(3)

    def __unicode__(self):
        info = u"Joint: name %s, pos %s, parent %d" % (self._name, str(self._pos), self._parent_id)
        return info

    __str__ = __unicode__
    __repr__= __unicode__

    def dof(self):
        return self._dof

    def expdof(self):
        return self._expdof

    def local_pos(self):
        return self._pos

    def set_pose(self, pose):
        assert(pose.shape == self.pose.shape)
        self.pose = pose
        # calculate quat
        self.quat = self.pose2quat(pose)

    def set_vel(self, vel):
        assert(vel.shape == self.vel.shape)
        self.omg = self.vel2omg(vel)

    def get_quat(self):
        return self.quat

    def get_omg(self):
        return self.omg

    def get_parent(self):
        return self._parent_id

    def get_a_min(self):
        std = (self.limhigh - self.limlow) / 2
        mean = (self.limhigh + self.limlow) / 2
        a_min = mean - std*2
        return a_min

    def get_a_max(self):
        std = (self.limhigh - self.limlow) / 2
        mean = (self.limhigh + self.limlow) / 2
        a_max = mean + std*2
        return a_max

class Body:
    def __init__(self, info):
        self._info = info
        self._pos = np.array([info["AttachX"], info["AttachY"], info["AttachZ"]])
        self._name = info["Name"]
        self._joint_id = info["ID"]

    def get_joint_id(self):
        return self._joint_id

    def get_local_pos(self):
        return self._pos

    def get_body_name(self):
        return self._name

    def get_body_shape(self):
        return self._info["Shape"]

    def get_body_size(self):
        ar = np.array([self._info["Param0"], self._info["Param1"], self._info["Param2"]])
        shape = self.get_body_shape()
        if shape == "sphere":
            ar /= 2.0
        if shape == "capsule":
            ar[0] /= 2.0
            ar[2] /= 2.0
        return ar

class HumanoidSkeleton(object):
    """ Store humanoid skeleton information, can perform following functions:

            1) compute link positions and velocities when given character pose and vel
            2) construct character state when given root position and orintaion
            3) slerp between two character poses
            4) transfer between action and target poses

            __init__()
                 _calc_kin_info()
                 _calc_ctrl_info()

            set_pose()

            set_vel()

            build_state()

            slerp()

            pose representation:
                - root pos (x, y, z)
                - root rot (w, x, y, z) in world frame
                - joint rots theta or (w, z, y, z) in local frame, VelRel

            vel w/ padding representation:
                - root vel (vx, vy, vz)
                - root angular vel (wx, wy, wz) in world frame
                - joint angular vel w or (wx, wy, wz, 0) in local frame, VelRel

            vel w/o padding representation
                - root vel (vx, vy, vz)
                - root angular vel (wx, wy, wz) in world frame
                - joint angular vel w or (wx, wy, wz) in local frame, VelRel

            state
                - Y coordinate of origin pos
                - Root link's pos (x, y, z) in xyz coordinate
                - Root link's quat (W, X, Y, Z) in XYZ coordinate
                - Other links' pos (x, y, z) and quat (w, x, y, z) in xyz coordinate
                - Root link's vel (Vx, Vy, Vz) in XYZ coordinate
                - Root link's omega (Wx, Wy, Wz) in XYZ coordinate
                - Other links' vel (vx, vy, vz) and omg (wx, wy, wz) in xyz coordinate

                xyz's origin is set by origin_pos, or character's root joint
                xyz's rotation is calculated by rotating Y-axis to make X-axis to heading
                direction calculated by origin_rot or root joint's rotation

            action
                - no root pose or joint
                - angle-axis (\theta, nx, ny, nz) for spherical joints and \theta for
                    revolute joints

            action represented in exponential map
                - no root pose or joint
                - exponential map (\Omega_x, \Omega_y, \Omega_z) for spherical joints and
                    \theta for revolute joints
    """

    def __init__(self, char_file, ctrl_file):
        # NOTICE humanoid.txt and humanoid.urdf have different scales, urdf is 4x larger
        self._kin_core = KinematicCore.cKinematicCore(char_file)

        ### all private members
        self.joints = None
        self.bodys  = None
        self.dof    = None
        self.expdof = None

        # TODO move to function
        self.num_joints  = None
        self.joint_types = None
        self.joint_dof   = None
        self.pos_start   = None     # start position in pose/vel vec
        self.act_start   = None     # start position in action vec
        self.joint_w     = None

        # control configuration
        self.kp = None
        self.kd = None
        self.a_max = None
        self.a_min = None

        ### all private members
        self._calc_kin_info(char_file)
        self._calc_ctrl_info(ctrl_file)

    def _calc_kin_info(self, char_file):
        # handle character skeleton information
        with open(char_file, 'r') as f:
            json_data = json.load(f)
            joints_data = json_data["Skeleton"]["Joints"]
            bodys_data = json_data["BodyDefs"]

        # build for pose/vel transfer
        self.joints = list(map(lambda x: Joint(x),  joints_data))
        self.bodys  = list(map(lambda x: Body(x),   bodys_data))

        self.dof = sum(map(lambda x: x.dof(), self.joints)) + 3
        self.expdof = sum(map(lambda x: x.expdof(), self.joints))

        # build velocity padding - compressing
        self.comp2pad = [0, 1, 2]  # root vel
        p_off = 3
        for j in self.joints:
            if j.dof() == 1:
                self.comp2pad += [p_off]
                p_off += 1
            if j.dof() == 4:
                self.comp2pad += [p_off, p_off+1, p_off+2]
                p_off += 4

        self.comp2pad = np.array(self.comp2pad, dtype=np.int)

        # build joint infomation mat
        self.num_joints     = len(self.joints)
        self.joint_types    = list(map(lambda x: x._type,   self.joints))
        self.joint_w        = list(map(lambda x: x._w,      self.joints))
        self.joint_dof      = list(map(lambda x: x.dof(),   self.joints))
        self.joint_dof      = np.array(self.joint_dof, dtype=np.int)
        self.joint_expdof   = list(map(lambda x: x.expdof(),self.joints))
        self.joint_expdof   = np.array(self.joint_expdof, dtype=np.int)
        self.end_effectors  = list(filter(
                                                        lambda x: self.joints[x]._is_end_effector,
                                                        range(self.num_joints)))

        self.pos_start = np.zeros(self.num_joints, dtype=np.int)  # start offset for each joint in pos vector
        for i in range(1, self.num_joints):
            self.pos_start[i] = self.pos_start[i-1] + self.joint_dof[i-1]

        self.exp_start = np.zeros(self.num_joints, dtype=np.int)  # start offset for each joint in exp vector
        for i in range(1, self.num_joints):
            self.exp_start[i] = self.exp_start[i-1] + self.joint_expdof[i-1]

        self.pos_start += 3                     # offset for root pos
        self.act_start = self.pos_start - 7     # 7 dim for root pos and root joint
                                                                                # so the act_start[0] is meaningless

        self.joint_mat = np.array([
                                                                self.joint_types, self.joint_dof,
                                                                self.pos_start, self.act_start,
                                                            ], dtype=np.int)
        self.joint_mat = self.joint_mat.transpose()

    def _calc_ctrl_info(self, ctrl_file):
        """ Build Kp, Kd and torque_lim
        """
        # handle controller information
        with open(ctrl_file, 'r') as f:
            data = json.load(f)
            controller = data["PDControllers"]
            assert(len(data["PDControllers"]) == self.num_joints)

        self.kp = [0,0,0] # for root pos
        self.kd = [0,0,0] # for root pos
        self.torque_lim = [0, 0, 0] # for root pos
        for i in range(self.num_joints):
            if self.joint_types[i] is not JointType.FIXED:
                self.kp += [controller[i]["Kp"]] * self.joint_dof[i]
                self.kd += [controller[i]["Kd"]] * self.joint_dof[i]
                self.torque_lim += self.joints[i]._torque_lim

        self.kp = np.array(self.kp)
        self.kd = np.array(self.kd)
        self.torque_lim = np.array(self.torque_lim)

    def build_a_min(self):
        a_mins = [j.get_a_min() for j in self.joints]
        a_min = np.concatenate(a_mins)
        assert(a_min.size == self.expdof)
        return a_min

    def build_a_max(self):
        a_maxs = [j.get_a_max() for j in self.joints]
        a_max = np.concatenate(a_maxs)
        assert(a_max.size == self.expdof)
        return a_max

    def get_link_ids(self):
        return list(range(len(self.bodys)))

    def get_link_names(self):
        return list(map(lambda body: body.get_body_name(), self.bodys))

    def get_link_shapes(self):
        return list(map(lambda body: body.get_body_shape(), self.bodys))

    def get_link_sizes(self):
        return list(map(lambda body: body.get_body_size().tolist(), self.bodys))

    def set_pose(self, pose):
        """ Set character's pose

        Inputs:
            pose   np.array of float, should be equal to self.dof
        """
        assert(pose.size == self.dof)
        self._kin_core.setPose(pose)

    def set_vel(self, vel):
        """ Set character's velocity

            *NOTICE* set_pose should be already called

        Inputs:
            vel  numpy array of velocity, shoubld be equal to self.dof
        """
        assert(vel.size == self.dof)
        self._kin_core.setVel(vel)

    def set_heading_vec(self, head):
        """ Set heading vector

        """
        assert(len(head) == 3)
        self._kin_core.setHeadingVec(head)

    def build_state(self, origin_pos=None, origin_rot=None, root_global=True):
        """ Build character state

            Inputs:
                origin_pos      new coordinate's origin (x, y, z)
                origin_rot      reference rotation (w, x, y, z), then a heading direction
                                                is calculated as new coordinate's x-axis

            Outputs:
                state           np.array of float
        """
        if origin_pos is None and origin_rot is None:
            state = self._kin_core.buildState()
            state = np.array(state)
        elif origin_pos is not None and origin_rot is not None:
            state = self._kin_core.buildState(origin_pos, origin_rot, root_global)
        else:
            raise ValueError("not implemented yet")
        return state

    def build_state2(self):
        state = self._kin_core.buildState2()
        return np.array(state)

    def get_com_pos(self):
        """ Return CoM position

            *NOTICE* set_pose should be already called
        """
        return np.array(self._kin_core.getCoMPos())

    def get_com_vel(self):
        """ Return CoM velocity

            *NOTICE* set_pose and set_vel should be already called
        """
        return np.array(self._kin_core.getCoMVel())

    def slerp(self, pose0, pose1, t):
        """ slerp between two poses

            Inputs:
                pose0   np.array of float, start pose
                pose1   np.array of float, end pose
                t       float in [0, 1], interpolating parameter

            Outputs:
                pose_t  np.array of float, interpolated pose
        """

        assert(pose0.size == self.dof)
        assert(pose1.size == self.dof)

        pose_t = self._kin_core.slerp(pose0, pose1, t)
        pose_t = np.array(pose_t)

        return pose_t

    def computeVel(self, pose0, pose1, dt, padding=True):
        """ Compute velocity between two poses

            Inputs:
                pose0   np.array of float, start pose
                pose1   np.array of float, end pose
                dt      float, duraction between two poses

            Outputs:
                avg_vel np.array of float, vel (w/ or w/o padding)
        """
        assert(pose0.size == self.dof)
        assert(pose1.size == self.dof)

        avg_vel = np.zeros_like(pose0)

        root0 = pose0[:3]
        root1 = pose1[:3]
        avg_vel[:3] = (root1 - root0) / dt

        offset = 3

        # root angular velocity is in world coordinate
        dof = self.joint_dof[0]
        quat0 = Quaternion.fromWXYZ(pose0[offset:offset+dof])
        quat1 = Quaternion.fromWXYZ(pose1[offset:offset+dof])
        avg_vel[offset : offset+3] = computeAngVel(quat0, quat1, dt)
        offset += dof

        # other joints
        for i in range(1, self.num_joints):
            dof = self.joint_dof[i]
            if dof == 1:  # revolute
                theta0 = pose0[offset]
                theta1 = pose1[offset]
                avg_vel[offset] = (theta1 - theta0) / dt
            elif dof == 4:  # spherical
                quat0 = Quaternion.fromWXYZ(pose0[offset:offset+dof])
                quat1 = Quaternion.fromWXYZ(pose1[offset:offset+dof])
                avg_vel[offset : offset+3] = computeAngVelRel(quat0, quat1, dt)
            offset += dof

        if padding is False:
            avg_vel = avg_vel[self.comp2pad]

        return avg_vel

    def toLocalFrame(self, pose, vel, ori_pos=None, ori_rot=None):
        """ Convert pose and vel from world frame to local frame,
                the local frame heading direction is rotated x-axis

            Inputs:
                pose        np.array of float, character pose
                vel         np.array of float, character vel w/ padding
                ori_pos     np.array of float, 3 dim, position of local coordinate origin
                ori_rot     np.array of float, 4 dim, (w, x, y, z) quat of local coordinate orientation

            Outputs:
                local_pose
                local_vel
        """
        if ori_pos is None:
            ori_pos = pose[:3]
        if ori_rot is None:
            ori_rot = pose[3:7]

        # heading theta
        inv_ori_rot = self.buildHeadingTrans(ori_rot)

        local_pos = pose.copy()
        local_vel = vel.copy()

        # ignore y difference, because local cooridnate shares xoz plane with world
        local_pos[0] -= ori_pos[0]                      # root x pos
        local_pos[2] -= ori_pos[2]                      # root y pos

        ori_rot = Quaternion.fromWXYZ(ori_rot)
        ori_rot = inv_ori_rot.mul(ori_rot)
        local_pos[3:7] = ori_rot.pos_wxyz()              # root orientation

        local_vel[:3] = inv_ori_rot.rotate(vel[:3])      # root velocity
        local_vel[3:6] = inv_ori_rot.rotate(vel[3:6])    # root angular velocity

        return local_pos, local_vel

    def buildHeadingTrans(self, rot):
        """ Build the rotation that rotate to local coordinate
            rot     np.array of float, 4 dim, (w, x, y, z) quat of coordinate orientation
        """
        theta = self._kin_core.getHeadingTheta(rot)
        inv_rot = Quaternion.fromExpMap(np.array([0, -theta, 0]))
        return inv_rot

    def compressVel(self, vel):
        """ Squeeze velocity from padded to compressed

            Inputs:
                vel     np.array of float, vel w/ padding

            Outputs:
                vel_cmp np.array of float, vel w/o padding
        """
        vel_cmp = vel[self.comp2pad]
        return vel_cmp

    def padVel(self, vel):
        """ Pad velocity from compressed to padded

            Inputs:
                vel     np.array of float, vel w/o padding

            Outputs:
                vel_pad np.array of float, vel w/ padding
        """
        vel_pad = np.zeros(self.dof)
        vel_pad[self.comp2pad] = vel
        return vel_pad

    def action_to_targ_pose(self, action):
        """ Converte action to PD controller target pose

            Inputs:
                action      np.array of float, action which DeepMimicSim can take in

            Outputs:
                pose        np.array of float, pose of character
        """
        assert(action.size == self.dof - 7)

        targ_pose = np.zeros(self.dof)
        targ_pose[3] = 1
        for i in range(1, self.num_joints):
            dof = self.joint_dof[i]
            p_off = self.pos_start[i]
            a_off = self.act_start[i]
            if dof == 1:  # revolute
                targ_pose[p_off] = action[a_off]
            elif dof == 4:  # spherical
                angle = action[a_off]
                axis = action[a_off+1:a_off+4]
                quata = Quaternion.fromAngleAxis(angle, axis)
                targ_pose[p_off : p_off+4] = quata.wxyz()

        assert(np.isfinite(sum(targ_pose))), embed() #(action, targ_pose)
        return targ_pose

    def targ_pose_to_action(self, pose):
        """ Convert desired pose to action

            Inputs:
                pose        np.array of float, pose of character

            Outputs:
                action      np.array of float, action which DeepMimicSim can take in
        """
        assert(pose.size == self.dof)

        action = np.zeros(self.dof-7)
        for i in range(1, self.num_joints):
            dof = self.joint_dof[i]
            p_off = self.pos_start[i]
            a_off = self.act_start[i]
            if dof == 1:  # revolute
                action[a_off] = pose[p_off]
            elif dof == 4:  # spherical
                quata = Quaternion.fromWXYZ(pose[p_off:p_off+4])
                action[a_off : a_off+4] = quata.angaxis()

        assert(np.isfinite(sum(action))), embed() #(action, targ_pose)
        return action

    def action_as_offset(self, pose, action):
        """ Take action as offset of a given reference pose, return standard action

            Inputs:
                pose        np.array of float, pose of character
                action      np.array of float, offset action to the character

            Outputs:
                new_action  np.array of float, standard action which DeepMimicSim can take in
        """
        assert(pose.size == self.dof)
        assert(action.size == self.dof - 7)

        new_action = self._kin_core.actionAsOffset(pose, action.tolist())

        return new_action

    def exp_to_action(self, expmap):
        """ Convert action represented in exponential map to angle-axis action
                which deepmimic_sim can take in

            Inputs:
                expmap      np.array of float, action represented in exponential map

            Outputs:
                action      np.array of float, action for deepmimic environment
        """
        assert(expmap.size == self.expdof)
        action = np.zeros(self.dof-7)

        for i in range(1, self.num_joints):
            dof = self.joint_dof[i]
            e_off = self.exp_start[i]
            a_off = self.act_start[i]
            if dof == 1:  # revolute
                action[a_off] = expmap[e_off]
            elif dof == 4:  # spherical
                quata = Quaternion.fromExpMap(expmap[e_off:e_off+3])
                action[a_off:a_off+4] = quata.angaxis()

        assert(np.isfinite(sum(action))), embed() #(action, targ_pose)
        return action

    def exp_to_targ_pose(self, expmap, padding=True):
        """ Convert action represented in exponential map to angle-axis action
                which deepmimic_sim can take in

            Inputs:
                expmap      np.array of float, action represented in exponential map

            Outputs:
                pose        np.array of float, pose of character
        """
        assert(expmap.size == self.expdof)
        return self._kin_core.expMapToTargetPose(expmap.tolist(), padding)

    def exp_to_targ_pose_old(self, expmap):
        """ Convert action represented in exponential map to angle-axis action
                which deepmimic_sim can take in

            Inputs:
                expmap      np.array of float, action represented in exponential map

            Outputs:
                pose        np.array of float, pose of character
        """
        assert(expmap.size == self.expdof)
        pose = np.zeros(self.dof)

        for i in range(1, self.num_joints):
            dof = self.joint_dof[i]
            e_off = self.exp_start[i]
            p_off = self.pos_start[i]
            if dof == 1:  # revolute
                pose[p_off] = expmap[e_off]
            elif dof == 4:  # spherical
                quata = Quaternion.fromExpMap(expmap[e_off:e_off+3])
                pose[p_off:p_off+4] = quata.wxyz()

        assert(np.isfinite(sum(pose))), embed() #(action, targ_pose)
        return pose

    def targ_pose_to_exp(self, pose):
        """ Convert target pose to exponential map, used as initialization of
                Actor_FDM reference memory

            Inputs:
                pose        np.array of float, pose of character

            Outputs:
                exp         np.array of float, action represented in exponential map
        """
        assert(pose.size == self.dof)

        expmap = np.zeros(self.expdof)
        for i in range(1, self.num_joints):
            dof = self.joint_dof[i]
            p_off = self.pos_start[i]
            e_off = self.exp_start[i]
            if dof == 1:  # revolute
                expmap[e_off] = pose[p_off]
            elif dof == 4:  # spherical
                quata = Quaternion.fromWXYZ(pose[p_off:p_off+4])
                expmap[e_off : e_off+3] = quata.expmap()

        assert(np.isfinite(sum(expmap))), embed() #(action, targ_pose)
        return expmap

    def normalize_target_pose(self, pose):
        """
            normalize given pose to make quaternions norm as 1
        """
        assert(pose.size == self.dof)
        norm_pose = pose.copy()
        for i in range(self.num_joints):
            dof = self.joint_dof[i]
            p_off = self.pos_start[i]
            if dof == 4:  # spherical
                quata = Quaternion.fromWXYZ(pose[p_off:p_off+4])
                norm_pose[p_off : p_off+4] = quata.pos_wxyz()

        assert(np.isfinite(sum(norm_pose))), embed() #(action, targ_pose)
        return norm_pose

    def pose_wxyz_to_xyzw(self, pose):
        pose_new = pose.copy()

        p_off = 3
        quat = pose[p_off:p_off+4]
        pose_new[p_off:p_off+3] = quat[1:]
        pose_new[p_off+3] = quat[0]

        for i in range(1, self.num_joints):
            dof = self.joint_dof[i]
            p_off = self.pos_start[i]
            if dof == 4:
                quat = pose[p_off:p_off+4]
                pose_new[p_off:p_off+3] = quat[1:]
                pose_new[p_off+3] = quat[0]

        return pose_new

    def get_reward(self, pose0, vel0, pose1, vel1):
        assert(len(pose0) == self.dof), "pose 0 size missmatch"
        assert(len(pose1) == self.dof), "pose 1 size missmatch"
        assert(len(vel0) == self.dof), "vel 0 size missmatch"
        assert(len(vel1) == self.dof), "vel 1 size missmatch"
        return self._kin_core.calcReward(pose0, vel0, pose1, vel1)

    def get_reward2(self, pose0, vel0, pose1, vel1):
        assert(len(pose0) == self.dof), "pose 0 size missmatch"
        assert(len(pose1) == self.dof), "pose 1 size missmatch"
        assert(len(vel0) == self.dof), "vel 0 size missmatch"
        assert(len(vel1) == self.dof), "vel 1 size missmatch"
        return self._kin_core.calcReward2(pose0, vel0, pose1, vel1)

    def get_err_vec(self, pose0, vel0, pose1, vel1):
        assert(len(pose0) == self.dof), "pose 0 size missmatch"
        assert(len(pose1) == self.dof), "pose 1 size missmatch"
        assert(len(vel0) == self.dof), "vel 0 size missmatch"
        assert(len(vel1) == self.dof), "vel 1 size missmatch"
        self._kin_core.calcReward(pose0, vel0, pose1, vel1)
        return np.array(self._kin_core.getErrorVec())

    def get_sub_rewards(self, pose0, vel0, pose1, vel1):
        err_vec = self.get_err_vec(pose0, vel0, pose1, vel1)
        scale = np.array([2, 0.1, 40, 5, 10])
        sub_rewards = np.exp(- scale * err_vec)
        return sub_rewards

    def lowest_height(self, pose):
        return self._kin_core.lowestHeight(pose)

    def disturb_pose(self, pose, noise):
        """ uniform disturb pose

            pose
            noise  float, deviation angle in radius
        """
        new_pose = pose.copy()

        noise3d = noise / np.sqrt(3)
        new_pose[:3] += uniform(-noise3d, noise3d, 3)

        for jstar, jdof in zip(self.pos_start, self.joint_dof):
            if jdof == 1:
                new_pose[jstar] += uniform(-noise, noise)
            elif jdof == 4:
                noise_quat = Quaternion.fromExpMap(uniform(-noise3d, noise3d, 3))
                ori_quat = Quaternion.fromWXYZ(pose[jstar:jstar+4])
                new_quat = noise_quat.mul(ori_quat)
                new_pose[jstar:jstar+4] = new_quat.wxyz()
            elif jdof == 0:
                pass
            else:
                assert(False and "not support jdof other than 1 and 4")

        return new_pose

    def disturb_vel(self, vel, noise):
        """ uniform disturb vel

            vel
            noise  float, deviation angle in radius
        """
        new_vel = vel.copy()

        noise3d = noise / np.sqrt(3)
        new_vel[:3] += uniform(-noise3d, noise3d, 3)

        for jstar, jdof in zip(self.pos_start, self.joint_dof):
            if jdof == 1:
                new_vel[jstar] += uniform(-noise, noise)
            elif jdof == 4:
                new_vel[jstar:jstar+3] += uniform(-noise3d, noise3d, 3)
            elif jdof == 0:
                pass
            else:
                assert(False and "not support jdof other than 1 and 4")

        return new_vel

    def get_joint_global_pos(self, id):
        return list(self._kin_core.getJointPos(id))

    def get_joint_global_quat(self, id):
        return list(self._kin_core.getJointQuat(id))

    def get_body_global_pos(self, id):
        return list(self._kin_core.getBodyPos(id))

    def get_feature(self):
        return self._kin_core.getFeature()

    def inv_feature(self, pos):
        return self._kin_core.invFeature(pos)


# -*- coding: utf-8 -*-
from os.path import dirname, exists, join

import matplotlib.pyplot as plt
import numpy as np
from astropy.coordinates import AltAz, EarthLocation
from exoorbit.orbit import Orbit

from exoplanet_transit_snr.snr_estimate import (
    calculate_cohen_d_for_dataset,
    coadd_cross_correlation,
    load_data,
    run_cross_correlation,
)
from exoplanet_transit_snr.stats import gaussfit
from exoplanet_transit_snr.stellardb import StellarDb

# define the names of the star and planet
# as well as the datasets within the datasets folder
star, planet = "WASP-107", "b"
datasets = {50: "WASP-107b_SNR50", 100: "WASP-107b_SNR100", 200: "WASP-107b_SNR200"}

# Load the nominal data for this star and planet from simbad/nasa exoplanet archive
sdb = StellarDb()
star = sdb.get(star)
planet = star.planets[planet]
orbit = Orbit(star, planet)
telescope = EarthLocation.of_site("Paranal")

# Define the +- range of the radial velocity points,
# and the density of the sampling
rv_range = 200
rv_step = 0.25

for snr in [200]:
    # Where to find the data, might need to be adjusted
    data_dir = join(dirname(__file__), "../datasets", datasets[snr], "Spectrum_00")
    # load the data from the fits files, returns several objects
    data = load_data(data_dir, load=True)
    wave, flux, uncs, times, segments = data

    # Determine orbital parameters
    phi = (times - planet.time_of_transit) / planet.period
    phi = phi.to_value(1)
    phi = phi % 1
    ingress = (-planet.transit_duration / 2 / planet.period).to_value(1) % 1
    egress = (planet.transit_duration / 2 / planet.period).to_value(1) % 1
    in_transit = (phi >= ingress) | (phi <= egress)
    out_transit = ~in_transit

    altaz = star.coordinates.transform_to(AltAz(obstime=times, location=telescope))
    airmass = altaz.secz.value

    rv = orbit.radial_velocity_planet(times)
    rv_star = -star.coordinates.radial_velocity_correction(
        obstime=times, location=telescope
    )

    # TODO: Create initial stellar spectrum guess
    # sme = SME_Structure()
    # sme.wave = wave[low:upp].to_value(u.AA)
    # sme.spec = flux[low:upp]
    # # We need the specific intensities, not the combined spectrum
    # sme.vrad_flag = "whole"
    # sme.vrad = -11
    # sme.vrad_limit = 20
    # sme.cscale_flag = "none"
    # sme.cscale_type = "match"
    # # Set stellar parameters
    # sme.teff = star.teff.to_value(u.K)
    # sme.logg = star.logg.to_value(u.one)
    # sme.abund = Abund(star.monh.to_value(u.one), "asplund2009")
    # sme.vmic = 1
    # sme.vmac = 0
    # sme.vsini = 0
    # # Define atmosphere
    # sme.atmo.source = "marcs2012.sav"
    # sme.atmo.method = "grid"
    # # Set linelist
    # valdfile = join(dirname(__file__), "ltt1445A.lin")
    # sme.linelist = ValdFile(valdfile)
    # # sme.linelist = sme.linelist.trim(sme.wave[0, 0], sme.wave[0, -1], 1000)
    # # Set Mu to the specific positions we need
    # d = orbit.projected_radius(times)
    # r = orbit.planet.radius
    # R = orbit.star.radius
    # mu = np.sqrt(1 - (d / R) ** 2)

    # sme_fname = f"sme_intensities_{star.name}_{planet.name}_snr{snr}.npz"
    # try:
    #     sme_data = np.load(sme_fname)
    #     wave_int, flux_int, cont_int = sme_data["wave"], sme_data["flux"], sme_data["cont"]
    #     wave_sme, flux_sme = sme_data["wave_sme"], sme_data["flux_sme"]
    #     sme.mu = mu[np.isfinite(mu)]
    # except (FileNotFoundError, KeyError):
    #     sme = synthesize_spectrum(sme)
    #     flux_sme = np.copy(sme.synth.ravel())
    #     wave_sme = np.copy(sme.wave.ravel())

    #     sme.specific_intensities_only = True
    #     sme.mu = mu[np.isfinite(mu)]
    #     wave_int, flux_int, cont_int = synthesize_spectrum(sme)
    #     wave_int, flux_int, cont_int = wave_int[0], flux_int[0], cont_int[0]

    #     np.savez(
    #         sme_fname,
    #         wave=wave_int,
    #         flux=flux_int,
    #         cont=cont_int,
    #         wave_sme=wave_sme,
    #         flux_sme=flux_sme,
    #     )

    # str_wave, str_flux = wave_sme, flux_sme

    # Run the cross correlation to the next neighbour
    cc_data, rv_array = run_cross_correlation(
        data,
        nsysrem=1,
        rv_range=rv_range,
        rv_step=rv_step,
        load=False,
        data_dir=data_dir,
        rv_star=rv_star,
        rv_planet=rv,
        airmass=airmass,
    )

    cc_data = cc_data["10"]
    cc_data_coadd, cc_it, cc_oot = coadd_cross_correlation(
        cc_data, rv, rv_array, times, planet, data_dir=data_dir, load=False
    )

    # Normalize by the oot signal
    cc_it /= cc_oot

    # Fit a gaussian
    p0 = [np.max(cc_data_coadd) - np.min(cc_data_coadd), 0, 10, np.min(cc_data_coadd)]
    gauss, pval = gaussfit(rv_array, cc_data_coadd, p0=p0)

    # Plot the results
    plt.plot(rv_array, cc_data_coadd, label="all observations")
    plt.plot(rv_array, gauss, label="best fit gaussian")
    plt.plot(rv_array, cc_it, label="in-transit")
    plt.plot(rv_array, cc_oot, label="out-of-transit")
    plt.legend()
    plt.title(f"{star.name} {planet.name} SNR{snr}")
    plt.xlabel(r"$\Delta$RV [km/s]")
    plt.ylabel("CCF")
    plt.show()


#!/usr/bin/env python
# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import ctypes
import time
from functools import reduce
import numpy
import h5py
from pyscf import lib
from pyscf.lib import logger
from pyscf import ao2mo
from pyscf.ao2mo import _ao2mo
from pyscf.ao2mo import outcore

# least memory requirements:
# nmo  ncore  ncas  outcore  incore
# 200  40     16    0.8GB    3.7 GB (_eri 1.6GB intermediates 1.3G)
# 250  50     16    1.7GB    8.2 GB (_eri 3.9GB intermediates 2.6G)
# 300  60     16    3.1GB    16.8GB (_eri 8.1GB intermediates 5.6G)
# 400  80     16    8.5GB    53  GB (_eri 25.6GB intermediates 19G)
# 500  100    16    19 GB
# 600  120    16    37 GB
# 750  150    16    85 GB

libmcscf = lib.load_library('libmcscf')

def trans_e1_incore(eri_ao, mo, ncore, ncas):
    nmo = mo.shape[1]
    nocc = ncore + ncas
    eri1 = ao2mo.incore.half_e1(eri_ao, (mo,mo[:,:nocc]), compact=False)
    eri1 = eri1.reshape(nmo,nocc,-1)

    klppshape = (0, nmo, 0, nmo)
    klpashape = (0, nmo, ncore, nocc)
    aapp = numpy.empty((ncas,ncas,nmo,nmo))
    for i in range(ncas):
        _ao2mo.nr_e2(eri1[ncore+i,ncore:nocc], mo, klppshape,
                      aosym='s4', mosym='s1', out=aapp[i])
    ppaa = lib.transpose(aapp.reshape(ncas*ncas,-1)).reshape(nmo,nmo,ncas,ncas)
    aapp = None

    papa = numpy.empty((nmo,ncas,nmo,ncas))
    for i in range(nmo):
        _ao2mo.nr_e2(eri1[i,ncore:nocc], mo, klpashape,
                      aosym='s4', mosym='s1', out=papa[i])

    pp = numpy.empty((nmo,nmo))
    j_cp = numpy.zeros((ncore,nmo))
    k_pc = numpy.zeros((nmo,ncore))
    for i in range(ncore):
        _ao2mo.nr_e2(eri1[i,i:i+1], mo, klppshape, aosym='s4', mosym='s1', out=pp)
        j_cp[i] = pp.diagonal()
    j_pc = j_cp.T.copy()

    pp = numpy.empty((ncore,ncore))
    for i in range(nmo):
        klshape = (i, i+1, 0, ncore)
        _ao2mo.nr_e2(eri1[i,:ncore], mo, klshape, aosym='s4', mosym='s1', out=pp)
        k_pc[i] = pp.diagonal()
    return j_pc, k_pc, ppaa, papa

# level = 1: ppaa, papa and jpc, kpc
# level > 1: ppaa, papa only.  It affects accuracy of hdiag
def trans_e1_outcore(mol, mo, ncore, ncas, erifile,
                     max_memory=None, level=1, verbose=logger.WARN):
    time0 = (time.clock(), time.time())
    log = logger.new_logger(mol, verbose)
    log.debug1('trans_e1_outcore level %d  max_memory %d', level, max_memory)
    nao, nmo = mo.shape
    nao_pair = nao*(nao+1)//2
    nocc = ncore + ncas

    faapp_buf = lib.H5TmpFile()
    if isinstance(erifile, h5py.Group):
        feri = erifile
    else:
        feri = lib.H5TmpFile(erifile, 'w')

    mo_c = numpy.asarray(mo, order='C')
    mo = numpy.asarray(mo, order='F')
    pashape = (0, nmo, ncore, nocc)
    papa_buf = numpy.zeros((nao,ncas,nmo*ncas))
    j_pc = numpy.zeros((nmo,ncore))
    k_pc = numpy.zeros((nmo,ncore))

    mem_words = int(max(2000,max_memory-papa_buf.nbytes/1e6)*1e6/8)
    aobuflen = mem_words//(nao_pair+nocc*nmo) + 1
    ao_loc = numpy.array(mol.ao_loc_nr(), dtype=numpy.int32)
    shranges = outcore.guess_shell_ranges(mol, True, aobuflen, None, ao_loc)
    intor = mol._add_suffix('int2e')
    ao2mopt = _ao2mo.AO2MOpt(mol, intor,
                             'CVHFnr_schwarz_cond', 'CVHFsetnr_direct_scf')
    nstep = len(shranges)
    maxbuflen = max([x[2] for x in shranges])
    log.debug('mem_words %.8g MB, maxbuflen = %d', mem_words*8/1e6, maxbuflen)
    bufs1 = numpy.empty((maxbuflen, nao_pair))
    bufs2 = numpy.empty((maxbuflen, nmo*ncas))
    if level == 1:
        bufs3 = numpy.empty((maxbuflen, nao*ncore))
        log.debug('mem cache %.8g MB',
                  (bufs1.nbytes+bufs2.nbytes+bufs3.nbytes)/1e6)
    else:
        log.debug('mem cache %.8g MB', (bufs1.nbytes+bufs2.nbytes)/1e6)
    ti0 = log.timer('Initializing trans_e1_outcore', *time0)

    # fmmm, ftrans, fdrv for level 1
    fmmm = libmcscf.AO2MOmmm_ket_nr_s2
    ftrans = libmcscf.AO2MOtranse1_nr_s4
    fdrv = libmcscf.AO2MOnr_e2_drv
    for istep,sh_range in enumerate(shranges):
        log.debug('[%d/%d], AO [%d:%d], len(buf) = %d',
                  istep+1, nstep, *sh_range)
        buf = bufs1[:sh_range[2]]
        _ao2mo.nr_e1fill(intor, sh_range,
                         mol._atm, mol._bas, mol._env, 's4', 1, ao2mopt, buf)
        if log.verbose >= logger.DEBUG1:
            ti1 = log.timer('AO integrals buffer', *ti0)
        bufpa = bufs2[:sh_range[2]]
        _ao2mo.nr_e1(buf, mo, pashape, 's4', 's1', out=bufpa)
# jc_pp, kc_pp
        if level == 1: # ppaa, papa and vhf, jcp, kcp
            if log.verbose >= logger.DEBUG1:
                ti1 = log.timer('buffer-pa', *ti1)
            buf1 = bufs3[:sh_range[2]]
            fdrv(ftrans, fmmm,
                 buf1.ctypes.data_as(ctypes.c_void_p),
                 buf.ctypes.data_as(ctypes.c_void_p),
                 mo.ctypes.data_as(ctypes.c_void_p),
                 ctypes.c_int(sh_range[2]), ctypes.c_int(nao),
                 (ctypes.c_int*4)(0, nao, 0, ncore),
                 ctypes.POINTER(ctypes.c_void_p)(), ctypes.c_int(0))
            p0 = 0
            for ij in range(sh_range[0], sh_range[1]):
                i,j = lib.index_tril_to_pair(ij)
                i0 = ao_loc[i]
                j0 = ao_loc[j]
                i1 = ao_loc[i+1]
                j1 = ao_loc[j+1]
                di = i1 - i0
                dj = j1 - j0
                if i == j:
                    dij = di * (di+1) // 2
                    buf = numpy.empty((di,di,nao*ncore))
                    idx = numpy.tril_indices(di)
                    buf[idx] = buf1[p0:p0+dij]
                    buf[idx[1],idx[0]] = buf1[p0:p0+dij]
                    buf = buf.reshape(di,di,nao,ncore)
                    mo1 = mo_c[i0:i1]
                    tmp = numpy.einsum('uvpc,pc->uvc', buf, mo[:,:ncore])
                    tmp = lib.dot(mo1.T, tmp.reshape(di,-1))
                    j_pc += numpy.einsum('vp,pvc->pc', mo1, tmp.reshape(nmo,di,ncore))
                    tmp = numpy.einsum('uvpc,uc->vcp', buf, mo1[:,:ncore])
                    tmp = lib.dot(tmp.reshape(-1,nmo), mo).reshape(di,ncore,nmo)
                    k_pc += numpy.einsum('vp,vcp->pc', mo1, tmp)
                else:
                    dij = di * dj
                    buf = buf1[p0:p0+dij].reshape(di,dj,nao,ncore)
                    mo1 = mo_c[i0:i1]
                    mo2 = mo_c[j0:j1]
                    tmp = numpy.einsum('uvpc,pc->uvc', buf, mo[:,:ncore])
                    tmp = lib.dot(mo1.T, tmp.reshape(di,-1))
                    j_pc += numpy.einsum('vp,pvc->pc',
                                         mo2, tmp.reshape(nmo,dj,ncore)) * 2
                    tmp = numpy.einsum('uvpc,uc->vcp', buf, mo1[:,:ncore])
                    tmp = lib.dot(tmp.reshape(-1,nmo), mo).reshape(dj,ncore,nmo)
                    k_pc += numpy.einsum('vp,vcp->pc', mo2, tmp)
                    tmp = numpy.einsum('uvpc,vc->ucp', buf, mo2[:,:ncore])
                    tmp = lib.dot(tmp.reshape(-1,nmo), mo).reshape(di,ncore,nmo)
                    k_pc += numpy.einsum('up,ucp->pc', mo1, tmp)
                p0 += dij
            if log.verbose >= logger.DEBUG1:
                ti1 = log.timer('j_cp and k_cp', *ti1)

        if log.verbose >= logger.DEBUG1:
            ti1 = log.timer('half transformation of the buffer', *ti1)

# ppaa, papa
        faapp_buf[str(istep)] = \
                bufpa.reshape(sh_range[2],nmo,ncas)[:,ncore:nocc].reshape(-1,ncas**2).T
        p0 = 0
        for ij in range(sh_range[0], sh_range[1]):
            i,j = lib.index_tril_to_pair(ij)
            i0 = ao_loc[i]
            j0 = ao_loc[j]
            i1 = ao_loc[i+1]
            j1 = ao_loc[j+1]
            di = i1 - i0
            dj = j1 - j0
            if i == j:
                dij = di * (di+1) // 2
                buf1 = numpy.empty((di,di,nmo*ncas))
                idx = numpy.tril_indices(di)
                buf1[idx] = bufpa[p0:p0+dij]
                buf1[idx[1],idx[0]] = bufpa[p0:p0+dij]
            else:
                dij = di * dj
                buf1 = bufpa[p0:p0+dij].reshape(di,dj,-1)
                mo1 = mo[j0:j1,ncore:nocc].copy()
                for i in range(di):
                    lib.dot(mo1.T, buf1[i], 1, papa_buf[i0+i], 1)
            mo1 = mo[i0:i1,ncore:nocc].copy()
            buf1 = lib.dot(mo1.T, buf1.reshape(di,-1))
            papa_buf[j0:j1] += buf1.reshape(ncas,dj,-1).transpose(1,0,2)
            p0 += dij
        if log.verbose >= logger.DEBUG1:
            ti1 = log.timer('ppaa and papa buffer', *ti1)

        ti0 = log.timer('gen AO/transform MO [%d/%d]'%(istep+1,nstep), *ti0)
    buf = buf1 = bufpa = None
    bufs1 = bufs2 = bufs3 = None
    time1 = log.timer('mc_ao2mo pass 1', *time0)

    log.debug1('Half transformation done. Current memory %d',
               lib.current_memory()[0])

    nblk = int(max(8, min(nmo, (max_memory*1e6/8-papa_buf.size)/(ncas**2*nmo))))
    log.debug1('nblk for papa = %d', nblk)
    dset = feri.create_dataset('papa', (nmo,ncas,nmo,ncas), 'f8')
    for i0, i1 in prange(0, nmo, nblk):
        tmp = lib.dot(mo[:,i0:i1].T, papa_buf.reshape(nao,-1))
        dset[i0:i1] = tmp.reshape(i1-i0,ncas,nmo,ncas)
    papa_buf = tmp = None
    time1 = log.timer('papa pass 2', *time1)

    tmp = numpy.empty((ncas**2,nao_pair))
    p0 = 0
    for istep, sh_range in enumerate(shranges):
        tmp[:,p0:p0+sh_range[2]] = faapp_buf[str(istep)]
        p0 += sh_range[2]
    nblk = int(max(8, min(nmo, (max_memory*1e6/8-tmp.size)/(ncas**2*nmo)-1)))
    log.debug1('nblk for ppaa = %d', nblk)
    dset = feri.create_dataset('ppaa', (nmo,nmo,ncas,ncas), 'f8')
    for i0, i1 in prange(0, nmo, nblk):
        tmp1 = _ao2mo.nr_e2(tmp, mo, (i0,i1,0,nmo), 's4', 's1', ao_loc=ao_loc)
        tmp1 = tmp1.reshape(ncas,ncas,i1-i0,nmo)
        for j in range(i1-i0):
            dset[i0+j] = tmp1[:,:,j].transpose(2,0,1)
    tmp = tmp1 = None
    time1 = log.timer('ppaa pass 2', *time1)

    time0 = log.timer('mc_ao2mo', *time0)
    return j_pc, k_pc

# level = 1: ppaa, papa and vhf, jpc, kpc
# level = 2: ppaa, papa, vhf,  jpc=0, kpc=0
class _ERIS(object):
    def __init__(self, casscf, mo, method='incore', level=1):
        mol = casscf.mol
        nao, nmo = mo.shape
        ncore = casscf.ncore
        ncas = casscf.ncas

        dm_core = numpy.dot(mo[:,:ncore], mo[:,:ncore].T)
        vj, vk = casscf._scf.get_jk(mol, dm_core)
        self.vhf_c = reduce(numpy.dot, (mo.T, vj*2-vk, mo))

        mem_incore, mem_outcore, mem_basic = _mem_usage(ncore, ncas, nmo)
        mem_now = lib.current_memory()[0]
        eri = casscf._scf._eri
        if (method == 'incore' and eri is not None and
            (mem_incore+mem_now < casscf.max_memory*.9) or
            mol.incore_anyway):
            if eri is None:
                eri = mol.intor('int2e', aosym='s8')
            self.j_pc, self.k_pc, self.ppaa, self.papa = \
                    trans_e1_incore(eri, mo, ncore, ncas)
        else:
            log = logger.Logger(casscf.stdout, casscf.verbose)
            self.feri = lib.H5TmpFile()
            max_memory = max(3000, casscf.max_memory*.9-mem_now)
            if max_memory < mem_basic:
                log.warn('Calculation needs %d MB memory, over CASSCF.max_memory (%d MB) limit',
                         (mem_basic+mem_now)/.9, casscf.max_memory)
            self.j_pc, self.k_pc = \
                    trans_e1_outcore(mol, mo, ncore, ncas, self.feri,
                                     max_memory=max_memory,
                                     level=level, verbose=log)
            self.ppaa = self.feri['ppaa']
            self.papa = self.feri['papa']

def _mem_usage(ncore, ncas, nmo):
    outcore = basic = ncas**2*nmo**2*2 * 8/1e6
    incore = outcore + (ncore+ncas)*nmo**3*4/1e6
    return incore, outcore, basic

def prange(start, end, step):
    for i in range(start, end, step):
        yield i, min(i+step, end)

if __name__ == '__main__':
    from pyscf import scf
    from pyscf import gto
    from pyscf.mcscf import mc1step

    mol = gto.Mole()
    mol.verbose = 0
    mol.output = None#"out_h2o"
    mol.atom = [
        ['O', ( 0., 0.    , 0.   )],
        ['H', ( 0., -0.757, 0.587)],
        ['H', ( 0., 0.757 , 0.587)],]
    mol.basis = {'H': 'cc-pvtz',
                 'O': 'cc-pvtz',}
    mol.build()

    m = scf.RHF(mol)
    ehf = m.scf()

    mc = mc1step.CASSCF(m, 6, 4)
    mc.verbose = 4
    mo = m.mo_coeff.copy()

    eris0 = _ERIS(mc, mo, 'incore')
    eris1 = _ERIS(mc, mo, 'outcore')
    eris2 = _ERIS(mc, mo, 'outcore', level=1)
    eris3 = _ERIS(mc, mo, 'outcore', level=2)
    print('vhf_c', numpy.allclose(eris0.vhf_c, eris1.vhf_c))
    print('j_pc ', numpy.allclose(eris0.j_pc , eris1.j_pc ))
    print('k_pc ', numpy.allclose(eris0.k_pc , eris1.k_pc ))
    print('ppaa ', numpy.allclose(eris0.ppaa , eris1.ppaa ))
    print('papa ', numpy.allclose(eris0.papa , eris1.papa ))

    print('vhf_c', numpy.allclose(eris0.vhf_c, eris2.vhf_c))
    print('j_pc ', numpy.allclose(eris0.j_pc , eris2.j_pc ))
    print('k_pc ', numpy.allclose(eris0.k_pc , eris2.k_pc ))
    print('ppaa ', numpy.allclose(eris0.ppaa , eris2.ppaa ))
    print('papa ', numpy.allclose(eris0.papa , eris2.papa ))

    print('vhf_c', numpy.allclose(eris0.vhf_c, eris3.vhf_c))
    print('ppaa ', numpy.allclose(eris0.ppaa , eris3.ppaa ))
    print('papa ', numpy.allclose(eris0.papa , eris3.papa ))

    ncore = mc.ncore
    ncas = mc.ncas
    nocc = ncore + ncas
    nmo = mo.shape[1]
    eri = ao2mo.incore.full(m._eri, mo, compact=False).reshape((nmo,)*4)
    aaap = numpy.array(eri[ncore:nocc,ncore:nocc,ncore:nocc,:])
    ppaa = numpy.array(eri[:,:,ncore:nocc,ncore:nocc])
    papa = numpy.array(eri[:,ncore:nocc,:,ncore:nocc])
    jc_pp = numpy.einsum('iipq->ipq', eri[:ncore,:ncore,:,:])
    kc_pp = numpy.einsum('ipqi->ipq', eri[:ncore,:,:,:ncore])
    vhf_c = numpy.einsum('cij->ij', jc_pp)*2 - numpy.einsum('cij->ij', kc_pp)
    j_pc = numpy.einsum('ijj->ji', jc_pp)
    k_pc = numpy.einsum('ijj->ji', kc_pp)

    print('vhf_c', numpy.allclose(vhf_c, eris1.vhf_c))
    print('j_pc ', numpy.allclose(j_pc, eris1.j_pc))
    print('k_pc ', numpy.allclose(k_pc, eris1.k_pc))
    print('ppaa ', numpy.allclose(ppaa , eris0.ppaa ))
    print('papa ', numpy.allclose(papa , eris0.papa ))


#!/usr/bin/env python
# Computes the matrix of similarities between structures in a xyz file
# by first getting SOAP descriptors for all environments, finding the best
# match between environments using the Hungarian algorithm, and finally
# summing up the environment distances.
# Supports periodic systems, matching between structures with different
# atom number and kinds, and sports the infrastructure for introducing an
# alchemical similarity kernel to match different atomic species

# import sys, os, pickle
import sys, os
import cPickle as pickle
import gc
from lap.lap import best_pairs, best_cost, lcm_best_cost
from lap.perm import xperm, mcperm, rematch
import numpy as np
from environments import environ, alchemy, envk
import quippy
__all__ = [ "structk", "structure" ]

class structure:
   def __init__(self, salchem=None):
      self.env={}
      self.species={}
      self.zspecies = []
      self.atz = []
      self.nenv=0
      self.alchem=salchem
      if self.alchem is None: self.alchem=alchemy()
      self.globenv = None

   def getnz(self, sp):
      if sp in self.species:
         return self.species[sp]
      else: return 0

   def getatomenv(self, i):
      if i>=len(self.atz):
          raise IndexError("Trying to access atom past structure size")
      k=0
      lsp = {}
      for z in self.atz:
         if z in lsp: lsp[z]+=1
         else: lsp[z] = 0
         if i==k:
             return self.env[z][lsp[z]]
         k+=1

   def getenv(self, sp, i):
      if sp in self.env and i<len(self.env[sp]):
         return self.env[sp][i]
      else:
         return environ(self.nmax,self.lmax,self.alchem,sp)  # missing atoms environments just returned as isolated species!

   def ismissing(self, sp, i):
      if sp in self.species and i<self.species[sp]:
         return False
      else: return True

   def parse(self, fat, coff=5.0, cotw=0.5, nmax=4, lmax=3, gs=0.5, cw=1.0, nocenter=[], noatom=[], unsoap=False, kit=None, soapdump=None):
      """ Takes a frame in the QUIPPY format and computes a list of its environments. """

      # removes atoms that are to be ignored
      at = fat.copy()
      nol = []
      for s in range(1,at.z.size+1):
         if at.z[s] in noatom: nol.append(s)
      if len(nol)>0: at.remove_atoms(nol)

      self.nmax = nmax
      self.lmax = lmax
      self.atz = at.z.copy()
      self.species = {}
      for z in at.z:
         if z in self.species: self.species[z]+=1
         else: self.species[z] = 1

      self.zspecies = self.species.keys();
      self.zspecies.sort();
      lspecies = 'n_species='+str(len(self.zspecies))+' species_Z={ '
      for z in self.zspecies: lspecies = lspecies + str(z) + ' '
      lspecies = lspecies + '}'

      at.set_cutoff(coff);
      at.calc_connect();

      self.nenv = 0
      if not soapdump is None:
         soapdump.write("####### SOAP VECTOR FRAME ######\n")
      for sp in self.species:

         if sp in nocenter:
            self.species[sp]=0
            continue # Option to skip some environments

         # first computes the descriptors of species that are present
         if not soapdump is None: sys.stderr.write("SOAP STRING:    "+"soap central_reference_all_species=F central_weight="+str(cw)+"  covariance_sigma0=0.0 atom_sigma="+str(gs)+" cutoff="+str(coff)+" cutoff_transition_width="+str(cotw)+" n_max="+str(nmax)+" l_max="+str(lmax)+' '+lspecies+' Z='+str(sp)+"\n")
         desc = quippy.descriptors.Descriptor("soap central_reference_all_species=F "+("normalise=F" if unsoap else "")+" central_weight="+str(cw)+"  covariance_sigma0=0.0 atom_sigma="+str(gs)+" cutoff="+str(coff)+" cutoff_transition_width="+str(cotw)+" n_max="+str(nmax)+" l_max="+str(lmax)+' '+lspecies+' Z='+str(sp) )
         try:
            psp = desc.calc(at)["descriptor"]
         except TypeError:
            print("Interface change in QUIP/GAP. Update your code first.")

         if not soapdump is None:
            soapdump.write("Specie %d - %d atoms\n"% (sp,len(psp)))
            for p in psp:
                np.savetxt(soapdump,[p])

         # now repartitions soaps in environment descriptors
         lenv = []
         for p in psp:
            nenv = environ(nmax, lmax, self.alchem)
            nenv.convert(sp, self.zspecies, p, unsoap)
            lenv.append(nenv)
         self.env[sp] = lenv
         self.nenv += self.species[sp]

      # adds kit data
      if kit is None: kit = {}

      for sp in kit:
         if not sp in self.species:
            self.species[sp]=0
            self.env[sp] = []
         for k in range(self.species[sp], kit[sp]):
            self.env[sp].append(environ(self.nmax,self.lmax,self.alchem,sp))
            self.nenv+=1
         self.species[sp] = kit[sp]

      self.zspecies = self.species.keys()
      self.zspecies.sort()

      # also compute the global (flattened) fingerprint
      self.globenv = environ(nmax, lmax, self.alchem)

      for k, se in self.env.items():
         for e in se:
            self.globenv.add(e)
      # divides by the number of atoms in the structure
      for sij in self.globenv.soaps:  self.globenv.soaps[sij]*=1.0/self.nenv
      # self.globenv.normalize()  #if needed, normalization will be done later on.....

def gcd(a,b):
   if (b>a): a,b = b, a

   while (b):  a, b = b, a%b

   return a

def lcm(a,b):
   return a*b/gcd(b,a)

#def gstructk(strucA, strucB, alchem=alchemy(), peratom=False):
#
#   return envk(strucA.globenv, strucB.globenv, alchem)

def structk(strucA, strucB, alchem=alchemy(), peratom=False, mode="match", fout=None, peps=0.0, gamma=1.0, zeta=1.0, xspecies=False):
   # computes the SOAP similarity KERNEL between two structures by combining atom-centered kernels
   # possible kernel modes include:
   #   average :  scalar product between averaged kernels
   #   match:     best-match hungarian kernel
   #   permanent: average over all permutations

   # average kernel. quick & easy!
   if mode=="fastavg":
       genvA=strucA.globenv
       genvB=strucB.globenv
       return envk(genvA, genvB, alchem)**zeta, 0
   elif mode=="fastspecies":
       # for now, only implement standard Kronecker alchemy
       senvB = environ(strucB.nmax, strucB.lmax, strucB.alchem)
       kk = 0
       for za in strucA.zspecies:
         if not za in strucB.zspecies: continue
         senvA = environ(strucA.nmax, strucA.lmax, strucA.alchem)
         for ia in xrange(strucA.getnz(za)):
            senvA.add(strucA.getenv(za, ia))
         senvB = environ(strucB.nmax, strucB.lmax, strucB.alchem)
         for ib in xrange(strucB.getnz(za)):
            senvB.add(strucB.getenv(za, ib))
         kk += envk(senvA, senvB, alchem)**zeta

       kk/=strucA.nenv*strucB.nenv
       return kk,0

       #  for zb, nzb in nspeciesB:
       #         for ib in xrange(nzb):
       #            return envk(genvA, genvB, alchem), 0

   nenv = 0

   if peratom: # replicate structures to match structures of different peratomity
      # we do not check for compatibility at this stage, just assume that the
      # matching will be done somehow (otherwise it would be exceedingly hard to manage in case of non-standard alchemy)
      nspeciesA = []
      nspeciesB = []
      for z in strucA.zspecies:
         nspeciesA.append( (z, strucA.getnz(z)) )
      for z in strucB.zspecies:
         nspeciesB.append( (z, strucB.getnz(z)) )
      nenv=nenvA = strucA.nenv
      nenvB = strucB.nenv
   else:
      # top up missing atoms with isolated environments
      # first checks which atoms are present
      zspecies = sorted(list(set(strucB.zspecies+strucA.zspecies)))
      nspecies = []
      for z in zspecies:
         nz = max(strucA.getnz(z),strucB.getnz(z))
         nspecies.append((z,nz))
         nenv += nz
      nenvA = nenvB = nenv
      nspeciesA = nspeciesB = nspecies

   np.set_printoptions(linewidth=500,precision=4)

   kk = np.zeros((nenvA,nenvB),float)
   ika = 0
   ikb = 0
   for za, nza in nspeciesA:
      for ia in xrange(nza):
         envA = strucA.getenv(za, ia)
         ikb = 0
         for zb, nzb in nspeciesB:
            for ib in xrange(nzb):
               envB = strucB.getenv(zb, ib)
               if alchem.mu > 0 and (strucA.ismissing(za, ia) ^ strucB.ismissing(zb, ib)):
                   # includes a penalty dependent on "mu", in a way that is consistent with the definition of kernel distance
                   kk[ika,ikb] = exp(-alchem.mu)
               else:
                  if za == zb or not xspecies:  #uncomment to zero out kernels between different species
                    kk[ika,ikb] = envk(envA, envB, alchem)**zeta
                  else: kk[ika,ikb] = 0
               ikb+=1
         ika+=1
   aidx = {}
   ika=0
   for za, nza in nspeciesA:
      aidx[za] = range(ika,ika+nza)
      ika+=nza
   ikb=0
   bidx = {}
   for zb, nzb in nspeciesB:
      bidx[zb] = range(ikb,ikb+nzb)
      ikb+=nzb

   if fout != None:
      # prints out similarity information for the environment pairs
      fout.write("# atomic species in the molecules (possibly topped up with dummy isolated atoms): \n")
      for za, nza in nspeciesA:
         for ia in xrange(nza): fout.write(" %d " % (za) )
      fout.write("\n");
      for zb, nzb in nspeciesB:
         for ib in xrange(nzb): fout.write(" %d " % (zb) )
      fout.write("\n");

      fout.write("# environment kernel matrix: \n")
      for r in kk:
         for e in r:
            fout.write("%20.14e " % (e) )
         fout.write("\n")
      #fout.write("# environment kernel eigenvalues: \n")
      #ev = np.linalg.eigvals(kk)
      #for e in ev:
      #    fout.write("(%8.4e,%8.4e) " % (e.real,e.imag) )
      #fout.write("\n");

   # Now we have the matrix of scalar products.
   # We can first find the optimal scalar product kernel
   # we must find the maximum "cost"
   if mode == "match":
        if peratom and nenvA != nenvB:
            nenv = lcm(nenvA, nenvB)
            hun = lcm_best_cost(1-kk)
        else:
            hun=best_cost(1.0-kk)
        cost = 1-hun/nenv
   elif mode == "permanent":
        # there is no place to hide: cross-species environments are not necessarily zero
        if peps>0: cost = mcperm(kk, peps)
        else: cost = xperm(kk)

        cost = cost/np.math.factorial(nenv)/nenv
   elif mode == "rematch":
       cost=rematch(kk, gamma, 1e-6)  # hard-coded residual error for regularized gamma
       # print cost, kk.sum()/(nenv*nenv), envk(strucA.globenv, strucB.globenv, alchem)
   elif mode == "average":
       cost = kk.sum()/(nenvA*nenvB)
       # print 'elem: {}'.format(kk.sum())
       # print 'elem norm: {}'.format(cost)
       # print 'avg norm: {}'.format((nenvA*nenvB))

   else: raise ValueError("Unknown global fingerprint mode ", mode)

   return cost,kk

class structurelist(list):
    def __init__(self, basedir="tmpstructures"):
        self.basedir=basedir
        # create the folder if it is not there
        if not os.path.exists(basedir):os.makedirs(basedir)
        self.count=0

    def exists(self, index):
        # return true if the file associated with index exists, false otherwise
        f=self.basedir+'/sl_'+str(index)+'.dat'
        return os.path.isfile(f)
    # @profile
    def append(self, element):
        #pickle the element for later use
        ind=self.count
        f=self.basedir+'/sl_'+str(ind)+'.dat'
        file = open(f,"wb")
        gc.disable()
        pickle.dump(element, file,protocol=pickle.HIGHEST_PROTOCOL) # HIGHEST_PROTOCOL is 2 in py 2.7
        file.close()
        gc.enable()
        self.count+=1

    # @profile
    def __getitem__(self, index):
        f = self.basedir+'/sl_'+str(index)+'.dat'
        try:
            file = open(f,"rb")
        except IOError:
            raise IOError("Cannot load descriptors for index %d" % (index) )
        gc.disable()
        l = pickle.load(file)
        file.close()
        gc.enable()
        return l


"""Contains the classes that deal with the normal mode representation.

Copyright (C) 2013, Joshua More and Michele Ceriotti

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program. If not, see <http.//www.gnu.org/licenses/>.

Deals with the normal mode transformation, including the complications
introduced by PA-CMD when the bead masses are rescaled. Also deals with
the change in the dynamics introduced by this mass-scaling, and has its
own functions to calculate the kinetic energy, and the exact propagator
in the normal mode representation under the ring polymer Hamiltonian.

Classes:
   NormalModes: Deals with the normal mode transformation in RPMD and PA-CMD.
"""

import numpy as np
from ipi.utils.depend import *
from ipi.utils import units
from ipi.utils import nmtransform
from ipi.utils.messages import verbosity, warning, info

__all__ = [ "NormalModes" ]

class NormalModes(dobject):
   """ A helper class to manipulate the path NM.

   Normal-modes transformation, determination of path frequencies,
   dynamical mass matrix change, etc.

   Attributes:
      natoms: The number of atoms.
      nbeads: The number of beads.
      beads: The beads object for which the normal mode transformation should
         be done.
      ensemble: The ensemble object, specifying the temperature to hold the
         system to.
      transform: A nm_trans object that contains the functions that are
         required for the normal mode transformation.

   Depend objects:
      mode: A string specifying how the bead masses are chosen.
      transform_method: A string specifying how to do the normal mode
         transformation.
      nm_freqs: An array that specifies how the normal mode frequencies
         of the ring polymers are to be calculated, and thus how the
         bead masses should be chosen.
      qnm: The bead positions in the normal mode representation. Depends on
         beads.q.
      pnm: The bead momenta in the normal mode representation. Depends on
         beads.p.
      omegan: The effective vibrational frequency for the interaction
         between the replicas. Depends on the simulation temperature.
      omegan2: omegan**2.
      omegak: The normal mode frequencies for the free ring polymer.
         Depends on omegan.
      prop_pq: An array holding the exact normal mode propagator for the
         free ring polymer, using mass scaled coordinates.
         See J. Chem. Phys. 133, 124101 (2010). Depends on the bead masses
         and the timestep.
      nm_factor: An array of dynamical mass factors associated with each of
         the normal modes. Depends on nm_freqs and mode.
      dynm3: An array that gives the dynamical masses of individual atoms in the
         normal modes representation. Depends on nm_factor and beads.m3.
      dynomegak: The scaled vibrational frequencies. Depends on nm_factor and
         omegak.
      kins: A list of the kinetic energy for each normal mode, as
         calculated in the normal mode representation, using the
         dynamical mass factors. Depends on beads.sm3, beads.p and nm_factor.
      kin: The total kinetic energy, as calculated in the normal mode
         representation, using the dynamical mass factors.
      kstress: The kinetic stress tensor, as calculated in the normal mode
         representation, using the dynamical mass factors. Depends on
         beads.sm3, beads.p and nm_factor.
   """

   def __init__(self, mode="rpmd", transform_method="fft", freqs=None):
      """Initializes NormalModes.

      Sets the options for the normal mode transform.

      Args:
         mode: A string specifying how to calculate the bead masses.
         transform_method: A string specifying how to do the normal mode
            transformation.
         freqs: A list of data used to calculate the dynamical mass factors.
      """

      if freqs is None:
         freqs = []
      dset(self,"mode",   depend_value(name='mode', value=mode))
      dset(self,"transform_method",
         depend_value(name='transform_method', value=transform_method))
      dset(self,"nm_freqs",
         depend_array(name="nm_freqs",value=np.asarray(freqs, float) ) )

   def bind(self, beads, ensemble):
      """ Initializes the normal modes object and binds to beads and ensemble.

      Do all the work down here as we need a full-formed necklace and ensemble
      to know how this should be done.

      Args:
         beads: A beads object to be bound.
         ensemble: An ensemble object to be bound.
      """

      self.nbeads = beads.nbeads
      self.natoms = beads.natoms

      # stores a reference to the bound beads and ensemble objects
      self.beads = beads
      self.ensemble = ensemble

      # sets up what's necessary to perform nm transformation.
      if self.transform_method == "fft":
         self.transform = nmtransform.nm_fft(nbeads=self.nbeads, natoms=self.natoms)
      elif self.transform_method == "matrix":
         self.transform = nmtransform.nm_trans(nbeads=self.nbeads)

      # creates arrays to store normal modes representation of the path.
      # must do a lot of piping to create "ex post" a synchronization between the beads and the nm
      sync_q = synchronizer()
      sync_p = synchronizer()
      dset(self,"qnm",
         depend_array(name="qnm",
            value=np.zeros((self.nbeads,3*self.natoms), float),
               func={"q": (lambda : self.transform.b2nm(depstrip(self.beads.q)) ) },
                  synchro=sync_q ) )
      dset(self,"pnm",
         depend_array(name="pnm",
            value=np.zeros((self.nbeads,3*self.natoms), float),
               func={"p": (lambda : self.transform.b2nm(depstrip(self.beads.p)) ) },
                  synchro=sync_p ) )

      # must overwrite the functions
      dget(self.beads, "q")._func = { "qnm": (lambda : self.transform.nm2b(depstrip(self.qnm)) )  }
      dget(self.beads, "p")._func = { "pnm": (lambda : self.transform.nm2b(depstrip(self.pnm)) )  }
      dget(self.beads, "q").add_synchro(sync_q)
      dget(self.beads, "p").add_synchro(sync_p)

      # also within the "atomic" interface to beads
      for b in range(self.nbeads):
         dget(self.beads._blist[b],"q")._func = { "qnm": (lambda : self.transform.nm2b(depstrip(self.qnm)) )  }
         dget(self.beads._blist[b],"p")._func = { "pnm": (lambda : self.transform.nm2b(depstrip(self.pnm)) )  }
         dget(self.beads._blist[b],"q").add_synchro(sync_q)
         dget(self.beads._blist[b],"p").add_synchro(sync_p)

      # finally, we mark the beads as those containing the set positions
      dget(self.beads, "q").update_man()
      dget(self.beads, "p").update_man()

      # create path-frequencies related properties
      dset(self,"omegan",
         depend_value(name='omegan', func=self.get_omegan,
            dependencies=[dget(self.ensemble,"temp")]) )
      dset(self,"omegan2", depend_value(name='omegan2',func=self.get_omegan2,
            dependencies=[dget(self,"omegan")]) )
      dset(self,"omegak", depend_array(name='omegak',
         value=np.zeros(self.beads.nbeads,float),
            func=self.get_omegak, dependencies=[dget(self,"omegan")]) )

      # sets up "dynamical" masses -- mass-scalings to give the correct RPMD/CMD dynamics
      dset(self,"nm_factor", depend_array(name="nmm",
         value=np.zeros(self.nbeads, float), func=self.get_nmm,
            dependencies=[dget(self,"nm_freqs"), dget(self,"mode") ]) )
      dset(self,"dynm3", depend_array(name="dm3",
         value=np.zeros((self.nbeads,3*self.natoms), float),func=self.get_dynm3,
            dependencies=[dget(self,"nm_factor"), dget(self.beads, "m3")] ) )
      dset(self,"dynomegak", depend_array(name="dynomegak",
         value=np.zeros(self.nbeads, float), func=self.get_dynwk,
            dependencies=[dget(self,"nm_factor"), dget(self,"omegak") ]) )

      dset(self,"prop_pq",
         depend_array(name='prop_pq',value=np.zeros((self.beads.nbeads,2,2)),
            func=self.get_prop_pq,
               dependencies=[dget(self,"omegak"), dget(self,"nm_factor"), dget(self.ensemble,"dt")]) )

      # if the mass matrix is not the RPMD one, the MD kinetic energy can't be
      # obtained in the bead representation because the masses are all mixed up
      dset(self,"kins",
         depend_array(name="kins",value=np.zeros(self.nbeads, float),
            func=self.get_kins,
               dependencies=[dget(self,"pnm"), dget(self.beads,"sm3"), dget(self, "nm_factor") ] ))
      dset(self,"kin",
         depend_value(name="kin", func=self.get_kin,
            dependencies=[dget(self,"kins")] ))
      dset(self,"kstress",
         depend_array(name="kstress",value=np.zeros((3,3), float),
            func=self.get_kstress,
               dependencies=[dget(self,"pnm"), dget(self.beads,"sm3"), dget(self, "nm_factor") ] ))

   def get_omegan(self):
      """Returns the effective vibrational frequency for the interaction
      between replicas.
      """

      return self.ensemble.temp*self.nbeads*units.Constants.kb/units.Constants.hbar

   def get_omegan2(self):
      """Returns omegan**2."""

      return self.omegan**2

   def get_omegak(self):
      """Gets the normal mode frequencies.

      Returns:
         A list of the normal mode frequencies for the free ring polymer.
         The first element is the centroid frequency (0.0).
      """

      return 2*self.omegan*np.array([np.sin(k*np.pi/self.nbeads) for k in range(self.nbeads)])

   def get_dynwk(self):
      """Gets the dynamical normal mode frequencies.

      Returns:
         A list of the scaled normal mode frequencies for the free ring polymer.
         The first element is the centroid frequency (0.0).
      """

      return self.omegak/np.sqrt(self.nm_factor)

   def get_prop_pq(self):
      """Gets the normal mode propagator matrix.

      Note the special treatment for the centroid normal mode, which is
      propagated using the standard velocity Verlet algorithm as required.
      Note that both the normal mode positions and momenta are propagated
      using this matrix.

      Returns:
         An array of the form (nbeads, 2, 2). Each 2*2 array prop_pq[i,:,:]
         gives the exact propagator for the i-th normal mode of the
         ring polymer.
      """

      dt = self.ensemble.dt
      pqk = np.zeros((self.nbeads,2,2), float)
      pqk[0] = np.array([[1,0], [dt,1]])

      for b in range(1, self.nbeads):
         sk = np.sqrt(self.nm_factor[b]) # NOTE THAT THE PROPAGATOR USES MASS-SCALED MOMENTA!

         dtomegak = self.omegak[b]*dt/sk
         c = np.cos(dtomegak)
         s = np.sin(dtomegak)
         pqk[b,0,0] = c
         pqk[b,1,1] = c
         pqk[b,0,1] = -s*self.omegak[b]*sk
         pqk[b,1,0] = s/(self.omegak[b]*sk)

      return pqk

   def get_nmm(self):
      """Returns dynamical mass factors, i.e. the scaling of normal mode
      masses that determine the path dynamics (but not statics)."""

      # also checks that the frequencies and the mode given in init are
      # consistent with the beads and ensemble

      dmf = np.zeros(self.nbeads,float)
      dmf[:] = 1.0
      if self.mode == "rpmd":
         if len(self.nm_freqs) > 0:
            warning("nm.frequencies will be ignored for RPMD mode.", verbosity.low)
      elif self.mode == "manual":
         if len(self.nm_freqs) != self.nbeads-1:
            raise ValueError("Manual path mode requires (nbeads-1) frequencies, one for each internal mode of the path.")
         for b in range(1, self.nbeads):
            sk = self.omegak[b]/self.nm_freqs[b-1]
            dmf[b] = sk**2
      elif self.mode == "pa-cmd":
         if len(self.nm_freqs) > 1:
            warning("Only the first element in nm.frequencies will be considered for PA-CMD mode.", verbosity.low)
         if len(self.nm_freqs) == 0:
            raise ValueError("PA-CMD mode requires the target frequency of all the internal modes.")
         for b in range(1, self.nbeads):
            sk = self.omegak[b]/self.nm_freqs[0]
            info(" ".join(["NM FACTOR", str(b), str(sk), str(self.omegak[b]), str(self.nm_freqs[0])]), verbosity.medium)
            dmf[b] = sk**2
      elif self.mode == "wmax-cmd":
         if len(self.nm_freqs) > 2:
            warning("Only the first two element in nm.frequencies will be considered for WMAX-CMD mode.", verbosity.low)
         if len(self.nm_freqs) < 2:
            raise ValueError("WMAX-CMD mode requires [wmax, wtarget]. The normal modes will be scaled such that the first internal mode is at frequency wtarget and all the normal modes coincide at frequency wmax.")
         wmax = self.nm_freqs[0]
         wt = self.nm_freqs[1]
         for b in range(1, self.nbeads):
            sk = 1.0/np.sqrt((wt)**2*(1+(wmax/self.omegak[1])**2)/(wmax**2+(self.omegak[b])**2))
            dmf[b] = sk**2

      return dmf

   def get_dynm3(self):
      """Returns an array with the dynamical masses of individual atoms in the normal modes representation."""

      dm3 = np.zeros(self.beads.m3.shape,float)
      for b in range(self.nbeads):
         dm3[b] = self.beads.m3[b]*self.nm_factor[b]

      return dm3

   def free_qstep(self):
      """Exact normal mode propagator for the free ring polymer.

      Note that the propagator works in mass scaled coordinates, so that the
      propagator matrix can be determined independently from the particular
      atom masses, and so the same propagator will work for all the atoms in
      the system. All the ring polymers are propagated at the same time by a
      matrix multiplication.

      Also note that the centroid coordinate is propagated in qcstep, so is
      not altered here.
      """

      if self.nbeads == 1:
         pass
      else:
         pq = np.zeros((2,self.natoms*3),float)
         sm = depstrip(self.beads.sm3)[0]
         prop_pq = depstrip(self.prop_pq)
         for k in range(1,self.nbeads):
            pq[0,:] = depstrip(self.pnm)[k]/sm
            pq[1,:] = depstrip(self.qnm)[k]*sm
            pq = np.dot(prop_pq[k],pq)
            self.qnm[k] = pq[1,:]/sm
            self.pnm[k] = pq[0,:]*sm

   def get_kins(self):
      """Gets the MD kinetic energy for all the normal modes.

      Returns:
         A list of the kinetic energy for each NM.
      """

      kmd = np.zeros(self.nbeads,float)
      sm = depstrip(self.beads.sm3[0])
      pnm = depstrip(self.pnm)
      nmf = depstrip(self.nm_factor)

      # computes the MD ke in the normal modes representation, to properly account for CMD mass scaling
      for b in range(self.nbeads):
         sp = pnm[b]/sm                      # mass-scaled momentum of b-th NM
         kmd[b] = np.dot(sp,sp)*0.5/nmf[b]   # include the partially adiabatic CMD mass scaling

      return kmd

   def get_kin(self):
      """Gets the total MD kinetic energy.

      Note that this does not correspond to the quantum kinetic energy estimate
      for the system.

      Returns:
         The sum of the kinetic energy of each NM in the path.
      """

      return self.kins.sum()

   def get_kstress(self):
      """Calculates the total MD kinetic stress tensor.

      Note that this does not correspond to the quantum kinetic stress tensor
      estimate for the system.

      Returns:
         The sum of the MD kinetic stress tensor contributions from each NM.
      """

      kmd = np.zeros((3,3),float)
      sm = depstrip(self.beads.sm3[0])
      pnm = depstrip(self.pnm)
      nmf = depstrip(self.nm_factor)

      for b in range(self.nbeads):
         sp = pnm[b]/sm  # mass-scaled momentum of b-th NM

         for i in range(3):
            for j in range(3):
               # computes the outer product of the p of various normal modes
               # singling out Cartesian components to build the tensor
               # also takes care of the possibility of having non-RPMD masses
               kmd[i,j] += np.dot(sp[i:3*self.natoms:3],sp[j:3*self.natoms:3])/nmf[b]

      return kmd