from trajectory import *
import sys
#from abc import ABC, abstractmethod
import numpy as np

from controller import Supervisor,Node
from gym import spaces
from gym.utils import seeding
import math

from agent import Agent

class robotisImitationEnv():
    def __init__(self, timestep=None, filepath="trajectory_symmetry.bin",actionpath="motion_symmetry.bin"):
        # z ---- go forward , y ---- up axis, x --- left axis
        self.supervisor = Supervisor()
        self.trajectory = np.fromfile(filepath).reshape((-1, 35))
        print(self.trajectory.shape)
        self.ref_action = np.fromfile(actionpath).reshape((-1, 20))
        print(self.ref_action.shape)
        self.rotation = np.array([[-0.9999999989699976, 3.0327660728272706e-06, -4.5285835280544456e-05, -0.21857917589470194],
[-0.9999999991019061, 3.367143945265502e-06, -4.2247486701443e-05, -0.21860266807156137],
[-0.9999998580559774, 0.00017000032127180363, 0.0005049632817574539, -0.22321864518970114],
[-0.9998194462326855, 0.001922925556407621, 0.018904425203387986, -0.2218460974394333],
[-0.9930211732134864, 0.01019052890403352, 0.11749511764480679, -0.2182865210768348],
[-0.9800772159269385, 0.017511353607993694, 0.19784338077306768, -0.21933789915629365],
[-0.9564550329478053, 0.028429927686300743, 0.2904918400929093, -0.22404976217865255],
[-0.9205124028564287, 0.04300157342800492, 0.38833977502980577, -0.23342608796493147],
[-0.8728922817243131, 0.060112987575936014, 0.48419571789806015, -0.24850459969493885],
[-0.8662680311951013, 0.06645400533803839, 0.4951399431513224, -0.25997894774004116],
[-0.8656229176946545, 0.07286883697382376, 0.4953656194770442, -0.26252673039876995],
[-0.8764280736200449, 0.07710372993451749, 0.47531973091892804, -0.25852054178131084],
[-0.9003241884316204, 0.07854707368721672, 0.4280732565112178, -0.24992776536994296],
[-0.936497723154842, 0.07619162929951007, 0.3422964360762219, -0.23824116195343267],
[-0.9679308416354844, 0.07052119648462457, 0.24111542185683563, -0.23015530530944023],
[-0.9934397530887789, 0.05758377152786723, 0.09880063886196679, -0.2238437237812888],
[-0.9977519339652996, 0.03823035023022427, -0.05504106276025859, -0.22353331965586368],
[-0.9781336101031687, 0.013923676686779081, -0.20751089613334586, -0.2295286595677816]])
        
        self._seed()
        if timestep is None:
            self.timestep = int(self.supervisor.getBasicTimeStep())
        else:
            self.timestep = timestep
        self.robot = self.supervisor.getFromDef("ROBOT")
        self.leftfoot = self.supervisor.getFromDef("ROBOT.LeftFoot")
        self.rightfoot = self.supervisor.getFromDef("ROBOT.RightFoot")
        self.GetLeftFootPosition = self.leftfoot.getField("translation")
        self.GetRightFootPosition = self.rightfoot.getField("translation")
        self.motorNames = ['ShoulderR', 'ShoulderL', 'ArmUpperR', 'ArmUpperL','ArmLowerR', 'ArmLowerL', 'PelvYR', 'PelvYL',
                       'PelvR', 'PelvL', 'LegUpperR', 'LegUpperL', 'LegLowerR', 'LegLowerL', 'AnkleR', 'AnkleL','FootR', 'FootL', 'Neck', 'Head']
        self.controlled_motorNames = ['ShoulderR', 'ShoulderL', 'PelvR', 'PelvL', 'LegUpperR', 'LegUpperL', 'LegLowerR', 'LegLowerL', 'AnkleR', 'AnkleL','FootR', 'FootL']
        self.uncontrolled_motorNames = ['ArmUpperR', 'ArmUpperL','ArmLowerR', 'ArmLowerL', 'PelvYR', 'PelvYL', 'Neck', 'Head']
        self.motors = []
        self.positionSensors = []
        self.old_motor_position = None

        for i in range(0, len(self.motorNames)):
            self.positionSensors.append(self.supervisor.getPositionSensor(self.motorNames[i] + 'S'))
            self.positionSensors[i].enable(self.timestep)

        for j in range(0, len(self.motorNames)):
            self.motors.append(self.supervisor.getMotor(self.motorNames[j]))

        self.translation_field = self.robot.getField("translation")
        #print(self.translation_field)
        self.rotation_field = self.robot.getField("rotation")

        self.observationDim = 71
        high = np.array([np.inf] * 71)
        self.observation_space = spaces.Box(-high, high, dtype=np.float32)
        ah = np.array([3.14, 2.85, 0.1, 0.02, 0.45, 1.68, 1.5, 0.03, 1.38, 1.22, 0.2, 0.1])
        al = np.array([-3.14, -3.14, -0.02, -0.1, -1.77, -0.5, -0.02, -1.5, -1.24, -1.39, -0.1, -0.2])
        self.action_high = np.array([3.14, 2.85, 2.3, 0.77, 1.16, 1.63, 0.66, 2.5, 0.1, 0.02, 0.45, 1.68, 1.50, 0.03, 1.38, 1.22, 0.2, 0.1, 1.81, 0.94])
        self.action_low = np.array([-3.14, -3.14, -0.68, -2.25, -1.65, -1.18, -2.42, -0.69, -0.02, -0.1, -1.77, -0.5, -0.02, -1.50, -1.24, -1.39,  -0.1, -0.2, -1.81, -0.36])
        #leglowerLS -2.25--1.5,leglowerRS 2.25-1.5,PelvR', 'PelvL'  among -0.5-0.5 'FootR', 'FootL -0.2-0.3 -0.3-0.2 
        self.action_space = spaces.Box(al, ah, dtype=np.float32)
        
        self.com_pos = np.array([0, 0, 0])
        self.com_pos_his = np.array([0, 0, 0])
        self.single_action_num = 1
        
        self.time = 0
        self.phase = 0
        self.iter = 0
        self.counter = 0
        self.time_limit = 4800
        self.bug = False

    def _seed(self,seed=9):
        self.np_random,seed=seeding.np_random(seed)
        print(seed)
        return [seed]

    def _seed(self,seed=9):
        self.np_random,seed=seeding.np_random(seed)
        print(seed)
        return [seed]

    def reset(self):
        #init = [-0.8179579106376571, 0.6989832003767069, -0.30190405558037775, 0.29744219697048097, 0.4952370439753212, -0.49969957190596326, 4.025495172224157e-06, -4.025840085326092e-06, 0.0033553072962370094, -0.003344802532738738, -0.6112416889513878, 0.6112416840086581, 0.8997598716180006, -0.8997598740004444, 0.5071369804878852, -0.5071369796685542, 0.013384907840742132, -0.013384907796730895, 4.6439830495366814e-08, 0.16805433609950954]
        if self.bug:
        #if self.test:
            self.supervisor.simulationReset()
            self.bug = False
            print('reset')
        #self.supervisor.simulationReset()
        self.phase = random.randint(0, 17)
        self.time = 0
        self.counter = 0
        self.iter = 0
        ref_obs = self.get_kin_state() 
        init_motor_pos = ref_obs[15:]
        init_pos = ref_obs[0:3]
        init_rotation = self.rotation[self.phase]
        #qpos = self.get_kin_state()
        #self.translation_filed.setSFVec3f(init_pos)
        #print(self.phase)
        
        #self.translation_field.setSFVec3f(list([0,0.35,1]))
        #self.rotation_field.setSFRotation(list([0,1,0,0.4]))
        for i in range(0, len(self.motorNames)):
            self.motors[i].setPosition(init_motor_pos[i])
        cnt = 0
        while self.supervisor.step(self.timestep) != -1:
            cnt = cnt + 1
            if cnt == 10:
                break
        #print('reset ok')
        self.translation_field.setSFVec3f(list(init_pos))        
        self.rotation_field.setSFRotation(list(init_rotation))
        self.supervisor.simulationResetPhysics()
        #self.phase = 0
        
        return self.get_observations()

    def step(self, action):
        cnt = 0
        self.stepSimulation(action)
        # self.supervisor.step(self.timestep)
        while self.supervisor.step(self.timestep) != -1:
            cnt = cnt + 1
            if cnt == self.single_action_num:
                break
        self.time += 1
        self.iter += 1
        if self.time%4==0:
            self.phase += 1
            self.iter = 0

        if (self.phase >= 18):
            self.phase = 0
            self.iter = 0
            self.counter +=1

        next_observations = self.get_observations()
        if self.bug:
            rewards = 0
            done = True
            info = np.array([])
        else:
            rewards = self.get_reward(action)
            done = self.is_done()
            info = self.get_info()
            
        
        self.com_pos_his = np.array(self.com_pos)
        return next_observations, rewards, done, info


    def stepSimulation(self, action):
        ref_act = self.get_ref_action()
        #target = action + ref_act
        #cancel action2 and action3, have a try
        if self.phase < 9 :
            act = [action[0],action[1],0,0,0,0,0,0,action[2],action[3],action[4],action[5],action[6],action[7],action[8],action[9], action[10], action[11], 0, 0]
            target = act + ref_act
        else:
            mirror_action = [-action[1],-action[0],0,0,0,0,0,0,-action[3],-action[2],-action[5],-action[4],-action[7],-action[6],-action[9],-action[8], -action[11], -action[10], 0, 0]
            target = mirror_action + ref_act


        for i in range(0, len(self.motorNames)):
            target[i] = np.clip(target[i], self.action_low[i] + 0.01, self.action_high[i] - 0.01)
            # print(action[i].item())
            self.motors[i].setPosition(target[i].item())

    def get_reward(self,action):
        sim_obs = self.get_sim_observations()
        ref_obs = self.get_kin_state() 
        joint_penalty = 0
        weight = [0.05, 0.05, 0.09, 0.09, 0.09, 0.09, 0.09, 0.09, 0.09, 0.09, 0.09, 0.09]
        joint_index = [15, 16, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32]
        for i in range(12):
            error = weight[i] * (ref_obs[joint_index[i]]-sim_obs[joint_index[i]])**2
            joint_penalty += error*30 

        com_penalty = (ref_obs[13] - sim_obs[13])**2 + (0 - sim_obs[12])**2
        orientation_penalty = (ref_obs[3] - sim_obs[3])**2 + (ref_obs[4] - sim_obs[4])**2 + (ref_obs[5] - sim_obs[5])**2
        imitation_reward = 0.6*np.exp(-joint_penalty)+0.3*np.exp(-com_penalty)+0.1*np.exp(-orientation_penalty)
        sim_reward = 50 * min((self.com_pos[2] - self.com_pos_his[2]),0.02)
        #vel = self.com_pos[2] - self.com_pos_his[2]
        #vel_error = 100000*(vel - 0.0025)**2
        #sim_reward = np.exp(-vel_error)
        #print(sim_reward)
        AxisAngle = self.rotation_field.getSFRotation()
        orientation = self.quaternionFromAxisAngle(AxisAngle)
        euler = self.GetEulerFromOrientation(orientation[3], orientation[0], orientation[1], orientation[2])
        #print(euler)
        if (euler[1]>20 or euler[1]<-20):
            sim_reward -=1
        """
        sim_reward+=np.exp(-abs(euler[1])/30)-1
        """
        total_reward = 0.62*imitation_reward+0.38*sim_reward
        """
        if self.is_done():
            total_reward = 0
        else:
            total_reward += 0.2
        """
        return total_reward


    def is_done(self):
        # while (True):
        #     print(self.accelerometer.getValues())
        AxisAngle = self.rotation_field.getSFRotation()
        orientation = self.quaternionFromAxisAngle(AxisAngle)
        euler = self.GetEulerFromOrientation(orientation[3], orientation[0], orientation[1], orientation[2])
        # print(euler)
        if self.time >= self.time_limit:
            return True

        if (self.robot.getCenterOfMass()[1] < 0.13):
            return True
        if (euler[0] > 45 or euler[0] < -45):
            return True
            # print('True')
        elif (euler[2] > 45 or euler[2] < -45):
            return True
            # print('True')
        else:
            return False

    def get_info(self):
        return np.array([])

    def get_observations(self):
        sim_obs = self.get_sim_observations()
        ref_obs = self.get_kin_state() 
        #0,1,2-pos, 3,4,5-euler, 6,7,8,9,10,11-velocity, 12,13,14-com 15..34-motor
        phase = np.array([self.phase])  
        
        #0-8 left leg 9-17 right
        if self.phase < 9:
            return np.concatenate([sim_obs, ref_obs, phase])
        else:
            ref_obs_copy = np.zeros(35)
            ref_obs_copy[0:15] = np.copy(ref_obs[0:15])

            ref_obs_copy[15] = -np.copy(ref_obs[16])
            ref_obs_copy[16] = -np.copy(ref_obs[15])
            ref_obs_copy[17:23] = np.copy(ref_obs[17:23])

            ref_obs_copy[23] = -np.copy(ref_obs[24])
            ref_obs_copy[24] = -np.copy(ref_obs[23])

            ref_obs_copy[25] = -np.copy(ref_obs[26])
            ref_obs_copy[26] = -np.copy(ref_obs[25])

            ref_obs_copy[27] = -np.copy(ref_obs[28])
            ref_obs_copy[28] = -np.copy(ref_obs[27])

            ref_obs_copy[29] = -np.copy(ref_obs[30])
            ref_obs_copy[30] = -np.copy(ref_obs[29])

            ref_obs_copy[31] = -np.copy(ref_obs[32])
            ref_obs_copy[32] = -np.copy(ref_obs[31])

            ref_obs_copy[33] = np.copy(ref_obs[33])
            ref_obs_copy[34] = np.copy(ref_obs[34])

            sim_obs_copy = np.zeros(35)
            sim_obs_copy[0:15] = np.copy(sim_obs[0:15])

            sim_obs_copy[15] = -np.copy(sim_obs[16])
            sim_obs_copy[16] = -np.copy(sim_obs[15])
            sim_obs_copy[17:23] = np.copy(sim_obs[17:23])

            sim_obs_copy[23] = -np.copy(sim_obs[24])
            sim_obs_copy[24] = -np.copy(sim_obs[23])

            sim_obs_copy[25] = -np.copy(sim_obs[26])
            sim_obs_copy[26] = -np.copy(sim_obs[25])

            sim_obs_copy[27] = -np.copy(sim_obs[28])
            sim_obs_copy[28] = -np.copy(sim_obs[27])

            sim_obs_copy[29] = -np.copy(sim_obs[30])
            sim_obs_copy[30] = -np.copy(sim_obs[29])

            sim_obs_copy[31] = -np.copy(sim_obs[32])
            sim_obs_copy[32] = -np.copy(sim_obs[31])

            sim_obs_copy[33] = np.copy(sim_obs[33])
            sim_obs_copy[34] = np.copy(sim_obs[34])
   
        # I think it is this not next ,so that they are combined as state as network input to get the delta action means: we need delta action to close the ref state
        return np.concatenate([sim_obs_copy, ref_obs_copy, phase])

    def get_sim_observations(self):
        observation = []
        motor_position = []
        # 7+6+20+20=53
        # x,y,z pos; x, y, z, w quaternion
        base_position = self.translation_field.getSFVec3f()
        base_orientation = self.rotation_field.getSFRotation()
        orientation = self.quaternionFromAxisAngle(base_orientation)
        base_orientation = self.GetEulerFromOrientation(orientation[3], orientation[0], orientation[1], orientation[2])
        observation.append(base_position[0])
        observation.append(base_position[1])
        observation.append(base_position[2])
        observation.append(base_orientation[0])
        observation.append(base_orientation[1])
        observation.append(base_orientation[2])

        # linear velocity--0,1,2; angular velocity--3,4,5
        base_velocity = self.robot.getVelocity()
        observation.append(base_velocity[0])
        observation.append(base_velocity[1])
        observation.append(base_velocity[2])
        observation.append(base_velocity[3])
        observation.append(base_velocity[4])
        observation.append(base_velocity[5])

        self.com_pos = self.robot.getCenterOfMass()
        observation.append(self.com_pos[0])
        observation.append(self.com_pos[1])
        observation.append(self.com_pos[2])

        for i in range(0, len(self.motorNames)):
            pos = self.positionSensors[i].getValue()
            motor_position.append(pos)
            observation.append(pos)
            if pos > self.action_high[i]+1 or pos < self.action_low[i]-1:
                self.bug = True
                print(i,pos)
                motor_position[-1] = 0
                observation[-1] = 0
                if i != len(self.motorNames)-1:
                    for j in range(i+1,len(self.motorNames)):
                        motor_position.append(0)
                        observation.append(0)
                break
            else:
                self.bug = False

        if self.old_motor_position is None:
            self.old_motor_position = motor_position
        

        self.old_motor_position = motor_position

        return np.array(observation)

    def get_ref_action(self):
        ref_act_now = self.ref_action[self.phase]
        if self.phase == 17:
            next_phase = 0
        else:
            next_phase = self.phase+1
        ref_act_next = self.ref_action[next_phase]
        ref_act = ((ref_act_next-ref_act_now)/4)*self.iter + ref_act_now
        return ref_act

    def get_kin_state(self):
        if self.counter==0:
            obs = np.copy(self.trajectory[self.phase])
            obs += (np.copy(self.trajectory[self.phase+1])-np.copy(self.trajectory[self.phase]))*self.iter/4
        elif self.counter==1:
            obs = np.copy(self.trajectory[self.phase+18])
            obs += (np.copy(self.trajectory[self.phase+19])-np.copy(self.trajectory[self.phase+18]))*self.iter/4
        else:
            obs = np.copy(self.trajectory[self.phase+18])
            obs += (np.copy(self.trajectory[self.phase+19])-np.copy(self.trajectory[self.phase+18]))*self.iter/4

            obs[2] += (self.trajectory[53, 2]- self.trajectory[35, 2])* (self.counter-1)
            obs[14] += (self.trajectory[53, 14]- self.trajectory[35, 14])* (self.counter-1)
        
        return obs

    def get_kin_next_state(self):
        phase = self.phase + 1
        if phase >= 18:
            phase = 0
            counter = self.counter+1
        if counter==0:
            obs = np.copy(self.trajectory[phase])
        elif counter==1:
            obs = np.copy(self.trajectory[phase+18])
        else:
            obs = np.copy(self.trajectory[phase+18])
            obs[2] += (self.trajectory[53, 2]- self.trajectory[35, 2])* (self.counter-2)
            obs[14] += (self.trajectory[53, 14]- self.trajectory[35, 14])* (self.counter-2)
        return obs

    def quaternionFromAxisAngle(self, rotation):
        axis = [rotation[0], rotation[1], rotation[2]]
        angle = [rotation[3]]
        theta = angle[0] / 2
        w = np.cos(np.array(theta))
        # print(w)
        x = axis[0] * np.sin(np.array(theta))
        y = axis[1] * np.sin(np.array(theta))
        z = axis[2] * np.sin(np.array(theta))
        quaternion = [x, y, z, w]
        return quaternion

    def GetEulerFromOrientation(self, w, x, y, z):
        r = math.atan2(2 * (w * x + y * z), 1 - 2 * (x * x + y * y))
        p = math.asin(2 * (w * y - z * x))
        y = math.atan2(2 * (w * z + x * y), 1 - 2 * (z * z + y * y))
        angleR = r * 180 / math.pi
        angleP = p * 180 / math.pi
        angleY = y * 180 / math.pi
        return [angleR, angleP, angleY]

    def close(self):
        pass

if __name__ == '__main__':
    env = robotisImitationEnv(timestep=None)
    agent = Agent(env, gamma = 0.995, lam = 0.98, kl_targ = 0.003, batch_size = 20, hid1_mult = 10, policy_logvar = -1.0,load_model = False)
    #env.is_done()
    # while env.supervisor.step(int(env.timestep)) != -1:
    #     print(env.accelerometer.getValues())
    agent.run_policy(episodes=5)
    episode = 0
    num_episodes = 200000
    batch_size = 20
    while episode < num_episodes:
        print(episode)
        trajectories = agent.run_policy(episodes=batch_size)
        episode += len(trajectories)
        agent.add_value(trajectories)  # add estimated values to episodes
        agent.add_disc_sum_rew(trajectories)  # calculated discounted sum of Rs
        agent.add_gae(trajectories)  # calculate advantage
        # concatenate all episodes into single NumPy arrays
        observes, actions, advantages, disc_sum_rew = agent.build_train_set(trajectories)
        # add various stats to training log:
        agent.log_batch_stats(observes, actions, advantages, disc_sum_rew, episode)
        agent.policy.update(observes, actions, advantages, agent.logger)  # update policy
        agent.val_func.fit(observes, disc_sum_rew, agent.logger)  # update value function
        agent.logger.write(display=True)  # write logger results to file and stdout
        if episode%10==0:
            agent.save()
    agent.logger.close()
    agent.policy.close_sess()
    agent.val_func.close_sess()