#!/usr/bin/env python

'''
    Construct SoS using integrated orbits.
'''

import numpy as np
from scipy.interpolate import interp1d
from scipy.optimize import brentq

def sos(orb, axis=0, direction='+'):

    ''' Construct SoS for integrated orbits. '''

    # interpolator,
    k_intp = interp1d(np.arange(orb.shape[0]), orb[:, axis], \
                      copy=True, assume_sorted=True)

    # find zeros,
    sign_t = 1.0 if direction == '+' else -1.0
    i_zeros = np.arange(orb.shape[0] - 1)[ \
              (orb[1:, axis] * orb[:-1, axis] < 0.) \
            * (np.sign(orb[1:, axis + 2]) * sign_t > 0.)]
    t_zeros = np.array([brentq(k_intp, zi, zi + 1) for zi in i_zeros])

    # interpolate other axes,
    w_intp = interp1d(np.arange(orb.shape[0]), orb, \
                      axis=0, copy=False, assume_sorted=True)

    # return surface if sectuibm
    return w_intp(t_zeros)


#!/usr/bin/env python
# coding=utf-8

__author__ = 'Minglong Li'
__date__ = 'Nov. 2017'

import numpy as np
import random
import math
import copy
import sys
''' A Python program using MCTS-UC (Monte Carlo Tree Search with useful cycles) for multi-robot patrol.'''

ROBOT_NUM = 2;
#MCTS scalar.  Larger scalar will increase exploitation, smaller will increase exploration.
SCALAR= math.sqrt(2.0)
GLOBAL_ID = 0 #维护一个全局ＩＤ，每个ｎｏｄｅ有且仅有一个独立ＩＤ，包括cyclic_node.

#robot_init_pose: robot_0:[1, 0]; roobt_1:[2, 2]
#create a map, 0 represents 'free', 1 represents 'obstacle'.
class Map():
    def __init__(self):
        self.map = [[0, 0, 1, 0, 0],
                    [0, 0, 1, 0, 0],
                    [1, 0, 0, 0, 1],
                    [0, 1, 0, 1, 0],
                    [0, 0, 0, 0, 0]]
        self.horizontal_length = 5
        self.vertical_length = 5

    def obstacles_in_map(self):
        temp_obstacles = []
        for i in range(self.horizontal_length):
            for j in range(self.vertical_length):
                if(self.map[i][j] == 1):
                    temp_obstacles.append([i,j])
        return temp_obstacles

    def random_pose_without_obstacles(self):
        obstacles = self.obstacles_in_map()
        while(True):
            temp_x = random.randint(0, self.horizontal_length-1)
            temp_y = random.randint(0, self.vertical_length-1)
            obstacle_flag = False
            for i in range(len(obstacles)):
                if(obstacles[i] == [temp_x, temp_y]):
                    obstacle_flag = True
                    break
            if(obstacle_flag == False):
                return [temp_x, temp_y]

    def change_scale(self):
        pass

    def add_obstacle(self):
        pass

class DynamicIntruder():
    def __init__(self):
        self.time_e = 0
        self.time_step = 0
        pass #TODO: version 0.0 of mcts_patrol, we use a stationary intruder.

    def update_time_step(self):
        pass

    def time_e(self):
        pass

    def time_e(self):
        pass

    def time_p(self):
        pass

    def pose(self, time):
        pass

class StationaryIntruder():
    def __init__(self):
        self.time_e = random.randint(2,10)#intruder出现的时刻，从第3个时间步开始计算，０，１时刻，ｉｎｔｒｕｄｅｒ机器人肯定不会出现。
        #从直觉上讲，time_p越大，在一定budget内，捕捉到intruder的概率越大，作者最小设置为９
        self.time_p = 9#intruder停留的时间为10个时间步。
        self.time_step = self.time_e
        self.map = Map()
        self.pose = self.map.random_pose_without_obstacles()

    def update_time_step(self):
        self.time_step = self.time_step + 1

    def update_initial_pose(self):
        self.pose = self.map.random_pose_without_obstacles()

    def get_time_e(self):
        return self.time_e

    def get_time_p(self):
        return self.time_p

    def get_pose(self):
        return self.pose

class Node():
    def __init__(self, actions = None, poses = None, parent = None, time_step = None, map = None, node_ID = None):
        self.visits = 1#初始化的时候就是为１的，只要ｃｈｉｌｄ被添加上，ｖｉｓｉｔｓ就是１．
        self.reward = 0.0
        self.children = []
        self.actions = actions
        self.poses = poses
        self.parent = parent
        self.time_step = time_step
        self.is_cyclic_arm = False
        self.cyclic_head_ID = None
        self.node_ID = node_ID

        if(time_step == None):
            time_step = 0#default parameter
        self.map = map
        if(self.map == None):
            self.map = Map()#default parameter

    #已经检查过的child
    def add_child(self, child):
        global GLOBAL_ID
        temp_child = copy.deepcopy(child)
        temp_child.node_ID = GLOBAL_ID
        self.children.append(temp_child)
        GLOBAL_ID += 1

    #根据当前节点的actions和poses算出child_poses,　child_actions经过传参得到,再算出grand_son，检查合法性。
    def check_and_add_child_actions(self,child_actions):
        global ROBOT_NUM
        global GLOBAL_ID
        #temp_actions = []#temp_actions是个一维数组，维度是robot_num
        child_poses = []#child_poses是个二维数组,第一个维度是robot_num
        #子节点的位置是由当前节点的动作和位置决定的，子节点对应的各个机器人动作可以指定。

        child_poses = NEXT_POSES(copy.deepcopy(self.actions), copy.deepcopy(self.poses))
        '''
        for(i = 0; i < ROBOT_NUM; i+=1):
            j = self.actions[i]#up, down, left, right
            if(j == 0):#up
                child_poses[i][0] = self.poses[0]
                child_poses[i][1] = self.poses[1] - 1
            if(j == 1):#down
                child_poses[i][0] = self.poses[0]
                child_poses[i][1] = self.poses[1] + 1
            if(j == 2):#left
                child_poses[i][0] = self.poses[0] - 1
                child_poses[i][1] = self.poses[1]
            if(j == 3):#right
                child_poses[i][0] = self.poses[0] + 1
                child_poses[i][1] = self.poses[1]
        '''
        #因为self节点的上一层节点在add_child的时候，会检测碰撞，所以，当前节点的poses不会有碰撞。
        #但是，根据child_poses和指定的child_actions加入grandson节点的poses，就会可能碰撞
        grandson_poses = []

        grandson_poses = NEXT_POSES(copy.deepcopy(child_actions), copy.deepcopy(child_poses))
        '''
        for(i = 0; i < ROBOT_NUM; i+=1):
            j = child_actions[i]
            if(j == 0):#up
                grandson_poses[i][0] = child.poses[0]
                grandson_poses[i][1] = child_poses[1] - 1
            if(j == 1):#down
                grandson_poses[i][0] = child_poses[0]
                grandson_poses[i][1] = child_poses[1] + 1
            if(j == 2):#left
                grandson_poses[i][0] = child_poses[0] - 1
                grandson_poses[i][1] = child_poses[1]
            if(j == 3):#right
                grandson_poses[i][0] = child_poses[0] + 1
                grandson_poses[i][1] = child_poses[1]
        '''
        #child节点的位置已经是无碰撞的了，需要用grandson的位置判断是否加入child节点。
        bump_flag = self.check_bump(copy.deepcopy(grandson_poses))
        if(not bump_flag):
            temp_child_actions = copy.deepcopy(child_actions)
            temp_child_poses = copy.deepcopy(child_poses)
            child = Node(actions = temp_child_actions, poses = temp_child_poses, parent = self, time_step = self.time_step + 1)
            child.node_ID = GLOBAL_ID
            self.children.append(child)
            GLOBAL_ID +=1

    def check_and_add_child_actions_poses(self, child_actions, child_poses):
        global GLOBAL_ID
        grandson_poses = NEXT_POSES(child_actions, child_poses)
        bump_flag = self.check_bump(grandson_poses)
        if(not bump_flag):
            temp_child_actions = copy.deepcopy(child_actions)
            temp_child_poses = copy.deepcopy(child_poses)
            child = Node(actions = temp_child_actions, poses = temp_child_poses, parent = self, time_step = self.time_step + 1)
            child.node_ID = GLOBAL_ID
            self.children.append(child)
            GLOBAL_ID += 1

    def check_bump(self, poses):
        global ROBOT_NUM
        #bump_flag = False#bump into obstacles, or bump into boundings
        for i in range(ROBOT_NUM):
            if(poses[i][0] == -1 or poses[i][0] == self.map.horizontal_length):#bump into left or right boundings
                #bump_flag = True
                #break
                return True
            if(poses[i][1] == -1 or poses[i][1] == self.map.vertical_length):#bump into up or down boundings
                #bump_flag = True
                #break
                return True
            else:
                obstacles = self.map.obstacles_in_map()#a list storing the coordinates of the obstacles
                for j in range(len(obstacles)):
                    if poses[i] == obstacles[j]:
                    #if(poses[i][0] == obstacles[j][0] and poses[i][1] == obstacles[j][1]):
                        #bump_flag == True
                        #break
                        return True
                #break
        #return bump_flag
        return False

    def update(self,reward):
        self.reward+=reward
        self.visits+=1

    def ucb(self):
        visits_parent = self.parent.visits
        visits_child = self.visits
        w_pi = self.reward
        return w_pi + SCALAR * math.sqrt(math.log(visits_parent/visits_child))

    def create_cyclic_sibling_arm(self, cyclic_head_ID):
        global GLOBAL_ID
        temp_parent = self.parent
        new_node = copy.deepcopy(self)
        new_node.cyclic_head_ID = cyclic_head_ID
        new_node.is_cyclic_arm = True
        new_node.visits = 1
        new_node.node_ID = GLOBAL_ID
        temp_parent.children.append(new_node)
        GLOBAL_ID += 1

    def get_is_cyclic_arm(self):
        return self.is_cyclic_arm

    def node_equal(self, other):
        self_node_ID = self.node_ID
        other_node_ID = other.node_ID

        if(self_node_ID == other_node_ID):
            return True
        else:
            return False

    def is_fully_expanded(self):
        #To be verified: Version 0 of mucts_patrol, all of the nodes are fully expanded
        pass

#global functions
#UCT search
#def SELECT(budget, root):
#TODO:关于深度拷贝的检查，看到了这里。
def SELECT(root):
    #pass
    selected_node = Node()
    global_ucb = 0
    temp_ucb = 0
    temp_children = root.children #'children' is a list
    #print "244: SELECT 1"
    while(len(temp_children) != 0):
        #print "246: SELECT while"
        for i in range(len(temp_children)):
            '''
            print "248: SELECT for"
            print "temp_ucb: "
            print temp_children[i].parent.visits
            print temp_children[i].visits
            print temp_children[i].reward
            '''
            temp_ucb = temp_children[i].ucb()
            #print temp_ucb
            if(global_ucb < temp_ucb or global_ucb == temp_ucb):
                #print "SELECT if"
                global_ucb = temp_ucb
                selected_node = temp_children[i]
        #如果始终找不到比当前ｕｃｂ还小的，那么temp_children应该就是不发生变化的？
        #print "SELECT 2"
        temp_children = selected_node.children#TODO:还是这里存在问题，一直下去肯定是一个Ｎｏｎｅ。２１点４２分。
        #len(temp_children)　等于０的时候不就跳出了吗？此时Ｓｅｌｅｃｔｅｄ应该不是０啊。
        global_ucb = 0#TODO:这里对吗？很有可能是个雷啊。
        #print "len(temp_children):"
        #print len(temp_children)
    return selected_node

def ITERATIVE_LOOP(temp_node, cycles, temp_actions):
    global ROBOT_NUM
    cycles -= 1
    index = ROBOT_NUM - cycles -1

    if(cycles==-1):
        return
    else:
        if(cycles == 0):
            for i in range(4):
                temp_actions[index] = i
                temp_node.check_and_add_child_actions(child_actions = temp_actions);
        else:
            for i in range(4):
                temp_actions[index] = i
                ITERATIVE_LOOP(temp_node, cycles, temp_actions)

#fully expand the node
def EXPAND(node):
    global ROBOT_NUM
    temp_actions = []
    for i in range(ROBOT_NUM):
        temp_actions.append(-1)
    ITERATIVE_LOOP(node, ROBOT_NUM, temp_actions)
    i = random.randint(0, len(node.children)-1)
    return node.children[i]

def NEXT_POSES(actions, poses):
    child_poses = []
    for index in range(ROBOT_NUM):
        child_poses.append([0,0])
    #print child_poses
    #print child_poses[0][0]
    #print actions[0]
    #print poses[0]
    for i in range(ROBOT_NUM):
        j = actions[i]#up, down, left, right
        if(j == 0):#up
            child_poses[i][0] = poses[i][0]
            child_poses[i][1] = poses[i][1] + 1
        if(j == 1):#down
            child_poses[i][0] = poses[i][0]
            child_poses[i][1] = poses[i][1] - 1
        if(j == 2):#left
            child_poses[i][0] = poses[i][0] - 1
            child_poses[i][1] = poses[i][1]
        if(j == 3):#right
            child_poses[i][0] = poses[i][0] + 1
            child_poses[i][1] = poses[i][1]
    return child_poses

def ITERATIVE_LOOP_ROOT_CHILD_ACTION(root, cycles, child_poses, child_actions):
    global ROBOT_NUM
    cycles -= 1
    index = ROBOT_NUM - cycles - 1

    if cycles == -1:
        return
    else:
        if cycles == 0:
            for i in range(4):
                child_actions[index] = i
                #if bump, not add. if not bump, add.
                root.check_and_add_child_actions_poses(child_actions, child_poses)
        else:
            for i in range(4):
                child_actions[index] = i
                ITERATIVE_LOOP_ROOT_CHILD_ACTION(root, cycles, child_poses, child_actions)

#root节点的初始化当中，action为空，poses是初始化指定的值。
def ITERATIVE_LOOP_ROOT_ACTION(root, cycles, root_actions):
    global ROBOT_NUM
    cycles -= 1
    index = ROBOT_NUM - cycles - 1
    count_num = 0
    #print "cycles in ITERATIVE_LOOP_ROOT_ACTION: "
    #print cycles
    #print "index in ITERATIVE_LOOP_ROOT_ACTION: "
    #print index
    if(cycles == -1):
        return
    else:
        if(cycles == 0):
            for i in range(4):
                #print "count_num:"
                #print count_num
                #print "index in if part: "
                #print index
                count_num +=1
                root_actions[index] = i
                #print "root_actions in if part:"
                print root_actions
                child_poses = NEXT_POSES(root_actions, root.poses)
                #print "child_poses"
                #print child_poses
                bump_flag = root.check_bump(child_poses)
                #print "bump_flag"
                #print bump_flag
                if bump_flag == True:
                    pass#检查合法性，如果不合法，continue。
                elif(bump_flag == False):
                    child_actions = []
                    for i in range(ROBOT_NUM):
                        child_actions.append(-1)
                    ITERATIVE_LOOP_ROOT_CHILD_ACTION(root=root, cycles=ROBOT_NUM, child_poses=child_poses,child_actions = child_actions)
        else:
            for i in range(4):
                #print "index in else part: "
                #print index
                root_actions[index] = i
                #print "root_actions in else part: "
                #print root_actions
                ITERATIVE_LOOP_ROOT_ACTION(root, cycles, root_actions)

def EXPAND_ROOT(root):
    global ROBOT_NUM
    root_actions = []
    for i in range(ROBOT_NUM):
        root_actions.append(-1)
    ITERATIVE_LOOP_ROOT_ACTION(root, ROBOT_NUM, root_actions)
    i = random.randint(0, len(root.children)-1)
    return root.children[i]

'''
#TODO:根据loop的逻辑写两个ITERATIVE LOOP
def loop(cycles):
    for a0 in range(4):
        for a1 in range(4):
            for a2 in range(4):
   　    　　　　    root_actions[0] = a0
   　    　　　　    root_actions[1] = a1
   　    　　　　    root_actions[2] = a2
   　    　　　　    root_poses = [[1,1],[2,2]]
   　    　　　　    child_poses = function(root_actions, root_poses)
                #对于root的话，CHILD_POSES在这里也要检查吗？

                for b0 in range(4):
                    for b1 in range(4):
                        for b2 in range(4):
                       　    child_actions[0] = a0
               　    　　　　    child_actions[1] = a1
               　    　　　　    child_actions[2] = a2

               　    　　　　    grandson_poses = function2(child_actions, child_poses);
               　    　　　　    bump_flag = check_bump(grandson_poses)

               　    　　　　    if bump_flag == False
               　    　　　　        add_child(child_poses, child_actions)

    for a0 in range(4):
       　root_actions[0] = a0
        for a1 in range(4):
           　root_actions[1] = a1
            for a2 in range(4):
                root_actions[2] = a2
                add_child(root_actions)

                print root_actions
'''

def RANDOM_ROLLOUT_FOR_ONE_STEP(actions, poses):
    global ROBOT_NUM
    map = Map()
    current_actions = actions
    current_poses = poses

    next_actions = []
    next_poses = []

    obstacles =  map.obstacles_in_map()

    for i in range(ROBOT_NUM):
        if(current_actions[i] == 0):#up
            next_poses.append( [current_poses[i][0], current_poses[i][1]-1] )
        if(current_actions[i] == 1):#down
            next_poses.append( [current_poses[i][0], current_poses[i][1]+1] )
        if(current_actions[i] == 2):#left
            next_poses.append( [current_poses[i][0]-1, current_poses[i][1]] )
        if(current_actions[i] == 3):#right
            next_poses.append( [current_poses[i][0]+1, current_poses[i][1]] )

    for i in range(ROBOT_NUM):
        while(True):
            temp_poses = [0,0]
            random_action = random.randint(0, 3)
            if(random_action == 0):#up
                temp_poses[0] = next_poses[i][0]
                temp_poses[1] = next_poses[i][1] + 1
            if(random_action == 1):#down
                temp_poses[0] = next_poses[i][0]
                temp_poses[1] = next_poses[i][1] - 1
            if(random_action == 2):#left
                temp_poses[0] = next_poses[i][0] - 1
                temp_poses[1] = next_poses[i][1]
            if(random_action == 3):#right
                temp_poses[0] = next_poses[i][0] + 1
                temp_poses[1] = next_poses[i][1]
            obstacle_flag = False
            for j in range(len(obstacles)):
                if(temp_poses == obstacles[j]):
                    obstacle_flag = True
                    break
            if(obstacle_flag == False):
                next_actions.append(random_action)
                break
    return next_actions, next_poses

def PERFORM_CYCLIC_ACTIONS(node, intruder):
    time_e = intruder.get_time_e()
    time_p = intruder.get_time_p()
    intruder_pose = intruder.get_pose()

    #status of the node to be evaluated
    temp_time_step = node.time_step
    #temp_poses = node.poses

    #temp_node = node.parent
    temp_node = node
    head_mark_node_ID = node.cyclic_head_ID
    #指针赋值＃TODO: TO BE REVISED>
    head_mark_node = node
    #arm_mark_node = node

    poses_buffer = []
    poses_buffer.append(temp_poses)
    while(not temp_node.node_equal(head_mark_node)):
        poses_buffer.append(temp_node.poses)
        temp_node = temp_node.parent

    index = 0
    index_increase_flag = True
    while(True):
        win_flag = False
        for i in range(len(poses_buffer[index])):
            if poses_buffer[index][i] == intruder_pose:
                win_flag = True
                break
        if temp_time_step < time_e + time_p:
            if win_flag == True:
                return 1
        elif temp_time_step == time_e + time_p:
            if win_flag == True:
                return 1
            elif win_flag == False:
                return 0

        if(index_increase_flag):
            index +=1
            temp_time_step+=1
            if index == len(poses_buffer):
                index-=1
                temp_time_step-=1
                index_increase_flag = False
        elif(not index_increase_flag):
            index -= 1
            temp_time_step+=1
            if index == -1:
                index+=2
                index_increase_flag = True

def ROLLOUT(node, intruder):
    #有两个前提：
    #第一个前提是当前ｎｏｄｅ上层的ｎｏｄｅ一定不是终止节点，ｓｅｌｅｃｔ和ｒｏｌｌｏｕｔ的都不是终止节点，只有当前ｎｏｄｅ和之后的ｎｏｄｅ可能是终止节点。
    #第二个前提是当前ｎｏｄｅ是经过ｅｘｐａｎｄ选出来的唯一一个待评估节点。
    #status of the intruder
    time_e = intruder.get_time_e()
    time_p = intruder.get_time_p()
    intruder_pose = intruder.get_pose()
    #status of the node to be evaluated
    temp_time_step = node.time_step
    temp_actions = node.actions
    temp_poses = node.poses

    next_actions = temp_actions
    next_poses = temp_poses

    #print "rollout 1"

    if(node.get_is_cyclic_arm() == True):
        #print "rollout 2"
        return PERFORM_CYCLIC_ACTIONS(node, intruder)
    else:
        #print "rollout 3"
        while(True):
            #print "rollout 4"
            win_flag = False
            for i in range(len(next_poses)):
                if(next_poses[i] == intruder_pose):
                    win_flag = True
                    break
            #if the intruder completing intruding or capture the intruder
            if(temp_time_step < time_e + time_p):
                if win_flag == True:
                    return 1
                elif win_flag == False:
                    #侵入完成时间也没到，机器人也没找着，继续搜索,TODO:一个idea，是否把找到的时间也算在内呢？
                    temp_time_step += 1
                    next_actions, next_poses = RANDOM_ROLLOUT_FOR_ONE_STEP(next_actions, next_poses)
            elif(temp_time_step == time_e + time_p):#入侵完成时间到了
                if(win_flag == True):
                    return 1
                elif(win_flag == False):
                    return 0

def ONE_EQUAL_IN_LIST_DEL(list_0, list_1):
    equal_flag = False
    for i in range(len(list_0)):
        for j in range(len(list_1)):
            if(list_0[i] == list_1[j]):
                del list_0[i]
                del list_1[j]
                equal_flag = True
                return list_0, list_1
    if(equal_flag == False):
        list_0[0] = "none"
        return list_0, list_1

#只需要node_0和node_1当中的所有的机器人的位置都有一一相等的，不用顺序相等。
def STATUS_EQUAL(node_0, node_1):
    global ROBOT_NUM
    poses_0 = []
    poses_1 = []
    for i in range(ROBOT_NUM):
        poses_0.append(node_0.poses[i])
        poses_1.append(node_1.poses[i])
    #poses_0 = node_0.poses
    #poses_1 = node_1.poses
    if(len(poses_0) != len(poses_1)):
        return False
    else:
        equal_flag = False
        while(True):
            poses_0, poses_1 = ONE_EQUAL_IN_LIST_DEL(poses_0, poses_1)
            if(len(poses_0) == 0 and len(poses_1) == 0):
                return True
            if(poses_0[0] == "none"):
                return False
                #TODO:对应的ｉ个机器人的概念在这里没有考虑，索引号没有加进来。

def BACK_PROPAGATION(leaf_node, reward):
    leaf_node.update(reward)
    temp_node = leaf_node.parent

    while(True):
        if(STATUS_EQUAL(temp_node, leaf_node)):
            temp_node_ID = temp_node.node_ID
            leaf_node.create_cyclic_sibling_arm(cyclic_head_ID = temp_node_ID)#it is an arm
        temp_node.update(reward)
        temp_node = temp_node.parent
        if temp_node == None:
            break

def IS_INTRUDER_CAPTURED(node, intruder):
    temp_poses = node.poses;
    intruder_pose = intruder.pose
    win_flag = False
    for i in range(len(temp_poses)):
        if(temp_poses[i] == intruder_pose):
            win_flag = True
            break
    if win_flag:
        return True
    else:
        return False

if __name__=="__main__":
    map = Map()
    intruder = StationaryIntruder()
    BUDGET = 100000;
    reward = 0.0;
    #初始的机器人的位置是(1, 0), (1, 1)
    root = Node(actions = None, poses = [[1,0],[1,1]], parent = None, time_step = 0, map = map)
    root.node_ID = -1
    print "Hello 1"
    EXPAND_ROOT(root)

    #print len(root.children)
    print len(root.children)

    cycle_count = 0
    print "Hello 2"
    while(cycle_count < BUDGET):
        print "Hello 3"
        leaf_node = SELECT(root)#TODO:这里需要判断node是不是时间终止节点。
        '''
        print "639, leaf_node.actions: "
        print leaf_node.actions
        print "641, leaf_node.poses: "
        print leaf_node.poses
        '''
        print "Hello 4"
        if leaf_node.time_step == intruder.time_e + intruder.time_p:
            print "Hello 5"
            if IS_INTRUDER_CAPTURED(leaf_node, intruder) == True:
                reward = 1
            else:
                reward = 0
            BACK_PROPAGATION(leaf_node, reward)
        elif leaf_node.time_step < intruder.time_e + intruder.time_p:
            print "Hello 6"
            if leaf_node.get_is_cyclic_arm():
                print "Hello 7"
                reward = ROLLOUT(leaf_node, intruder)
                BACK_PROPAGATION(leaf_node, reward)
            else:

                print "Hello 8"

                print "664, leaf_node.actions: "
                print leaf_node.actions
                print "666, leaf_node.poses: "
                print leaf_node.poses

                added_node = EXPAND(leaf_node)#TODO:测试过程中，这一步极其花费时间。？？？？？？？？
                #TODO：leaf_node是不是一定不是cyclic arm？？？？？？？？？？？
                print len(leaf_node.children)#只有12的是肯定不对的。
                print "Hello 9"
                reward = ROLLOUT(added_node, intruder)
                print "Hello 10"
                BACK_PROPAGATION(added_node, reward)
                print "Hello 11"
        elif leaf_node.time_step > intruder.time_e + intruder.time_p:
            print "程序异常，异常点标号为１"
        intruder.update_initial_pose()
        print "cycle_count :"
        print cycle_count
        cycle_count += 1


# Generate adversarial examples for EAT

import numpy as np
import tensorflow as tf

def gen_examples_fgs(discriminator, train_data, config):
    x = train_data[0]
    x = 2.0*x - 1.0

    epsilon = config['epsilon']
    d_out = discriminator(x)
    y = tf.stop_gradient(tf.argmax(d_out, 1))
    d_loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(
        logits=d_out, labels=tf.one_hot(y,10)))

    grad_x, = tf.gradients(d_loss, x)
    x_adv = tf.stop_gradient(x + epsilon*tf.sign(grad_x))
    x_adv = tf.clip_by_value(x_adv, -1.0, 1.0)

    x_adv = (x_adv+1.0)/2.0

    return x_adv

# least likely class
def gen_examples_ll(discriminator, train_data, config):
    x = train_data[0]
    x = 2.0*x - 1.0
    epsilon = config['epsilon']

    d_out = discriminator(x)
    yhat = tf.stop_gradient(tf.argmin(d_out, 1))

    d_loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(
        logits=d_out, labels=tf.one_hot(yhat,10)))

    grad_x, = tf.gradients(d_loss, x)
    x_adv = tf.stop_gradient(x - epsilon*tf.sign(grad_x))
    x_adv = tf.clip_by_value(x_adv, -1.0, 1.0)

    x_adv = (x_adv+1.0)/2.0

    return x_adv

def gen_examples_pgd(discriminator, train_data, config):
    x0 = train_data[0]

    x0 = 2.0*x0 - 1.0
    y = train_data[1]
    epsilon = config['epsilon']
    class_num = config['class_num']
    pgd_iter = config['pgd_iter']

    step_size = epsilon*0.25

    # randomize
    x = x0 + tf.random_uniform(x0.shape, -epsilon, epsilon)
    for i in range(pgd_iter):
        d_out = discriminator(x)
        d_loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(
            logits=d_out, labels=tf.one_hot(y,class_num)))

        grad_x, = tf.gradients(d_loss, x)
        x = tf.stop_gradient(x + step_size*tf.sign(grad_x))
        x = tf.clip_by_value(x, x0 - epsilon, x0 + epsilon)
        x = tf.clip_by_value(x, -1.0, 1.0)

    x_adv = (x+1.0)/2.0

    return x_adv

# one model def for one model filename
def gen_adv_examples(discriminator, ckpt_filename, train_data, test_data, config):

    batch_size = config['batch_size']
    x_fgs = gen_examples_fgs(discriminator, train_data, config)
    x_pgd = gen_examples_pgd(discriminator, train_data, config)
    x_ll = gen_examples_ll(discriminator, train_data, config)

    var_list = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope='discriminator')
    saver = tf.train.Saver(var_list)

    global_step = tf.Variable(0, trainable=False)
    global_step_op = global_step.assign_add(1)

    time_diff = tf.placeholder(tf.float32)
    Wall_clock_time = tf.Variable(0., trainable=False)
    update_Wall_op = Wall_clock_time.assign_add(time_diff)

    #with sv.managed_session() as sess:
    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())

        coord = tf.train.Coordinator()
        threads = tf.train.start_queue_runners(sess=sess, coord=coord)

        saver.restore(sess, ckpt_filename)

        train_size = config['train_size']
        samples_fgs, samples_pgd, samples_ll, x, y = sess.run([x_fgs, x_pgd, x_ll, train_data[0], train_data[1]])

        num_steps = int(train_size/batch_size)
        my_dim = list(samples_fgs.shape)
        my_dim = [my_dim[0]*4*num_steps] + my_dim[1:]
        examples = np.zeros(my_dim)
        labels = np.zeros(4*num_steps*batch_size)
        i=0
        examples[(4*i*batch_size):((4*i+1)*batch_size),:,:,:] = samples_fgs
        examples[((4*i+1)*batch_size):((4*i+2)*batch_size),:,:,:] = samples_pgd
        examples[((4*i+2)*batch_size):((4*i+3)*batch_size),:,:,:] = samples_ll
        examples[((4*i+3)*batch_size):((4*i+4)*batch_size),:,:,:] = x

        labels[(4*i*batch_size):((4*i+1)*batch_size)] = y
        labels[((4*i+1)*batch_size):((4*i+2)*batch_size)] = y
        labels[((4*i+2)*batch_size):((4*i+3)*batch_size)] = y
        labels[((4*i+3)*batch_size):((4*i+4)*batch_size)] = y

        for i in range(1, num_steps):
            samples_fgs, samples_pgd, samples_ll, x, y = sess.run([x_fgs, x_pgd, x_ll, train_data[0], train_data[1]])

            examples[(4*i*batch_size):((4*i+1)*batch_size),:,:,:] = samples_fgs
            examples[((4*i+1)*batch_size):((4*i+2)*batch_size),:,:,:] = samples_pgd
            examples[((4*i+2)*batch_size):((4*i+3)*batch_size),:,:,:] = samples_ll
            examples[((4*i+3)*batch_size):((4*i+4)*batch_size),:,:,:] = x

            labels[(4*i*batch_size):((4*i+1)*batch_size)] = y
            labels[((4*i+1)*batch_size):((4*i+2)*batch_size)] = y
            labels[((4*i+2)*batch_size):((4*i+3)*batch_size)] = y
            labels[((4*i+3)*batch_size):((4*i+4)*batch_size)] = y

        # test data
        test_size = config['test_size']
        num_test_steps = int(test_size/batch_size)
        x_test, y_test = sess.run([test_data[0], test_data[1]])
        test_dim = list(x_test.shape)
        test_dim = [num_test_steps*batch_size]+test_dim[1:]
        test_examples = np.zeros(test_dim)
        test_labels = np.zeros(num_test_steps*batch_size)
        i = 0
        test_examples[(i*batch_size):((i+1)*batch_size),:,:,:] = x_test
        test_labels[(i*batch_size):((i+1)*batch_size)] = y_test

        for i in range(1, num_test_steps):
            x_test, y_test = sess.run([test_data[0], test_data[1]])
            test_examples[(i*batch_size):((i+1)*batch_size),:,:,:] = x_test
            test_labels[(i*batch_size):((i+1)*batch_size)] = y_test

        coord.request_stop()
        #sess.run(model.queue.close(cancel_pending_enqueues=True))
        coord.join(threads)

    # save numpy examples
    np.savez(config['sample_file'], examples=examples, labels=labels)
    np.savez(config['test_sample_file'], examples=test_examples, labels=test_labels)


#!/usr/bin/env python
# -*- coding: utf-8 -*

"""
decomp module
"""

__author__ = 'Dr. Janus Juul Eriksen, University of Bristol, UK'
__license__ = 'MIT'
__version__ = '0.8'
__maintainer__ = 'Dr. Janus Juul Eriksen'
__email__ = 'janus.eriksen@bristol.ac.uk'
__status__ = 'Development'

import numpy as np
from pyscf import gto
from typing import List, Dict, Union, Any

# component keys
COMP_KEYS = ['coul', 'exch', 'kin', 'solvent', 'nuc_att_glob', 'nuc_att_loc', 'nuc_att', 'xc', 'el', 'struct']

class DecompCls(object):
        """
        this class contains all decomp attributes
        """
        def __init__(self, loc: str = '', pop: str = 'mulliken', \
                     part = 'atoms', ndo: bool = False, multiproc: bool = False, \
                     gauge_origin: Union[List[Any], np.ndarray] = np.zeros(3), \
                     prop: str = 'energy', write: str = '', verbose: int = 0) -> None:
                """
                init molecule attributes
                """
                # set system defaults
                self.loc = loc
                self.pop = pop
                self.part = part
                self.ndo = ndo
                self.multiproc = multiproc
                self.gauge_origin = gauge_origin
                self.prop = prop
                self.write = write
                self.verbose = verbose
                # set internal defaults
                self.res: Dict[str, np.ndarray] = {comp_key: None for comp_key in COMP_KEYS}
                self.charge_atom: np.ndarray = None
                self.dist: np.ndarray = None
                self.weights: np.ndarray = None
                self.centres: np.ndarray = None

def sanity_check(mol: gto.Mole, decomp: DecompCls) -> None:
        """
        this function performs sanity checks of decomp attributes
        """
        # localization procedure
        assert decomp.loc in ['', 'fb', 'pm', 'ibo-2', 'ibo-4'], \
            'invalid localization procedure. valid choices: none (default), `fb`, `pm`, `ibo-2`, and `ibo-4`'
        # population scheme
        assert decomp.pop in ['mulliken', 'iao'], \
            'invalid population scheme. valid choices: `mulliken` (default) or `iao`'
        # partitioning
        assert decomp.part in ['atoms', 'eda', 'orbitals'], \
            'invalid partitioning. valid choices: `atoms` (default), `eda`, or `orbitals`'
        # ndo decomposition
        assert isinstance(decomp.ndo, bool), \
            'invalid ndo argument. must be a bool'
        # multiprocessing
        assert isinstance(decomp.multiproc, bool), \
            'invalid multiprocessing argument. must be a bool'
        # gauge origin
        assert isinstance(decomp.gauge_origin, (list, np.ndarray)), \
            'invalid gauge origin. must be a list or numpy array of ints/floats'
        # property
        assert decomp.prop in ['energy', 'dipole'], \
            'invalid property. valid choices: `energy` (default) and `dipole`'
        # write
        assert isinstance(decomp.write, str), \
            'invalid write format argument. must be a str'
        assert decomp.write in ['', 'cube', 'numpy'], \
            'invalid write format. valid choices: `cube` and `numpy`'
        # verbosity
        assert isinstance(decomp.verbose, int), \
            'invalid verbosity. valid choices: 0 <= `verbose` (default: 0)'
        assert 0 <= decomp.verbose, \
            'invalid verbosity. valid choices: 0 <= `verbose` (default: 0)'


""" Tools for parsing Line List data
"""
from __future__ import print_function, absolute_import, division, unicode_literals

import numpy as np
import os, imp, glob, pdb, gzip, sys
if not sys.version_info[0] > 2:
    import codecs
    open = codecs.open

from astropy import units as u
from astropy.units.quantity import Quantity
from astropy.io import fits, ascii
from astropy.table import Column, Table, vstack

from ..abund import roman, ions
from ..abund.elements import ELEMENTS

lt_path = imp.find_module('linetools')[1]

# TODO
# Ingest AGN lines
# Add Ej, Ek, Ex for emission lines (specially Balmer, Paschen and Brackett)

#
def line_data(nrows=1):
    """ Defines the dict (and/or Table) for spectral line Data

    Parameters
    ----------
    nrows : int, optional
      Number of rows in Table [default = 1]

    Notes
    -----
    Group definition:
       *    0: None
       *    1: "All" ISM (intended to be all atomic lines ever observed)
       *    2: Strong ISM
       *    4: HI Lyman series
       *    8: H2
       *   16: CO
       *   32: EUV
       *   64: Galaxy Emission
       *  128: Galaxy Absorption
       *  256: AGN
       *  512: ??
       * 1024: User1 (Reserved)
       * 2048: User2 (Reserved)
    """
    ldict = {
        'name': ' '*20,       # Name
        'wrest': 0.*u.AA,     # Rest Wavelength (Quantity)
        'f':  0.,             # Oscillator strength
        'Ej': 0./u.cm,        # Energy of lower level (relative to ground state)
        'Ek': 0./u.cm,        # Energy of upper level (relative to ground state)
        'Ex': 0./u.cm,        # Excitation energy (cm^-1)
        'A': 0./u.s,          # Einstein coefficient
        'gj': 0,              # Lower statistical weight (2J+1)
        'gk': 0,              # Upper statistical weight (2J+1)
        'gamma': 0./u.s,      # Sum of A
        'nj': 0,              # Orbital level of lower state (or vibrational level)
        'nk': 0,              # Orbital level of upper state (or vibrational level)
        'Jj': 0.,             # Tot ang mom (z projection) of lower state (or rotation level)
        'Jk': 0.,             # Tot ang mom (z projection) of upper state (or rotation level)
        'el': 0,              # Electronic transition (2=Lyman (B-X), 3=Werner (C-X))
        'Z': 0,               # Atomic number (for atoms)
        'Am': 0,              # Mass number (often written as "A"; only used for D)
        'ion': 0,             # Ionic state (1=Neutral)
        'mol': ' '*10,        # Molecular name (H2, HD, CO, C13O)
        'Ref': ' '*50,        # References
        'group': 0            # Flag for grouping
        }

    # Table
    clms = []
    for key in ldict.keys():
        if type(ldict[key]) is Quantity:
            clm = Column( ([ldict[key].value]*nrows), name=key)
            clm.unit = ldict[key].unit
        else:
            clm = Column( [ldict[key]]*nrows, name=key)
        # Append
        clms.append(clm)

    # make it a masked Table so we can deal with Galaxy
    # emission and ISM absorption simultaneously by masking
    # out what does not make sense in one case or the other
    tbl = Table(clms, masked=True)
    # import pdb
    # pdb.set_trace()

    return ldict, tbl

def read_sets(infil=None):
    """ Read sets file

    Parameters
    ----------
    infil : str, optional
      Set file
    """
    if infil is None:
        fils = glob.glob(lt_path+'/lists/sets/llist_v*')
        fils.sort()
        infil = fils[-1] # Should grab the lateset
    # Read
    print('read_sets: Using set file -- \n  {:s}'.format(infil))
    set_data = ascii.read(infil, format='fixed_width')

    # Return
    return set_data

def read_euv():
    """ read additional EUV lines

    Returns
    -------
    Table of EUV lines
    """
    EUV_fil = lt_path + '/data/lines/EUV_lines.ascii'
    print('linetools.lists.parse: Reading linelist --- \n   {:s}'.format(EUV_fil))
    data = Table.read(EUV_fil, format='ascii', guess=False, comment=';',delimiter='|')

    # Units
    data['wrest'].unit = u.AA

    # Return
    return data

def read_H2():
    """ Simple def to read H2 data

    Returns
    -------
    Table of H2 lines

    References
    ----------
    * Abgrall et al. 1993, A&AS, 101, 323
    * Abgrall et al. 1993, A&AS, 101, 273

    Kindly provide by co-author E. Roueff to JC Howk to JXP

    """
    H2_fil = lt_path + '/data/lines/H2_resonance.ascii'
    print('linetools.lists.parse: Reading linelist --- \n   {:s}'.format(H2_fil))
    data = Table.read(H2_fil, format='ascii', guess=False, comment=';')

    # Units
    data['wrest'].unit = u.AA
    data['gamma'].unit = 1./u.s

    # Rename some columns
    data.rename_column('Jp', 'Jj')
    data.rename_column('Jpp', 'Jk')
    data.rename_column('np', 'nj')
    data.rename_column('npp', 'nk')

    # Molecule column
    cmol = Column(['H2']*len(data), name='mol')
    data.add_column(cmol)

    # Group
    cgroup = Column(np.ones(len(data),dtype='int')*(2**3), name='group')
    data.add_column(cgroup)

    # Return
    return data

def read_CO():
    """ Simple def to read CO UV data

    Generated by JXP with some great effort.  See GRB 080607 paper.

    Returns
    -------
    Table of CO lines
    """
    CO_fil = lt_path + '/data/lines/CO_UV.ascii'
    print('linetools.lists.parse: Reading linelist --- \n   {:s}'.format(CO_fil))
    data = ascii.read(CO_fil)

    # Units

    # Rename some columns
    data.rename_column('Jp', 'Jj')
    data.rename_column('Jpp', 'Jk')
    data.rename_column('np', 'nj')
    data.rename_column('npp', 'nk')
    data.rename_column('iso', 'Am') # Isotope
    data.rename_column('wave', 'wrest')

    data['wrest'].unit = u.AA

    # Fvalues
    data['fv'] = 10.**data['fv']
    data.rename_column('fv', 'f')

    # Molecule column
    cmol = Column(['CO']*len(data), name='mol')
    data.add_column(cmol)

    # Group
    cgroup = Column(np.ones(len(data),dtype='int')*(2**4), name='group')
    data.add_column(cgroup)

    # Return
    return data

def read_verner94():
    """ Read Verner1994 Table
    """
    # Read
    verner94 = lt_path + '/data/lines/verner94_tab6.fits'
    print(
        'linetools.lists.parse: Reading linelist --- \n   {:s}'.format(
            verner94))
    tbl_6 = Table.read(verner94)

    # Deal with bad unit
    tbl_6['lambda'].unit = u.AA
    tbl_6 = Table(tbl_6)

    # My table
    ldict, data = line_data(nrows=len(tbl_6))

    # Fill
    data['wrest'] = tbl_6['lambda']
    data['f'] = tbl_6['Fik']
    data['gj'] = tbl_6['Gi']
    data['gk'] = tbl_6['Gk']
    data['Z'] = tbl_6['Z']
    data['ion'] = tbl_6['Z'] - tbl_6['N'] + 1
    for ii,row in enumerate(tbl_6):
        data[ii]['name'] = (
            row['Species'][0:2].strip() + row['Species'][2:].strip() +
            ' {:d}'.format(int(row['lambda'])))
    # name
    names = []
    for row in data:
        ionnm = ions.ion_name((row['Z'], row['ion']))
        names.append('{:s} {:d}'.format(ionnm, int(row['wrest'])))
    data['name'] = names
    #  Finish
    data['group'] = 1
    data['Ref'] = 'Verner1994'
    data['mol'] = ''

    # Return
    return data

def read_forbidden():
    """ read galaxy emission lines (forbidden)

    There may be more here: https://github.com/moustakas/impro/blob/master/pro/hiiregions/im_getmatrix.pro

    Returns
    -------
    Table of forbidden lines
    """
    forb_fil = lt_path + '/data/lines/galaxy_forbidden.ascii'
    print('linetools.lists.parse: Reading linelist --- \n   {:s}'.format(forb_fil))
    aux = Table.read(forb_fil, format='ascii')

    # My table
    ldict, data = line_data(nrows=len(aux))

    # load values using convention names
    data['wrest'] = aux['wave']
    data['wrest'].unit = u.AA
    data['Z'] = aux['Z']
    data['ion'] = aux['ion']
    for ii, row in enumerate(data):
        row['name'] = aux['name'][ii].replace('_', ' ')
    data['Ref'] = 'DESI_NIST_JM'

    # mask the galaxy data using default mask_keys
    data = mask_gal(data)

    # Return
    return data

def read_recomb():
    """ read galaxy emission lines (recombination)

    Returns
    -------
    Table of recombination lines
    """
    recomb_fil = lt_path + '/data/lines/galaxy_recomb.ascii'
    print('linetools.lists.parse: Reading linelist --- \n   {:s}'.format(recomb_fil))
    aux = Table.read(recomb_fil, format='ascii')

    # My table
    ldict, data = line_data(nrows=len(aux))

    # load values using convention names
    data['wrest'] = aux['wave']
    data['wrest'].unit = u.AA
    data['Z'] = aux['Z']
    data['ion'] = aux['ion']
    for ii, row in enumerate(data):
        row['name'] = aux[ii]['name'].replace('_', ' ')
    data['Ref'] = 'DESI_NIST_JM'

    # mask the galaxy data using default mask_keys
    data = mask_gal(data)

    # Return
    return data

def read_galabs():
    """ read galaxy absorption lines

    Returns
    -------
    Table of recombination lines
    """
    galabs_fil = lt_path + '/data/lines/galaxy_abs.ascii'
    print('linetools.lists.parse: Reading linelist --- \n   {:s}'.format(galabs_fil))
    aux = Table.read(galabs_fil, format='ascii')

    # My table
    ldict, data = line_data(nrows=len(aux))

    # load values using convention names
    data['wrest'] = aux['wave']
    data['wrest'].unit = u.AA
    data['Z'] = aux['Z']
    data['ion'] = aux['ion']
    for ii, row in enumerate(data):
        row['name'] = aux[ii]['name'].replace('_', ' ')
    data['Ref'] = 'JXP_DK_Unknown'

    # mask the galaxy data using default mask_keys
    data = mask_gal(data)

    # Return
    return data

def mask_gal(data, mask_keys=None):
    """Masks linelist attributes for all galaxy lines

    Parameters
    ----------
    data : Table (masked)
        The original table to mask columns for
    mask_keys : list of strings, optional
        List of column names to be masked if given
        Otherwise it uses the default:

    Returns
    -------
    data_masked : Table (masked)
        The masked version of `data`
    """
    # check input
    if not isinstance(data, (Table)):
        raise RuntimeError('The input table has to be astropy Table')

    if data.masked is not True:
        raise RuntimeError('The input Table has to be masked.')

    # set default keys to mask
    if mask_keys is None:
        mask_keys = ['A', 'el', 'nj', 'nk','group','Ek','f','mol',
                     'Ej','Am','Ex','Jj','Jk','gk','gj','gamma']

    for key in mask_keys:
        data[key].mask = True

    return data

def parse_verner96(orig=False, write=False):
    """Parse tables from Verner, Verner, & Ferland (1996, Atomic Data and Nuclear Data Tables, Vol. 64, p.1)

    Parameters
    ----------
    orig : bool, optional
      Use original code to parse the ASCII file
      Else, read from a FITS file

    Returns
    -------
    data : Table
      Atomic data
    """
    # Look for FITS
    fitsf = lt_path + '/data/lines/verner96_tab1.fits.gz'
    verner96_tab1 = glob.glob(fitsf)

    if (len(verner96_tab1) > 0) & (not orig):
        print('linetools.lists.parse: Reading linelist --- \n   {:s}'.format(
            verner96_tab1[0]))
        data = Table.read(verner96_tab1[0])
    else:
        # File
        verner96_tab1 = lt_path + '/data/lines/verner96_tab1.txt'
        # Read
        with open(verner96_tab1) as f:
            lines = f.readlines()
        # Grab the 'good' ones
        gdlines = [iline.strip() for iline in lines if len(iline.strip()) > 113]
        ldict, data = line_data(nrows=len(gdlines))
        # Loop
        for kk, line in enumerate(gdlines):
            # Z, ion
            data[kk]['Z'] = ELEMENTS[line[0:2].strip()].number
            data[kk]['ion'] = int(line[2:4].strip())
            # wrest
            data[kk]['wrest'] = float(line[47:56].strip())
            # name
            ionnm = ions.ion_name((data[kk]['Z'], data[kk]['ion']))
            data[kk]['name'] = '{:s} {:d}'.format(ionnm,
                    int(data[kk]['wrest']))
            # Ej, Ek
            data[kk]['Ej'] = float(line[59:73].strip())
            data[kk]['Ek'] = float(line[73:89].strip())
            # gj, gk
            data[kk]['gj'] = int(line[89:92].strip())
            data[kk]['gk'] = int(line[92:95].strip())
            # Ak
            data[kk]['A'] = float(line[95:103].strip())
            # f
            data[kk]['f'] = float(line[104:112].strip())
        # Update
        data['Ref'] = 'Verner1996'

        # Write
        if write:
            outfil = lt_path + '/data/lines/verner96_tab1.fits'
            data.write(outfil,overwrite=True)
            print('parse_verner96: Wrote {:s}'.format(outfil))
            # Compress and delete
            print('Now compressing...')
            with open(outfil) as src:
                with gzip.open(outfil+'.gz', 'wb') as dst:
                    dst.writelines(src)
            os.unlink(outfil)

    # Return
    return data

def parse_morton00(orig=False):
    """Parse tables from Morton 2000, ApJS, 130, 403

    Parameters
    ----------
    orig : bool, optional
      Use original code to parse the ASCII file

    Returns
    -------
    data : Table
      Atomic data
    """
    # Look for FITS
    fitsf = lt_path + '/data/lines/morton00_table2.fits.gz'
    morton00_tab2 = glob.glob(fitsf)

    if (len(morton00_tab2) > 0) & (not orig):
        print('linetools.lists.parse: Reading linelist --- \n   {:s}'.format(
            morton00_tab2[0]))
        data = Table.read(morton00_tab2[0])
    else:
        # File
        morton00_tab2 = lt_path + '/data/lines/morton00_table2.dat'
        # Call
        data = parse_morton03(orig=True, tab_fil=morton00_tab2)
        # Update
        data['Ref'] = 'Morton2000'

    # Return
    return data

#
def parse_morton03(orig=False, tab_fil=None, HIcombine=True):
    """Parse tables from Morton 2003, ApJS, 149, 205

    Parameters
    ----------
    orig : bool, optional
      Use original code to parse the ASCII file
    tab_fil : str, optional
      Filename to use.  Default = /data/lines/morton03_table2.dat
    HIcombine : bool, optional
      Combine doublet for HI [True]

    Returns
    -------
    data : Table
      Atomic data
    """
    # Look for FITS
    fitsf = lt_path + '/data/lines/morton03_table2.fits.gz'
    morton03_tab2 = glob.glob(fitsf)

    if (len(morton03_tab2) > 0) & (not orig):
        print('linetools.lists.parse: Reading linelist --- \n   {:s}'.format(
            morton03_tab2[0]))
        data = Table.read(morton03_tab2[0])
    else:
        ## Read Table 2
        if tab_fil is None:
            morton03_tab2 = lt_path + '/data/lines/morton03_table2.dat'
        else:
            morton03_tab2 = tab_fil
        print(
            'linetools.lists.parse: Reading linelist --- \n   {:s}'.format(
                morton03_tab2))
        f = open(morton03_tab2, 'r', encoding="ISO-8859-1")
        lines = f.readlines()
        f.close()

        ## Find Elements and Ions
        elmi = []
        elmZ = []
        elmc = []
        ioni = []
        isoi = []
        ionv = []
        for kk,line in enumerate(lines):
            #print('kk = {:d}'.format(kk))
            try: # Deals with bad Byte in Morton00
                tmp = ('Z = ' in line) & ('A =' in line)
            except UnicodeDecodeError:
                tmp = False
            if tmp:
                # Grab Z
                ipos = line.find('Z = ')
                elmZ.append(int(line[ipos+4:ipos+7]))
                ipos2 = line.find('= ')
                elmc.append(line[ipos2+2:ipos].strip())
                #xdb.set_trace()
                # Line index
                elmi.append(kk)

            # ISOTOPE and ION
            try: # Deals with bad Byte in Morton00
                tmp2 = ( (('I ' in line[0:13]) | ('V ' in line[0:13]))
                    & (line[0:3] not in ['IOD','VAN']) & (line[0:2] != 'I ') )
            except UnicodeDecodeError:
                tmp2 = False
            if tmp2:
                # Grab ion
                ipos = line[0:10].find(' ')
                if ipos > 4:
                    ipos3 = line[0:10].find('I')
                    iionv = line[ipos3:ipos]
                else:
                    iionv = line[ipos:6].strip()
                if (len(iionv) == 0) | (iionv == '5s') | (iionv == 'B I'):
                    pdb.set_trace()
                ionv.append(iionv)
                if iionv == 'Z =':
                    pdb.set_trace()

                # Line index
                ioni.append(kk)

                # Deal with Isotope
                if line[0] in ['0','1','2','3','4','5','6','7','8','9']:
                    # Skip ArI !
                    if 'Ar I' in line:
                        pass
                    else:
                        isoi.append(kk)
                # Deuterium
                if line[0] == 'D':
                    Dline = kk

        #pdb.set_trace()
        ## Initialize table
        ldict, tbl = line_data(nrows=len(lines))

        ## Parse lines with UV rest wavelength
        count = 0
        for kk,line in enumerate(lines):
            try:
                tmp = line[23] == '.'
            except IndexError:
                pass
            else:
                if tmp: # UV wavelength?
                    # Parse

                    # Ion/Isotope
                    if kk > np.max(ioni):
                        gdi = len(ioni)-1
                    else:
                        gdi = np.where( (kk > np.array(ioni)) & (kk < np.roll(np.array(ioni),-1)))[0]
                        if len(gdi) != 1:
                            pdb.set_trace()
                            raise ValueError('Uh oh ion')
                        else:
                            gdi = gdi[0]
                    if ioni[gdi] in isoi: # Isotope
                        continue
                    # Ion
                    tbl[count]['ion'] = roman.fromRoman(ionv[gdi])

                    # Wavelength
                    tbl[count]['wrest'] = float(line[19:28]) #* u.AA
                    # Z
                    gdZ = np.where( (kk > np.array(elmi)) & (kk < np.roll(np.array(elmi),-1)))[0]
                    if len(gdZ) != 1:
                        if kk > np.max(elmi):
                            gdZ = len(elmi)-1
                        else:
                            #xdb.set_trace()
                            raise ValueError('Uh oh elm')
                    else:
                        gdZ = gdZ[0]
                    tbl[count]['Z'] = elmZ[gdZ]
                    # Name
                    tbl[count]['name'] = elmc[gdZ]+ionv[gdi]+' {:d}'.format(
                        int(tbl[count]['wrest']))
                    # Isotope (Atomic number)
                    if ioni[gdi] == Dline:
                        tbl[count]['Am'] = 2
                        tbl[count]['name'] = 'D'+ionv[gdi]+' {:d}'.format(
                            int(tbl[count]['wrest']))
                    # f
                    try:
                        tbl[count]['f'] = float(line[79:89])
                    except ValueError:
                        continue # Skip ones without f-value
                    # Ej, Ek
                    tbl[count]['Ej'] = float(line[29:38]) #/ u.cm
                    tbl[count]['Ek'] = float(line[40:50]) #/ u.cm
                    # A
                    try:
                        tbl[count]['A'] = float(line[59:68]) #/ u.s
                    except ValueError:
                        pass
                    # gamma
                    try:
                        tbl[count]['gamma'] = float(line[69:79]) #/ u.s
                    except ValueError:
                        pass
                    # gl, gu
                    tbl[count]['gj'] = int(line[52:54])
                    tbl[count]['gk'] = int(line[56:58])

                    # Only use combined HI lines
                    if HIcombine:
                        if ((tbl[count]['Z'] == 1) & (tbl[count]['ion']==1)
                            & (tbl[count]['gk'] != 6)):
                            #print('Skipping HI line {:g}'.format(tbl[count]['wrest']))
                            continue
                    # Ex
                    #all_dict[count]['Ex'] = 0.  # Zero out units (for Table)

                    # Increment
                    count += 1
        # Trim
        data = tbl[0:count]

        # Last
        data['group'] = 1
        data['Ref'] = 'Morton2003'
        data['mol'] = ''

    # Return
    #pdb.set_trace()
    return data

def mktab_morton03(do_this=False, outfil=None, fits=True):
    """Used to generate a VO or FITS Table for the Morton2003 paper

    Only intended for builder usage (1.5Mb file; gzip FITS is 119kb)

    Parameters
    ----------
    do_this : bool, optional
      Set to True to actually do this. Default=False
    outfil : str, optional
      Name of output file.  Defaults to a given value
    fits :  bool, optional
      Generate a FITS file?  Default=True
    """
    if not do_this:
        print('mktab_morton03: It is very unlikely you want to do this')
        print('mktab_morton03: Returning...')
        return

    # Read Morton2003 ASCII file
    m03 = parse_morton03(orig=True)

    # Write
    if fits:
        if outfil is None:
            outfil = lt_path + '/data/lines/morton03_table2.fits'
        m03.write(outfil,overwrite=True)
    else:
        if outfil is None:
            outfil = lt_path + '/data/lines/morton03_table2.vot'
        m03.write(outfil, format='votable', overwrite=True)
    print('mktab_morton03: Wrote {:s}'.format(outfil))
    # Compress and delete
    print('mktab_morton03: Now compressing...')
    with open(outfil) as src:
        with gzip.open(outfil+'.gz', 'wb') as dst:
            dst.writelines(src)
    os.unlink(outfil)

def mktab_morton00(do_this=False, outfil=None):
    """Used to generate a FITS Table for the Morton2000 paper

    Only intended for builder usage

    Parameters
    ----------
    do_this : bool, optional
      Set to True to actually do this. Default=False
    outfil : str, optional
      Name of output file.  Defaults to a given value
    """
    if not do_this:
        print('mktab_morton00: It is very unlikely you want to do this')
        print('mktab_morton00: Returning...')
        return

    # Read Morton2003
    m00 = parse_morton00(orig=True)

    # Write
    if outfil is None:
        outfil = lt_path + '/data/lines/morton00_table2.fits'
    m00.write(outfil, overwrite=True)
    print('mktab_morton00: Wrote {:s}'.format(outfil))
    #
    print('mktab_morton03: Now compressing...')
    with open(outfil) as src:
        with gzip.open(outfil+'.gz', 'wb') as dst:
            dst.writelines(src)
    os.unlink(outfil)

def grab_galaxy_linelists(do_this=False):
    """ Pulls galaxy emission line lists from DESI project

    Specifically, desisim
    Writes to hard-drive
    Only run if you are building

    Parameters
    ----------
    do_this : bool, optional
      Set to True to actually do this. Default=False

    """
    if not do_this:
        print('mktab_morton00: It is very unlikely you want to do this')
        print('mktab_morton00: Returning...')
        return

    try:
        # For Python 3.0 and later
        from urllib.request import urlopen
    except ImportError:
        # Fall back to Python 2's urllib2
        from urllib2 import urlopen

    # Forbidden
    url = 'https://raw.githubusercontent.com/desihub/desisim/master/data/forbidden_lines.dat'
    f = urlopen(url)
    tab_fil = lt_path+'/data/lines/galaxy_forbidden.ascii'
    print('Writing {:s}'.format(tab_fil))
    with open(tab_fil, "wb") as code:
        code.write(f.read())

    # Recombination
    url = 'https://raw.githubusercontent.com/desihub/desisim/master/data/recombination_lines.dat'
    f = urlopen(url)
    tab_fil = lt_path+'/data/lines/galaxy_recomb.ascii'
    print('Writing {:s}'.format(tab_fil))
    with open(tab_fil, "wb") as code:
        code.write(f.read())

def update_fval(table, verbose=False):
    """Update f-values from the literature

    Primarily for modifying lines in the ISM lists (e.g. Morton2003)

    Parameters
    ----------
    table : Table
      Data to be updated
    verbose : bool, optional

    Returns
    -------
    table : Table
      Updated table.
      Note: This return is required to handle the vstack, i.e.
      as opposed to modifying the input table in place.
    """
    # Shectman et al. 1998, ApJ, 504, 921
    #   Morton2003 cites this but uses a different f-value
    imn = np.argmin(np.abs(table['wrest']-1526.707))
    table['f'][imn] = 0.127

    # Howk 2000 (using Weise 2002 as in Morton for FeII 1142,1143,1144)
    howk00_fil = lt_path + '/data/lines/howk00_table1.ascii'
    howk00 = ascii.read(howk00_fil, comment='#')

    # Dress up
    howk00['wrest'].unit = u.AA

    fval = []
    fsig = []
    for row in howk00:
        ipos1 = row['fval_sig'].find('(')
        ipos2 = row['fval_sig'].find(')')
        #
        fval.append(float(row['fval_sig'][0:ipos1]))
        fsig.append(float(row['fval_sig'][ipos1+1:ipos2]))
    # Add columns
    howk00.add_column(Column(np.array(fval), name='f'))
    howk00.add_column(Column(np.array(fsig), name='fsig')) # Error in last decimal

    # Now, finally, update
    for row in howk00:
        mt = np.where( (np.abs(table['wrest']-row['wrest']*u.AA) < 1e-3*u.AA) &
            (table['Z'] == 26) & (table['ion'] == 2))[0]
        if len(mt) == 0:
            if verbose:
                print('update_fval: Line {:g} not in your table.'.format(row['wrest']))
        elif len(mt) == 1:
            table['f'][mt[0]] = row['f']
        else:
            raise ValueError('Uh oh')

    ## ##
    # Lines without f-value but of interest

    # AsII
    mn = np.min(np.abs(table['wrest']-1355.934))  # In Morton2000
    if mn > 0.05:  # Ang
        _, new_row = line_data()
        new_row['Z'] = 33
        new_row['ion'] = 2
        new_row['wrest'] = 1355.934
        new_row['name'] = 'AsII 1355'
        new_row['f'].mask = True
        # Stack
        table = vstack([table, new_row])

    return table

def update_gamma(table):
    """Update/add-in gamma values

    Parameters
    ----------
    table : Table
      Data to be updated
    verbose : bool, optional
    """
    # HI - Morton doesn't give these for the combined HI lines (sensible)
    #  Nor does he give them for the lines beyond Ly-d (not sure why)
    try:
        HI = np.where((table['Z']==1) & (table['ion']==1))[0]
    except KeyError: # Molecules
        pass
    else:
        if len(HI) > 0:
            # Same kludge as in atom.dat of VPFIT for higher order lines
            table['gamma'] = table['A']
            # More accurate for stronger lines (pulled from Morton) [all in units of s^-1]
            gdict = {1215.670: 6.265E+08, 1025.7222: 1.897E+08, # From Morton
                972.5367: 8.127E+07, 949.7430: 4.204E+07, 937.8034: 2.450E+07}
            for key in gdict.keys():
                mt = np.where( (np.abs(table['wrest']-key) < 1e-4))[0]
                if len(mt) > 0:
                    table['gamma'][mt[0]] = gdict[key]

def update_wrest(table, verbose=True):
    """Update wrest values (and Ej,Ek)

    Parameters
    ----------
    table : Table
      Data to be updated
    verbose : bool, optional
    """

    '''
    # TiII line (Morton 2003 vs Weise 2001)
    # Went back to Morton 2003.  If you go to Weise, you have
    #  to expunge the Verner94 row
    mt = np.where( (np.abs(table['wrest']-1910.9538*u.AA) < 1e-3*u.AA))[0]
    table['wrest'][mt[0]] = 1910.938 * u.AA
    table['Ek'][mt[0]] = 52330.33 / u.cm
    '''
#

def _write_ref_ISM_table():
    """ Write a reference table enabling faster I/O for ISM-related lists

    For developer use only.

    Note that after running this, you need to manually copy the table
    produced to linetools/data/lines/ISM_table.fits inside the github
    repository, and then check it in.
    """
    from linetools.lists.linelist import LineList

    ism = LineList('ISM', use_ISM_table=False)
    strong = LineList('Strong', use_ISM_table=False)
    euv = LineList('EUV', use_ISM_table=False)
    hi = LineList('HI', use_ISM_table=False)

    # need a Table to write
    tab = ism._data.copy()
    tab.sort(('wrest'))

    # Using np.in1d doesn't work for some reason. Do it the long way
    cond = []
    for table in (strong, euv, hi):
        igood = []
        for row in Table(table._data):
            ind = tab['wrest'].searchsorted(row['wrest'])
            dw = abs(tab['wrest'][ind] - row['wrest'])
            if abs(tab['wrest'][ind+1] - row['wrest']) < dw:
                ind = ind + 1
            rism = tab[ind]
            # check this is the right row.
            if all(row[k] == rism[k] for k in hi._data.colnames if
                   not hasattr(row[k], 'mask') or not row[k].mask):
                igood.append(ind)
            else:
                raise RuntimeError('No match found!')

        cond.append(np.zeros(len(tab), dtype=bool))
        cond[-1][np.array(igood)] = True

    col1 = Column(cond[0], name='is_Strong')
    col2 = Column(cond[1], name='is_EUV')
    col3 = Column(cond[2], name='is_HI')
    tab.add_columns([col1, col2, col3])

    tab.write('ISM_table.fits', overwrite=True)


# -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2014-2016 GEM Foundation
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake. If not, see <http://www.gnu.org/licenses/>.

"""
Module exports :class:`SomervilleEtAl2001NSHMP2008`.
"""
from __future__ import division

import numpy as np

from openquake.hazardlib.gsim.base import CoeffsTable, GMPE
from openquake.hazardlib.gsim.utils import clip_mean
from openquake.hazardlib import const
from openquake.hazardlib.imt import PGA, SA

class SomervilleEtAl2001NSHMP2008(GMPE):
    """
    Implements GMPE developed by P. Somerville, N. Collins, N. Abrahamson,
    R. Graves, and C. Saika and documented in "GROUND MOTION ATTENUATION
    RELATIONS FOR THE CENTRAL AND EASTERN UNITED STATES" (Final report, June
    30, 2001: Report to U.S. Geological Survey for award 99HQGR0098). This GMPE
    is used by the National Seismic Hazard Mapping Project (NSHMP) for the 2008
    US hazard model.

    Document available at:
    http://earthquake.usgs.gov/hazards/products/conterminous/2002/99HQGR0098.pdf

    This class replicates the algorithm for the Somerville et. al. 2001 GMPE as
    coded in the subroutine ``getSomer`` in the ``hazgridXnga2.f``
    Fortran code available at:
    http://earthquake.usgs.gov/hazards/products/conterminous/2008/software/

    Coefficients are given for the B/C site conditions.
    """
    #: Supported tectonic region type is stable continental crust,
    #: given that the equations have been derived for central and eastern
    #: north America
    DEFINED_FOR_TECTONIC_REGION_TYPE = const.TRT.STABLE_CONTINENTAL

    #: Supported intensity measure types are spectral acceleration,
    #: and peak ground acceleration
    DEFINED_FOR_INTENSITY_MEASURE_TYPES = set([
        PGA,
        SA
    ])

    #: Supported intensity measure component is the geometric mean of
    #two : horizontal components
    #:attr:`~openquake.hazardlib.const.IMC.AVERAGE_HORIZONTAL`,
    DEFINED_FOR_INTENSITY_MEASURE_COMPONENT = const.IMC.AVERAGE_HORIZONTAL

    #: Supported standard deviation type is only total.
    DEFINED_FOR_STANDARD_DEVIATION_TYPES = set([
        const.StdDev.TOTAL
    ])

    #: No site parameters required
    REQUIRES_SITES_PARAMETERS = set()

    #: Required rupture parameter is only magnitude (Mw).
    REQUIRES_RUPTURE_PARAMETERS = set(('mag', ))

    #: Required distance measure is rjb
    REQUIRES_DISTANCES = set(('rjb', ))

    def get_mean_and_stddevs(self, sites, rup, dists, imt, stddev_types):
        """
        See :meth:`superclass method
        <.base.GroundShakingIntensityModel.get_mean_and_stddevs>`
        for spec of input and result values.
        """
        assert all(stddev_type in self.DEFINED_FOR_STANDARD_DEVIATION_TYPES
                   for stddev_type in stddev_types)

        C = self.COEFFS[imt]

        mean = self._compute_mean(C, rup.mag, dists.rjb)
        mean = clip_mean(imt, mean)

        stddevs = self._compute_stddevs(C, dists.rjb.size, stddev_types)

        return mean, stddevs

    def _compute_mean(self, C, mag, rjb):
        """
        Compute and return mean value (table 8, page 8)
        """
        d1 = np.sqrt(50. ** 2 + 6. ** 2)
        d = np.sqrt(rjb ** 2 + 6 ** 2)

        mean = np.zeros_like(rjb)

        mean += (
            C['a1'] + C['a2'] * (mag - 6.4) +
            C['a7'] * (8.5 - mag) ** 2
        )

        idx = rjb < 50.
        mean[idx] += (
            C['a3'] * np.log(d[idx]) +
            C['a4'] * (mag - 6.4) * np.log(d[idx]) +
            C['a5'] * rjb[idx]
        )

        idx = rjb >=50.
        mean[idx] += (
            C['a3'] * np.log(d1) +
            C['a4'] * (mag - 6.4) * np.log(d[idx]) +
            C['a5'] * rjb[idx] + C['a6'] * (np.log(d[idx]) - np.log(d1))
        )

        return mean

    def _compute_stddevs(self, C, num_sites, stddev_types):
        """
        Return total standard deviation.
        """
        stddevs = []
        for _ in stddev_types:
            stddevs.append(np.zeros(num_sites) + C['sigma'])

        return stddevs

    #: Coefficient table obtained from coefficient arrays (a1, a2, a3, a4,
    #: a5, a6, a7, sig0) defined in subroutine getSomer in hazgridXnga2.f
    COEFFS = CoeffsTable(sa_damping=5, table="""\
    IMT    a1      a2      a3        a4       a5           a6         a7         sigma
    pga    0.658   0.805  -0.679     0.0861  -0.00498     -0.477      0.0        0.587
    0.1    1.442   0.805  -0.679     0.0861  -0.00498     -0.477      0.0        0.595
    0.2    1.358   0.805  -0.679     0.0861  -0.00498     -0.477      0.0        0.611
    0.3    1.2353  0.805  -0.67023   0.0861  -0.0048045   -0.523792  -0.030298   0.6057
    0.5    0.8532  0.805  -0.671792  0.0861  -0.00442189  -0.605213  -0.0640237  0.6242
    1.0   -0.0143  0.805  -0.696     0.0861  -0.00362     -0.755     -0.102      0.693
    2.0   -0.9497  0.805  -0.728     0.0861  -0.00221     -0.946     -0.140      0.824
    """)


#!/usr/bin/python
# author: Charlotte Bunne

# imports
import numpy as np
import matplotlib.pyplot as plt
import torch
import pickle
import os
from time import time
from torchvision import datasets, transforms
from torchvision.utils import save_image

# internal imports
from model.utils import *
from model.model_cnn import Generator, Adversary
from model.model_cnn import weights_init_generator, weights_init_adversary
from model.loss import gwnorm_distance, loss_total_variation, loss_procrustes

# get arguments
args = get_args()

# system preferences
seed = np.random.randint(100)
torch.set_default_dtype(torch.double)
np.random.seed(seed)
torch.manual_seed(seed)

# settings
batch_size = 256
z_dim = 100
lr = 0.0002
ngen = 3
beta = args.beta
lam = 0.5
niter = 10
epsilon = 0.005
num_epochs = args.num_epochs
cuda = args.cuda
channels = args.n_channels
id = args.id

model = 'gwgan_{}_eps_{}_tv_{}_procrustes_{}_ngen_{}_channels_{}_{}' \
        .format(args.data, epsilon, lam, beta, ngen, channels, id)
save_fig_path = 'out_' + model
if not os.path.exists(save_fig_path):
    os.makedirs(save_fig_path)

# data import
if args.data == 'mnist':
    dataloader = torch.utils.data.DataLoader(
        datasets.MNIST('./data/mnist', train=True, download=True,
                       transform=transforms.Compose([
                           transforms.Resize(32),
                           transforms.ToTensor(),
                           transforms.Normalize((0.5, 0.5, 0.5),
                                                (0.5, 0.5, 0.5))])),
        batch_size=batch_size, drop_last=True, shuffle=True)
elif args.data == 'fmnist':
    dataloader = torch.utils.data.DataLoader(
        datasets.FashionMNIST('./data/fmnist', train=True, download=True,
                              transform=transforms.Compose([
                                  transforms.Resize(32),
                                  transforms.ToTensor(),
                                  transforms.Normalize((0.5, 0.5, 0.5),
                                                       (0.5, 0.5, 0.5))])),
        batch_size=batch_size, drop_last=True, shuffle=True)
elif args.data == 'cifar_gray':
    dataloader = torch.utils.data.DataLoader(
        datasets.CIFAR10('./data/cifar10', train=True, download=True,
                         transform=transforms.Compose([
                            # transform RGB to grayscale
                            transforms.Grayscale(num_output_channels=1),
                            transforms.ToTensor(),
                            transforms.Normalize((0.5, 0.5, 0.5),
                                                 (0.5, 0.5, 0.5))])),
        batch_size=batch_size, drop_last=True, shuffle=True)
elif args.data == 'cifar':
    dataloader = torch.utils.data.DataLoader(
        datasets.CIFAR10('./data/cifar10', train=True, download=True,
                         transform=transforms.Compose([
                            transforms.ToTensor(),
                            transforms.Normalize((0.5, 0.5, 0.5),
                                                 (0.5, 0.5, 0.5))])),
        batch_size=batch_size, drop_last=True, shuffle=True)
else:
    raise NotImplementedError('dataset does not exist or not integrated.')

# print example images
save_image(next(iter(dataloader))[0][:25],
           os.path.join(save_fig_path, 'real.pdf'), nrow=5, normalize=True)

# define networks and parameters
generator = Generator(output_dim=channels)
adversary = Adversary(input_dim=channels)

# weight initialisation
generator.apply(weights_init_generator)
adversary.apply(weights_init_adversary)

if cuda:
    generator = generator.cuda()
    adversary = adversary.cuda()

# create optimizer
g_optimizer = torch.optim.Adam(generator.parameters(), lr, betas=(0.5, 0.99))
# zero gradients
generator.zero_grad()

c_optimizer = torch.optim.Adam(adversary.parameters(), lr, betas=(0.5, 0.99))
# zero gradients
adversary.zero_grad()

# sample for plotting
num_test_samples = batch_size
z_ex = torch.randn(num_test_samples, z_dim)
if cuda:
    z_ex = z_ex.cuda()

loss_history = list()
loss_tv = list()
loss_orth = list()
loss_og = 0
is_hist = list()

for epoch in range(num_epochs):
    t0 = time()

    for it, (image, _) in enumerate(dataloader):
        train_c = ((it + 1) % (ngen + 1) == 0)

        x = image.double()
        if cuda:
            x = x.cuda()

        # sample random number z from Z
        z = torch.randn(image.shape[0], z_dim)

        if cuda:
            z = z.cuda()

        if train_c:
            for q in generator.parameters():
                q.requires_grad = False
            for p in adversary.parameters():
                p.requires_grad = True
        else:
            for q in generator.parameters():
                q.requires_grad = True
            for p in adversary.parameters():
                p.requires_grad = False

        # result generator
        g = generator.forward(z)

        # result adversary
        f_x = adversary.forward(x)
        f_g = adversary.forward(g)

        # compute inner distances
        D_g = get_inner_distances(f_g, metric='euclidean', concat=False)
        D_x = get_inner_distances(f_x, metric='euclidean', concat=False)

        # distance matrix normalisation
        D_x_norm = normalise_matrices(D_x)
        D_g_norm = normalise_matrices(D_g)

        # compute normalized gromov-wasserstein distance
        loss, T = gwnorm_distance((D_x, D_x_norm), (D_g, D_g_norm),
                                  epsilon, niter, loss_fun='square_loss',
                                  coupling=True, cuda=cuda)

        if train_c:
            # train adversary
            loss_og = loss_procrustes(f_x, x.view(x.shape[0], -1), cuda)
            loss_to = -loss + beta * loss_og
            loss_to.backward()

            # parameter updates
            c_optimizer.step()
            # zero gradients
            reset_grad(generator, adversary)

        else:
            # train generator
            loss_t = loss_total_variation(g)
            loss_to = loss + lam * loss_t
            loss_to.backward()

            # parameter updates
            g_optimizer.step()
            # zero gradients
            reset_grad(generator, adversary)

    # plotting
    # get generator example
    g_ex = generator.forward(z_ex)
    g_plot = g_ex.cpu().detach()

    # plot result
    save_image(g_plot.data[:25],
               os.path.join(save_fig_path, 'g_%d.pdf' % epoch),
               nrow=5, normalize=True)

    fig1, ax = plt.subplots(1, 3, figsize=(15, 5))
    ax0 = ax[0].imshow(T.cpu().detach().numpy(), cmap='RdBu_r')
    colorbar(ax0)
    ax1 = ax[1].imshow(D_x.cpu().detach().numpy(), cmap='Blues')
    colorbar(ax1)
    ax2 = ax[2].imshow(D_g.cpu().detach().numpy(), cmap='Blues')
    colorbar(ax2)
    ax[0].set_title(r'$T$')
    ax[1].set_title(r'inner distances of $D$')
    ax[2].set_title(r'inner distances of $G$')
    plt.tight_layout(h_pad=1)
    fig1.savefig(os.path.join(save_fig_path, '{}_ccc.pdf'.format(
            str(epoch).zfill(3))), bbox_inches='tight')

    loss_history.append(loss)
    loss_tv.append(loss_t)
    loss_orth.append(loss_og)
    plt.close('all')

# plot loss history
fig2 = plt.figure(figsize=(2.4, 2))
ax2 = fig2.add_subplot(111)
ax2.plot(loss_history, 'k.')
ax2.set_xlabel('Iterations')
ax2.set_ylabel(r'$\overline{GW}_\epsilon$ Loss')
plt.tight_layout()
plt.grid()
fig2.savefig(save_fig_path + '/loss_history.pdf')

fig3 = plt.figure(figsize=(2.4, 2))
ax3 = fig3.add_subplot(111)
ax3.plot(loss_tv, 'k.')
ax3.set_xlabel('Iterations')
ax3.set_ylabel(r'Total Variation Loss')
plt.tight_layout()
plt.grid()
fig3.savefig(save_fig_path + '/loss_tv.pdf')

fig4 = plt.figure(figsize=(2.4, 2))
ax4 = fig4.add_subplot(111)
ax4.plot(loss_orth, 'k.')
ax4.set_xlabel('Iterations')
ax4.set_ylabel(r'$R_\beta(f_\omega(X), X)$ Loss')
plt.tight_layout()
plt.grid()
fig4.savefig(save_fig_path + '/loss_orth.pdf')