import numpy as np
import imutils
import pickle
import cv2
import os
from EmoPy.EmoPy.src.fermodel import FERModel
def recognize(image,#BGR format
			  pathDetector='CloudAndGridREC/face_detection_model',
			  pathEmbModel='CloudAndGridREC/openface_nn4.small2.v1.t7',
			  pathRecognizer='CloudAndGridREC/output/recognizer.pickle',
			  pathLabelEnc='CloudAndGridREC/output/le.pickle',
			  conf=0.5):

	student_names=[]
	student_emotions=[]
	# load our serialized face detector from disk
	print("[INFO] loading face detector...")
	protoPath = os.path.sep.join([pathDetector, "deploy.prototxt"])
	modelPath = os.path.sep.join([pathDetector,
		"res10_300x300_ssd_iter_140000.caffemodel"])
	detector = cv2.dnn.readNetFromCaffe(protoPath, modelPath)
	# load our serialized face embedding model from disk
	print("[INFO] loading face recognizer...")
	embedder = cv2.dnn.readNetFromTorch(pathEmbModel)
	# load the actual face recognition model along with the label encoder
	recognizer = pickle.load(open(pathRecognizer, "rb"))
	le = pickle.load(open(pathLabelEnc, "rb"))

	# load the image, resize it to have a width of 600 pixels (while
	# maintaining the aspect ratio), and then grab the image dimensions
	#image = cv2.imread(pathImage)
	image = imutils.resize(image, width=600)
	(h, w) = image.shape[:2]
	# construct a blob from the image
	imageBlob = cv2.dnn.blobFromImage(
		cv2.resize(image, (300, 300)), 1.0, (300, 300),
		(104.0, 177.0, 123.0), swapRB=False, crop=False)
	# apply OpenCV's deep learning-based face detector to localize
	# faces in the input image
	detector.setInput(imageBlob)
	detections = detector.forward()

	# loop over the detections
	for i in range(0, detections.shape[2]):
		# extract the confidence (i.e., probability) associated with the
		# prediction
		confidence = detections[0, 0, i, 2]
		# filter out weak detections
		if confidence > conf:
			# compute the (x, y)-coordinates of the bounding box for the
			# face
			box = detections[0, 0, i, 3:7] * np.array([w, h, w, h])
			(startX, startY, endX, endY) = box.astype("int")
			# extract the face ROI
			face = image[startY:endY, startX:endX]
			(fH, fW) = face.shape[:2]
			# ensure the face width and height are sufficiently large
			if fW < 20 or fH < 20:
				continue

			# construct a blob for the face ROI, then pass the blob
			# through our face embedding model to obtain the 128-d
			# quantification of the face

			face_emo, face_emo_index = recognize_emo(face)
			faceBlob = cv2.dnn.blobFromImage(face, 1.0 / 255, (96, 96),
				(0, 0, 0), swapRB=True, crop=False)
			embedder.setInput(faceBlob)
			vec = embedder.forward()
			# perform classification to recognize the face
			preds = recognizer.predict_proba(vec)[0]
			j = np.argmax(preds)
			proba = preds[j]
			name = le.classes_[j]
			student_names.append(name)
			student_emotions.append(face_emo)
			# draw the bounding box of the face along with the associated
			# probability
			text = "{}: {:.2f}% is {}".format(name, proba * 100, face_emo)
			y = startY - 10 if startY - 10 > 10 else startY + 10
			cv2.rectangle(image, (startX, startY), (endX, endY),
			(0, 0, 255), 2)
			cv2.putText(image, text, (startX, y),
				cv2.FONT_HERSHEY_SIMPLEX, 0.45, (0, 0, 255), 2)
	# show the output image
	print("[INFO] Recognition ended")
	return image, student_names ,student_emotions

def recognize_emo(image):
	target_emotions=set(['calm', 'disgust', 'surprise'])
	fermodel = FERModel(target_emotions=target_emotions)
	return fermodel.predict_from_ndarray(image)
	return "happiness", "3"


# http://stackoverflow.com/questions/4189766/range-with-step-of-type-float
# floating point range function
import sys, os

# for openCV
sys.path.append('/usr/local/lib/python3.5/site-packages/')

USE_OPENCV = 1

import math
import marlincodes
import sensorhead
import json
import numpy as np

if USE_OPENCV:
    import cv2

# floating point range
def xfrange(start, stop, step):
    while start < stop:
        yield start
        start += step

# simple point
class Point:
    def __init__(self, x, y):
        self.x = x
        self.y = y

# general class for scanner
# all units are in mm and deg celsius
class CapScan:
    def __init__(self, name, sensorSerial, machineSerial):

        self.name = name

        # set true for random fake data, used for testing
        self.faked = False

        self.sensorSerial = sensorSerial
        self.machineSerial = machineSerial

        ## coarse scan
        # self.stepSize = 10
        # self.start = Point(0, 0)
        # self.end = Point(120, 120)

        ##detail scan
        self.stepSize = .5
        self.start = Point(0, 0)
        self.end = Point(170, 120)

        self.averages = 20

        # just do 2d scans for now
        self.z = 0

        self.toolPath = []
        self.makeToolPath()

        self.resultsPath = 'results'

        # static results size so that we can scan out of order if needed
        self.results = self.makeResultsArray()

    # generates the tool path for scanning based on the bounding box (start, end)
    def makeToolPath(self):
        for yFloat in xfrange(self.start.y, self.end.y, self.stepSize * 2):
            # x is inner because moving y more causes cooling
            # scan up then back down again to save tool time
            # scan up
            for xFloat in xfrange(self.start.x, self.end.x, self.stepSize):
                self.toolPath += [Point(xFloat, yFloat)]
            yFloat += self.stepSize
            # scan back down
            for xFloat in reversed(list(xfrange(self.start.x, self.end.x, self.stepSize))):
                self.toolPath += [Point(xFloat, yFloat)]

    # prints out tool path for debugging
    def printToolPath(self):
        for p in self.toolPath:
            print('x:{x} - y:{y}'.format(x=p.x, y=p.y))

    # generate a results array based on scanning area, needed so we can scan out of order
    def makeResultsArray(self):
        numCols = math.ceil((self.end.y - self.start.y) / self.stepSize)
        numRows = math.ceil((self.end.x - self.start.x) / self.stepSize)
        resultArr = [[{} for a in range(numCols)] for b in range(numRows)]
        return resultArr

    # takes a float real life location in mm and converts it to results array
    def getResultsCoords(self, point):
        x = math.floor((point.x - self.start.x) / self.stepSize)
        y = math.floor((point.y - self.start.y) / self.stepSize)
        return x, y

    # iterates through toolpath and takes readings at each point
    def doScan(self):
        for point in self.toolPath:
            marlincodes.moveToPos(self.machineSerial, point.x, point.y, self.z)
            resx, resy = self.getResultsCoords(point)
            # self.results[resx][resy] = sensorhead.getAveragedCinPf(self.sensorSerial, averages=self.averages,
            #                                                        faked=self.faked)
            print('x:{x}, y:{y}'.format(x=resx, y=resy))
            self.sensorSerial.processStream()  # clear out any values accumulated while moving
            self.sensorSerial.processStream(averages=self.averages)
            self.results[resx][resy] = self.sensorSerial.getValuesAsDict()
            print(self.results[resx][resy])
        self.saveResults()

    def getValsAsDict(self):
        return self.sensorSerial

    def getResultsFilePath(self, optionalPath=''):
        if not optionalPath:
            optionalPath = self.resultsPath
        return os.path.join(optionalPath, self.name + '.json')

    def saveResults(self):
        f = open(self.getResultsFilePath(), 'w')
        json.dump(self.results, f)
        f.close()

    def loadResults(self):
        with open(self.getResultsFilePath()) as data_file:
            self.results = json.load(data_file)

    # sets the 2d results array from 1D array (assumed in (toolpath) time order)
    def setResultsFrom1D(self, arr):
        p = Point(0, 0)
        count = 0
        for yFloat in xfrange(self.start.y, self.end.y, self.stepSize * 2):
            # scan up the X axis
            for xFloat in xfrange(self.start.x, self.end.x, self.stepSize):
                p.x = xFloat
                p.y = yFloat
                indx, indy = self.getResultsCoords(p)
                self.results[indx][indy] = arr[count]
                count += 1
            yFloat += self.stepSize
            # scan back down the X axis at + y step
            for xFloat in reversed(list(xfrange(self.start.x, self.end.x, self.stepSize))):
                self.toolPath += [Point(xFloat, yFloat)]
                p.x = xFloat
                p.y = yFloat
                indx, indy = self.getResultsCoords(p)
                self.results[indx][indy] = arr[count]
                count += 1

    # returns a 1D array from the 2D results, order is given by the toolpath
    def getResultsAs1D(self):
        resArr = []
        p = Point(0, 0)
        for yFloat in xfrange(self.start.y, self.end.y, self.stepSize * 2):
            # scan up the X axis
            for xFloat in xfrange(self.start.x, self.end.x, self.stepSize):
                p.x = xFloat
                p.y = yFloat
                indx, indy = self.getResultsCoords(p)
                resArr += [self.results[indx][indy]]
            yFloat += self.stepSize
            # scan back down the X axis at + y step
            for xFloat in reversed(list(xfrange(self.start.x, self.end.x, self.stepSize))):
                self.toolPath += [Point(xFloat, yFloat)]
                p.x = xFloat
                p.y = yFloat
                indx, indy = self.getResultsCoords(p)
                resArr += [self.results[indx][indy]]
        return resArr

    def getDims(self):
        xdim = len(self.results)
        ydim = len(self.results[0])
        return [xdim, ydim]

    def minMaxChannel(self, channelName):
        minDiff = 9999
        maxDiff = -9999
        for x in range(len(self.results)):
            for y in range(len(self.results[0])):
                v = self.results[x][y][channelName]
                if v < minDiff:
                    minDiff = v
                if v > maxDiff:
                    maxDiff = v
        return [minDiff, maxDiff]

    # image channel 0=blue, 1=green, 2=red
    def getImageChannel(self, channelName, imageChannel=0, ignoreTop=0, ignoreBottom=0):
        xdim, ydim = self.getDims()
        blank_image = np.zeros((xdim, ydim, 3), np.uint8)
        minDiff, maxDiff = self.minMaxChannel(channelName)
        # set image values
        for x in range(len(self.results)):
            for y in range(len(self.results[0])):
                v = self.results[x][y][channelName]
                # clip out top and bottom outilers
                if (v < (maxDiff - ignoreTop)):
                    if (v > (minDiff + ignoreBottom)):
                        # normalize on min/max
                        v = v - (minDiff + ignoreBottom)
                        v *= 255 / ((maxDiff - ignoreBottom) - (minDiff + ignoreTop))
                        # set blue channel of BGR image, [1] is Green
                        blank_image[x][y][imageChannel] = v
        return blank_image

    def viewMultiChannel(self):
        if USE_OPENCV:

            dims = self.getDims()

            # offsets from capacitive center
            # distance sensor:
            #   24.08 mm back
            #   19.907 mm to the right
            distanceSensorXoffset = 24.08
            distanceSensorYoffset = 19.907

            # temp sensor:
            #   26.62 mm back
            #   13.333 mm to the left
            tempSensorXoffset = 24.08
            tempSensorYoffset = 13.333

            # get the largest offset and round it up to an int for pixels
            maxOffset = int(24.08 * self.stepSize + 1)

            capChan = self.getImageChannel('capReading', imageChannel=0)
            distanceChan = self.getImageChannel('distance', imageChannel=1)
            tempChan = self.getImageChannel('objectTemp', imageChannel=2)

            result = np.zeros((dims[0] + maxOffset * 3, dims[1] + maxOffset * 3, 3), np.uint8)

            # copy cap in
            result[maxOffset:maxOffset + capChan.shape[0], maxOffset:maxOffset + capChan.shape[1], 0] = capChan[::0]
            result[0:0 + distanceChan.shape[0], 20:20 + distanceChan.shape[1], 0] = distanceChan[::1]
            result[0:0 + tempChan.shape[0], 13:13 + tempChan.shape[1], 0] = tempChan[::2]

            resultLArge = self.scaleCVMat(result)
            cv2.imshow('', resultLArge)
            cv2.waitKey(0)  # 0ms wait means wait for key press
            self.saveAsPNG(result, name=self.name + 'multi-chan')

        else:
            print('need opencv to visualise results')

    # quick method to view preliminary results
    def showResults(self, valueToShow='capReading', filter=True, ignoreTop=.000, ignoreBottom=.000):
        if USE_OPENCV:
            # these values can be used to remove outlier readings
            # ignoreTop = .000
            # ignoreBottom = .000

            xdim = len(self.results)
            ydim = len(self.results[0])
            blank_image = np.zeros((xdim, ydim, 3), np.uint8)
            minDiff = 9999
            maxDiff = -9999

            rowAverages = [0]

            # get averages for filtering
            if filter:
                for y in range(len(self.results[0])):
                    sum = 0
                    for x in range(len(self.results)):
                        sum += self.results[x][y][valueToShow]
                    rowAverages[-1] = sum / len(self.results[0])
                    rowAverages.append(0)

            # get min/max values with optional row filtering
            for x in range(len(self.results)):
                for y in range(len(self.results[0])):
                    if filter:
                        v = self.results[x][y][valueToShow] - rowAverages[y]
                    else:
                        v = self.results[x][y][valueToShow]

                    if v < minDiff:
                        minDiff = v
                    if v > maxDiff:
                        maxDiff = v
            print('min: ' + str(minDiff))
            print('max: ' + str(maxDiff))

            # set image values
            for x in range(len(self.results)):
                for y in range(len(self.results[0])):
                    if filter:
                        v = self.results[x][y][valueToShow] - rowAverages[y]
                    else:
                        v = self.results[x][y][valueToShow]
                    # clip out top and bottom outilers
                    if (v < (maxDiff - ignoreTop)):
                        if (v > (minDiff + ignoreBottom)):
                            # normalize on min/max
                            v = v - (minDiff + ignoreBottom)
                            v *= 255 / ((maxDiff - ignoreBottom) - (minDiff + ignoreTop))
                            # set blue channel of BGR image, [1] is Green
                            blank_image[x][y][0] = v

            blankLarge = self.scaleCVMat(blank_image)
            cv2.imshow('', blankLarge)
            cv2.waitKey(0)  # 0ms wait means wait for key press
            self.saveAsPNG(blankLarge)
        else:
            print('No opencv enabled')

    # scales a CV mat using factor
    def scaleCVMat(self, matToScale, factor=10):
        if USE_OPENCV:
            # no sense in scaling beyond screen size
            if factor * len(matToScale) > 600:
                factor = 600.0 / len(matToScale)
            return cv2.resize(matToScale, (0, 0), fx=factor, fy=factor, interpolation=cv2.INTER_NEAREST)

    # given an openCV mat save as png in path
    def saveAsPNG(self, anOCVMat, optionalPath='', name=''):
        if USE_OPENCV:
            # allow optional saving elsewhere
            if not optionalPath:
                optionalPath = self.resultsPath
            if not name:
                name = self.name
            cv2.imwrite(os.path.join(optionalPath, name + '.png'), anOCVMat)


import http.client
import os
import tempfile
import wave
import numpy as np
import json
import time

def transcribeAudio(path_to_audio_file, samplerate=16000):
    headers = {"Content-type": "audio/wav; codec=\"audio/pcm\"; samplerate=" + str(samplerate)}

    with open(path_to_audio_file, 'rb') as audio_file:
        response = ""
        try:

            body = audio_file.read()
            # Connect to server to recognize the wave binary

            conn = http.client.HTTPConnection(os.environ['ASR_SERVER'], int(os.environ['GSTREAM_PORT']))

            conn.request("POST", "/client/dynamic/recognize",
                         body, headers)
            response = conn.getresponse().read().decode("UTF-8")
            conn.close()
        finally:
            audio_file.close()

        return response

def main(audio_chunk):
    with tempfile.NamedTemporaryFile(mode='w+b', suffix='.wav') as fp:
        succes = False
        while not succes:
            try:
                sample_rate = 16000
                wf = wave.open(fp, 'w')
                wf.setnchannels(1)
                wf.setsampwidth(2)
                wf.setframerate(sample_rate)
                wf.writeframes(np.frombuffer(audio_chunk, dtype=np.uint8))
                wf.close()
                res = transcribeAudio(fp.name)
                #print(res, flush=True)
                res = json.loads(res)
                if res:
                    return res['hypotheses'][0]['utterance']
                    succes = True
            except Exception as e:
                print('trying again', flush=True)
                time.sleep(0.5)


import numpy as np
import rospy
from rosgraph_msgs.msg import Log
from sensor_msgs.msg import JointState
from std_msgs.msg import String

class HeadControl(object):
    def __init__(self, ns):
        self.ns = ns
        _ = rospy.Subscriber("/rosout", Log, self._cb_headinfo, queue_size=10)

        self._pub_head_scan = rospy.Publisher(ns + "/head_control/scan_command", String)
        self._pub_head_joints = rospy.Publisher(ns + "/head_control/set_joint_states", JointState)
        self._pub_head_offset = rospy.Publisher(ns + "/head_control/set_joint_states_offset", JointState)
        self.is_head_done = False

    def _cb_headinfo(self, msg):
        if msg.name == "/op3_demo_opc":
            if msg.msg == "Head movement is finished.":
                self.is_head_done = True

    def set_head(self, pan, tilt, offset=False, blocking=False):
        self.is_head_done = False
        self.check_module("head_control_module")
        msg = JointState()
        msg.name = ["head_pan", "head_tilt"]
        msg.position = [np.pi * (pan / 180.0),
                        np.pi * (tilt / 180.0)]
        if offset:
            self._pub_head_offset.publish(msg)
        else:
            self._pub_head_joints.publish(msg)

        if blocking:
            while not self.is_head_done:
                rospy.sleep(0.01)

    def stop_head_scan(self):
        self._pub_head_scan.publish("stop")

    def head_scan(self):
        self._pub_head_scan.publish("scan")


"""Load Image, predict intentional framing"""
import os
import glob
import conf
import urllib
import operator
import numpy as np
import pandas as pd
from xml.dom import minidom

from keras.preprocessing import image
from keras.models import model_from_json
from keras.applications.inception_v3 import preprocess_input

classes = {'art': 0,
	'candid': 1,
	'landscape': 2,
	'macro': 3,
	'media_capture': 4,
	'portrait': 5,
	'product_presentation': 6,
	'setting': 7,
	'social_event_private': 8,
	'social_event_public': 9,
	'structure': 10}
classes = sorted(classes.items(), key=operator.itemgetter(1))

def prepare_model():
	"""prepare pre-trained model for intents classification."""
	def load_model():
		"""load model"""
		with open(conf.MODEL_PATH, 'r') as f:
			loaded_model_json = f.read()
			return model_from_json(loaded_model_json)

	def load_weights(model):
		"""Load mode weights."""
		model.load_weights(conf.WEIGHT_PATH)
		return model

	def compile_model(model):
		"""Compile current model."""
		model.compile(optimizer='rmsprop',
			loss='categorical_crossentropy',
			metrics=['accuracy'])
		return model

	return compile_model(load_weights(load_model()))

def make_prediction(url, img_id, model, classes, from_path=None):
	"""Download store in temp, make prediction, remove."""
	def download_img(url, img_id):
		"""Download image."""
		if from_path:
			return from_path
		img_path = ''.join([TEMP_PATH, img_id, '.jpg'])
		urllib.urlretrieve(url, img_path)
		return img_path

	def pre_process(img_path):
		"""Preprocess image"""
		img = image.load_img(img_path, target_size=(224, 224))
		x = image.img_to_array(img)
		x = np.expand_dims(x, axis=0)
		x = preprocess_input(x)
		return x

	def predict(model, x, classes):
		"""Predict class."""
		preds = model.predict(x)
		prob = 0
		label = ''
		for pred, c in zip(preds[0], classes):
			if pred > prob:
				prob = pred
				label = c
		return prob, label

	img_path = download_img(url, img_id)
	img = pre_process(img_path)
	prob, label = predict(model, img, classes)
	return prob, label


import gym
from average_functions import DNN
import random
import numpy as np
from collections import deque
from datetime import datetime
from collections import defaultdict
from Environment import Environment
from AgentUtility import AgentUtility
import pickle
import sys

"""
AutoNet Agent implementation using DDQN

usage: python3 train_autonet.py <learning_rate> <epsilon_step_size> <replay_buffer_size> <episodes> <optional: training_save_file_name>

The neural network is implemented in average_function.py
"""
class AutoNet:
    def __init__(self, gamma=0.995, learning_rate=0.001, epsilon_step_size=125, replay_buffer_size=50000):
        self.env = Environment()
        self.model = DNN(20,128,2,len(self.env.action_space), learning_rate)
        self.epsilon = 1.0
        self.min_exploration = 5
        self.replay_buffer_size = replay_buffer_size
        self.exploration_set = deque(maxlen=(self.replay_buffer_size))
        self.exploration_ct = 0
        self.gamma = gamma
        self.replay_start_marker = 0
        self.transer_weight_count = 5
        self.replay_time = 1
        self.moving_avg = deque(maxlen=100)
        self.learning_rate = learning_rate
        self.epsilon_step_size = epsilon_step_size
        self.epsilon_stat = defaultdict(int)
        print ("Number of action space is: {0}, learning_rate is: {1}, gamma is {2}, epsilon_step_size is {3} and replay_buffer_size is {4}".format(len(self.env.action_space), self.learning_rate, self.gamma, self.epsilon_step_size, self.replay_buffer_size))

    """
    Method to render the environment
    """
    def render(self):
        pass

    """
    Method to train the agent. The number of epsiodes to run is controlled by episodes argument.

    Logic:
        At the begining of the episode, the environment is first reet. For each iteration, the getAction method is called
        to get the action. getAction method returns the action as per epsilon - greedy policy.
        Each time the agent takes an action, it receives reward and the trainsition state information. The infomration
        is stored in the replay buffer in a form of tuple - (x,action,reward,x_new,done).

        The agent calls replay_batch at the end of an episode, provided the experience collected in the replay buffer
        is greater than the minimum exploration count as configured. The default is 3000. The batch size is provided as 1024.

        Epsilon has starting value of 1.0. The epsilon is reduced by 0.05 every "step_size" number of episodes.
        Once the value reaches 0.1. It remains so until 3000th episode. At 3001st episode, epsilon is set to constant value of 0.05

        The agent keeps track of best score above 199 and is saved to save_path. The file that is available at the end of the training
        will be from the best moving average score above 199 as seen by the agent.
    """
    def act(self, episodes=1, render=False, save_path=None):
        outfile = open('training_reward_avg_per_episode_{0}_{1}_{2}_{3}.txt'.format(self.learning_rate, self.gamma, self.epsilon_step_size, self.replay_buffer_size),'w')
        logfile = open('training_{0}_{1}_{2}_{3}.log'.format(self.learning_rate, self.gamma, self.epsilon_step_size, self.replay_buffer_size),'w')
        best_score = 199
        import time
        print("starting epispdes")
        time.sleep(10)
        for ep in range(episodes):
            print("Episode: ", ep)
            agent_utility = AgentUtility('../data/amazon_reviews_us_Mobile_Electronics_v1_00.tsv', 'review_body', 'star_rating', sep='\t')
            print("Agent utility loaded")
            time.sleep(10)
            self.env.set_action_utility(agent_utility)
            index = 0
            x = self.env.reset_stats()
            rv = 0
            while index<5:
                print("index: ",index)
                index += 1
                if render:
                    self.render()
                action = self.getAction(x)
                print(action)
                x_new, reward, done = self.env.step(action)
                rv += reward
                self.exploration_set.append((x,action,reward,x_new,done))
                x = x_new
                self.exploration_ct += 1
                pickle.dump(self.exploration_set,open("../autonetdata/autonetdata.pkl","wb"))
                if done:
                    break
            self.moving_avg.append(rv)
            test_rw = np.average(self.moving_avg)
            #logfile.write('{4}:  Ran episode {0}, scored {1}, epsilon is {2}, average score {3}\n'.format(ep, rv, self.epsilon, test_rw, str(datetime.now())))
            print('{4}:  Ran episode {0}, scored {1}, epsilon is {2}, average score {3}'.format(ep, rv, self.epsilon, test_rw, str(datetime.now())))
            outfile.write("{0},{1},{2}\n".format(ep,np.average(self.moving_avg),rv))
            if self.exploration_ct > self.min_exploration and ep % self.replay_time == 0:
                self.replay_batch(ep, 3)
                self.model.save('./../models/'+str(ep)+"_"+str(index))

            #self.epsilon_stat[self.epsilon] += 1
            if ep % self.epsilon_step_size == 0 and ep != 0 :
                self.epsilon = (self.epsilon - 0.05 if self.epsilon > 0.1 else 0.1)
            if ep > 3000 and self.epsilon == 0.1:
                self.epsilon = 0.05
            if test_rw > best_score and save_path:
                #self.save("lunar_learned_weights_{0}_{1}_{2}".format(ep, test_rw, self.learning_rate))
                self.save(save_path)
                best_score = test_rw
                #if test_rw >= 200:
                #    break
        outfile.close()
        logfile.close()
        print ('reward is: {0} at iteration {1}'.format(self.test_online(), episodes))
        print ('The weight saved for best moving averate of {0}'.format(best_score))

    """
    This methods recieves the batch_size as the input parameter. It samples those many number of records from the replay buffer.
    For each record in the sample set, following operations are performed -
        1. Call predict on the model (online network) using x
        2. Call predict on the model (online network) using x_new
        3. Call predict on the old_model (target network) using x_new
        4. Choose the best action from step 2, which is the index value of the highest value in the array
        5. Use the index number from 4 to select the q value from step 2
        6. Calcualte the q estimate using target equation - reward if the state is terminal else (reward + gamma* (q value from step 5))
        7. Set the q value to the array returned on step 1 to the index indentified by "action" within the sample record
        8. perform gradient descent

    The weigts are copied to the target network from online network every few episodes as configued in  self.transer_weight_count
    """
    def replay_online(self, ep_st, batch_size=32):
        indexes = np.random.choice(self.min_exploration, size=batch_size)
        batch = [self.exploration_set[itr] for itr in indexes]
        if ep_st % self.transer_weight_count == 0:
            self.model.clone()
        for x,action,reward,x_new,done in batch:
            q_x = self.model.predict(x.reshape(1,8))[0]
            qa = self.model.predict(x_new.reshape(1,8))[0]
            qas_old = self.model.old_model.predict(x_new.reshape(1,8))[0]
            q = reward if done else (reward + self.gamma * qas_old[np.argmax(qa)])
            q_x[action] = q
            self.model.train((x.reshape(1,8),q_x.reshape(1,4)),epoch=1)

    """
    This methods recieves the batch_size as the input parameter. It samples those many number of records from the replay buffer.
    On entire batch, following operations are performed -
        1. Call predict on the model (online network) using x
        2. Call predict on the model (online network) using x_new
        3. Call predict on the old_model (target network) using x_new
        4. Choose the best action from step 2, which is the index value of the highest value in the array
        5. Use the index number from 4 to select the q value from step 2
        6. Calcualte the q estimate using target equation - reward if the state is terminal else (reward + gamma* (q value from step 5))
        7. Set the q value to the array returned on step 1 to the index indentified by "action" within the sample record
        8. perform gradient descent on the whole batch using 32 samples at a time for gradient update.

    The weigts are copied to the target network from online network every few episodes as configued in  self.transer_weight_count
    """
    def replay_batch(self, ep_st, batch_size=2):
        indexes = np.random.choice(len(self.exploration_set), size=batch_size)
        batch = [self.exploration_set[itr] for itr in indexes]
        qvals = np.zeros((batch_size,20)).tolist()
        #qactions = np.zeros((batch_size,4))
        ind = 0
        if ep_st % self.transer_weight_count == 0:
            self.model.clone()
        x_new_batch = np.zeros((batch_size, 20)).tolist()
        x_batch = np.zeros((batch_size, 20)).tolist()
        for i in range(len(batch)):
            x,action,reward,x_new,done = batch[i]
            x_new_batch[i] = x_new
            x_batch[i] = x
        q_x = self.model.batch_predict(x_batch)
        qas = self.model.batch_predict(x_new_batch)
        qas_old = self.model.batch_predict(x_new_batch, False)
        for i in range(len(batch)):
            x,action,reward,x_new,done = batch[i]
            q = reward if done else reward + self.gamma * qas_old[i][np.argmax(qas[i])] # Choose best action from new model, but, pick the estimate for the action from stationary model
            q_x[i][self.model.output_index[action]] = q
            qvals[ind] = x
            ind+=1
        self.model.train((qvals,q_x),epoch=1, batch_size=32)

    """
    Returns action as per epsilon greedy policy if random_enabled is True. If
    random_enabled is false, uses only the online network to determine the action.
    """
    def getAction(self, x, random_enabled=True):
        a = 0
        if (random.random() < self.epsilon) and random_enabled:
            a = self.env.sample_actionset()
        else:
            qa = self.model.predict(" ".join(x))
            a = np.argmax(qa[0])
        return a

    def save(self, path):
        self.model.save(path)

    """
    Runs bench mark on the weights provided by path variable. By default it runs
     100 iterations and returns rewards from all the episodes
    """
    def test(self, path=None, iteration=100):
        self.env = gym.make('AutoNet-v2')
        self.env = gym.wrappers.Monitor(self.env, directory="temp")
        self.model.load(path)

        for it in range(iteration):
            x = self.env.reset()
            while True:
                #self.render()
                action = self.getAction(x, False)
                x_new, reward, done, _ = self.env.step(action)
                x = x_new
                if done:
                    break
        return self.env.get_episode_rewards()

    """
    Runs bench mark on the weights currently with in the model. By default it runs
     100 iterations and returns rewards from all the episodes
    """
    def test_online(self, iteration=100):
        rewards = []
        for it in range(iteration):
            x = self.env.reset()
            rew = 0
            while True:
                #self.render()
                action = self.getAction(x, False)
                x_new, reward, done, _ = self.env.step(action)
                x = x_new
                rew += reward
                if done:
                    break
            rewards.append(rew)
        return np.average(rewards)

def main():
    lander = AutoNet(learning_rate=float(0.001), epsilon_step_size=int(125), replay_buffer_size=int(50000), )
    lander.act(episodes=int(100000), save_path="lunar_learned_weights")
    #lander.save(sys.argv[5] if len(sys.argv) == 6 else "lunar_learned_weights")
    print (lander.epsilon_stat)

if __name__ == '__main__':
    main()


"""module for loading annotations from .mat files
created by SongAnnotationGUI:
https://github.com/yardencsGitHub/BirdSongBout/tree/master/helpers/GUI
"""
import os
import pathlib
from typing import ClassVar, List

import attr
import numpy as np
import scipy.io

import crowsetta
from crowsetta.typing import PathLike

def _cast_to_arr(val):
    """helper function that casts single elements to 1-d numpy arrays"""
    if type(val) == int or type(val) == float:
        # this happens when there's only one syllable in the file
        # with only one corresponding label
        return np.asarray([val])  # so make it a one-element list
    elif type(val) == np.ndarray:
        # this should happen whenever there's more than one label
        return val
    else:
        # something unexpected happened
        raise TypeError(
            f"Type {type(val)} not recognized when converting annotations to arrays."
        )

VALID_AUDIO_FORMATS = ['wav']

def _recursive_stem(path_str):
    """helper function that 'recursively' removes file extensions
    to recover name of an audio file from the name of an array file

    i.e. bird1_122213_1534.wav.mat -> i.e. bird1_122213_1534.wav
    and i.e. bird1_122213_1534.cbin.not.mat -> i.e. bird1_122213_1534.cbin

    adapted from ``vak`` library
    """
    name = pathlib.Path(path_str).name
    stem, ext = os.path.splitext(name)
    ext = ext.replace('.', '')
    while ext not in VALID_AUDIO_FORMATS:
        new_stem, ext = os.path.splitext(stem)
        ext = ext.replace('.', '')
        if new_stem == stem:
            raise ValueError(
                f'unable to compute stem of {path_str}'
            )
        else:
            stem = new_stem
    return stem

@crowsetta.interface.SeqLike.register
@attr.define
class SongAnnotationGUI:
    """Class that represents annotations
    from .mat files
    created by SongAnnotationGUI:
    https://github.com/yardencsGitHub/BirdSongBout/tree/master/helpers/GUI

    Attributes
    ----------
    name: str
        Shorthand name for annotation format: ``'yarden'``.
    ext: str
        Extension of files in annotation format: ``'.mat'``.
    annotations : numpy.ndarray
        ``numpy`` record array where each record is an annotation.
    audio_paths : numpy.ndarray
        ``numpy`` array where each element is a path to an audio file.
        Same length as ``annotations``. Each element in ``annotations``
        is the annotation for the corresponding path in ``audio_paths``.
    annot_path : str, pathlib.Path
        Path to .mat file from which annotations were loaded.
    """
    name: ClassVar[str] = 'yarden'
    ext: ClassVar[str] = '.mat'

    annotations: np.ndarray = attr.field(eq=attr.cmp_using(eq=np.array_equal))
    audio_paths: np.ndarray = attr.field(eq=attr.cmp_using(eq=np.array_equal))
    annot_path: pathlib.Path = attr.field(converter=pathlib.Path)

    @classmethod
    def from_file(cls,
                  annot_path: PathLike) -> 'Self':
        """load annotations from .mat files
        created by SongAnnotationGUI:
        https://github.com/yardencsGitHub/BirdSongBout/tree/master/helpers/GUI

        Parameters
        ----------
        annot_path: str, pathlib.Path
            Path to .mat file with annotations.
        """
        annot_path = pathlib.Path(annot_path)
        crowsetta.validation.validate_ext(annot_path, extension=cls.ext)

        # annotation structure loads as a Python dictionary with two keys
        # one maps to a list of filenames,
        # and the other to a Numpy record array,
        # where each element is the annotation
        # corresponding to the filename at the same index in the list.
        annot_mat = scipy.io.loadmat(annot_path, squeeze_me=True)
        audio_paths = annot_mat['keys']
        annotations = annot_mat['elements']
        if len(audio_paths) != len(annotations):
            raise ValueError(
                f'list of filenames and list of annotations in {annot_path} do not have the same length'
            )

        return cls(annotations=annotations,
                   audio_paths=audio_paths,
                   annot_path=annot_path)

    def to_seq(self,
               round_times: bool = True,
               decimals: int = 3
               ) -> List[crowsetta.Sequence]:
        """Convert this .not.mat annotation to a ``crowsetta.Sequence``.

        Parameters
        ----------
        round_times : bool
            If True, round times of onsets and offsets.
            Default is True.
        decimals : int
            Number of decimals places to round floating point numbers to.
            Only meaningful if round_times is True.
            Default is 3, so that times are rounded to milliseconds.

        Returns
        -------
        seqs : list
            of ``crowsetta.Sequence``, one for each element in ``annotations``.

        Notes
        -----
        The ``round_times`` and ``decimals`` arguments are provided
        to reduce differences across platforms
        due to floating point error, e.g. when loading annotation files
        and then sending them to a csv file,
        the result should be the same on Windows and Linux.
        """
        seqs = []
        for annotation in self.annotations:
            # below, .tolist() does not actually create a list,
            # instead gets ndarray out of a zero-length ndarray of dtype=object.
            # This is just weirdness that results from loading complicated data
            # structure in .mat file.
            seq_dict = {}
            seq_dict['onsets_s'] = annotation['segFileStartTimes'].tolist()
            seq_dict['offsets_s'] = annotation['segFileEndTimes'].tolist()
            seq_dict['labels'] = annotation['segType'].tolist()
            # cast all to numpy arrays
            seq_dict = dict((k, _cast_to_arr(seq_dict[k]))
                            for k in ['onsets_s', 'offsets_s', 'labels'])
            # after casting 'labels' to array, convert all values to string
            seq_dict['labels'] = np.asarray(
                [str(label) for label in seq_dict['labels']]
            )

            samp_freq = annotation['fs'].tolist()
            seq_dict['onset_samples'] = np.round(seq_dict['onsets_s'] * samp_freq).astype(int)
            seq_dict['offset_samples'] = np.round(seq_dict['offsets_s'] * samp_freq).astype(int)

            if round_times:
                seq_dict['onsets_s'] = np.around(seq_dict['onsets_s'], decimals=decimals)
                seq_dict['offsets_s'] = np.around(seq_dict['offsets_s'], decimals=decimals)

            seq = crowsetta.Sequence.from_dict(seq_dict)
            seqs.append(seq)

        return seqs

    def to_annot(self,
                 round_times: bool = True,
                 decimals: int = 3
                 ) -> List[crowsetta.Annotation]:
        """Convert this .not.mat annotation to a ``crowsetta.Annotation``.

        Parameters
        ----------
        round_times : bool
            If True, round times of onsets and offsets.
            Default is True.
        decimals : int
            Number of decimals places to round floating point numbers to.
            Only meaningful if round_times is True.
            Default is 3, so that times are rounded to milliseconds.

        Returns
        -------
        annots : list
            of ``crowsetta.Annotation``

        Notes
        -----
        The ``round_times`` and ``decimals`` arguments are provided
        to reduce differences across platforms
        due to floating point error, e.g. when loading annotation files
        and then sending them to a csv file,
        the result should be the same on Windows and Linux.
        """
        seqs = self.to_seq(round_times=round_times, decimals=decimals)
        annots = []
        for audio_path, seq in zip(self.audio_paths, seqs):
            annots.append(
                crowsetta.Annotation(annot_path=self.annot_path,
                                     notated_path=audio_path,
                                     seq=seq)
            )
        return annots


import numpy as np
import os
from . import mmdpy_texture
from . import mmdpy_type
from . import pmxpy_type
from typing import Any, cast, Union

# #### #### PMD Adjuster #### ####
def adjust(pmx_data: pmxpy_type.pmxpyType) -> Union[None, mmdpy_type.mmdpyTypeModel]:

    adjust_data = mmdpy_type.mmdpyTypeModel()

    # #### #### Vertex #### ####
    # 複数形は加工データ
    for vv in pmx_data.vertex:
        v = mmdpy_type.mmdpyTypeVertex()
        v.ver = np.asarray(vv.ver)
        v.nor = np.asarray(vv.nor)
        v.uv = np.asarray(vv.uv)
        v.bone_id = np.asarray(vv.bone_id)
        v.bone_weight = np.asarray(vv.bone_weight)
        v.edge = vv.edge

        adjust_data.vertices.append(v)

    # #### #### Face #### ####
    adjust_data.faces = [list(f) for f in pmx_data.face]

    # #### #### Material #### ####
    m_top = 0
    for mm in pmx_data.material:
        m = mmdpy_type.mmdpyTypeMaterial()
        m.diffuse = np.asarray(mm.diffuse)
        m.alpha = 1.00  # cast(float, mm.specular_scale)
        m.specularity = np.asarray(mm.specular_color)
        m.specular_color = np.asarray(mm.specular_color)
        m.mirror_color = np.asarray(mm.ambient_color)
        m.toon_index = mm.texture_index
        m.edge_size = mm.edge_size
        m.face_vert_count = mm.face_vert_count

        m.top = m_top
        m.size = m.face_vert_count
        m_top += m.size

        m.texture = None
        m.texture_name = mm.texrure_name
        if len(m.texture_name) > 0:
            if "*" in str(m.texture_name):
                m.texture_name = m.texture_name[:str(m.texture_name).find("*")]
            texture_path = os.path.join(pmx_data.directory, cast(str, m.texture_name))
            m.texture = mmdpy_texture.mmdpyTexture(texture_path)
        m.both_side_flag = ((mm.bit_flag & 0x01) != 0x00)
        m.color = m.diffuse
        if m.texture is not None:
            m.color[3] = 0.00

        adjust_data.materials.append(m)

    # #### #### Bone #### ####
    for i, bb in enumerate(pmx_data.bone):
        b = mmdpy_type.mmdpyTypeBone()
        b.id = i
        b.parent_id = bb.parent_id
        b.position = np.asarray(bb.position)
        b.name = bb.name.replace("\x00", "")
        b.weight = 1.00
        b.rotatable_control = lambda b, axis, rot: b.rot(axis, rot * b.weight)

        b.ik = None
        if bb.ik:
            ik = mmdpy_type.mmdpyTypeIK()
            ik.bone_index = int(i)
            ik.bone_effect_index = bb.ik.target_bone_index
            ik.iteration = bb.ik.iteration
            ik.child_bones_index = [c.bone_id for c in bb.ik.child_bone_id]
            b.ik = ik

        adjust_data.bones.append(b)

    # # #### #### IK #### ####
    # for pmd_ik in pmx_data.ik:
    #     mmd_ik = mmdpy_type.mmdpyTypeIK()

    #     mmd_ik.bone_index = pmd_ik.bone_id
    #     mmd_ik.bone_effect_index = pmd_ik.bone_to
    #     mmd_ik.length = pmd_ik.length
    #     mmd_ik.iteration = pmd_ik.iteration
    #     mmd_ik.weight = np.asarray([pmd_ik.weight, 0, 0, 0])
    #     mmd_ik.child_bones_index = pmd_ik.child_bone_id
    #     for c in mmd_ik.child_bones_index:
    #         adjust_data.bones[c].weight = pmd_ik.weight
    #         adjust_data.bones[c].rotatable_control = lambda b, axis, rot: b.rot(axis, rot * b.weight)

    #     adjust_data.iks.append(mmd_ik)

    # ひざの処理
    def knee_control(b: Any, axis: np.ndarray, rot: np.ndarray) -> None:
        v = np.matmul(b.rot(axis, rot, update_flag=False), np.array([0, 0, 1, 1]))
        if v[1] > 1e-4:
            return
        b.rot(axis, rot)

    for bone in adjust_data.bones:
        if "ひざ" in bone.name:
            bone.rotatable_control = knee_control

    # # #### #### Physics #### ####
    adjust_data.physics_flag = False
    # adjust_data.physics_flag = pmd_data.physics_flag
    # if pmd_data.physics_flag and pmd_data.physics is not None:
    #     adjust_data.physics = mmdpy_type.mmdpyTypePhysics()
    #     adjust_data.physics.body = []
    #     adjust_data.physics.joint = []

    #     if pmd_data.physics.body is not None:
    #         for pmd_body in pmd_data.physics.body:
    #             body = mmdpy_type.mmdpyTypePhysicsBody()

    #             bone_position = pmd_data.bone[pmd_body.bone_id].position
    #             body_pos = np.array(bone_position) + np.array(pmd_body.pos)
    #             # body_pos = np.array(pmd_body.pos)

    #             body.bone_id = pmd_body.bone_id
    #             body.name = pmd_body.name
    #             body.group_id = pmd_body.group_id
    #             body.group_mask = pmd_body.group_mask
    #             body.type_id = pmd_body.type_id
    #             body.sizes = np.asarray([pmd_body.sizes.x, pmd_body.sizes.y, pmd_body.sizes.z])
    #             body.pos = np.asarray(body_pos)
    #             body.rot = np.asarray(pmd_body.rot)
    #             body.mass = pmd_body.mass
    #             body.pos_dim = pmd_body.pos_dim
    #             body.rot_dim = pmd_body.rot_dim
    #             body.recoil = pmd_body.recoil
    #             body.friction = pmd_body.friction
    #             body.calc = pmd_body.calc

    #             adjust_data.physics.body.append(body)

    #     if pmd_data.physics.joint is not None:
    #         for pmd_joint in pmd_data.physics.joint:
    #             joint = mmdpy_type.mmdpyTypePhysicsJoint()

    #             joint.name = pmd_joint.name
    #             joint.rigidbody_a = pmd_joint.rigidbody_a
    #             joint.rigidbody_b = pmd_joint.rigidbody_b

    #             joint.pos = np.asarray(pmd_joint.pos)
    #             joint.rot = np.asarray(pmd_joint.rot)
    #             joint.constrain_pos = [np.asarray(x) for x in pmd_joint.constrain_pos]
    #             joint.spring_pos = np.asarray(pmd_joint.spring_pos)
    #             joint.spring_rot = np.asarray(pmd_joint.spring_rot)

    #             adjust_data.physics.joint.append(joint)

    return adjust_data


#coding=utf-8
#性别识别

import cv2
from keras.models import load_model
import numpy as np
import chineseText

img = cv2.imread("img/gather.png")
face_classifier = cv2.CascadeClassifier(
    "d:\Python36\Lib\site-packages\opencv-master\data\haarcascades\haarcascade_frontalface_default.xml"
)
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
faces = face_classifier.detectMultiScale(
    gray, scaleFactor=1.2, minNeighbors=3, minSize=(140, 140))

gender_classifier = load_model(
    "classifier/gender_models/simple_CNN.81-0.96.hdf5")
gender_labels = {0: '女', 1: '男'}
color = (255, 255, 255)

for (x, y, w, h) in faces:
    face = img[(y - 60):(y + h + 60), (x - 30):(x + w + 30)]
    face = cv2.resize(face, (48, 48))
    face = np.expand_dims(face, 0)
    face = face / 255.0
    gender_label_arg = np.argmax(gender_classifier.predict(face))
    gender = gender_labels[gender_label_arg]
    cv2.rectangle(img, (x, y), (x + h, y + w), color, 2)
    img = chineseText.cv2ImgAddText(img, gender, x + h, y, color, 30)

cv2.imshow("Image", img)
cv2.waitKey(0)
cv2.destroyAllWindows()


# -*- coding: utf-8 -*-
"""
General CNN player

Convolutional neural network that solves OpenAI gym environments
(https://gym.openai.com) from raw pixel data via Q-learning

Download the code, assign to game_name the name of environment you wish to
run, and let the script learn how to solve it.
Note the code only works for environments with discrete action space and
3-dimensional observation space representing the raw pixel data of the screen.

https://github.com/FlankMe/general-CNN-player
@author: Riccardo Rossi
"""

# Choice of the game and definition of the goal
game_name = 'Pong-v0'
MAX_EPISODES = 2500
CONSECUTIVE_EPISODES = 100   # Number of trials' rewards to average for solving
IS_RECORDING = True

# Fine-tuning the EPSILON_DECAY parameters will lead to better results for
# some environments and worse for others. As this code is a go at a
# general player, it is neater to treat it as a global constant
EPSILON_DECAY = 0.99

# Import basic libraries
import numpy as np

"""
Main loop for the game's initialization and runs
"""
def main():

    # Import gym and launch the game
    import gym
    env = gym.make(game_name)

    assert isinstance(env.action_space, gym.spaces.discrete.Discrete), (
        'env.action_space is not Discrete and is currently unsupported')
    assert isinstance(env.observation_space, gym.spaces.box.Box), (
        'env.observation_space is not continuous and is currently unsupported')
    assert len(env.observation_space.shape) == 3, (
        '''CNN not required for this environment
        Download and run general-gym-player from my github page
        https://github.com/FlankMe/general-gym-player''')
    if IS_RECORDING:
        env.monitor.start('results-' + game_name, force=True)

    # Create a gym instance and initialize the agent
    agent = Agent(env.action_space.n,
                  env.observation_space.shape)
    reward, done, = 0.0, False

    # Start the game loop
    for episode in range(1, MAX_EPISODES + 1):
        obs, done = env.reset(), False
        action = agent.act(obs, reward, done, episode)

        while not done:

            # Un-comment to show the game on screen
            #env.render()

            # Decide next action and feed the decision to the environment
            obs, reward, done, _ = env.step(action)
            action = agent.act(obs, reward, done, episode)

    # Save info and shut activities
    env.close()
    agent.close()
    if IS_RECORDING:
        env.monitor.close()

"""
The general game player implements the Q-learning method with minibatches
"""
class Agent():

    def __init__(self, n_actions, obs_space):

        # Initialization of useful variables and constants
        self._N_ACTIONS = n_actions
        self._nn = FeedForwardNeuralNetwork(n_actions, obs_space)
        self._STATE_FRAMES = self._nn._STATE_FRAMES
        self._close = lambda : self._nn.close()

        # Hyperparameters of the training
        self._DISCOUNT_FACTOR = 0.99    # discount of future rewards
        self._TRAINING_PER_STAGE = 1
        self._MINIBATCH_SIZE = 32
        self._REPLAY_MEMORY = 20000

        # Exploration/exploitations parameters
        self._epsilon = 1.
        self._EPSILON_DECAY = EPSILON_DECAY
        self._EPISODES_PURE_EXPLORATION = 10
        self._MIN_EPSILON = 0.1

        # Define useful variables
        self._total_reward, self._list_rewards = 0.0, []
        self._last_action = np.zeros(self._N_ACTIONS)
        self._previous_observations, self._last_state = [], None
        self._LAST_STATE_IDX = 0
        self._ACTION_IDX = 1
        self._REWARD_IDX = 2
        self._CURR_STATE_IDX = 3
        self._TERMINAL_IDX = 4

    def act(self, obs, reward, done, episode):

        self._total_reward += reward

        if done:
            self._list_rewards.append(self._total_reward)
            average = np.mean(self._list_rewards[-CONSECUTIVE_EPISODES:])
            print ('Episode', episode, 'Reward', self._total_reward,
                   'Average Reward', round(average, 2))
            self._total_reward = 0.0
            self._epsilon = max(self._epsilon * self._EPSILON_DECAY,
                                self._MIN_EPSILON)

        # Compress the input image into a pre-set format
        compressedImage = self._nn._compressImage(obs)

        # If the CNN has no last state, fill it by using the current state,
        # choose a random action, and return the action to the game
        if self._last_state is None:
            self._last_state = compressedImage.copy()
            for _ in range(self._STATE_FRAMES - 1):
                self._last_state = np.append(self._last_state,
                                             compressedImage,
                                             axis=3)
            value_per_action = self._nn.predict(self._last_state)
            chosen_action_index = np.argmax(value_per_action)
            self._last_action = np.zeros(self._N_ACTIONS)
            self._last_action[chosen_action_index] = 1
            return (chosen_action_index)

        # Update the current state
        # current_state is made by (STATE_FRAMES) reduced images
        current_state = np.append(compressedImage,
                                  self._last_state[:,:,:,:-3],
                                  axis=3)

        # Store the last transition
        new_observation = [0 for _ in range(5)]
        new_observation[self._LAST_STATE_IDX] = self._last_state.copy()
        new_observation[self._ACTION_IDX] = self._last_action.copy()
        new_observation[self._REWARD_IDX] = reward
        new_observation[self._CURR_STATE_IDX] = current_state.copy()
        new_observation[self._TERMINAL_IDX] = done
        self._previous_observations.append(new_observation)
        self._last_state = current_state.copy()

        # If the memory is full, pop the oldest stored transition
        while len(self._previous_observations) >= self._REPLAY_MEMORY:
            self._previous_observations.pop(0)

        # Only train and decide after enough episodes of random play
        if episode > self._EPISODES_PURE_EXPLORATION:

            for _ in range(self._TRAINING_PER_STAGE):
                self._train()

            # Chose the next action with an epsilon-greedy approach
            if np.random.random() > self._epsilon:
                value_per_action = self._nn.predict(self._last_state)
                chosen_action_index = np.argmax(value_per_action)
            else:
                chosen_action_index = np.random.randint(0, self._N_ACTIONS)

        else:
            chosen_action_index = np.random.randint(0, self._N_ACTIONS)

        next_action_vector = np.zeros([self._N_ACTIONS])
        next_action_vector[chosen_action_index] = 1
        self._last_action = next_action_vector

        return (chosen_action_index)

    def _train(self):

        # Sample a mini_batch to train on
        permutations = np.random.permutation(
            len(self._previous_observations))[:self._MINIBATCH_SIZE]
        previous_states = np.concatenate(
            [self._previous_observations[i][self._LAST_STATE_IDX]
            for i in permutations],
            axis=0)
        actions = np.concatenate(
            [[self._previous_observations[i][self._ACTION_IDX]]
            for i in permutations],
            axis=0)
        rewards = np.array(
            [self._previous_observations[i][self._REWARD_IDX]
            for i in permutations]).astype('float')
        current_states = np.concatenate(
            [self._previous_observations[i][self._CURR_STATE_IDX]
            for i in permutations],
            axis=0)
        done = np.array(
            [self._previous_observations[i][self._TERMINAL_IDX]
            for i in permutations]).astype('bool')

        # Calculates the value of the current_states (per action)
        valueCurrentstates = self._nn.predict(current_states)

        # Calculate the empirical target value for the previous_states
        valuePreviousstates = rewards.copy()
        valuePreviousstates += ((1. - done) *
                                self._DISCOUNT_FACTOR *
                                valueCurrentstates.max(axis=1))

        # Run a training step
        self._nn.fit(previous_states,
                          actions,
                          valuePreviousstates)

"""
Plain Feed Forward Neural Network
The chosen activation function is the Leaky ReLU function
"""
class FeedForwardNeuralNetwork:

    def __init__(self, n_actions, obs_space):

        # NN variables
        self._generateNetwork(n_actions, obs_space)
        self._previous_observations = []
        self._PARAMETERS_FILE_PATH = 'Parameters_CNN.ckpt'

    def _generateNetwork(self, n_actions, obs_space):
        """
        The network is implemented in TensorFlow
        Change this method if you wish to use a different library
        """

        import tensorflow as tf
        self._ALPHA = 1e-3              # learning rate
        RESIZED_SCREEN = 84
        self._STATE_FRAMES = 3         # states/images used for taking a decision

        # Graph for compressing the input image
        x, y, z = obs_space
        self._image_input_layer = tf.placeholder("float",
            [None, x, y, z])
        image_step_size_x = int(np.ceil(float(x / RESIZED_SCREEN)))
        image_step_size_y = int(np.ceil(float(y / RESIZED_SCREEN)))
        extra_pad_x = RESIZED_SCREEN - int(x / image_step_size_x)
        extra_pad_y = RESIZED_SCREEN - int(y / image_step_size_y)
        self._image_output_layer = tf.nn.max_pool(
                self._image_input_layer,
                ksize=[1, image_step_size_x, image_step_size_y, 1],
                strides=[1, image_step_size_x, image_step_size_y, 1],
                padding="VALID")

        # Function for compressing (and reshaping) the image
        self._compressImage = lambda obs : np.pad(
            self._session.run(
                self._image_output_layer,
                feed_dict={self._image_input_layer: np.array([obs])})/255.0,
            ((0,0), (0,extra_pad_x), (0,extra_pad_y), (0,0)),
            mode='constant')

        CONVOLUTION_FILTER_VECTOR = [6, 6, 4]
        CONVOLUTION_STRIDE_VECTOR = [3, 3, 2]
        CONVOLUTION_KERNEL_VECTOR = [16, 16, 36]
        CONVOLUTION_INPUT_VECTOR = ([z * self._STATE_FRAMES] +
                                    CONVOLUTION_KERNEL_VECTOR[:-1])
        FEED_FWD_VECTOR = [(3**2) * CONVOLUTION_KERNEL_VECTOR[-1], 64,
                           n_actions]

        # The chosen activation function is the Leaky ReLU function
        self._activation = lambda x : tf.maximum(0.01*x, x)

        # Initialization parameters
        INITIALIZATION_STDDEV = 0.1
        INITIALIZATION_MEAN = 0.00
        INITIALIZATION_BIAS = -0.001

        # Convolutional layers
        self._input_layer = tf.placeholder("float",
                                           [None,
                                            RESIZED_SCREEN,
                                            RESIZED_SCREEN,
                                            z * self._STATE_FRAMES])
        self._convolutional_weights = []
        self._convolutional_bias = []
        self._hidden_convolutional_layer = [self._input_layer]

        for i in range(len(CONVOLUTION_FILTER_VECTOR)):
            self._convolutional_weights.append(tf.Variable(tf.truncated_normal(
                [CONVOLUTION_FILTER_VECTOR[i],
                 CONVOLUTION_FILTER_VECTOR[i],
                 CONVOLUTION_INPUT_VECTOR[i],
                 CONVOLUTION_KERNEL_VECTOR[i]],
                mean=INITIALIZATION_MEAN,
                stddev=INITIALIZATION_STDDEV)))
            self._convolutional_bias.append(tf.Variable(tf.constant(
                INITIALIZATION_BIAS,
                shape=[CONVOLUTION_KERNEL_VECTOR[i]])))
            self._hidden_convolutional_layer.append(
                self._activation(tf.nn.conv2d(
                                    self._hidden_convolutional_layer[i],
                                    self._convolutional_weights[i],
                                    strides=[1,
                                             CONVOLUTION_STRIDE_VECTOR[i],
                                             CONVOLUTION_STRIDE_VECTOR[i],
                                             1],
                                    padding="VALID")
                                + self._convolutional_bias[i]))

        # Feed forward layers
        self._hidden_activation_layer = [tf.reshape(
            self._hidden_convolutional_layer[-1],
            [-1, FEED_FWD_VECTOR[0]])]
        self._feed_forward_weights = []
        self._feed_forward_bias = []

        for i in range(len(FEED_FWD_VECTOR) - 2):
            self._feed_forward_weights.append(tf.Variable(tf.truncated_normal(
                [FEED_FWD_VECTOR[i],
                 FEED_FWD_VECTOR[i+1]],
                mean=INITIALIZATION_MEAN,
                stddev=INITIALIZATION_STDDEV)))
            self._feed_forward_bias.append(tf.Variable(tf.constant(
                INITIALIZATION_BIAS, shape=[FEED_FWD_VECTOR[i+1]])))
            self._hidden_activation_layer.append(self._activation(
                    tf.matmul(self._hidden_activation_layer[i],
                              self._feed_forward_weights[i])
                    + self._feed_forward_bias[i])
                    )

        # The calculation of the state-action value function does not
        # require the neurons' activation function
        self._feed_forward_weights.append(tf.Variable(tf.truncated_normal(
            [FEED_FWD_VECTOR[-2],
             FEED_FWD_VECTOR[-1]],
            mean=INITIALIZATION_MEAN,
            stddev=INITIALIZATION_STDDEV)))
        self._feed_forward_bias.append(tf.Variable(tf.constant(
            INITIALIZATION_BIAS,
            shape=[FEED_FWD_VECTOR[-1]])))
        self._state_value_layer = (tf.matmul(self._hidden_activation_layer[-1],
                                             self._feed_forward_weights[-1])
                                    + self._feed_forward_bias[-1])

        # Define the logic of the optimization
        self._action = tf.placeholder("float", [None, n_actions])
        self._target = tf.placeholder("float", [None])
        self._action_value_vector = tf.reduce_sum(tf.mul(
            self._state_value_layer, self._action), reduction_indices=1)
        self._cost = tf.reduce_sum(tf.square(
            self._target - self._action_value_vector))
        self._alpha = tf.placeholder('float')
        self._train_operation = tf.train.AdamOptimizer(
            self._alpha).minimize(self._cost)
        self._session = tf.Session()

        operation_intizializer = tf.initialize_all_variables()
        self._saver = tf.train.Saver()

        try:
            self._saver.restore(self._session, self._PARAMETERS_FILE_PATH)
            print ('Calibrated parameters SUCCESSFULLY LOADED.',
                   flush=True)
        except:
            self._session.run(operation_intizializer)
            print ('It was not possible to load calibrated parameters.',
                   flush=True)

        # Definition of feed_forward and optimization functions
        self._feedFwd = lambda state : self._session.run(
                            self._state_value_layer,
                            feed_dict={self._input_layer: state})

        self._backProp = lambda valueStates, actions, valueTarget : (
            self._session.run(self._train_operation,
            feed_dict={self._input_layer: valueStates,
                       self._action: actions,
                       self._target: valueTarget,
                       self._alpha : self._ALPHA}))

    def close(self):

        # If training, save the RAM memory to file
        if self._is_training:
            self._saver.save(self._session, self._PARAMETERS_FILE_PATH)

        # Close the session and clear TensorFlow's graphs
        from tensorflow.python.framework import ops
        ops.reset_default_graph()
        self._session.close()

    def predict(self, state):
        return(self._feedFwd(state))

    def fit(self, valueStates, actions, valueTarget):
        self._backProp(valueStates, actions, valueTarget)

if __name__=="__main__":
   main()


import moderngl
import numpy as np
from typing import Tuple

Color = Tuple[float, float, float, float]
Point2 = Tuple[float, float]

_SIGNAL_VERTEX_SHADER = '''
    #version 330

    uniform int Start;
    uniform float XScale;
    uniform mat4 MainRect;

    in int in_vert;

    void main() {
        float x = 2. * (Start + gl_VertexID) * XScale - 1.;
        float y = in_vert / 32768.;
        gl_Position = vec4(x, y, 0.0, 1.0) * MainRect;
    }
'''

_BASIC_FRAGMENT_SHADER = '''
    #version 330

    uniform vec4 Color;

    out vec4 f_color;

    void main() {
        f_color = Color;
    }
'''

_BASIC_VERTEX_SHADER = '''
    #version 330
    out int inst;
    void main() {
        inst = gl_InstanceID;
    }
'''

_LINE_GEOMETRY_SHADER = '''
    #version 330
    layout (points) in;
    layout (line_strip, max_vertices = 2) out;
    uniform mat4x2 LineMatrix;
    uniform int Segments;
    uniform mat4 MainRect;
    void main() {
        float start = float(2 * gl_PrimitiveIDIn) / Segments;
        gl_Position = vec2(start, 1) * LineMatrix * MainRect;
        EmitVertex();
        float end = float(2 * gl_PrimitiveIDIn + 1) / Segments;
        gl_Position = vec2(end, 1) * LineMatrix * MainRect;
        EmitVertex();
        EndPrimitive();
    }
'''

_GRID_GEOMETRY_SHADER = '''
    #version 330
    layout (points) in;
    layout (line_strip, max_vertices = 2) out;
    in int inst[1];
    uniform mat4 MainRect;
    uniform int HCount;
    uniform int VCount;
    uniform int Segments;
    void vline(float n, float m) {
        float x = 0.;
        if (HCount != 0)
            x = 2 * n / HCount - 1;
        float y1 = 4 * m / Segments - 1;
        float y2 = (4 * m + 2) / Segments - 1;
        gl_Position = vec4(x, y1, 0, 1) * MainRect;
        EmitVertex();
        gl_Position = vec4(x, y2, 0, 1) * MainRect;
        EmitVertex();
        EndPrimitive();
    }
    void hline(float n, float m) {
        float y = 0;
        if (VCount != 0)
            y = 2 * n / VCount - 1;
        float x1 = 4 * m / Segments - 1;
        float x2 = (4 * m + 2) / Segments - 1;
        gl_Position = vec4(x1, y, 0, 1) * MainRect;
        EmitVertex();
        gl_Position = vec4(x2, y, 0, 1) * MainRect;
        EmitVertex();
        EndPrimitive();
    }
    void main() {
        if (inst[0] <= HCount)
            vline(inst[0], gl_PrimitiveIDIn);
        else
            hline(inst[0] - HCount - 1, gl_PrimitiveIDIn);
    }
'''

_TEXTURE_VERTEX_SHADER = '''
    #version 330

    uniform mat4 MainRect;
    uniform mat2 TexScale;
    uniform int TexOffset;
    uniform int Height;

    in vec2 in_vert;
    in vec2 in_texcoord;

    out vec2 v_texcoord;

    void main() {
        v_texcoord = (in_texcoord - vec2(0, float(TexOffset) / Height))
                     * TexScale;
        gl_Position = vec4(in_vert[0], in_vert[1], 0.0, 1.0) * MainRect;
    }
'''

_TEXTURE_FRAGMENT_SHADER = '''
    #version 330

    uniform sampler2D Texture;
    uniform int Palette;

    in vec2 v_texcoord;
    out vec4 f_color;

    vec4 palette0(float v) {
        return vec4(v*v, v*v*v, v*(1-v)+v*v*v, 1.);
    }

    vec4 palette1(float v) {
        return vec4(v, v, v*v*v*v, v);
    }

    vec4 palette2(float v) {
        return vec4(v, 0, 0, v);
    }

    void main() {
        float val = texture(Texture, v_texcoord)[0];
        if (Palette == 2) {
            f_color = palette2(val);
        }
        else if (Palette == 1) {
            f_color = palette1(val);
        }
        else {
            f_color = palette0(val);
        }
        if (f_color.a < 0.1)
            discard;
    }
'''

class VertexArrayWrapper:

    def __init__(self, ctx: moderngl.Context, program: moderngl.Program,
                 contents: list, main_rect: np.ndarray):
        self.vao = ctx.vertex_array(program, contents)
        self.program['MainRect'] = tuple(main_rect.flatten())
        self.mode = None

    @property
    def program(self) -> 'moderngl.Program':
        return self.vao.program

    def render(self, mode=None, vertices=-1, **kargs):
        r_mode = mode if mode is not None else self.mode
        self.vao.render(mode=r_mode, vertices=vertices, **kargs)

class SignalVA(VertexArrayWrapper):
    VERTEX_WIDTH = 2

    def __init__(self, ctx: moderngl.Context, main_rect: np.ndarray,
                 max_count: int):
        program = ctx.program(
            vertex_shader=_SIGNAL_VERTEX_SHADER,
            fragment_shader=_BASIC_FRAGMENT_SHADER)
        self.buffer = ctx.buffer(
            reserve=max_count * self.VERTEX_WIDTH)
        super().__init__(ctx, program, [(self.buffer, 'i2', 'in_vert')],
                         main_rect)
        self.mode = moderngl.LINE_STRIP

    def render_data(self, points: np.ndarray, count: int, scale: int,
                    start: int, color: Color):
        self.buffer.orphan()
        self.buffer.write(points)
        self.program['Color'] = color
        self.program['XScale'] = 1 / (scale - 1)
        self.program['Start'] = start
        super().render(vertices=count)

class GridVA(VertexArrayWrapper):
    # Grid color
    GRID_COLOR = (0.4, 0.4, 0.4, 1.0)
    # Color for frame, center lines and ticks
    FRAME_COLOR = (0.8, 0.8, 0.8, 1.0)
    # Number of segments of grid lines
    GRID_SEGMENTS = 101

    def __init__(self, ctx: moderngl.Context, main_rect: np.ndarray,
                 horizontal_div: int, vertical_div: int):
        program = ctx.program(
            vertex_shader=_BASIC_VERTEX_SHADER,
            geometry_shader=_GRID_GEOMETRY_SHADER,
            fragment_shader=_BASIC_FRAGMENT_SHADER)
        super().__init__(ctx, program, [], main_rect)
        self.mode = moderngl.POINTS
        self.horizontal_div = horizontal_div
        self.vertical_div = vertical_div

    def render_fragment(self, hcount: int, vcount: int, segments: int,
                        color: Color, ticks: bool = False):
        self.program['HCount'] = hcount
        self.program['VCount'] = vcount
        self.program['Color'] = color
        self.program['Segments'] = segments
        segs = segments // 2 + 1 if not ticks else 1
        super().render(vertices=segs, instances=hcount + vcount + 2)
        # super().render(vertices=segs, instances=1)

    def render_grid(self):
        # Render basic grid
        self.render_fragment(self.horizontal_div, self.vertical_div,
                             self.GRID_SEGMENTS, self.GRID_COLOR)
        # Render central horizontal and vertical lines
        self.render_fragment(0, 0, self.GRID_SEGMENTS, self.FRAME_COLOR)

    def render_frame(self):
        # Render frame
        self.render_fragment(1, 1, 1, self.FRAME_COLOR)
        # Render ticks
        self.render_fragment(self.horizontal_div * 5, self.vertical_div * 5,
                             200, self.FRAME_COLOR, ticks=1)

class LineVA(VertexArrayWrapper):

    def __init__(self, ctx: moderngl.Context, main_rect: np.ndarray):
        program = ctx.program(
            vertex_shader=_BASIC_VERTEX_SHADER,
            geometry_shader=_LINE_GEOMETRY_SHADER,
            fragment_shader=_BASIC_FRAGMENT_SHADER)
        super().__init__(ctx, program, [], main_rect)
        self.mode = moderngl.POINTS

    def render_line(self, start: Point2, end: Point2, color: Color,
                    segments: int):
        self.program['LineMatrix'] = (end[0] - start[0], start[0],
                                      end[1] - start[1], start[1],
                                      0,                 0,
                                      0,                 1)
        self.program['Color'] = color
        self.program['Segments'] = segments
        super().render(vertices=segments // 2 + 1)

    def render_hline(self, y: float, color: Color, segments: int):
        self.render_line((-1, y), (1, y), color, segments)

    def render_vline(self, x: float, color: Color, segments: int):
        self.render_line((x, -1), (x, 1), color, segments)

class TextureVA(VertexArrayWrapper):

    def __init__(self, ctx: moderngl.Context, main_rect: np.ndarray,
                 width: int, height: int,
                 palette: int = 0):
        program = ctx.program(
            vertex_shader=_TEXTURE_VERTEX_SHADER,
            fragment_shader=_TEXTURE_FRAGMENT_SHADER)
        self.width = width
        self.height = height
        self.buffer = ctx.buffer(
            np.array([
                -1.0, -1.0,  0, 0,  # lower left
                -1.0,  1.0,  0, 1,  # upper left
                1.0, -1.0,  1, 0,  # lower right
                1.0,  1.0,  1, 1,  # upper right
            ], dtype=np.float32))
        super().__init__(ctx, program,
                         [(self.buffer, '2f4 2f4', 'in_vert', 'in_texcoord')],
                         main_rect)
        self.program['Height'] = height
        self.program['TexScale'] = (1, 0, 0, 1)
        self.program['Palette'] = palette
        self.mode = moderngl.TRIANGLE_STRIP
        self.texture = ctx.texture((width, height), 1, dtype='f4')
        # self.texture.filter = moderngl.NEAREST, moderngl.NEAREST
        self.texture.swizzle = 'R001'
        self.texture_buffer = ctx.buffer(reserve=width * height * 4)
        self.position = 0

    def render(self):
        self.texture.use(location=0)
        super().render()

    def write_data(self, data: np.ndarray):
        self.texture_buffer.write(data, offset=(
                self.height - self.position - 1)
                * self.width * 4)
        self.texture.write(self.texture_buffer)
        self.program['TexOffset'] = self.position
        self.position = (self.position + 1) % self.height

    def render_data(self, data: np.ndarray):
        self.write_data(data)
        self.render()

    def set_scale(self, scale_x: float):
        self.program['TexScale'] = (1 / scale_x, 0, 0, 1)

    def clear(self):
        self.texture_buffer.clear()


"""
legacy test for the coddiwomple.openmm package.
"""
# Import package, test suite, and other packages as needed
import numpy as np
import sys
import os
from simtk import openmm, unit
from openmmtools.testsystems import HarmonicOscillator
from coddiwomple.tests.utils import get_harmonic_testsystem
from coddiwomple.tests.utils import HarmonicAlchemicalState

# from coddiwomple.openmm.integrators import *
# from coddiwomple.openmm.propagators import *
# from coddiwomple.openmm.reporters import *
# from coddiwomple.openmm.states import *
# from coddiwomple.openmm.utils import *
# from pkg_resources import resource_filename
# import tempfile
import pickle
import tqdm
import mdtraj as md

def test_ais_propagator():
    """
    Test the OpenMM Annealed Importance Sampling Propagator and its subclasses
    """
    from perses.annihilation.lambda_protocol import RelativeAlchemicalState
    traj, factory, alchemical_functions = load_prereqs()
    num_frames = traj.n_frames
    pdf_state = OpenMMPDFState(system = factory.hybrid_system, alchemical_composability = RelativeAlchemicalState, pressure=None)
    n_steps = 10

    ais_integrator = OMMLIAIS(
                             alchemical_functions,
                             n_steps,
                             temperature=300.0 * unit.kelvin,
                             collision_rate=1.0 / unit.picoseconds,
                             timestep=1.0 * unit.femtoseconds,
                             splitting="P I H N S V R O R V",
                             constraint_tolerance=1e-6)

    ais_propagator = OMMAISP(openmm_pdf_state = pdf_state, integrator = ais_integrator, record_state_work_interval = 1, reassign_velocities=True)
    frames = np.random.choice(range(num_frames), 10)

    for num_trajectories in tqdm.trange(len(frames)):
        positions = traj.xyz[frames[num_trajectories]] * unit.nanometers
        particle_state = OpenMMParticleState(positions=positions, box_vectors=traj.unitcell_vectors[frames[num_trajectories]]*unit.nanometers)
        _, return_dict = ais_propagator.apply(particle_state, n_steps=n_steps, apply_pdf_to_context = True, reset_integrator = True)
        print(return_dict)

def test_smc_mcmc_resampler():
    """
    Test coddiwomple.openmm.coddiwomple.mcmc_smc_resampler
    """
    from coddiwomple.openmm.coddiwomple import mcmc_smc_resampler
    traj, factory, alchemical_functions = load_prereqs()
    default_functions = {'lambda_sterics_core':
                     lambda x: x,
                     'lambda_electrostatics_core':
                     lambda x: x,
                     'lambda_sterics_insert':
                     lambda x: 2.0 * x if x < 0.5 else 1.0,
                     'lambda_sterics_delete':
                     lambda x: 0.0 if x < 0.5 else 2.0 * (x - 0.5),
                     'lambda_electrostatics_insert':
                     lambda x: 0.0 if x < 0.5 else 2.0 * (x - 0.5),
                     'lambda_electrostatics_delete':
                     lambda x: 2.0 * x if x < 0.5 else 1.0,
                     'lambda_bonds':
                     lambda x: x,
                     'lambda_angles':
                     lambda x: x,
                     'lambda_torsions':
                     lambda x: x
                     }

    lambda_sequence = [{key: default_functions[key](q) for key in default_functions.keys()} for q in np.linspace(0,1,10)]
    _traj_filename = resource_filename('coddiwomple', f"/data/perses_data/fluorobenzene_zero_endstate_cache/vacuum.cache.pdb")

    particles = mcmc_smc_resampler(system = factory.hybrid_system,
                           endstate_cache_filename = _traj_filename,
                           directory_name = None,
                           trajectory_prefix = None,
                           md_topology = factory.hybrid_topology,
                           number_of_particles = 10,
                           parameter_sequence = lambda_sequence)

    assert all(particle.iteration == 9 for particle in particles)
    reduced_free_energy = 1.33
    tolerance = 0.2
    average_free_energy = np.mean([particle.cumulative_work for particle in particles])
    upper_threshold, lower_threshold = (reduced_free_energy + tolerance)/reduced_free_energy, (reduced_free_energy - tolerance)/reduced_free_energy
    assert average_free_energy / reduced_free_energy < upper_threshold and average_free_energy / reduced_free_energy > lower_threshold
    return particles


import torch
import torchvision
from torch import nn
from torch.autograd import Variable
import torch.nn.functional as F
from torchvision.utils import save_image
from torchvision.models import vgg16
import math
def weights_init_normal(m):
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        torch.nn.init.normal_(m.weight.data, 0.0, 0.02)
    elif classname.find('BatchNorm') != -1:
        torch.nn.init.normal_(m.weight.data, 1.0, 0.02)
        torch.nn.init.constant_(m.bias.data, 0.0)

class Generator(nn.Module):
    def __init__(self, scale_factor):
        upsample_block_num = int(math.log(scale_factor, 2))

        super(Generator, self).__init__()
        self.block1 = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=9, padding=4),
            nn.PReLU()
        )
        self.block2 = ResidualBlock(64)
        self.block3 = ResidualBlock(64)
        self.block4 = ResidualBlock(64)
        self.block5 = ResidualBlock(64)
        self.block6 = ResidualBlock(64)
        self.block7 = nn.Sequential(
            nn.Conv2d(64, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64)
        )
        block8 = [UpsampleBLock(64, 2) for _ in range(upsample_block_num)]
        block8.append(nn.Conv2d(64, 3, kernel_size=9, padding=4))
        self.block8 = nn.Sequential(*block8)

    def forward(self, x):
        block1 = self.block1(x)
        block2 = self.block2(block1)
        block3 = self.block3(block2)
        block4 = self.block4(block3)
        block5 = self.block5(block4)
        block6 = self.block6(block5)
        block7 = self.block7(block6)
        block8 = self.block8(block1 + block7)

        return (F.tanh(block8) + 1) / 2

class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()
        self.net = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=3, padding=1),
            nn.LeakyReLU(0.2),

            nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(64),
            nn.LeakyReLU(0.2),

            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.BatchNorm2d(128),
            nn.LeakyReLU(0.2),

            nn.Conv2d(128, 128, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(128),
            nn.LeakyReLU(0.2),

            nn.Conv2d(128, 256, kernel_size=3, padding=1),
            nn.BatchNorm2d(256),
            nn.LeakyReLU(0.2),

            nn.Conv2d(256, 256, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(256),
            nn.LeakyReLU(0.2),

            nn.Conv2d(256, 512, kernel_size=3, padding=1),
            nn.BatchNorm2d(512),
            nn.LeakyReLU(0.2),

            nn.Conv2d(512, 512, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(512),
            nn.LeakyReLU(0.2),

            nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(512, 1024, kernel_size=1),
            nn.LeakyReLU(0.2),
            nn.Conv2d(1024, 1, kernel_size=1)
        )

    def forward(self, x):
        batch_size = x.size(0)
        return F.sigmoid(self.net(x).view(batch_size))

class ResidualBlock(nn.Module):
    def __init__(self, channels):
        super(ResidualBlock, self).__init__()
        self.conv1 = nn.Conv2d(channels, channels, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(channels)
        self.prelu = nn.PReLU()
        self.conv2 = nn.Conv2d(channels, channels, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm2d(channels)

    def forward(self, x):
        residual = self.conv1(x)
        residual = self.bn1(residual)
        residual = self.prelu(residual)
        residual = self.conv2(residual)
        residual = self.bn2(residual)

        return x + residual

class UpsampleBLock(nn.Module):
    def __init__(self, in_channels, up_scale):
        super(UpsampleBLock, self).__init__()
        self.conv = nn.Conv2d(in_channels, in_channels * up_scale ** 2, kernel_size=3, padding=1)
        self.pixel_shuffle = nn.PixelShuffle(up_scale)
        self.prelu = nn.PReLU()

    def forward(self, x):
        x = self.conv(x)
        x = self.pixel_shuffle(x)
        x = self.prelu(x)
        return x
#   WGAN

import os
import pickle
import glob
import numpy as np
import torch
import torchvision
import torchvision.transforms as transforms
from PIL import Image
from torch.utils.data import Dataset
# HR_train_data_path = 'C:/Users/User/Desktop/GAN_data_set/imagenet64_train/Imagenet64_train'
HR_train_data_path = './SRGAN_training_data/*'
test_data_path = './SRGAN_test_data/*'

batch_size = 64
input_channels = 3
hr_height = 128
lr_height = 32
n_critic = 1
n_critic_D = 1
clip_value = 0.02
epochs = 100
num_epochs = 201
def unpickle(file):
    with open(file, 'rb') as fo:
        dict = pickle.load(fo)
    return dict

class ImageDataset(Dataset):
    def __init__(self, imgs, lr_transforms=None, hr_transforms=None):
        self.lr_transform = transforms.Compose(lr_transforms)
        self.hr_transform = transforms.Compose(hr_transforms)

        self.files = imgs

    def __getitem__(self, index):
        img =  Image.fromarray(self.files[index].astype('uint8'), 'RGB')
        img_lr = self.lr_transform(img)
        img_hr = self.hr_transform(img)

        return {'lr': img_lr, 'hr': img_hr}

    def __len__(self):
        return len(self.files)

class TestImageDataset(Dataset):
    def __init__(self, imgs, lr_transforms=None, hr_transforms=None):
        self.lr_transform = transforms.Compose(lr_transforms)
        self.hr_transform = transforms.Compose(hr_transforms)

        self.files = imgs

    def __getitem__(self, index):
        img =  Image.fromarray(self.files[index].astype('uint8'), 'RGB')
        img_lr = self.lr_transform(img)
        img_hr = self.hr_transform(img)

        return {'lr': img_lr, 'hr': img_hr}

    def __len__(self):
        return len(self.files)
def load_databatch(data_folder, idx):
    data_file_HR = os.path.join(HR_train_data_path, 'train_data_batch_')
    d_HR = unpickle(data_file_HR + str(idx))
    x_HR = d_HR['data'][:4000]

    data_size = x_HR.shape[0]

    hr_height2 = hr_height * hr_height
    x_HR = np.dstack((x_HR[:, :hr_height2], x_HR[:, hr_height2:2*hr_height2], x_HR[:, 2*hr_height2:]))
    x_HR = x_HR.reshape((x_HR.shape[0], hr_height, hr_height, 3))

    lr_transforms = [
                    transforms.Resize((hr_height//4, hr_height//4), Image.BICUBIC),
                    transforms.ToTensor() ]

    hr_transforms = [
                    transforms.Resize((hr_height, hr_height), Image.BICUBIC),
                    transforms.ToTensor() ]

    train_loader = torch.utils.data.DataLoader(ImageDataset(x_HR, lr_transforms=lr_transforms, hr_transforms=hr_transforms),
                        batch_size=batch_size, shuffle=True)
    return train_loader

def load_jpg(data_folder,batch_size = batch_size, shuffle=True):
#     print(data_folder)
    image_list = []
    print("start loading")
    for filename in glob.glob(data_folder): #assuming gif
        im = Image.open(filename)
        im = im.resize((hr_height, hr_height), Image.BICUBIC)
        im = np.array(im)
        if im.shape == (hr_height ,hr_height ,3):
            image_list.append(im)

    image = np.asarray(image_list)

    lr_transforms = [
                    transforms.Resize((lr_height, lr_height), Image.BICUBIC),
                    transforms.ToTensor()
                    ]

    hr_transforms = [
                    transforms.Resize((hr_height, hr_height), Image.BICUBIC),
                    transforms.ToTensor()
                    ]

    train_loader = torch.utils.data.DataLoader(TestImageDataset(image, lr_transforms=lr_transforms, hr_transforms=hr_transforms),
                        batch_size=batch_size, shuffle=shuffle)

    return train_loader

dataloader = load_jpg(HR_train_data_path)
test_data_set = load_jpg(test_data_path,shuffle = False)
# print(dataloader)
print("end")

# Loss function
class GeneratorLoss(nn.Module):
    def __init__(self):
        super(GeneratorLoss, self).__init__()
        vgg = vgg16(pretrained=True)
        loss_network = nn.Sequential(*list(vgg.features)[:31]).eval()
        for param in loss_network.parameters():
            param.requires_grad = False
        self.loss_network = loss_network
        self.mse_loss = nn.MSELoss()
        self.tv_loss = TVLoss()

    def forward(self, out_labels, out_images, target_images):
        # Adversarial Loss
        adversarial_loss = torch.mean(1 - out_labels)
        # Perception Loss
        perception_loss = self.mse_loss(self.loss_network(out_images), self.loss_network(target_images))
        # Image Loss
        image_loss = self.mse_loss(out_images, target_images)
        # TV Loss
        tv_loss = self.tv_loss(out_images)
        return image_loss + 0.001 * adversarial_loss + 0.006 * perception_loss + 2e-8 * tv_loss
class TVLoss(nn.Module):
    def __init__(self, tv_loss_weight=1):
        super(TVLoss, self).__init__()
        self.tv_loss_weight = tv_loss_weight

    def forward(self, x):
        batch_size = x.size()[0]
        h_x = x.size()[2]
        w_x = x.size()[3]
        count_h = self.tensor_size(x[:, :, 1:, :])
        count_w = self.tensor_size(x[:, :, :, 1:])
        h_tv = torch.pow((x[:, :, 1:, :] - x[:, :, :h_x - 1, :]), 2).sum()
        w_tv = torch.pow((x[:, :, :, 1:] - x[:, :, :, :w_x - 1]), 2).sum()
        return self.tv_loss_weight * 2 * (h_tv / count_h + w_tv / count_w) / batch_size

    @staticmethod
    def tensor_size(t):
        return t.size()[1] * t.size()[2] * t.size()[3]
# def generator_loss(logits_fake):
#     size = logits_fake.shape[0]
#     true_labels = Variable(torch.ones(size, 1)).float().cuda()
#     loss = torch.mean(logits_fake)
#     return loss
# def adversarial_loss(logits_real, logits_fake): # 判别器的 loss
#     size = logits_real.shape[0]
#     true_labels = Variable(torch.ones(size, 1)).float().cuda()
#     false_labels = Variable(torch.zeros(size, 1)).float().cuda()
#     loss = torch.nn.L1Loss(logits_real, true_labels) + torch.nn.L1Loss(logits_fake, false_labels)
#     return loss

generator_loss = GeneratorLoss().cuda()
discriminator_loss = torch.nn.MSELoss()
content_loss = nn.MSELoss(size_average = True)

#Optimizer

def psnr_cal(img1, img2):
    mse = np.mean( (img1 - img2) ** 2 )

    PIXEL_MAX = 1.0
    return 20 * np.log10(PIXEL_MAX / np.sqrt(mse))

cuda = True if torch.cuda.is_available() else False
Tensor = torch.cuda.FloatTensor if cuda else torch.Tensor
input_lr = Tensor(batch_size, input_channels, lr_height, lr_height).cuda()
input_hr = Tensor(batch_size, input_channels, hr_height, hr_height).cuda()
valid = Variable(torch.ones(batch_size, 1)).float().cuda()
fake = Variable(torch.zeros(batch_size, 1)).float().cuda()

# def debug_memory():
#     import collections, gc, torch
#     tensors = collections.Counter((str(o.device), o.dtype, tuple(o.shape))
#                                   for o in gc.get_objects()
#                                   if torch.is_tensor(o))
#     for line in sorted(tensors.items()):
#         print('{}\t{}'.format(*line))
# debug_memory()

def train_a_gan(D_net, G_net, D_optimizer, G_optimizer, generator_loss, sample_interval = 50, checkpoint_interval = 10):
    iter_count = 0
    current_batch = 0
    index = 0
    torch.cuda.empty_cache()
    for epoch in range(epochs, num_epochs):
        for i, img in enumerate(dataloader):
            if i == len(dataloader)-1:
                break;
            # model input
            imgs_lr = Variable(input_lr.copy_(img['lr'])).cuda()
            imgs_hr = Variable(input_hr.copy_(img['hr'])).cuda()
#             print(imgs_lr.size())

            fake_images = G_net(imgs_lr).detach()
            # Train Discriminator
#             if i %  n_critic_D == 0 :
            D_optimizer.zero_grad()

            #loss of real and fake image
            # Total loss

            # GAN

            # loss_D = discriminator_loss(D_net(imgs_hr), valid) + discriminator_loss(D_net(fake_images), fake)

            # WGAN
            loss_D = 1 - D_net(imgs_hr).mean() + D_net(fake_images).mean() # 判别器的 loss

            # RGAN
            # loss_D = -torch.mean(torch.log(torch.sigmoid(-Relative)))  # 判别器的 loss

            loss_D.backward()
            D_optimizer.step()

            # train generators
#             if i % n_critic == 0:
            G_optimizer.zero_grad()
            # Generate Hr from low resolution input
            # adversarial loss
            gen_img =  G_net(imgs_lr)
            # GAN
            # gen_error = generator_loss(D_net(fake_images), valid)
            # WGAN
#             gen_error = -torch.mean(D_net(gen_img))
#             content_error = content_loss(gen_img , imgs_hr)

            # VGG_feature_loss
#                 gen_features = feature_extractor(fake_images)
#                 real_features = Variable(feature_extractor(imgs_hr).data, requires_grad=False)
#                 VGG_loss = content_loss(real_features , gen_features)
    #             print(imgs_hr.size())

            #Total loss
            # GAN
            real_out = D_net(imgs_hr).mean()
            fake_out = D_net(fake_images).mean()
            loss_G = generator_loss(fake_out, gen_img, imgs_hr)

            # WGAN
            # loss_G = 1e-2*gen_error + content_error

            # RGAN
            # Relative = D_net(fake_images)-D_net(imgs_hr)

            loss_G.backward()
            G_optimizer.step()

            print("[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f]" %
                                                            (epoch, num_epochs, i, len(dataloader),
                                                            loss_D.item(), loss_G.item()))

#             print(fake_images.data.cpu().numpy())
            # Save image sample turn to next batch
            if i % sample_interval == 0 :
                index = index+1
                psnr_val = 0
                for i , test_img in enumerate(test_data_set):
                    if i >= 1:
                        break;
#                     print(test_img)
                    test_imgs_lr = Variable(input_lr.copy_(test_img['lr'])).detach()
                    test_imgs_hr = Variable(input_hr.copy_(test_img['hr'])).detach()
                    test_fake_images = G_net(test_imgs_lr)

                    psnr_val += psnr_cal(test_fake_images.data.cpu().numpy(),test_imgs_hr.data.cpu().numpy())
                    print("[PSNR: %f]"%psnr_val)
#                     save_image(test_imgs_lr,"C:/Users/User/Desktop/Deep_Learning/GAN_images/TESTIMAGE'{0}'.png".format(i+(epoch)*batch_size), normalize=True)
                save_image(torch.cat((test_fake_images.data, test_imgs_hr.data), -2),
                                "./SRGAN_save_image/'{0}'.png".format(i+(epoch)*len(dataloader)), normalize=True)

#             if i == (len(dataloader)-1) :
#                 break;

        if checkpoint_interval != -1 and epoch % checkpoint_interval == 0:
            # Save model checkpoints
            torch.save(G_net.state_dict(), './SRGAN_save_model/generator_%d-RGAN.pth' % epoch)
            torch.save(D_net.state_dict(), './SRGAN_save_model/discriminator_%d-RGAN.pth' % epoch)

D = Discriminator().cuda()
G = Generator(4).cuda()
def get_optimizer_G(net):
    optimizer = torch.optim.Adam(net.parameters(), lr = 1e-4)
    return optimizer
def get_optimizer_D(net):
    optimizer = torch.optim.Adam(net.parameters(), lr = 1e-4)
    return optimizer
D_optim = get_optimizer_D(D)
G_optim = get_optimizer_G(G)

if epochs != 0:
    # Load pretrained models
    D.apply(weights_init_normal)
#    D.load_state_dict(torch.load('./SRGAN_save_model/discriminator_%d-RGAN.pth'% epochs))
    G.load_state_dict(torch.load('./SRGAN_save_model/generator_%d-RGAN.pth'% epochs))
else:
    # Initialize weights
    D.apply(weights_init_normal)
    G.apply(weights_init_normal)
train_a_gan(D, G, D_optim, G_optim, generator_loss)