import numpy as np
import tensorflow as tf
# Helper Function------------------------------------------------------------------------------------------------------------
# Copies one set of variables to another.
# Used to set worker network parameters to those of global network.

def dlrelu(x, alpha=0.1):
  return tf.nn.relu(x) - alpha * tf.nn.relu(0.05-x) - (1 - alpha) *  tf.nn.relu(x-0.95) 

class RunningStats:

    def __init__(self):
        self.n = 0
        self.old_m = 0
        self.new_m = 0
        self.old_s = 0
        self.new_s = 0

    def clear(self):
        self.n = 0

    def push(self, x):
        self.n += 1

        if self.n == 1:
            self.old_m = self.new_m = x
            self.old_s = 0
        else:
            self.new_m = self.old_m + (x - self.old_m) / self.n
            self.new_s = self.old_s + (x - self.old_m) * (x - self.new_m)

            self.old_m = self.new_m
            self.old_s = self.new_s

    def mean(self):
        return self.new_m if self.n else 0.0

    def variance(self):
        return self.new_s / (self.n - 1) if self.n > 1 else 0.0

    def standard_deviation(self):
        return np.sqrt(self.variance())
        
    def normalize(self,x):
        self.push(x)
        return (x - self.mean()) / (self.standard_deviation()+1e-3) if self.n > 1 else x



# process state (the last 3 entires are obstacle info which should not be processed)
def process_state(s,s1,center=True,diff=0):
    s = np.asarray(s)
    s1 = np.asarray(s1)
    s_14 = (s1[22:36]-s[22:36]) / 0.01
    s_3 = (s1[38:]-s[38:]) / 0.01
    s = np.hstack((s1[:36],s_14,s1[36:],s_3))
    if diff == 0:
        s[-6:] = 0.0 # if diff = 0, then manully turn off all obstacles

    if center:
      # transform into all relative quantities
      x_pos = [1,22,24,26,28,30,32,34]
      y_pos = [i+1 for i in x_pos]
      for i in x_pos:
          s[i] -= s[18]
      for j in y_pos:
          s[j] -= s[19]
      
      x_vs = [i+14 for i in x_pos]
      x_vs[0] = 4
      y_vs = [i+1 for i in x_vs]
      for i in x_vs:
          s[i] -= s[20]
      for j in y_vs:
          s[j] -= s[21]
      # transform cm as origin
      s[18:22] = 0.0
        
    return s
        
def n_step_transition(episode_buffer,n_step,gamma):
    _,_,_,s1,done = episode_buffer[-1]
    s,action,_,_,_ = episode_buffer[-1-n_step]
    r = 0
    for i in range(n_step):
      r += episode_buffer[-1-n_step+i][2]*gamma**i
    return [s,action,r,s1,done]

def engineered_action(seed):
    test = np.ones(18)*0.05
    if seed < 0.5:
        test[0] = 0.3
        test[3] = 0.8
        test[4] = 0.5
        test[6] = 0.3
        test[8] = 0.8
        test[9] = 0.3
        test[11] = 0.5
        test[14] = 0.3
        test[17] = 0.5
    else:
        test[9] = 0.3
        test[12] = 0.8
        test[13] = 0.5
        test[15] = 0.3
        test[17] = 0.8
        test[0] = 0.3
        test[2] = 0.5
        test[3] = 0.3
        test[8] = 0.5
            
    return test

# [Hacked] the memory might always be leaking, here's a solution #58
# https://github.com/stanfordnmbl/osim-rl/issues/58 
# separate process that holds a separate RunEnv instance.
# This has to be done since RunEnv() in the same process result in interleaved running of simulations.

import opensim as osim
from osim.http.client import Client
from osim.env import *

import multiprocessing
from multiprocessing import Process, Pipe

def standalone_headless_isolated(conn,vis,seed,diff):
    e = RunEnv(visualize=vis)
    while True:
        try:
            msg = conn.recv()

            # messages should be tuples,
            # msg[0] should be string

            if msg[0] == 'reset':
                o = e.reset(difficulty=diff,seed=seed)
                conn.send(o)
            elif msg[0] == 'step':
                ordi = e.step(msg[1])
                conn.send(ordi)
            else:
                conn.close()
                del e
                return
        except:
            conn.close()
            del e
            raise

# class that manages the interprocess communication and expose itself as a RunEnv.
class ei: # Environment Instance
    def __init__(self,vis,seed,diff):
        self.pc, self.cc = Pipe()
        self.p = Process(
            target = standalone_headless_isolated,
            args=(self.cc,vis,seed,diff,)
        )
        self.p.daemon = True
        self.p.start()

    def reset(self):
        self.pc.send(('reset',))
        return self.pc.recv()

    def step(self,actions):
        self.pc.send(('step',actions,))
        try:
            return self.pc.recv()
        except :  
            print('Error in recv()')
            raise

    def __del__(self):
        self.pc.send(('exit',))
        #print('(ei)waiting for join...')
        self.p.join()
	try:
	    del self.pc
	    del self.cc
	    del self.p
	except:
	    raise