'''
===============================================================================
Library for easy control of NIRSPEC.

Notes:

Dependencies:______________________________________________
  Custom:
  - ktl         -- Library to use Keck Keywords
  - get_path    -- Function used to get the path where data has to be saved.

  Native:
  - numpy       --
  - time        --
  - fits        --

===============================================================================
Author:         Jacques-Robert Delorme (jdelorme@caltech.edu)
Origin Date:    01/02/2019
'''
############################### Import Library ################################

## Math Library
import numpy as np
## Time library
import time
## import system library
import sys
## Import library used to manipulate .fits file
from astropy.io import fits
## Import OS information
import os
## Import medfilt. Function use to remove the bad pixel
from scipy.signal import medfilt as medfilt
## KTL = Library to use Keck Keywords
import ktl
## Library used to plot graphics and show images
import matplotlib.pyplot as plt

## Function use to get the path where data should be saved
from Organization import get_path
## Sylvain shmlib (good one)
import sce_shmlib

##
from photutils import aperture_photometry
from photutils import CircularAperture
from photutils import CircularAnnulus
from photutils import DAOStarFinder
##
from astropy.stats import sigma_clipped_stats

shm_im_0   = sce_shmlib.shm('/tmp/Scam_Raw_Image.im.shm')
shm_im_1   = sce_shmlib.shm('/tmp/Scam_Reduced_Image.im.shm')
shm_im_2   = sce_shmlib.shm('/tmp/Scam_Cleaned_Image.im.shm')
shm_im_3   = sce_shmlib.shm('/tmp/Scam_Slit_Mask.im.shm')
#shm_spec_0 = sce_shmlib.shm('/tmp/Spec_Raw_Image.im.shm')
#shm_spec_1 = sce_shmlib.shm('/tmp/Spec_Reduced_Image.im.shm')

#################################### Scam #####################################

class Scam_cmds():

    ## Take image with Scam:
    def TakeImage(self, Coadds = '', tint = '', Save = False, filename = '', overwrite = True):
        # Check if instrument is connected
        if ktl.read('nscam','status') != 'Idle':
            print('Error in Scam.py: The instrument is not connected.')
        else:
            Coadds_ini   = self.get('coadds')
            tint_ini     = np.round(self.get('tint')/656.)*656
            filename_ini = self.get('lastfile')
            # Change the number of Coadds (if requested)
            if Coadds != '':
                if Coadds_ini != Coadds:
                    if Coadds >= 20 :
                        Coadds = 20
                        print('Too many Coadds (>20)')
                    ktl.write('nscam','coadds', Coadds)
            self.Coadds = self.get('coadds')
            # Change the integration time (if requested)
            if tint != '':
                if tint_ini != tint:
                    if tint >= 300000 :
                        tint = 300000
                        print('Exposure time too long (>5 min).')
                    ktl.write('nscam','itime', tint)
            self.tint = self.get('tint')
            # The Flag is used to make sure we do not have bad frames
            Flag = False
            while not Flag:
                # Check if instrument is ready
                while ktl.read('nscam','ready') != 'Yes': time.sleep(0.01)
                # Save time at the beginning of the exposure
                self.timeIni = time.time()
                # Take an image
                ktl.write('nscam','go',1)
                # Wait the end of the exposure
                while ktl.read('nscam','ready') != 'Yes': time.sleep(0.01)
                # Save time at the end of the exposure
                self.timeEnd = time.time()

                # Read the path + filename of the last image saved
                while self.get('lastfile') == filename_ini:
                    time.sleep(0.01)
                # Read the fits file saved in the nfiuserver
                hdu = fits.open(ktl.read('nscam','lastfile'))

                # Extract data and header
                im  = hdu[0].data
                hdr = hdu[0].header

                # Bad frames has weird values
                if np.max(np.abs(im)) < 40000.:
                    Flag = True
                else:
                    #print('Bad Frame!!!')
                    Flag = False

            # Return to inital number of coadds
            if Coadds_ini != self.Coadds:
                ktl.write('nscam','coadds', Coadds_ini)
            # Return to inital integration time
            if tint_ini != self.tint:
                ktl.write('nscam','itime', tint_ini)

            # Define the filename
            bkgd  = 'Tint_'
            bkgd += '%06d' %(np.round(self.tint,0))
            bkgd += 'ms_Coadds_' + '%02d' %(np.int(self.Coadds))
            self.bkgd_filename = bkgd
            # Get the path where calibration images are saved.
            Path, tmp = get_path('Scam/Calib')

            # Define the name of the file
            fullname = Path + bkgd + '.fits'

            if os.path.isfile(fullname):
                bkgd = fits.open(fullname)[0].data
                self.bkgd = True
            else:
                bkgd = 0.
                self.bkgd = False

            self.Slit_Mask = shm_im_3.get_data()

            self.im = np.zeros([3,np.shape(im)[0],np.shape(im)[1]])
            self.im[0,:,:] = im
            self.im[1,:,:] = im - bkgd
            self.im[2,:,:] = (np.abs(medfilt(self.im[1,:,:],3))+1)*self.Slit_Mask

            shm_im_0.set_data(self.im[0,:,:])
            shm_im_1.set_data(self.im[1,:,:])
            shm_im_2.set_data(self.im[2,:,:])

            # Update the header
            text_1 = 'Time at the beginning of the exposure.'
            text_2 = 'Time at the end of the exposure (Scam ready).'
            hdr.append(('','',''), end=True)
            hdr.append(('','- Timing Information   -',''), end=True)
            hdr.append(('','',''), end=True)
            hdr.append(('TimeIni', self.timeIni,text_1), end=True)
            hdr.append(('TimeEnd', self.timeEnd,text_2), end=True)
            hdr.append(('','',''), end=True)

            self.hdr = hdr
            # Save a local copy of the image on request
            if Save == True:
                # Get the path of the directory where the image has to be saved
                Path, tmp = get_path('Scam')
                # If no filename provided, use the one provide by get_path().
                if filename == '': filename = tmp
                # If the user provided a filename, verify if valid.
                else:
                    # Verify if filename is a string
                    if not isinstance(filename,(str)):
                        print('The filename must be a string.')

                    # Remove unexpected / at the beginning if needed
                    if filename[0] == '/': filename = filename[1:]

                    # Verify if filename contains only valid characters
                    tmp = set(filename)
                    if not tmp.intersection(set('!?/@#$%^&*()><}{+~`')) == set():
                        print('The filename contains unexpected character(s).')
                        return False
                # Full name of the image
                fullname = Path + filename + '.fits'
                # Create a Header Data Unit (HDU) based on the data to save.
                hdu = fits.PrimaryHDU(data = self.im, header = self.hdr)
                # Save the timestamps associated to the image
                hdu.writeto(fullname, overwrite = overwrite)
            # Return the image and the associated header
            return

    # =========================================================================
    def get(self, keyword):
        if keyword.lower() in ['coadds','coadd','coad','coads']:
            return np.double(ktl.read('nscam','coadds'))
        if keyword.lower() in ['itime','tint','exptime']:
            return np.double(ktl.read('nscam','itime'))
        if keyword.upper() in ktl.keywords('nscam'):
            return ktl.read('nscam',keyword)

    # =========================================================================
    def take_calibration_images(self, Coadds = '', tint = '', overwrite = True):
        ''' -------------------------------------------------------------------
        Description:
            Function to save Scam calibration images.
        Arguments:
            - nb_im	: (optional) Number of images per cube.
                        -> must be a number.
                        -> will be automatically converted as an int
                        -> min = 1.
                        -> max = 1000.
            - overwire	: (optional) Boolean to overwrite data

        Returns:
        Nominal Case:
            - path	: path and part of the name common to all files saved.
        In case of error:
            - False
        ------------------------------------------------------------------- '''
        # Get the path where images will be saved.
        Path, tmp = get_path('Scam/Calib')

        if tint   == '': tint   = self.get('tint'  )
        if Coadds == '': Coadds = self.get('Coadds')

        # Take the calibration images
        self.TakeImage(Coadds = Coadds, tint = tint)

        # Define the filename
        filename  = 'Tint_'
        filename += '%06d' %(np.round(self.tint,0))
        filename += 'ms_Coadds_' + '%02d' %(np.int(self.Coadds))

        # -- Save data
        # Define the name of the file
        fullname = Path + filename + '.fits'
        # Create a Header Data Unit (HDU) based on the data to save.
        hdu = fits.PrimaryHDU(data = self.im[0,:,:],header = self.hdr)
        # Save the image
        hdu.writeto(fullname, overwrite = overwrite)

        return

    # =========================================================================
    def take_all_calibration_images(self,):
        for i in [1,3,5]:
            for j in [1,2,4,8]:
                 self.take_calibration_images(Coadds = i, tint = 656*j)
        return

    # =========================================================================
    def slit_extraction(self,):

        self.Slit_Param = np.zeros([9])

        # Take an image with scam (long exposure)
        self.TakeImage(tint = 64*656)
        # Create a copy of the raw image acquired with Scam
        I = np.copy(self.im[0,:,:])
        # Reduce the image in order to extract the slit
        I[np.where(np.abs(I) > 1000)] = 0
        I[np.where(I != 0)] = 1
        # Determine position X and Y of the non-nul pixels
        X = np.where(I == 1)[0]
        Y = np.where(I == 1)[1]
        # Determine the number of non-nul pixels
        nb_pix = np.size(X)
        # Prepare a variable N. This variable will contain the patch number
        # associated to each non-nul pixel
        N = np.zeros([nb_pix])
        # Associate a patch number to each non-nul pixel.
        # If two pixels are next to each other, they got the same patch number.
        for i in np.arange(nb_pix):
            # Attribute a patch number to a non-nul pixel if it does not
            # already have one.
            if N[i] == 0: N[i] = np.max(N+1)
            for j in np.arange(np.size(X)):
                cdt_1 = (i != j)
                cdt_2 = (np.abs(X[i]-X[j]) < 2)
                cdt_3 = (np.abs(Y[i]-Y[j]) < 2)
                if cdt_1 and cdt_2 and cdt_3: N[j] = N[i]

        # Determine the number of independent patch
        nb_patch = np.int(np.max(N))
        # Determine the size of each patch
        patch_size = np.zeros([nb_patch])
        for i in np.arange(nb_patch): patch_size[i] = np.size(np.where(N == i))
        # Determine the list of non-nul pixels contain in the biggest patch
        Valid_pts = np.where(N == np.where(patch_size == max(patch_size))[0][0])[0]
        # Remove all non-null pixel not located in the bigest patch.
        I[:,:] *= 0
        I[X[Valid_pts],Y[Valid_pts]] = 1

        # If a null pixel is located iside the main patch consider it as a non
        # null pixel
        for i in np.arange(np.shape(I)[0])-2:
            for j in np.arange(np.shape(I)[1])-2:
                if I[i+1,j+1] == 0:
                    if np.sum(I[i:i+3,j:j+3]) >= 7:
                        I[i+1,j+1] = 1

        self.Slit_Mask = 1-I
        shm_im_3.set_data(self.Slit_Mask)

        # Direction and size of the slit
        # List of points contain into the slit
        list_of_point = np.where(I == 1)
        # Position (X & Y) of each point
        x = list_of_point[0]
        y = list_of_point[1]
        # Fit a curve to this cloud of points
        z = np.polyfit(x, y, 1)
        f = np.poly1d(z)

        # Determine number of point to extrapolate
        if np.int(np.max(x)-np.min(x)) > np.int(np.max(y)-np.min(y)):
            nb_point = np.int(np.max(x)-np.min(x))
        else:
            nb_point = np.int(np.max(y)-np.min(y))

        # Compute new x's and y's
        x_new = np.linspace(np.min(x),np.max(x),nb_point)
        y_new = f(x_new)

        # Compute the length of the slit in pixels
        Length  = np.sqrt((np.max(x)-np.min(x))**2 + (np.max(y)-np.min(y))**2 )
        # Compute the surfaceof the slit
        Surface = np.sum(I)
        # Compute the width of the slit
        Width   = Surface/Length

        fig_1 = plt.figure(num = 1)
        # Title of the fig
        plt.title('Image of NIRSPEC Slit')
        # Create a subplot for the image
        plt.subplot(111)
        plt.imshow(self.im[0,:,:],vmin=0, vmax=5000)
        # Modify the axis: one ticks every 128 pixels
        X_ticks = np.arange(0, 257, 64)
        plt.gca().set_xticks(X_ticks)
        plt.xticks(fontsize = 8)
        Y_ticks = np.arange(0, 257, 64)
        plt.gca().set_yticks(Y_ticks)
        plt.gca().invert_yaxis()
        plt.yticks(fontsize = 8)
        # Label of the axis
        plt.xlabel('Pixel', fontsize = 8)
        plt.ylabel('Pixel', fontsize = 8)
        # Adjust the size of the subplot to fit the color bar on the right side
        plt.subplots_adjust(bottom = 0.1, right=0.8, top=0.85)
        # Define the size of the colorbar
        cax = plt.axes([0.8,0.1,0.03,0.75])
        # Add a colorbar
        cbar = plt.colorbar(cax = cax)
        # Add the label
        cbar.set_label('Flux (ADU)', rotation=90,fontsize = 8)
        # Adjust the size of the tick and associated labels
        cbar.ax.tick_params(labelsize = 8, width = 1)

        ax      = fig_1.add_subplot(111)
        # Overplot the slit on the image
        plt.plot(y_new,x_new,'r',linewidth=0.5)
        # Print main information
        plt.text(10,25, 'Slit length = %05.2f pixels' %(Length) ,color='w',fontsize=10)
        plt.text(10,10, 'Slit width  = %05.2f pixels' %(Width)  ,color='w',fontsize=10)
        # Get the path where image must be saved
        Path, tmp = get_path('Scam/Slit')
        # Save the plot
        plt.savefig(Path + tmp + '.pdf' %(i), bbox_inches='tight', pad_inches=0.25, dpi=600)
        plt.show(fig_1)
        # Close the figure
        plt.close(fig_1)

        self.Slit_Param[0] = np.min(x)
        self.Slit_Param[1] = np.max(x)
        self.Slit_Param[2] = np.min(y)
        self.Slit_Param[3] = np.max(y)
        self.Slit_Param[4] = z[0]
        self.Slit_Param[5] = z[1]
        self.Slit_Param[6] = Length
        self.Slit_Param[7] = Width
        self.Slit_Param[8] = time.time()

        # --- Save Slit Mask
        # Create the name of the data
        fullname = Path + tmp + '_Slit_Mask.fits'
        # Create a Header Data Unit (HDU) based on the data to save.
        hdu = fits.PrimaryHDU(self.Slit_Mask)
        # Save the data
        hdu.writeto(fullname, overwrite = True)

        # --- Save Slit parameters
        # Create the name of the data
        fullname = Path + tmp + '_Slit_Parameters.fits'
        # Create a Header Data Unit (HDU) based on the data to save.
        hdu = fits.PrimaryHDU(self.Slit_Param)
        # Save the data
        hdu.writeto(fullname, overwrite = True)

    # =========================================================================
    def Photometry(self,ref_im = 'Last_Image', psf_fwhm = 3.):
        if isinstance(ref_im,(str)):
            if ref_im == 'Last_Image': ref_im = self.im[1,:,:]
        data = medfilt(ref_im,kernel_size = np.int(psf_fwhm))
        try:
            # Find source position and define apertures
            mean, median, std = sigma_clipped_stats(data, sigma = 5.0)
            # Find the psf position
            daofind = DAOStarFinder(fwhm=psf_fwhm, threshold=10.*std)
            # Remove background offset on the data
            sources = daofind(data - median)
            for col in sources.colnames:
                # for consistent table output
                sources[col].info.format = '%.8g'
            positions = (sources['xcentroid'], sources['ycentroid'])

            # Define apertures for encircled energy
            apertures = CircularAperture(positions, r=3*psf_fwhm)
            # Define apertures for the flux calibration
            # apertures_total = CircularAperture(positions, r=5*psf_fwhm)
            # Define apertures for the background
            annulus_apertures = CircularAnnulus(positions, r_in=10*psf_fwhm, r_out=30*psf_fwhm)

            apers      = [apertures,annulus_apertures]
            phot_table = aperture_photometry(data, apers)
            bkg_mean   = phot_table['aperture_sum_1']/annulus_apertures.area()
            bkg_sum    = bkg_mean * apertures.area()
            final_sum  = phot_table['aperture_sum_0'] - bkg_sum
            if not isinstance(final_sum[0],(int,float)): return [0,0,0]
            psf_X      = sources['xcentroid']
            psf_Y      = sources['ycentroid']
            if np.size(final_sum) > 1:
                flux_min = np.min(final_sum) -1
                Results = np.zeros([np.size(final_sum),3])
                for i in np.arange(np.size(final_sum)):
                   Results[i,0] = final_sum[np.where(final_sum == np.max(final_sum))]
                   Results[i,1] =     psf_X[np.where(final_sum == np.max(final_sum))]
                   Results[i,2] =     psf_Y[np.where(final_sum == np.max(final_sum))]
                   final_sum[np.where(final_sum == np.max(final_sum))] = flux_min
                return Results
            else:
                Results = np.zeros([1,3])
                Results[0,0],Results[0,1],Results[0,2] = final_sum, psf_X, psf_Y
                return Results
        except: np.zeros([1,3])

class Spec_cmds():

    ## Take image with Spec
    def TakeImage(self, Coadds = '', tint = '', Save = False, filename = '', overwrite = True):
        # Check if instrument is connected
        if ktl.read('nspec','status') != 'Idling':
            print('Error in Nirspec_cmds.py: Nirspec is probably not connected.')
        else:
            Coadds_ini   = self.get('coadds')
            tint_ini     = np.round(self.get('tint')/656.)*656
            filename_ini = self.get('lastfile')
            # Change the number of Coadds (if requested)
            if Coadds != '':
                if Coadds_ini != Coadds:
                    if Coadds >= 20 :
                        Coadds = 20
                        print('Too many Coadds (>20)')
                    ktl.write('nspec','coadds', Coadds)
            self.Coadds = self.get('coadds')
            # Change the integration time (if requested)
            if tint != '':
                if tint_ini != tint:
                    if tint >= 3600000 :
                        tint = 3600000
                        print('Exposure time too long (>1 hour).')
                    ktl.write('nspec','itime', tint)
            self.tint = self.get('tint')
            # The Flag is used to make sure we do not have bad frames
            Flag = False
            while not Flag:
                # Check if instrument is ready
                while ktl.read('nspec','ready') != '1': time.sleep(0.01)
                # Save time at the beginning of the exposure
                self.timeIni = time.time()
                # Take an image
                ktl.write('nspec','go',1)
                # Wait the end of the exposure
                while ktl.read('nspec','ready') != '1': time.sleep(0.01)
                # Save time at the end of the exposure
                self.timeEnd = time.time()

                # Read the path + filename of the last image saved
                while self.get('lastfile') == filename_ini:
                    time.sleep(0.01)
                # Read the fits file saved in the nfiuserver
                hdu = fits.open(ktl.read('nspec','lastfile'))

                # Extract data and header
                im  = hdu[0].data
                hdr = hdu[0].header

                # Bad frames has weird values
                if np.max(np.abs(im)) < 40000.:
                    Flag = True
                else:
                    #print('Bad Frame!!!')
                    Flag = False

            # Return to inital number of coadds
            if Coadds_ini != self.Coadds:
                ktl.write('nspec','coadds', Coadds_ini)
            # Return to inital integration time
            if tint_ini != self.tint:
                ktl.write('nspec','itime', tint_ini)

            # Define the filename
            bkgd  = 'Tint_'
            bkgd += '%06d' %(np.round(self.tint,0))
            bkgd += 'ms_Coadds_' + '%02d' %(np.int(self.Coadds))
            self.bkgd_filename = bkgd
            # Get the path where calibration images are saved.
            Path, tmp = get_path('Spec/Calib')

            # Define the name of the file
            fullname = Path + bkgd + '.fits'

            if os.path.isfile(fullname):
                bkgd = fits.open(fullname)[0].data
                self.bkgd = True
            else:
                bkgd = 0.
                self.bkgd = False

            self.im = np.zeros([2,np.shape(im)[0],np.shape(im)[1]])
            self.im[0,:,:] = im
            self.im[1,:,:] = im - bkgd

            shm_im_0.set_data(self.im[0,:,:])
            shm_im_1.set_data(self.im[1,:,:])

            # Update the header
            text_1 = 'Time at the beginning of the exposure.'
            text_2 = 'Time at the end of the exposure (Spec ready).'
            hdr.append(('','',''), end=True)
            hdr.append(('','- Timing Information   -',''), end=True)
            hdr.append(('','',''), end=True)
            hdr.append(('TimeIni', self.timeIni,text_1), end=True)
            hdr.append(('TimeEnd', self.timeEnd,text_2), end=True)
            hdr.append(('','',''), end=True)

            self.hdr = hdr
            # Save a local copy of the image on request
            if Save == True:
                # Get the path of the directory where the image has to be saved
                Path, tmp = get_path('Spec')
                # If no filename provided, use the one provide by get_path().
                if filename == '': filename = tmp
                # If the user provided a filename, verify if valid.
                else:
                    # Verify if filename is a string
                    if not isinstance(filename,(str)):
                        print('The filename must be a string.')

                    # Remove unexpected / at the beginning if needed
                    if filename[0] == '/': filename = filename[1:]

                    # Verify if filename contains only valid characters
                    tmp = set(filename)
                    if not tmp.intersection(set('!?/@#$%^&*()><}{+~`')) == set():
                        print('The filename contains unexpected character(s).')
                        return False
                # Full name of the image
                fullname = Path + filename + '.fits'
                # Create a Header Data Unit (HDU) based on the data to save.
                hdu = fits.PrimaryHDU(data = self.im, header = self.hdr)
                # Save the timestamps associated to the image
                hdu.writeto(fullname, overwrite = overwrite)
            # Return the image and the associated header
            return

    # =========================================================================
    def get(self, keyword):
        if keyword.lower() in ['coadds','coadd','coad','coads']:
            return np.double(ktl.read('nscam','coadds'))
        if keyword.lower() in ['itime','tint','exptime']:
            return np.double(ktl.read('nscam','itime'))
        if keyword.upper() in ktl.keywords('nscam'):
            return ktl.read('nscam',keyword)

    # =========================================================================
    def take_calibration_images(self, Coadds = '', tint = '', overwrite = True):
        # Get the path where images will be saved.
        Path, tmp = get_path('Scam/Calib')

        if tint   == '': tint   = self.get('tint'  )
        if Coadds == '': Coadds = self.get('Coadds')

        # Take the calibration images
        self.TakeImage(Coadds = Coadds, tint = tint)

        # Define the filename
        filename  = 'Tint_'
        filename += '%06d' %(np.round(self.tint,0))
        filename += 'ms_Coadds_' + '%02d' %(np.int(self.Coadds))

        # -- Save data
        # Define the name of the file
        fullname = Path + filename + '.fits'
        # Create a Header Data Unit (HDU) based on the data to save.
        hdu = fits.PrimaryHDU(data = self.im[0,:,:],header = self.hdr)
        # Save the image
        hdu.writeto(fullname, overwrite = overwrite)

        return


# -*- coding: utf-8 -*-
import gym
import numpy as np

import torch
import torch.optim as optim

from rllite import Base
from rllite.common import ActorCritic3,make_env,test_env2,compute_gae,SubprocVecEnv
from tensorboardX import SummaryWriter

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

class PPO(Base):
    def __init__(
            self,
            env_name = 'BipedalWalkerHardcore-v2',
            load_dir = './ckpt',
            log_dir = "./log",
            seed = 1,
            num_envs = 8,
            num_steps = 20,
            mini_batch_size = 5,
            ppo_epochs = 4,
            max_episode_steps = None,
            save_steps_num = 5000
            ):
        self.env_name = env_name
        self.load_dir = load_dir
        self.log_dir = log_dir
        self.seed = seed
        self.num_envs = num_envs
        self.num_steps = num_steps
        self.mini_batch_size = mini_batch_size
        self.ppo_epochs = ppo_epochs
        self.max_episode_steps = max_episode_steps
        self.save_steps_num = save_steps_num

        torch.manual_seed(self.seed)
        np.random.seed(self.seed)
        self.writer = SummaryWriter(log_dir=self.log_dir)

        self.env = gym.make(self.env_name)
        if self.max_episode_steps != None:
            self.env._max_episode_steps = self.max_episode_steps
        else:
            self.max_episode_steps = self.env._max_episode_steps

        self.envs = [make_env(self.env_name) for i in range(self.num_envs)]
        self.envs = SubprocVecEnv(self.envs)

        self.num_inputs  = self.envs.observation_space.shape[0]
        self.num_outputs = self.envs.action_space.shape[0]

        #Hyper params:
        self.hidden_size      = 256
        self.lr               = 3e-4

        self.model = ActorCritic3(self.num_inputs, self.num_outputs, self.hidden_size).to(device)
        try:
            self.load(directory=self.load_dir, filename=self.env_name)
            print('Load model successfully !')
        except:
            print('WARNING: No model to load !')

        self.optimizer = optim.Adam(self.model.parameters(), lr=self.lr)

        self.total_steps = 0
        self.state = self.envs.reset()

    def load(self, directory, filename):
        self.model.load(directory=self.load_dir, filename=self.env_name)

    def save(self, directory, filename):
        self.model.save(directory=self.load_dir, filename=self.env_name)

    def predict(self, state):
        state = torch.FloatTensor(state).unsqueeze(0).to(device)
        dist, _ = self.model(state)
        action = dist.sample()
        action = action.cpu().numpy()[0]
        return action

    def ppo_iter(self, states, actions, log_probs, returns, advantage):
        batch_size = states.size(0)
        for _ in range(batch_size // self.mini_batch_size):
            rand_ids = np.random.randint(0, batch_size, self.mini_batch_size)
            yield states[rand_ids, :], actions[rand_ids, :], log_probs[rand_ids, :], returns[rand_ids, :], advantage[rand_ids, :]

    def train_step(self, states, actions, log_probs, returns, advantages, clip_param=0.2):
        for _ in range(self.ppo_epochs):
            for state, action, old_log_probs, return_, advantage in self.ppo_iter(states, actions, log_probs, returns, advantages):
                dist, value = self.model(state)
                entropy = dist.entropy().mean()
                new_log_probs = dist.log_prob(action)

                ratio = (new_log_probs - old_log_probs).exp()
                surr1 = ratio * advantage
                surr2 = torch.clamp(ratio, 1.0 - clip_param, 1.0 + clip_param) * advantage

                actor_loss  = - torch.min(surr1, surr2).mean()
                critic_loss = (return_ - value).pow(2).mean()

                loss = 0.5 * critic_loss + actor_loss - 0.001 * entropy

                self.optimizer.zero_grad()
                loss.backward()
                self.optimizer.step()

    def learn(self, max_steps=1e7):
        while self.total_steps < max_steps:# and not self.early_stop:
            log_probs = []
            values    = []
            states    = []
            actions   = []
            rewards   = []
            masks     = []
            entropy = 0

            for _ in range(self.num_steps):
                state = torch.FloatTensor(self.state).to(device)
                dist, value = self.model(state)

                action = dist.sample()
                next_state, reward, done, _ = self.envs.step(action.cpu().numpy())

                log_prob = dist.log_prob(action)
                entropy += dist.entropy().mean()

                log_probs.append(log_prob)
                values.append(value)
                rewards.append(torch.FloatTensor(reward).unsqueeze(1).to(device))
                masks.append(torch.FloatTensor(1 - done).unsqueeze(1).to(device))

                states.append(state)
                actions.append(action)

                state = next_state
                self.total_steps += 1

                if self.total_steps % 100 == 0:
                    test_reward = np.mean([test_env2(self.env, self.model) for _ in range(10)])
                    self.writer.add_scalar('test_reward', test_reward, self.total_steps)

                if self.total_steps % self.save_steps_num == 0:
                    self.save(directory=self.load_dir, filename=self.env_name)

            next_state = torch.FloatTensor(next_state).to(device)
            _, next_value = self.model(next_state)
            returns = compute_gae(next_value, rewards, masks, values)

            returns   = torch.cat(returns).detach()
            log_probs = torch.cat(log_probs).detach()
            values    = torch.cat(values).detach()
            states    = torch.cat(states)
            actions   = torch.cat(actions)
            advantage = returns - values

            self.train_step(states, actions, log_probs, returns, advantage)

    def save_expert_traj(self, max_expert_num = 50000):
        #Saving trajectories for GAIL
        from itertools import count

        num_steps = 0
        expert_traj = []

        for i_episode in count():
            state = self.env.reset()
            done = False
            total_reward = 0

            while not done:
                state = torch.FloatTensor(state).unsqueeze(0).to(device)
                dist, _ = self.model(state)
                action = dist.sample().cpu().numpy()[0]
                next_state, reward, done, _ = self.env.step(action)
                state = next_state
                total_reward += reward
                expert_traj.append(np.hstack([state, action]))
                num_steps += 1

            print("episode:", i_episode, "reward:", total_reward)

            if num_steps >= max_expert_num:
                break

        expert_traj = np.stack(expert_traj)
        print(expert_traj.shape)
        np.save("expert_traj.npy", expert_traj)

if __name__ == '__main__':
    model = PPO()
    model.learn()
    #model.save_expert_traj(50)