'''
Artifact detection on MEG data.

TODO: Update doc.

For now
'''
import glob
import logging

import mne
import numpy as np
import pandas as pd

from .tools import hilbert

def annotate_blinks(raw, ch_mapping={'x': 'UADC002-3705', 'y': 'UADC003-3705', 'p': 'UADC004-3705'}):
    '''
    Detect blinks and annotate as bad blinks
    '''
    logging.info('Annotating blinks artifacts')
    x, y, p = eye_voltage2gaze(raw, ch_mapping=ch_mapping)
    xpos, ypos = x.ravel() / ppd(), y.ravel() / ppd()
    sc = saccade_detection(xpos, ypos, threshold=10, acc_thresh=2000, Hz=1200)
    blinks, scfilt = blink_detection(xpos, ypos, sc)
    if len(blinks) == 0:
        return None
    blink_onsets = raw.times[blinks[:, 0]]
    blink_durations = raw.times[blinks[:, 1]] - raw.times[blinks[:, 0]]
    return mne.Annotations(blink_onsets, blink_durations, 'bad blinks')

def annotate_muscle(raw, cutoff=10):
    logging.info('Annotating muscle artifacts')
    try:
        arts, z = detect_muscle(raw.copy(), cutoff=cutoff)
    except MemoryError:
        print('Memory Error detected:', raw)
        raise RuntimeError(
            'Memory error detected in annotate muscle: ' + raw.info['filename'])

    annotations = None
    if len(arts) > 0:
        annotations = mne.Annotations(arts[:, 0],
                                      arts[:, 1], 'bad muscle')
    return annotations, z

def annotate_cars(raw, cutoff=4.0, der_cutoff=7.):
    logging.info('Annotating car artifacts')
    arts, z, d = detect_cars(raw.copy(), cutoff=cutoff, der_cutoff=der_cutoff)
    annotations = None
    if len(arts) > 0:
        annotations = mne.Annotations(arts[:, 0], arts[:, 1], 'bad car')
    return annotations, z, d

def annotate_jumps(raw, cutoff=25, allowed_before_bad=np.inf):
    logging.info('Annotating jump artifacts')
    arts, z, jumps_per_channel = detect_jumps(raw.copy(), cutoff=cutoff)
    # Need to check for jumps_per_channel
    bads = [k for k, v in jumps_per_channel.items() if v >
            allowed_before_bad]
    arts = [arts[k] for k, v in jumps_per_channel.items() if v <=
            allowed_before_bad]

    if len(bads) > 0:
        if 'bads' in raw.info.keys():
            raw.info['bads'].extend(bads)
        else:
            raw.info['bads'] = bads
    a = []
    for k in arts:
        if len(k) is not 0:
            a.extend(k)
    arts = np.array(a)
    annotations = None
    try:
        if len(arts) > 0:
            annotations = mne.Annotations(arts[:, 0], arts[:, 1], 'bad jump')
    except IndexError:
        pass
    return raw, annotations, z, jumps_per_channel

def detect_cars(raw, cutoff=3.5, der_cutoff=5.0, frequency_band=(None, 1)):
    '''
    Detect cars artifacts on blocks and ignore intermediate data.

    This works analagously to the fieldtrip detect artifact routine.

    Data epochs that are already marked as bad by the annotations in raw will be
    excluded for computation of mean and std.
    '''

    logging.info('Detecting car events with cutoff %i' % cutoff)
    if not hasattr(raw, '_data'):
        logging.info('Loading data for car artifact detection')
        raw.load_data()
    raw.pick_channels([x for x in raw.ch_names if x.startswith('M')])
    raw.filter(l_freq=None, h_freq=1)
    hilb = hilbert(raw._data)
    del raw._data
    hilb = np.abs(hilb).astype(float)

    # Compute IGR and median
    Qs = np.percentile(hilb, [10, 50, 90], axis=1)
    IQR = Qs[2, :] - Qs[0, :]
    m = Qs[1, :]
    zh = ((((hilb - m[:, np.newaxis]) / IQR[:, np.newaxis]))**2).mean(0)

    # Normalize zh to have 80 between 0 an 1
    q80 = np.percentile(zh, [80])
    zh = zh / q80

    # Compute derivative of zh
    d = np.concatenate(([0], np.diff(zh)))
    # Normalize to have 80 between -1 and 1
    d = d / np.diff(np.percentile(d, [10, 80]))
    d[:int(raw.info['sfreq'])] = 0
    d[-int(raw.info['sfreq']):] = 0

    # Compute artifact borders
    art_borders = np.where(np.diff(np.concatenate([[0], zh > cutoff, [0]])))[0]
    artifacts = []
    for start, end in zip(art_borders[0::2], art_borders[1::2]):
        onset_t = max((start - 1) - int(2.5 * raw.info['sfreq']), 0)
        end_t = min((end) + int(2.5 * raw.info['sfreq']), len(zh))
        # Check for derivative
        if d[onset_t:end_t].min() < -der_cutoff and d[onset_t:end_t].max() > der_cutoff:
            onset_t /= raw.info['sfreq']
            end_t /= raw.info['sfreq']
            duration = end_t - onset_t
            artifacts.append((onset_t, duration))
    return np.array(artifacts), zh, d

def detect_muscle(raw, cutoff=10, frequency_band=(110, 140)):
    '''
    Detect muscle artifacts on blocks and ignore intermediate data.

    This works analagously to the fieldtrip detect artifact routine.

    Data epochs that are already marked as bad by the annotations in raw will be
    excluded for computation of mean and std.
    '''

    logging.info('Detecting muscle events with cutoff %i' % cutoff)
    if not hasattr(raw, '_data'):
        logging.info('Loading data for muscle artifact detection')
        raw.load_data()
    raw.pick_channels([x for x in raw.ch_names if x.startswith('M')])

    # filt = mne.filter.band_pass_filter(raw._data, raw.info['sfreq'],
    #                                   frequency_band[0], frequency_band[1],  method='iir',
    #                                   iir_params = dict(order=9, ftype='butter'),
    #                                   copy=False)
    filt = mne.filter.filter_data(raw._data, raw.info['sfreq'],
                                  l_freq=frequency_band[
                                      0], h_freq=frequency_band[1],  method='iir',
                                  iir_params=dict(order=9, ftype='butter'),
                                  copy=False)

    hilb = abs(hilbert(filt)).astype(float)
    del filt
    # Compute IGR and median
    Qs = np.percentile(hilb, [10, 50, 90], axis=1)
    IQR = Qs[2, :] - Qs[0, :]
    m = Qs[1, :]
    zh = ((((hilb - m[:, np.newaxis]) / IQR[:, np.newaxis]))**2).mean(0)

    # Normalize zh to have 80 between 0 an 1
    q80 = np.percentile(zh, [80])
    zh = zh / q80

    art_borders = np.where(np.diff(np.concatenate([[0], zh > cutoff, [0]])))[0]

    artifacts = []
    for start, end in zip(art_borders[0::2], art_borders[1::2]):
        artifacts.append((
                         ((start - 1) / raw.info['sfreq']) - 0.2,
                         ((end - start) / raw.info['sfreq']) + 0.2,
                         ))
    return np.array(artifacts), zh

def detect_jumps(raw, cutoff=25):
    '''
    Detect jumps by convolving with a jump detection filter.
    '''
    logging.info('Detecting muscle events with cutoff %i' % cutoff)
    if not hasattr(raw, '_data'):
        logging.info('Loading data for muscle artifact detection')
        raw.load_data()
    raw.pick_channels([x for x in raw.ch_names if x.startswith('M')])
    #filt = mne.filter.low_pass_filter(raw._data, raw.info['sfreq'], 1)
    filt = mne.filter.filter_data(
        raw._data, raw.info['sfreq'], l_freq=None, h_freq=1)

    jump_kernel = (np.array([1] * 50),
                   [0, 0],
                   np.array([-1] * 50))
    jump_kernel = np.concatenate(jump_kernel)
    filt = 0 * raw._data.copy()
    for i in range(filt.shape[0]):
        filt[i, :] = np.convolve(
            raw._data[i, :], jump_kernel, mode='same') - filt[i, :]
        filt[i, :int(len(jump_kernel) / 2)] = 0
        filt[i, -int(len(jump_kernel) / 2):] = 0

    # Compute IGR and median
    Qs = np.percentile(filt, [10, 50, 90], axis=1)
    IQR = Qs[2, :] - Qs[0, :]
    m = Qs[1, :]
    filt = (((filt - m[:, np.newaxis]) / IQR[:, np.newaxis]))**2
    # Need to keep information about channels here.

    artifacts = {}
    channel_count = {}
    for i, zh in enumerate(filt):
        artifacts[raw.ch_names[i]] = []
        art_borders = np.where(
            np.diff(np.concatenate([[0], zh > cutoff, [0]])))[0]
        channel_count[raw.ch_names[i]] = 0
        for start, end in zip(art_borders[0::2], art_borders[1::2]):
            artifacts[raw.ch_names[i]].append(
                ((start - 1) / raw.info['sfreq'], (end - start) / raw.info['sfreq']))
            channel_count[raw.ch_names[i]] += 1
    return artifacts, zh, channel_count

def eye_voltage2gaze(raw, ranges=(-5, 5), screen_x=(0, 1920),
                     screen_y=(0, 1080),
                     ch_mapping={'x': 'UADC002-3705', 'y': 'UADC003-3705', 'p': 'UADC004-3705'}):
    '''
    Convert analog output of EyeLink 1000+ to gaze coordinates.
    '''
    minvoltage, maxvoltage = ranges
    maxrange, minrange = 1., 0.
    screenright, screenleft = screen_x
    screenbottom, screentop = screen_y

    idx = np.where(np.array(raw.ch_names) == ch_mapping['x'])[0][0]
    R = (raw[idx, :][0] - minvoltage) / (maxvoltage - minvoltage)
    S = R * (maxrange - minrange) + minrange
    x = S * (screenright - screenleft + 1) + screenleft

    idy = np.where(np.array(raw.ch_names) == ch_mapping['y'])[0][0]
    R = (raw[idy, :][0] - minvoltage) / (maxvoltage - minvoltage)
    S = R * (maxrange - minrange) + minrange
    y = S * (screenbottom - screentop + 1) + screentop

    idp = np.where(np.array(raw.ch_names) == ch_mapping['p'])[0][0]
    p = raw[idp, :][0]
    return x, y, p

def eye_voltage2gaze_epochs(epochs, ranges=(-5, 5), screen_x=(0, 1920),
                            screen_y=(0, 1080),
                            ch_mapping={'x': 'UADC002-3705', 'y': 'UADC003-3705', 'p': 'UADC004-3705'}):
    '''
    Convert analog output of EyeLink 1000+ to gaze coordinates.
    '''
    minvoltage, maxvoltage = ranges
    maxrange, minrange = 1., 0.
    screenright, screenleft = screen_x
    screenbottom, screentop = screen_y

    idx = np.where(np.array(epochs.ch_names) == ch_mapping['x'])[0][0]
    R = (epochs._data[:, idx, :].squeeze() -
         minvoltage) / (maxvoltage - minvoltage)
    S = R * (maxrange - minrange) + minrange
    x = S * (screenright - screenleft + 1) + screenleft

    idy = np.where(np.array(epochs.ch_names) == ch_mapping['y'])[0][0]
    R = (epochs._data[:, idy, :].squeeze() -
         minvoltage) / (maxvoltage - minvoltage)
    S = R * (maxrange - minrange) + minrange
    y = S * (screenbottom - screentop + 1) + screentop

    idp = np.where(np.array(epochs.ch_names) == ch_mapping['p'])[0][0]
    p = epochs._data[:, idp:].squeeze()
    return x, y, p

velocity_window_size = 3

def get_velocity(x, y, Hz):
    '''
    Compute velocity of eye-movements.

    'x' and 'y' specify the x,y coordinates of gaze location. The function
    assumes that the values in x,y are sampled continously at a rate specified
    by 'Hz'.
    '''
    Hz = float(Hz)
    distance = ((np.diff(x) ** 2) +
                (np.diff(y) ** 2)) ** .5
    distance = np.hstack(([distance[0]], distance))
    win = np.ones((velocity_window_size)) / float(velocity_window_size)
    velocity = np.convolve(distance, win, mode='same')
    velocity = velocity / (velocity_window_size / Hz)
    acceleration = np.diff(velocity) / (1. / Hz)
    acceleration = np.abs(np.hstack(([acceleration[0]], acceleration)))
    return velocity, acceleration

def saccade_detection(x, y, Hz=1200, threshold=30,
                      acc_thresh=2000):
    '''
    Detect saccades in a stream of gaze location samples.

    Coordinates of x,y are assumed to be in degrees.

    Saccades are detect by a velocity/acceleration threshold approach.
    A saccade starts when a) the velocity is above threshold, b) the
    acceleration is above acc_thresh at least once during the interval
    defined by the velocity threshold.
    '''

    velocity, acceleration = get_velocity(x, y, float(Hz))
    saccades = (velocity > threshold)

    borders = np.where(np.diff(saccades.astype(int)))[0] + 1
    if velocity[1] > threshold:
        borders = np.hstack(([0], borders))

    saccade = 0 * np.ones(x.shape)

    saccade_times = []
    # Only count saccades when acceleration also surpasses threshold
    for i, (start, end) in enumerate(zip(borders[0::2], borders[1::2])):
        if np.sum(acceleration[start:end] > acc_thresh) >= 1:
            saccade[start:end] = 1
            saccade_times.append((start, end))

    return np.array(saccade_times)

def microssacade_detection(x, y, VFAC):
    '''
    Microsaccade detection a la Engbert et al.
    '''
    if len(x) < 5:
        return None
    dt = 1 / 1200.
    kernel = np.array([1., 1., 0., -1., -1.])
    vx = np.convolve(x, kernel, mode='same') / (6 * dt)
    vy = np.convolve(y, kernel, mode='same') / (6 * dt)
    msdx = np.sqrt(np.median((vx - np.median(vx))**2))
    msdy = np.sqrt(np.median((vy - np.median(vy))**2))
    radiusx = VFAC * msdx
    radiusy = VFAC * msdy
    test = (vx / radiusx)**2 + (vy / radiusy)**2
    borders = np.where(np.diff((test > 1).astype(int)))[0] + 1
    if test[0] > 1:
        borders = np.hstack(([0], borders))
    if test[-1] > 1:
        borders = np.hstack((borders, [len(x)]))

    borders = borders.reshape(len(borders) / 2, 2)
    return borders

def blink_detection(x, y, saccades):
    '''
    A blink is everything that is surrounded by two saccades and period in
    between where the eye is off screen.
    '''
    rm_sac = (saccades[:, 0] * 0).astype(bool)
    blinks = []
    skipnext = False
    for i, ((pss, pse), (nss, nse)) in enumerate(zip(saccades[:-1], saccades[1:])):
        if skipnext:
            skipnext = False
            continue
        xavg = x[pse:nss].mean()
        yavg = y[pse:nss].mean()

        if (xavg > 40) and (yavg > 20):
            rm_sac[i:i + 2] = True
            blinks.append((pss, nse))
            skip_next = True

    return np.array(blinks), saccades[~rm_sac, :]

def nan_bad_epochs(data, raw):
    '''
    Overwrite data with NANs for all epochs in the Nx3 artifacts matrix.
    '''
    Hz = raw.info['sfreq']
    if raw.annotations is not None:
        for start, duration, desc in zip(raw.annotations.onset,
                                         raw.annotations.duration,
                                         raw.annotations.description):
            if not 'bad' in desc:
                continue
            start = int(start * Hz)
            data[start:start + int(duration * Hz)] = nan
    return data

def ppd(vieweing_distance=62.5, screen_width=38.0, x_resolution=1450):
    '''
    Compute pixels per degree for the current setup.
    '''
    o = np.tan(0.5 * np.pi / 180) * vieweing_distance
    return 2 * o * x_resolution / screen_width

def combine_annotations(annotations):
    '''
    Add annotations to a raw object. Makes sure that old annotations are kept.
    '''
    if len(annotations) == 0:
        return mne.Annotations(np.array([0]), np.array([0]), 'dummy')
    elif len(annotations) == 1:
        return annotations[0]
    else:
        old = annotations[0]
        for new in annotations[1:]:
            if new is None:
                continue
            orig_time = None
            onsets = np.concatenate((old.onset, new.onset))
            duration = np.concatenate((old.duration, new.duration))
            descr = np.concatenate((old.description, new.description))
            if old.orig_time is not None:
                orig_time = np.concatenate((old.orig_time, new.orig_time))
            old = mne.Annotations(onsets, duration, descr, orig_time)
        return old

def concatenate_annotations(annotations, durations):
    '''
    Concatenat annotations for several consecutive raw objects that are to be
    merged. Each annotation must come with one duration for the raw object that
    it belongs to.

    durations is in s.

    TODO: This was because of a bug in pymne - which should have been fixed upstream by now.
    '''
    durations[0] = 0
    onsets = np.concatenate([a.onset + (li)
                             for a, li in zip(annotations, cumsum(durations))])
    durations = np.concatenate([a.duration
                                for a in annotations])
    description = np.concatenate([a.description
                                  for a in annotations])
    return mne.Annotations(onsets, durations, description)


import os, sys
sys.path.append(os.getcwd())

import time
import functools

import numpy as np
import tensorflow as tf
import scipy.misc

import tflib as lib
import tflib.ops.linear
import tflib.ops.conv2d
import tflib.ops.batchnorm
import tflib.ops.deconv2d
import tflib.save_images
import tflib.celebA_64x64
import tflib.small_imagenet
import tflib.ops.layernorm
import tflib.plot

FLAGS = tf.app.flags.FLAGS

# Configurations
tf.app.flags.DEFINE_string('mode', 'wgan-gp',
                           "loss function option. [wgan-gp | dcgan | wgan | lsgan]")
tf.app.flags.DEFINE_string('data_dir', 'data/celebA_64x64', "data directory")
tf.app.flags.DEFINE_string('train_dir', 'train', "image output direcotory")
tf.app.flags.DEFINE_string('summary_dir', 'summary', "tensorboard summary directory")
tf.app.flags.DEFINE_integer('max_runtime', 20, "maximum run time in min")
tf.app.flags.DEFINE_integer('max_iter', 500, "maximum mini-batch iterations")
tf.app.flags.DEFINE_float('LAMBDA', 10., "gradient penalty lambda parameter")
tf.app.flags.DEFINE_float('gen_l1_weight', 0.9, "weight of L1 difference in generator loss")
tf.app.flags.DEFINE_integer('architecture', 0, "index of architecture")

# Download 64x64 ImageNet at http://image-net.org/small/download.php and
# fill in the path to the extracted files here!
DATA_DIR = FLAGS.data_dir
SUMMARY_DIR = FLAGS.summary_dir
GEN_L1_WEIGHT = FLAGS.gen_l1_weight # Weighting factor for L1 difference in generator loss
TRAIN_DIR = FLAGS.train_dir # Directory to output image
MODE = FLAGS.mode # dcgan, wgan, wgan-gp, lsgan
ITERS = FLAGS.max_iter # How many iterations to train for
LAMBDA = FLAGS.LAMBDA # Gradient penalty lambda hyperparameter

if len(DATA_DIR) == 0:
    raise Exception('Please specify path to data directory in gan_64x64.py!')

DIM = 64 # Model dimensionality
K = 4 # How much to downsample
CRITIC_ITERS = 5 # How many iterations to train the critic for
N_GPUS = 1 # Number of GPUs
BATCH_SIZE = 16 # Batch size. Must be a multiple of N_GPUS
INPUT_DIM = 16*16*3 # Number of pixels in each input
OUTPUT_DIM = 64*64*3 # Number of pixels in each iamge
DELETE_TRAIN_DIR=True

lib.print_model_settings(locals().copy())

# create summary dir
if not tf.gfile.Exists(FLAGS.summary_dir):
    tf.gfile.MakeDirs(FLAGS.summary_dir)

# clean directory
if DELETE_TRAIN_DIR:
    if tf.gfile.Exists(FLAGS.train_dir):
        tf.gfile.DeleteRecursively(FLAGS.train_dir)
        tf.gfile.MakeDirs(FLAGS.train_dir)
    tf.gfile.MakeDirs(FLAGS.train_dir)

# architecture dictionary
def get_architectures():
    ARCHITECTURE_TABLE = {
        # Baseline (G: DCGAN, D: DCGAN)
        0: (DCGANGenerator, DCGANDiscriminator),

        # No BN and constant number of filts in G
        1: (WGANPaper_CrippledDCGANGenerator, DCGANDiscriminator),

        # 512-dim 4-layer ReLU MLP G
        2: (FCGenerator, DCGANDiscriminator),

        # No normalization anywhere
        3: (functools.partial(DCGANGenerator, bn=False),
            functools.partial(DCGANDiscriminator, bn=False)),

        # Gated multiplicative nonlinearities everywhere
        4: (MultiplicativeDCGANGenerator, MultiplicativeDCGANDiscriminator),

        # tanh nonlinearities everywhere
        5: (functools.partial(DCGANGenerator, bn=True, nonlinearity=tf.tanh),
            functools.partial(DCGANDiscriminator, bn=True, nonlinearity=tf.tanh)),

        # 101-layer ResNet G and D
        6: (ResnetGenerator, ResnetDiscriminator)
    }
    return ARCHITECTURE_TABLE

def GeneratorAndDiscriminator():
    """
    Choose which generator and discriminator architecture to use by
    uncommenting one of these lines.
    """
    table = get_architectures()
    if FLAGS.architecture <= len(table):
        return table[FLAGS.architecture]

    raise Exception('You must choose an architecture!')

DEVICES = ['/gpu:{}'.format(i) for i in range(N_GPUS)]

def LeakyReLU(x, alpha=0.2):
    return tf.maximum(alpha*x, x)

def ReLULayer(name, n_in, n_out, inputs):
    output = lib.ops.linear.Linear(name+'.Linear', n_in, n_out, inputs, initialization='he')
    return tf.nn.relu(output)

def LeakyReLULayer(name, n_in, n_out, inputs):
    output = lib.ops.linear.Linear(name+'.Linear', n_in, n_out, inputs, initialization='he')
    return LeakyReLU(output)

def Batchnorm(name, axes, inputs):
    if ('Discriminator' in name) and (MODE == 'wgan-gp'):
        if axes != [0,2,3]:
            raise Exception('Layernorm over non-standard axes is unsupported')
        return lib.ops.layernorm.Layernorm(name,[1,2,3],inputs)
    else:
        return lib.ops.batchnorm.Batchnorm(name,axes,inputs,fused=True)

def pixcnn_gated_nonlinearity(a, b):
    return tf.sigmoid(a) * tf.tanh(b)

def SubpixelConv2D(*args, **kwargs):
    kwargs['output_dim'] = 4*kwargs['output_dim']
    output = lib.ops.conv2d.Conv2D(*args, **kwargs)
    output = tf.transpose(output, [0,2,3,1])
    output = tf.depth_to_space(output, 2)
    output = tf.transpose(output, [0,3,1,2])
    return output

def ResidualBlock(name, input_dim, output_dim, filter_size, inputs, resample=None, he_init=True):
    """
    resample: None, 'down', or 'up'
    """
    if resample=='down':
        conv_shortcut = functools.partial(lib.ops.conv2d.Conv2D, stride=2)
        conv_1        = functools.partial(lib.ops.conv2d.Conv2D, input_dim=input_dim, output_dim=input_dim//2)
        conv_1b       = functools.partial(lib.ops.conv2d.Conv2D, input_dim=input_dim//2, output_dim=output_dim//2, stride=2)
        conv_2        = functools.partial(lib.ops.conv2d.Conv2D, input_dim=output_dim//2, output_dim=output_dim)
    elif resample=='up':
        conv_shortcut = SubpixelConv2D
        conv_1        = functools.partial(lib.ops.conv2d.Conv2D, input_dim=input_dim, output_dim=input_dim//2)
        conv_1b       = functools.partial(lib.ops.deconv2d.Deconv2D, input_dim=input_dim//2, output_dim=output_dim//2)
        conv_2        = functools.partial(lib.ops.conv2d.Conv2D, input_dim=output_dim//2, output_dim=output_dim)
    elif resample==None:
        conv_shortcut = lib.ops.conv2d.Conv2D
        conv_1        = functools.partial(lib.ops.conv2d.Conv2D, input_dim=input_dim,  output_dim=input_dim//2)
        conv_1b       = functools.partial(lib.ops.conv2d.Conv2D, input_dim=input_dim//2,  output_dim=output_dim/2)
        conv_2        = functools.partial(lib.ops.conv2d.Conv2D, input_dim=input_dim//2, output_dim=output_dim)

    else:
        raise Exception('invalid resample value')

    if output_dim==input_dim and resample==None:
        shortcut = inputs # Identity skip-connection
    else:
        shortcut = conv_shortcut(name+'.Shortcut', input_dim=input_dim, output_dim=output_dim, filter_size=1,
                                 he_init=False, biases=True, inputs=inputs)

    output = inputs
    output = tf.nn.relu(output)
    output = conv_1(name+'.Conv1', filter_size=1, inputs=output, he_init=he_init, weightnorm=False)
    output = tf.nn.relu(output)
    output = conv_1b(name+'.Conv1B', filter_size=filter_size, inputs=output, he_init=he_init, weightnorm=False)
    output = tf.nn.relu(output)
    output = conv_2(name+'.Conv2', filter_size=1, inputs=output, he_init=he_init, weightnorm=False, biases=False)
    output = Batchnorm(name+'.BN', [0,2,3], output)

    return shortcut + (0.3*output)

# ! Generators

def FCGenerator(n_samples, noise=None, FC_DIM=512, input_dim=INPUT_DIM):
    if noise is None:
        noise = tf.random_normal([n_samples, input_dim])

    output = ReLULayer('Generator.1', input_dim, FC_DIM, noise)
    output = ReLULayer('Generator.2', FC_DIM, FC_DIM, output)
    output = ReLULayer('Generator.3', FC_DIM, FC_DIM, output)
    output = ReLULayer('Generator.4', FC_DIM, FC_DIM, output)
    output = lib.ops.linear.Linear('Generator.Out', FC_DIM, OUTPUT_DIM, output)

    output = tf.tanh(output)

    return output

def DCGANGenerator(
        n_samples, noise=None, dim=DIM, input_dim=INPUT_DIM,
        k=K, bn=True, nonlinearity=tf.nn.relu):

    lib.ops.conv2d.set_weights_stdev(0.02)
    lib.ops.deconv2d.set_weights_stdev(0.02)
    lib.ops.linear.set_weights_stdev(0.02)

    if noise is None:
        noise = tf.random_normal([n_samples, input_dim])
        output = lib.ops.linear.Linear(
            'Generator.Input', 256, (dim//k)*(dim//k)*8*dim, noise)
        output = tf.reshape(output, [-1, 8*dim, dim//k, dim//k])
        if bn:
            output = Batchnorm('Generator.BN1', [0,2,3], output)
            output = nonlinearity(output)
    else:
        # downsampled data as input (noise)
        # input (noise) dimension [batchsize, 3*(dim/K)*(dim/K)]
        # decode twice to tensor of [batchsize, 8*dim, 4, 4]
        output = tf.reshape(noise, [-1, 3, dim//k, dim//k])
        output = tflib.ops.conv2d.Conv2D(
            'Generator.Encoder1.1', 3, 4*dim, 5, output, stride=2)
        if bn:
            output = Batchnorm('Generator.BN1.1', [0,2,3], output)
            output = nonlinearity(output)

        output = tflib.ops.conv2d.Conv2D(
            'Generator.Encode1.2', 4*dim, 8*dim, 5, output, stride=2)
        if bn:
            output = Batchnorm('Generator.BN1.2', [0, 2, 3], output)
            output = nonlinearity(output)

    output = lib.ops.deconv2d.Deconv2D('Generator.2', 8*dim, 4*dim, 5, output)
    if bn:
        output = Batchnorm('Generator.BN2', [0,2,3], output)
    output = nonlinearity(output)

    output = lib.ops.deconv2d.Deconv2D('Generator.3', 4*dim, 2*dim, 5, output)
    if bn:
        output = Batchnorm('Generator.BN3', [0,2,3], output)
    output = nonlinearity(output)

    output = lib.ops.deconv2d.Deconv2D('Generator.4', 2*dim, dim, 5, output)
    if bn:
        output = Batchnorm('Generator.BN4', [0,2,3], output)
    output = nonlinearity(output)

    output = lib.ops.deconv2d.Deconv2D('Generator.5', dim, 3, 5, output)
    output = tf.tanh(output)

    lib.ops.conv2d.unset_weights_stdev()
    lib.ops.deconv2d.unset_weights_stdev()
    lib.ops.linear.unset_weights_stdev()

    return tf.reshape(output, [-1, OUTPUT_DIM])

def WGANPaper_CrippledDCGANGenerator(
        n_samples, noise=None, dim=DIM, input_dim=INPUT_DIM):
    if noise is None:
        noise = tf.random_normal([n_samples, input_dim])

    output = lib.ops.linear.Linear('Generator.Input', input_dim, 4*4*dim, noise)
    output = tf.nn.relu(output)
    output = tf.reshape(output, [-1, dim, 4, 4])

    output = lib.ops.deconv2d.Deconv2D('Generator.2', dim, dim, 5, output)
    output = tf.nn.relu(output)

    output = lib.ops.deconv2d.Deconv2D('Generator.3', dim, dim, 5, output)
    output = tf.nn.relu(output)

    output = lib.ops.deconv2d.Deconv2D('Generator.4', dim, dim, 5, output)
    output = tf.nn.relu(output)

    output = lib.ops.deconv2d.Deconv2D('Generator.5', dim, 3, 5, output)
    output = tf.tanh(output)

    return tf.reshape(output, [-1, OUTPUT_DIM])

def ResnetGenerator(n_samples, noise=None, dim=DIM, input_dim=INPUT_DIM):
    if noise is None:
        noise = tf.random_normal([n_samples, input_dim])

    output = lib.ops.linear.Linear('Generator.Input', input_dim, 4*4*8*dim, noise)
    output = tf.reshape(output, [-1, 8*dim, 4, 4])

    for i in range(6):
        output = ResidualBlock('Generator.4x4_{}'.format(i), 8*dim, 8*dim, 3, output, resample=None)
    output = ResidualBlock('Generator.Up1', 8*dim, 4*dim, 3, output, resample='up')
    for i in range(6):
        output = ResidualBlock('Generator.8x8_{}'.format(i), 4*dim, 4*dim, 3, output, resample=None)
    output = ResidualBlock('Generator.Up2', 4*dim, 2*dim, 3, output, resample='up')
    for i in range(6):
        output = ResidualBlock('Generator.16x16_{}'.format(i), 2*dim, 2*dim, 3, output, resample=None)
    output = ResidualBlock('Generator.Up3', 2*dim, 1*dim, 3, output, resample='up')
    for i in range(6):
        output = ResidualBlock('Generator.32x32_{}'.format(i), 1*dim, 1*dim, 3, output, resample=None)
    output = ResidualBlock('Generator.Up4', 1*dim, dim//2, 3, output, resample='up')
    for i in range(5):
        output = ResidualBlock('Generator.64x64_{}'.format(i), dim/2, dim/2, 3, output, resample=None)

    output = lib.ops.conv2d.Conv2D('Generator.Out', dim//2, 3, 1, output, he_init=False)
    output = tf.tanh(output / 5.)

    return tf.reshape(output, [-1, OUTPUT_DIM])

def MultiplicativeDCGANGenerator(n_samples, noise=None, dim=DIM, bn=True, input_dim=INPUT_DIM):
    if noise is None:
        noise = tf.random_normal([n_samples, input_dim])

    output = lib.ops.linear.Linear('Generator.Input', input_dim, 4*4*8*dim*2, noise)
    output = tf.reshape(output, [-1, 8*dim*2, 4, 4])
    if bn:
        output = Batchnorm('Generator.BN1', [0,2,3], output)
    output = pixcnn_gated_nonlinearity(output[:,::2], output[:,1::2])

    output = lib.ops.deconv2d.Deconv2D('Generator.2', 8*dim, 4*dim*2, 5, output)
    if bn:
        output = Batchnorm('Generator.BN2', [0,2,3], output)
    output = pixcnn_gated_nonlinearity(output[:,::2], output[:,1::2])

    output = lib.ops.deconv2d.Deconv2D('Generator.3', 4*dim, 2*dim*2, 5, output)
    if bn:
        output = Batchnorm('Generator.BN3', [0,2,3], output)
    output = pixcnn_gated_nonlinearity(output[:,::2], output[:,1::2])

    output = lib.ops.deconv2d.Deconv2D('Generator.4', 2*dim, dim*2, 5, output)
    if bn:
        output = Batchnorm('Generator.BN4', [0,2,3], output)
    output = pixcnn_gated_nonlinearity(output[:,::2], output[:,1::2])

    output = lib.ops.deconv2d.Deconv2D('Generator.5', dim, 3, 5, output)
    output = tf.tanh(output)

    return tf.reshape(output, [-1, OUTPUT_DIM])

# ! Discriminators

def MultiplicativeDCGANDiscriminator(inputs, dim=DIM, bn=True):
    output = tf.reshape(inputs, [-1, 3, 64, 64])

    output = lib.ops.conv2d.Conv2D('Discriminator.1', 3, dim*2, 5, output, stride=2)
    output = pixcnn_gated_nonlinearity(output[:,::2], output[:,1::2])

    output = lib.ops.conv2d.Conv2D('Discriminator.2', dim, 2*dim*2, 5, output, stride=2)
    if bn:
        output = Batchnorm('Discriminator.BN2', [0,2,3], output)
    output = pixcnn_gated_nonlinearity(output[:,::2], output[:,1::2])

    output = lib.ops.conv2d.Conv2D('Discriminator.3', 2*dim, 4*dim*2, 5, output, stride=2)
    if bn:
        output = Batchnorm('Discriminator.BN3', [0,2,3], output)
    output = pixcnn_gated_nonlinearity(output[:,::2], output[:,1::2])

    output = lib.ops.conv2d.Conv2D('Discriminator.4', 4*dim, 8*dim*2, 5, output, stride=2)
    if bn:
        output = Batchnorm('Discriminator.BN4', [0,2,3], output)
    output = pixcnn_gated_nonlinearity(output[:,::2], output[:,1::2])

    output = tf.reshape(output, [-1, 4*4*8*dim])
    output = lib.ops.linear.Linear('Discriminator.Output', 4*4*8*dim, 1, output)

    return tf.reshape(output, [-1])

def ResnetDiscriminator(inputs, dim=DIM):
    output = tf.reshape(inputs, [-1, 3, 64, 64])
    output = lib.ops.conv2d.Conv2D('Discriminator.In', 3, dim//2, 1, output, he_init=False)

    for i in range(5):
        output = ResidualBlock('Discriminator.64x64_{}'.format(i), dim/2, dim/2, 3, output, resample=None)
    output = ResidualBlock('Discriminator.Down1', dim//2, dim*1, 3, output, resample='down')
    for i in range(6):
        output = ResidualBlock('Discriminator.32x32_{}'.format(i), dim*1, dim*1, 3, output, resample=None)
    output = ResidualBlock('Discriminator.Down2', dim*1, dim*2, 3, output, resample='down')
    for i in range(6):
        output = ResidualBlock('Discriminator.16x16_{}'.format(i), dim*2, dim*2, 3, output, resample=None)
    output = ResidualBlock('Discriminator.Down3', dim*2, dim*4, 3, output, resample='down')
    for i in range(6):
        output = ResidualBlock('Discriminator.8x8_{}'.format(i), dim*4, dim*4, 3, output, resample=None)
    output = ResidualBlock('Discriminator.Down4', dim*4, dim*8, 3, output, resample='down')
    for i in range(6):
        output = ResidualBlock('Discriminator.4x4_{}'.format(i), dim*8, dim*8, 3, output, resample=None)

    output = tf.reshape(output, [-1, 4*4*8*dim])
    output = lib.ops.linear.Linear('Discriminator.Output', 4*4*8*dim, 1, output)

    return tf.reshape(output / 5., [-1])

def FCDiscriminator(inputs, FC_DIM=512, n_layers=3):
    output = LeakyReLULayer('Discriminator.Input', OUTPUT_DIM, FC_DIM, inputs)
    for i in range(n_layers):
        output = LeakyReLULayer('Discriminator.{}'.format(i), FC_DIM, FC_DIM, output)
    output = lib.ops.linear.Linear('Discriminator.Out', FC_DIM, 1, output)

    return tf.reshape(output, [-1])

def DCGANDiscriminator(inputs, dim=DIM, bn=True, nonlinearity=LeakyReLU):
    output = tf.reshape(inputs, [-1, 3, 64, 64])

    lib.ops.conv2d.set_weights_stdev(0.02)
    lib.ops.deconv2d.set_weights_stdev(0.02)
    lib.ops.linear.set_weights_stdev(0.02)

    output = lib.ops.conv2d.Conv2D('Discriminator.1', 3, dim, 5, output, stride=2)
    output = nonlinearity(output)

    output = lib.ops.conv2d.Conv2D('Discriminator.2', dim, 2*dim, 5, output, stride=2)
    if bn:
        output = Batchnorm('Discriminator.BN2', [0,2,3], output)
    output = nonlinearity(output)

    output = lib.ops.conv2d.Conv2D('Discriminator.3', 2*dim, 4*dim, 5, output, stride=2)
    if bn:
        output = Batchnorm('Discriminator.BN3', [0,2,3], output)
    output = nonlinearity(output)

    output = lib.ops.conv2d.Conv2D('Discriminator.4', 4*dim, 8*dim, 5, output, stride=2)
    if bn:
        output = Batchnorm('Discriminator.BN4', [0,2,3], output)
    output = nonlinearity(output)

    output = tf.reshape(output, [-1, 4*4*8*dim])
    output = lib.ops.linear.Linear('Discriminator.Output', 4*4*8*dim, 1, output)

    lib.ops.conv2d.unset_weights_stdev()
    lib.ops.deconv2d.unset_weights_stdev()
    lib.ops.linear.unset_weights_stdev()

    return tf.reshape(output, [-1])

# kernel for downsampling
arr = np.zeros([K, K, 3, 3])
arr[:,:,0,0] = 1.0/(K*K)
arr[:,:,1,1] = 1.0/(K*K)
arr[:,:,2,2] = 1.0/(K*K)
_downsample_weight = tf.constant(arr, dtype=tf.float32)

def downsample(data, method='conv'):
    data = tf.reshape(data, [-1, 3, DIM, DIM])
    # BCHW -> BHWC
    data = tf.transpose(data, [0, 2, 3, 1])
    if method == 'conv':
        data = tf.nn.conv2d(data, _downsample_weight,
                            strides=[1, K, K, 1], padding='SAME')
    elif method == 'area':
        data = tf.image.resize_area(data, [DIM//K, DIM//K])
    # BHWC -> BCHW
    data = tf.transpose(data, [0, 3, 1, 2])
    data = tf.reshape(data, [-1, 3 * DIM//K * DIM//K])
    return data

Generator, Discriminator = GeneratorAndDiscriminator()

with tf.Session(config=tf.ConfigProto(allow_soft_placement=True)) as session:

    all_real_data_conv = tf.placeholder(tf.int32, shape=[BATCH_SIZE, 3, 64, 64])
    if tf.__version__.startswith('1.'):
        split_real_data_conv = tf.split(all_real_data_conv, len(DEVICES))
    else:
        split_real_data_conv = tf.split(0, len(DEVICES), all_real_data_conv)

    gen_l1_costs, gen_gan_costs = [], []
    gen_costs, disc_costs = [],[]

    for device_index, (device, real_data_conv) in enumerate(zip(DEVICES, split_real_data_conv)):
        with tf.device(device):
            real_data = 2*((tf.cast(real_data_conv, tf.float32)/255.)-.5)
            real_data = tf.reshape(real_data, [BATCH_SIZE//len(DEVICES), OUTPUT_DIM])
            # downsampled (by K) as generator input
            real_data_downsampled = downsample(real_data)
            fake_data = Generator(BATCH_SIZE//len(DEVICES), noise=real_data_downsampled)

            disc_real = Discriminator(real_data)
            disc_fake = Discriminator(fake_data)

            if MODE == 'wgan':
                gen_cost = tf.reduce_mean(disc_fake)
                disc_cost = tf.reduce_mean(disc_real) - tf.reduce_mean(disc_fake)

            elif MODE == 'wgan-gp':
                gen_cost = tf.reduce_mean(disc_fake)
                disc_cost = tf.reduce_mean(disc_real) - tf.reduce_mean(disc_fake)

                alpha = tf.random_uniform(
                    shape=[BATCH_SIZE//len(DEVICES),1],
                    minval=0.,
                    maxval=1.
                )
                differences = fake_data - real_data
                interpolates = real_data + (alpha*differences)
                gradients = tf.gradients(Discriminator(interpolates), [interpolates])[0]
                slopes = tf.sqrt(tf.reduce_sum(tf.square(gradients), reduction_indices=[1]))
                gradient_penalty = tf.reduce_mean((slopes-1.)**2)
                disc_cost += LAMBDA*gradient_penalty

            elif MODE == 'dcgan':
                try: # tf pre-1.0 (bottom) vs 1.0 (top)
                    gen_cost = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=disc_fake,
                                                                                labels=tf.ones_like(disc_fake)))
                    disc_cost =  tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=disc_fake,
                                                                                labels=tf.zeros_like(disc_fake)))
                    disc_cost += tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=disc_real,
                                                                                labels=tf.ones_like(disc_real)))
                except Exception as e:
                    gen_cost = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(disc_fake, tf.ones_like(disc_fake)))
                    disc_cost =  tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(disc_fake, tf.zeros_like(disc_fake)))
                    disc_cost += tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(disc_real, tf.ones_like(disc_real)))
                disc_cost /= 2.

            elif MODE == 'lsgan':
                gen_cost = tf.reduce_mean((disc_fake - 1)**2)
                disc_cost = (tf.reduce_mean((disc_real - 1)**2) + tf.reduce_mean((disc_fake - 0)**2))/2.

            else:
                raise Exception()

            # add L1 difference to penalty
            fake_data_downsampled = downsample(fake_data)
            gen_l1_cost = tf.reduce_mean(
                tf.abs(fake_data_downsampled - real_data_downsampled))

            gen_l1_costs.append(gen_l1_cost)
            gen_gan_costs.append(gen_cost)

            gen_cost = GEN_L1_WEIGHT * gen_l1_cost + (1 - GEN_L1_WEIGHT) * gen_cost

            gen_costs.append(gen_cost)
            disc_costs.append(disc_cost)

    gen_cost = tf.add_n(gen_costs) / len(DEVICES)
    disc_cost = tf.add_n(disc_costs) / len(DEVICES)
    gen_gan_cost = tf.add_n(gen_gan_costs) / len(DEVICES)
    gen_l1_cost = tf.add_n(gen_l1_costs) / len(DEVICES)
    tf.summary.scalar('gen gan loss', gen_gan_cost, collections=['scalars'])
    tf.summary.scalar('gen l1 diff', gen_l1_cost, collections=['scalars'])
    tf.summary.scalar('gen loss', gen_cost, collections=['scalars'])
    tf.summary.scalar('disc loss', disc_cost, collections=['scalars'])

    if MODE == 'wgan':
        gen_train_op = tf.train.RMSPropOptimizer(learning_rate=1e-4).minimize(
            gen_cost, var_list=lib.params_with_name('Generator'), colocate_gradients_with_ops=True)
        disc_train_op = tf.train.RMSPropOptimizer(learning_rate=1e-4).minimize(disc_cost,
                                             var_list=lib.params_with_name('Discriminator.'), colocate_gradients_with_ops=True)

        clip_ops = []
        for var in lib.params_with_name('Discriminator'):
            clip_bounds = [-.01, .01]
            clip_ops.append(tf.assign(var, tf.clip_by_value(var, clip_bounds[0], clip_bounds[1])))
        clip_disc_weights = tf.group(*clip_ops)

    elif MODE == 'wgan-gp':
        gen_train_op = tf.train.AdamOptimizer(
            learning_rate=1e-4, beta1=0.5, beta2=0.9).minimize(
                gen_cost,var_list=lib.params_with_name('Generator'), colocate_gradients_with_ops=True)
        disc_train_op = tf.train.AdamOptimizer(learning_rate=1e-4, beta1=0.5, beta2=0.9).minimize(disc_cost,
                                           var_list=lib.params_with_name('Discriminator.'), colocate_gradients_with_ops=True)

    elif MODE == 'dcgan':
        gen_train_op = tf.train.AdamOptimizer(learning_rate=2e-4, beta1=0.5).minimize(gen_cost,
                                          var_list=lib.params_with_name('Generator'), colocate_gradients_with_ops=True)
        disc_train_op = tf.train.AdamOptimizer(learning_rate=2e-4, beta1=0.5).minimize(disc_cost,
                                           var_list=lib.params_with_name('Discriminator.'), colocate_gradients_with_ops=True)

    elif MODE == 'lsgan':
        gen_train_op = tf.train.RMSPropOptimizer(learning_rate=1e-4).minimize(gen_cost,
                                             var_list=lib.params_with_name('Generator'), colocate_gradients_with_ops=True)
        disc_train_op = tf.train.RMSPropOptimizer(learning_rate=1e-4).minimize(disc_cost,
                                              var_list=lib.params_with_name('Discriminator.'), colocate_gradients_with_ops=True)

    else:
        raise Exception()

#     # For generating samples
#     fixed_noise = tf.constant(np.random.normal(size=(BATCH_SIZE, INPUT_DIM)).astype('float32'))
#     all_fixed_noise_samples = []
#     for device_index, device in enumerate(DEVICES):
#         n_samples = BATCH_SIZE // len(DEVICES)
#         all_fixed_noise_samples.append(Generator(n_samples, noise=fixed_noise[device_index*n_samples:(device_index+1)*n_samples]))
#     if tf.__version__.startswith('1.'):
#         all_fixed_noise_samples = tf.concat(all_fixed_noise_samples, axis=0)
#     else:
#         all_fixed_noise_samples = tf.concat(0, all_fixed_noise_samples)

#     def generate_image(iteration):
#         # add image to summary
#         samples_reshaped = tf.reshape(
#             all_fixed_noise_samples, (BATCH_SIZE, 3, DIM, DIM))
#         samples_reshaped = tf.transpose(samples_reshaped, [0, 2, 3, 1])
#         image_op = tf.summary.image(
#             'generator output', samples_reshaped)
#         image_summary = session.run(image_op)
#         summary_writer.add_summary(image_summary, iteration)

#         samples = session.run(all_fixed_noise_samples)
#         samples = ((samples+1.)*(255.99/2)).astype('int32')
#         lib.save_images.save_images(samples.reshape((BATCH_SIZE, 3, 64, 64)), 'samples_{}.png'.format(iteration))

    def generate_test_image(iteration, real_data, fake_data,  max_samples=10):
        feature = tf.reshape(real_data_downsampled, [-1, 3, DIM//K, DIM//K])
        # BCHW -> BHWC
        feature = (tf.transpose(feature, [0, 2, 3, 1]) + 1)/2.
        nearest = tf.image.resize_nearest_neighbor(feature, [DIM, DIM])
        nearest = tf.maximum(tf.minimum(nearest, 1.), 0.)
        bicubic = tf.image.resize_bicubic(feature, [DIM, DIM])
        bicubic = tf.maximum(tf.minimum(bicubic, 1.), 0.)
        fake_data = (tf.reshape(fake_data, [-1, 3, DIM, DIM]) + 1.)/2.
        fake_data = tf.transpose(fake_data, [0, 2, 3, 1])
        real_data = tf.reshape(real_data, [-1, 3, DIM, DIM])
        real_data = tf.transpose(real_data, [0, 2, 3, 1])
        real_data = (real_data + 1.) / 2.
        clipped = tf.maximum(tf.minimum(fake_data, 1.), 0.)
        image = tf.concat([nearest, bicubic, clipped, real_data], 2)

        feed_dict = {real_data_conv: test_data}
        image_col = tf.summary.image('generator output', image, max_samples)
        image_summary = session.run(image_col, feed_dict=feed_dict)
        summary_writer.add_summary(image_summary, iteration)

        image = image[0:max_samples,:,:,:]
        image = tf.concat([image[i,:,:,:] for i in range(max_samples)], 0)
        clipped = clipped[0:max_samples, :, :, :]
        clipped = tf.concat([clipped[i, :, :, :] for i in range(max_samples)], 1)

        image, clipped = session.run([image, clipped], feed_dict=feed_dict)

        filename_1 = 'batch%06d_image.png' % iteration
        filename_2 = 'batch%06d_row.png' % iteration
        filename_1 = os.path.join(TRAIN_DIR, filename_1)
        filename_2 = os.path.join(TRAIN_DIR, filename_2)
        scipy.misc.toimage(image, cmin=0., cmax=1.).save(filename_1)
        scipy.misc.toimage(clipped, cmin=0., cmax=1.).save(filename_2)
        print("Saved %s %s" % (filename_1, filename_2))

    # Dataset iterator and test set (for visualization)
    train_gen, test_data = lib.celebA_64x64.load(BATCH_SIZE, data_dir=DATA_DIR)
    #train_gen, dev_gen = lib.small_imagenet.load(BATCH_SIZE, data_dir=DATA_DIR)

    def inf_train_gen():
        while True:
            for (images,) in train_gen():
                yield images

    # Save a batch of ground-truth samples
    _x = next(inf_train_gen())
    _x_r = session.run(real_data, feed_dict={real_data_conv: _x})
    _x_r = ((_x_r+1.)*(255.99/2)).astype('int32')
    lib.save_images.save_images(_x_r.reshape((BATCH_SIZE, 3, 64, 64)), 'samples_groundtruth.png')

    # Train loop
    merged_scalars = tf.summary.merge_all(key='scalars')
    summary_writer = tf.summary.FileWriter(SUMMARY_DIR, session.graph)

    session.run(tf.global_variables_initializer())
    gen = inf_train_gen()
    all_start_time = time.time()
    for iteration in range(ITERS):
        start_time = time.time()
        # finish if run overtime
        total_elapsed = (start_time - all_start_time) / 60.
        if total_elapsed > FLAGS.max_runtime:
            break

        # Train generator
        if iteration > 0:
            _ = session.run(gen_train_op, feed_dict={all_real_data_conv: _data})

        # Train critic
        if (MODE == 'dcgan') or (MODE == 'lsgan'):
            disc_iters = 1
        else:
            disc_iters = CRITIC_ITERS
        for i in range(disc_iters):
            _data = next(gen)
            _disc_cost, _ = session.run([disc_cost, disc_train_op], feed_dict={all_real_data_conv: _data})
            if MODE == 'wgan':
                _ = session.run([clip_disc_weights])

        lib.plot.plot('train disc cost', _disc_cost)
        lib.plot.plot('time', time.time() - start_time)
        #print('iter={0} disc_loss={1:.3g} time={2:.2g}'.format(
        #    iteration, _disc_cost, time.time() - start_time))

        if iteration % 10 == 0:
            merged_summary = session.run(merged_scalars, feed_dict={all_real_data_conv: _data})
            summary_writer.add_summary(merged_summary, iteration)

        if iteration % 200 == 9:
            t = time.time()
            #dev_disc_costs = []
            #for (images,) in dev_gen():
            #    _dev_disc_cost = session.run(disc_cost, feed_dict={all_real_data_conv: _data})
            #    dev_disc_costs.append(_dev_disc_cost)
            #lib.plot.plot('dev disc cost', np.mean(dev_disc_costs))
            generate_test_image(iteration, real_data, fake_data)

        if (iteration < 5) or (iteration % 200 == 199):
            lib.plot.flush()

        lib.plot.tick()

if __name__ == '__main__':
    tf.app.run()


# Main entrance of GAIL
import numpy as np
import torch
import torch.nn.functional as F
import gym
import safety_gym
import time
import os.path as osp

import wandb

from neural_nets import Discriminator, ActorCritic, count_vars

from utils import BufferStudent, BufferTeacher
from utils import mpi_fork, proc_id, num_procs, EpochLogger,\
    average_gradients, sync_all_params, setup_pytorch_for_mpi, sync_params, mpi_avg_grads

def gail_penalized(env_fn, actor_critic=ActorCritic, ac_kwargs=dict(),
         disc=Discriminator,
         dc_kwargs=dict(), seed=0,
         episodes_per_epoch=40,
         epochs=500,
         gamma=0.99, lam=0.97,
         # Cost constraints / penalties:
         cost_lim=25,
         penalty_init=1.,
         penalty_lr=5e-3,
         clip_ratio=0.2,
         pi_lr=3e-3, vf_lr=3e-3, dc_lr=5e-4,
         train_v_iters=80, train_pi_iters=80, train_dc_iters=80,
         max_ep_len=1000, logger_kwargs=dict(), config_name = 'standard', save_freq=10):
    # W&B Logging
    wandb.login()

    composite_name = 'new_gail_penalized_' + config_name
    wandb.init(project="LearningCurves", group="GAIL Clone", name=composite_name)

    # Special function to avoid certain slowdowns from PyTorch + MPI combo.
    setup_pytorch_for_mpi()

    l_lam = 0  # balance two loss terms

    # Set up logger and save configuration
    logger = EpochLogger(**logger_kwargs)
    logger.save_config(locals())

    seed += 10000 * proc_id()
    torch.manual_seed(seed)
    np.random.seed(seed)

    # Instantiate environment
    env = env_fn()
    obs_dim = env.observation_space.shape
    act_dim = env.action_space.shape

    ac_kwargs['action_space'] = env.action_space

    # Models  # Create actor-critic and discriminator modules
    ac = actor_critic(input_dim=obs_dim[0], **ac_kwargs)
    discrim = disc(input_dim=obs_dim[0], **dc_kwargs)

    # Set up model saving
    logger.setup_pytorch_saver([ac, discrim])

    # Sync params across processes
    sync_params(ac)
    sync_params(discrim)

    # Load expert policy here
    expert = actor_critic(input_dim=obs_dim[0], **ac_kwargs)
    # expert_name = "expert_torch_save.pt"
    expert_name = "model.pt"
    # expert = torch.load(osp.join(logger_kwargs['output_dir'],'pyt_save' , expert_name))
    # expert = torch.load('/home/tyna/Documents/openai/research-project/data/anonymous-expert/anonymous-expert_s0/pyt_save/model.pt')
    expert = torch.load(
        '/home/tyna/Documents/openai/research-project/data/test-pen-ppo/test-pen-ppo_s0/pyt_save/model.pt')

    print('RUNNING GAIL')

    # Buffers
    local_episodes_per_epoch = int(episodes_per_epoch / num_procs())
    buff_s = BufferStudent(obs_dim[0], act_dim[0], local_episodes_per_epoch, max_ep_len)
    buff_t = BufferTeacher(obs_dim[0], act_dim[0], local_episodes_per_epoch, max_ep_len)

    # Count variables
    var_counts = tuple(count_vars(module) for module in [ac.pi, ac.v, discrim.pi])
    logger.log('\nNumber of parameters: \t pi: %d, \t v: %d, \t d: %d\n' % var_counts)

    # Optimizers
    pi_optimizer = torch.optim.Adam(ac.pi.parameters(), lr=pi_lr)
    vf_optimizer = torch.optim.Adam(ac.v.parameters(), lr=vf_lr)
    discrim_optimizer = torch.optim.Adam(discrim.pi.parameters(), lr=dc_lr)

    # # Parameters Sync
    # sync_all_params(ac.parameters())
    # sync_all_params(disc.parameters())

    # Set up function for computing PPO policy loss
    def compute_loss_pi(obs, act, adv, logp_old):
        # Policy loss # policy gradient term + entropy term
        # Policy loss with clipping (without clipping, loss_pi = -(logp*adv).mean()).
        # TODO: Think about removing clipping
        _, logp, _ = ac.pi(obs, act)
        ratio = torch.exp(logp - logp_old)
        clip_adv = torch.clamp(ratio, 1 - clip_ratio, 1 + clip_ratio) * adv
        loss_pi = -(torch.min(ratio * adv, clip_adv)).mean()

        return loss_pi

    def penalty_update(cur_penalty):
        cur_cost = logger.get_stats('EpCostS')[0]
        cur_rew = logger.get_stats('EpRetS')[0]

        # Penalty update
        cur_penalty = max(0, cur_penalty + penalty_lr * (cur_cost - cost_lim))
        return cur_penalty

    def update(e):
        obs_s, act, adv, ret, log_pi_old = [torch.Tensor(x) for x in buff_s.retrieve_all()]
        obs_t, _ = [torch.Tensor(x) for x in buff_t.retrieve_all()]

        # Policy
        _, logp, _ = ac.pi(obs_s, act)
        entropy = (-logp).mean()

        # Policy loss   # policy gradient term + entropy term
        # loss_pi = -(logp * adv).mean() - l_lam * entropy

        # Train policy
        if e > 10:
            # Train policy with multiple steps of gradient descent
            for _ in range(train_pi_iters):
                pi_optimizer.zero_grad()
                loss_pi = compute_loss_pi(obs, act, adv, ret)
                loss_pi.backward()
                mpi_avg_grads(ac.pi)
                pi_optimizer.step()

        # Value function
        v = ac.v(obs_s)
        v_l_old = F.mse_loss(v, ret)

        for _ in range(train_v_iters):
            v = ac.v(obs_s)
            v_loss = F.mse_loss(v, ret)

            # Value function train
            vf_optimizer.zero_grad()
            v_loss.backward()
            mpi_avg_grads(ac.v)  # average gradients across MPI processes
            vf_optimizer.step()

        # Discriminator
        gt1 = torch.ones(obs_s.size()[0], dtype=torch.int)
        gt2 = torch.zeros(obs_t.size()[0], dtype=torch.int)
        _, logp_student, _ = discrim(obs_s, gt=gt1)
        _, logp_teacher, _ = discrim(obs_t, gt=gt2)
        discrim_l_old = - logp_student.mean() - logp_teacher.mean()

        for _ in range(train_dc_iters):
            _, logp_student, _ = discrim(obs_s, gt=gt1)
            _, logp_teacher, _ = discrim(obs_t, gt=gt2)
            dc_loss = - logp_student.mean() - logp_teacher.mean()

            # Discriminator train
            discrim_optimizer.zero_grad()
            dc_loss.backward()
            # average_gradients(discrim_optimizer.param_groups)
            mpi_avg_grads(discrim.pi)
            discrim_optimizer.step()

        _, logp_student, _ = discrim(obs_s, gt=gt1)
        _, logp_teacher, _ = discrim(obs_t, gt=gt2)
        dc_loss_new = - logp_student.mean() - logp_teacher.mean()

        # Log the changes
        _, logp, _, v = ac(obs, act)
        entropy_new = (-logp).mean()
        pi_loss_new = -(logp * adv).mean() - l_lam * entropy
        v_loss_new = F.mse_loss(v, ret)
        kl = (log_pi_old - logp).mean()
        logger.store(
            # LossPi=loss_pi,
            LossV=v_l_old, LossDC=discrim_l_old,
                     # DeltaLossPi=(pi_loss_new - loss_pi),
                     DeltaLossV=(v_loss_new - v_l_old), DeltaLossDC=(dc_loss_new - discrim_l_old),
                     DeltaEnt=(entropy_new - entropy),
                     Entropy=entropy, KL=kl)

    start_time = time.time()
    o, r, sdr, d, ep_ret, ep_cost, ep_sdr, ep_len = env.reset(), 0, 0, False, 0, 0, 0, 0
    total_t = 0
    ep_len_t = 0

    # Initialize penalty
    cur_penalty = np.log(max(np.exp(penalty_init) - 1, 1e-8))

    for epoch in range(epochs):
        ac.eval()
        discrim.eval()
        # We recognize the probability term of index [0] correspond to the teacher's policy

        # Student's policy rollout
        for _ in range(local_episodes_per_epoch):
            for _ in range(max_ep_len):
                obs = torch.Tensor(o.reshape(1, -1))
                a, _, lopg_t, v_t = ac(obs)

                buff_s.store(o, a.detach().numpy(), r, sdr, v_t.item(), lopg_t.detach().numpy())
                logger.store(VVals=v_t)

                o, r, d, info = env.step(a.detach().numpy()[0])
                # print("INFO: ", info)
                c = info.get("cost")

                _, sdr, _ = discrim(torch.Tensor(o.reshape(1, -1)), gt=torch.Tensor([0]))
                if sdr < -4:  # Truncate rewards
                    sdr = -4

                ep_ret += r
                ep_cost += c
                ep_sdr += sdr
                ep_len += 1
                total_t += 1

                terminal = d or (ep_len == max_ep_len)
                if terminal:
                    buff_s.end_episode()
                    logger.store(EpRetS=ep_ret, EpCostS= ep_cost, EpLenS=ep_len, EpSdrS=ep_sdr)
                    print("Student Episode Return: \t", ep_ret)
                    o, r, sdr, d, ep_ret, ep_cost, ep_sdr, ep_len = env.reset(), 0, 0, False, 0, 0, 0, 0

        # Teacher's policy rollout
        for _ in range(local_episodes_per_epoch):
            for _ in range(max_ep_len):
                # obs =
                a, _, _, _ = expert(torch.Tensor(o.reshape(1, -1)))

                buff_t.store(o, a.detach().numpy(), r)

                o, r, d, info = env.step(a.detach().numpy()[0])
                c = info.get("cost")
                ep_ret += r
                ep_cost += c
                ep_len += 1
                total_t += 1

                terminal = d or (ep_len == max_ep_len)
                if terminal:
                    buff_t.end_episode()
                    logger.store(EpRetT=ep_ret, EpCostT=ep_cost, EpLenT=ep_len)
                    print("Teacher Episode Return: \t", ep_ret)
                    o, r, d, ep_ret, ep_cost, ep_len = env.reset(), 0, False, 0, 0, 0

        if (epoch % save_freq == 0) or (epoch == epochs - 1):
            logger.save_state({'env': env}, [ac, discrim], None)

        # Update
        ac.train()
        discrim.train()

        # update penalty
        cur_penalty = penalty_update(cur_penalty)

        # update networks
        update(epoch)

        # Log
        logger.log_tabular('Epoch', epoch)
        logger.log_tabular('EpRetS', average_only=True)
        logger.log_tabular('EpSdrS', average_only=True)
        logger.log_tabular('EpLenS', average_only=True)
        logger.log_tabular('EpRetT', average_only=True)
        logger.log_tabular('EpLenT', average_only=True)
        logger.log_tabular('VVals', with_min_and_max=True)
        logger.log_tabular('TotalEnvInteracts', total_t)
        # logger.log_tabular('LossPi', average_only=True)
        # logger.log_tabular('DeltaLossPi', average_only=True)
        logger.log_tabular('LossV', average_only=True)
        logger.log_tabular('DeltaLossV', average_only=True)
        logger.log_tabular('LossDC', average_only=True)
        logger.log_tabular('DeltaLossDC', average_only=True)
        # logger.log_tabular('Entropy', average_only=True)
        # logger.log_tabular('DeltaEnt', average_only=True)
        # logger.log_tabular('KL', average_only=True)
        logger.log_tabular('Time', time.time() - start_time)
        logger.dump_tabular()

if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument('--env', type=str, default='Safexp-PointGoal1-v0')
    parser.add_argument('--hid', type=int, default=64)
    parser.add_argument('--l', type=int, default=2)
    parser.add_argument('--gamma', type=float, default=0.99)
    parser.add_argument('--lam', type=float, default=0.97)
    parser.add_argument('--seed', '-s', type=int, default=0)
    parser.add_argument('--cpu', type=int, default=1)
    parser.add_argument('--episodes-per-epoch', type=int, default=5)
    # parser.add_argument('--episodes-per-epoch', type=int, default=40)
    parser.add_argument('--epochs', type=int, default=1000)
    parser.add_argument('--exp_name', type=str, default='valor-anonymous-expert')
    parser.add_argument('--con', type=int, default=5)
    args = parser.parse_args()

    mpi_fork(args.cpu)

    from utils import setup_logger_kwargs

    logger_kwargs = setup_logger_kwargs(args.exp_name, args.seed)

    gail_penalized(lambda: gym.make(args.env), actor_critic=ActorCritic, ac_kwargs=dict(hidden_dims=[args.hid] * args.l),
         disc=Discriminator, dc_kwargs=dict(hidden_dims=[args.hid] * args.l), gamma=args.gamma, lam=args.lam,
         seed=args.seed, episodes_per_epoch=args.episodes_per_epoch, epochs=args.epochs, logger_kwargs=logger_kwargs)


#!/usr/bin/env python
'''
mcu: Modeling and Crystallographic Utilities
Copyright (C) 2019 Hung Q. Pham. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

Email: Hung Q. Pham <pqh3.14@gmail.com>
'''

import numpy as np
from ..utils import plot
from ..vasp import const
from . import crystal_io
import matplotlib as mpl
import matplotlib.pyplot as plt

class main:
    def __init__(self,  seedname="outfile"):
        '''

        '''
        self.seedname = seedname

############ Plotting #################
    def get_band(self, phonon=False, gamma_correct=False, threshold=8.06554):
        '''Make a band from from
           NOTE: the proj_kpath is computed from the dk and nkp. Due to the round out error in f25, the computed high symmetric k-point coordinates won't be exactly the same as values obtained from *.BAND file.

           if phonon == True: the phonon band is read. Three small negative acoustic modes at Gamma can be removed by using the gamma_correct keyword and threshold 1 meV = 8.06554 cm^-1
        '''
        data = crystal_io.read_f25(self.seedname + ".f25")
        temp = []
        ihferm_list = []
        proj_kpath = []
        sym_kpoint_coor = [0.0]
        shift = 0
        for block in data:
            ihferm, type, nband, nkp, dk, efermi, eigenvals = block
            if type == 'BAND':
                temp.append(eigenvals.reshape(-1, nband))
                path = np.arange(nkp)*dk + shift
                proj_kpath.append(path)
                shift = path[-1] + dk
                sym_kpoint_coor.append(path[-1])

        if ihferm == 0 or ihferm == 2:
            band = np.float64([np.vstack(temp)])
            proj_kpath = np.hstack(proj_kpath)
        elif ihferm == 1 or ihferm == 3:
            nblock = len(temp) // 2
            band_up = np.vstack(temp[:nblock])
            band_down = np.vstack(temp[nblock:])
            band = np.float64([band_up, band_down])
            proj_kpath = np.hstack(proj_kpath[:nblock])
        sym_kpoint_coor = np.float64(sym_kpoint_coor)
        if phonon == True:
            if gamma_correct == True:
                # Should be checked first before using the correction to make sure it is really acoustic phonon modes
                nspin, nkpts, nband = band.shape
                band = band.flatten()
                imag_mode_idx = band < 0.0
                imag_mode = band[imag_mode_idx]
                imag_mode_idx2 = imag_mode > -threshold
                imag_mode[imag_mode_idx2] = 0.0
                band[imag_mode_idx] = imag_mode
                band = band.reshape(nspin, nkpts, nband)
        else:
            band = const.AUTOEV * band
            efermi = const.AUTOEV * efermi

        return band, proj_kpath, sym_kpoint_coor, efermi

    def get_bandgap(self):
        '''Get the bandgap'''

        band, proj_kpath, sym_kpoint_coor, efermi = self.get_band()
        nspin, nkpts, nbands = band.shape
        for spin in range(nspin):
            print('Spin:', spin)
            CBM = None
            for bandth in range(nbands):
                shifted_band = band[spin,:,bandth] - efermi
                if (shifted_band > 0.0).all() == True:
                    CBM = band[spin,:, bandth]
                    VBM = band[spin,:, bandth -1]
                    break
                elif ((shifted_band < 0.0).any() == True) and ((shifted_band > 0.0).any() == True):
                    print("This is a metal")
                    break

            if CBM is not None:
                vbm_idx = np.argmax(VBM)
                cbm_idx = np.argmin(CBM)
                bandgap = CBM[cbm_idx] - VBM[vbm_idx]
                direct = False
                if vbm_idx == cbm_idx: direct = True

                # TODO: the kpath_frac currently cannot be obtained form f25.
                # Other outputs are needed
                # kpath_frac = self.cp2k_io.kpath_frac[set_block]
                # print('  E(VBM) = %7.4f at k = [%6.4f,%6.4f,%6.4f]' % (VBM[vbm_idx],
                                                                # kpath_frac[vbm_idx,0], kpath_frac[vbm_idx,1], kpath_frac[vbm_idx,2]))
                # print('  E(CBM) = %7.4f at k = [%6.4f,%6.4f,%6.4f]' % (CBM[cbm_idx],
                                                                # kpath_frac[cbm_idx,0], kpath_frac[cbm_idx,1], kpath_frac[cbm_idx,2]))
                if direct == True:
                    print('  Direct bandgap   : %6.3f' % (bandgap))
                else:
                    print('  Indirect bandgap : %6.3f' % (bandgap))
                    gap1 = CBM[vbm_idx] - VBM[vbm_idx]
                    gap2 = CBM[cbm_idx] - VBM[cbm_idx]
                    direct_gap = min(gap1, gap2)
                    print('  Direct bandgap   : %6.3f' % (direct_gap))

    def _generate_band(self, efermi=None, spin=0, label=None):
        '''Processing/collecting the band data before the plotting function
        '''
        band, proj_kpath, sym_kpoint_coor, efermi_ = self.get_band()
        if efermi is None: efermi = efermi_
        band = band[spin] - efermi

        return band, proj_kpath, sym_kpoint_coor, label

    def plot_band(self, efermi=None, label=None, spin=0, save=False, band_color=['#007acc','#808080','#808080'],
                    figsize=(6,6), figname='BAND', xlim=None, ylim=[-6,6], fontsize=18, dpi=600, format='png'):
        '''Plot band structure

            Attribute:
                efermi          : a Fermi level or a list of Fermi levels
                spin            : 0  for spin unpolarized and LSORBIT = .TRUE.
                                  0 or 1 for spin polarized
                color           : a list of three color codes for band curves, high symmetric kpoint grid, and Fermi level

        '''
        assert isinstance(band_color,list)
        assert len(band_color) == 3
        plot.plot_band(self, efermi=efermi, spin=spin, save=save, band_color=band_color,
                figsize=figsize, figname=figname, xlim=xlim, ylim=ylim, fontsize=fontsize, dpi=dpi, format=format, label=label)

    def _generate_phononband(self, unit="CM", gamma_correct=False, threshold=8.06554, spin=0, label=None):
        '''Processing/collecting the phnon band (cm^-1) data before the plotting function
           Unit: CM = CM^-1, THZ = THz, MEV = meV
        '''
        band, proj_kpath, sym_kpoint_coor, efermi_ = self.get_band(phonon=True, gamma_correct=gamma_correct, threshold=threshold)
        if unit.lower() == "thz":
            band = const.CMTOTHZ * band
        elif unit.lower() == "mev":
            band = const.CMTOMEV * band

        return band[spin], proj_kpath, sym_kpoint_coor, label

    def plot_phononband(self, unit="CM", gamma_correct=False, threshold=8.06554, label=None, spin=0, save=False, band_color=['#007acc','#808080','#808080'],
                    figsize=(6,6), figname='PHONONBAND', xlim=None, ylim=None, fontsize=18, dpi=600, format='png'):
        '''Plot band structure

            Attribute:
                efermi          : a Fermi level or a list of Fermi levels
                spin            : 0  for spin unpolarized and LSORBIT = .TRUE.
                                  0 or 1 for spin polarized
                color           : a list of three color codes for band curves, high symmetric kpoint grid, and Fermi level

        '''
        assert isinstance(band_color,list)
        assert len(band_color) == 3
        plot.plot_phononband(self, unit=unit, gamma_correct=gamma_correct, threshold=threshold, spin=spin, save=save, band_color=band_color,
                figsize=figsize, figname=figname, xlim=xlim, ylim=ylim, fontsize=fontsize, dpi=dpi, format=format, label=label)


"""
Colloid_Setup contains background classes and methods to prepare a
model domain for colloid simulation. The user should not need to
import or call methods directly from this module.
"""

import numpy as np
from scipy import interpolate
from copy import copy
import h5py as H
import sys

class Hdf5Reader(object):
    """
    Hdf5 reader class to grab results from lattice Boltzmann model runs
    to parameterize the colloid simulation. Consider moving this class to
    the Colloid_IO module

    Parameters:
    ----------
    :param str HDF_name: lattice boltzmann hdf file name.

    :ivar np.ndarray imarray: binary image array defining model boundaries
    :ivar np.ndarray uarry: velocity array of [y, x]
    :ivar np.ndarray yu: y velocity array from lattice Boltzmann simulation
    :ivar np.ndarray xu: x velocity array from lattice Boltzmann simulation
    :ivar float mean_yu: mean velocity in the y direction
    :ivar float mean_xu: mean velocity in the x direction
    :ivar float velocity_factor: velocity dimensionalization factor
    """
    def __init__(self, hdf_name):
        hdf = H.File(hdf_name, 'r+')
        self.imarray = hdf['Binary_image'][()]
        self.uarray = hdf['results/uarray'][()]
        self.yu = hdf['results/uarray'][()][0]
        self.xu = hdf['results/uarray'][()][1]
        self.mean_xu = hdf['results/mean_ux'][()]
        self.mean_yu = hdf['results/mean_uy'][()]
        self.velocity_factor = hdf['results/velocity_factor'][()]
        hdf.close()

class GridArray(object):
    """
    Gridarray class creates arrays of distances from pore spaces, corrects for
    interpolation effects at pore boundaries, and creates vector arrays that are
    later used to give direction to forces.

    Parameters:
    ----------
    :param (np.array, bool) arr: Array of boolean porous media (segmented image array)
    :param float gridres: model resolution in meters
    :param int gridsplit: interpolation factor for refining grid mesh
    :param bool solid: solid phase boolean identifier, default=True
    """
    def __init__(self, arr, gridres, gridsplit, solid=True):
        self.yarr = np.copy(arr.T)
        self._vimgx, self._vimgy = self._create_vector_array(arr, solid)
        self.__gridx, self.__vector_x = self._arrx(arr, self._vimgx, gridres, gridsplit)
        self.__gridy, self.__vector_y = self._arry(self.yarr, self._vimgy, gridres, gridsplit)

    @property
    def gridx(self):
        """
        :return: (np.array, np.float) Array of distances from nearest solid phase in the x-direction
        """
        return copy(self.__gridx)

    @property
    def gridy(self):
        """
        :return: (np.array, np.float) Array of distances from nearest solid phase in the y-direction
        """
        return copy(self.__gridy)

    @property
    def vector_x(self):
        """
        :return: (np.array, np.float) Array of specific vector directions in the x-direction (-1 == left, 1 == right)
        """
        return copy(self.__vector_x)

    @property
    def vector_y(self):
        """
        :return: (np.array, np.float) Array of specific vector directions in the y-direction (-1 == down, 1 == up)
        """
        return copy(self.__vector_y)

    def _create_vector_array(self, img, solid):
        """
        creates an x-dir copy and y-dir copy of the image domain for vector
        directions to populate.
        """
        vimgx = np.copy(img)
        vimgy = np.copy(img.T)
        vimgx[vimgx == solid] = np.nan
        vimgy[vimgy == solid] = np.nan
        return vimgx, vimgy

    def _arrx(self, arr, vector_x, gridres, gridsplit):
        """
        Method that handles looping through array and sends it to distance_gridx
        """
        for line in range(len(arr)):
            arr[line], vector_x[line] = self._distance_gridx(arr[line], vector_x[line], gridres, gridsplit)
        return arr, vector_x

    def _arry(self, yarr, vector_y, gridres, gridsplit):
        """
        Method that handles looping through array and sends it to distance_gridy
        """
        for line in range(len(yarr)):
            ylen = len(yarr[0])
            yarr[line], vector_y[line] = self._distance_gridy(yarr[line], vector_y[line], gridres, ylen, gridsplit)
        return yarr.T, vector_y.T

    def _distance_gridx(self, line, vline, gridres, gridsplit):
        """
        Method uses linear interpolation to correct pore boundary space

        Follows by creating an array of pore boundaries and then counts the distance
        from the nearest solid. It also creates a a vector direction array that corresponds.
        """

        # Saner(?) method which does not raise warnings from Python 2.7.12
        # Forward compatable with Python 3.5.2
        for i in range(len(line)):
            if line[i] > 0.:
                vline[i] = 1
                line[i] = 1
            else:
                vline[i] = 0
                line[i] = 0
        np.array(line)

        # create an array of pore boundaries from binary system
        boundary= np.where(np.abs(np.diff(line)) >= 1)[0]
        try:
            for i in range(1, len(boundary), 2):
                rbound = boundary[i] + 1
                lbound = boundary[i-1]
                gap = rbound - lbound

                if gap % 2 == 0:
                    gap = gap//2

                    if gridres > 1e-9:
                        left = (np.arange(1, gap + 1) * gridres) - (gridres - 1e-9)
                    else:
                        left = np.arange(1, gap + 1) * gridres

                    right = left[::-1]
                    line[lbound:rbound] = np.append(left, right)

                    left = np.ones(gap) * -1
                    right = np.ones(gap)
                    vline[lbound:rbound] = np.append(left, right)

                else:
                    gap = gap//2

                    if gridres > 1e-9:
                        left = (np.arange(1, gap + 2) * gridres) - (gridres - 1e-9)
                        right = (np.arange(1, gap + 1) * gridres) - (gridres - 1e-9)
                    else:
                        left = np.arange(1, gap + 2) * gridres
                        right = np.arange(1, gap + 1) * gridres

                    line[lbound:rbound] = np.append(left, right)

                    left = np.ones(gap + 1) * -1
                    right = np.ones(gap)
                    vline[lbound:rbound] = np.append(left, right)

        except IndexError:
            print('volume does not percolate')
            sys.exit()
        return line, vline

    def _distance_gridy(self, line, vline, gridres, ylen, gridsplit):
        '''
        Method uses linear interpolation to correct pore boundary space

        Follows by creating an array of pore boundaries and then counts the distance
        from the nearest solid. It also creates a a vector direction array that corresponds.
        '''

        for i in range(len(line)):
            if line[i] > 0:
                vline[i] = 1
                line[i] = 1
            else:
                vline[i] = 0
                line[i] = 0
        np.array(line)

        # create an array of pore boundaries from binary system
        boundary = np.where(np.abs(np.diff(line)) >= 1)[0]

        if len(boundary) > 0:
            rbound = boundary[0] + 1
            lbound = 0

            if gridres > 1e-9:
                # Use this statement to enforce nm scale DLVO @ boundary
                top = (np.arange(rbound, lbound, -1) * gridres) - (gridres - 1e-9)
            else:
                top = np.arange(rbound, lbound, -1) * gridres

            line[lbound:rbound] = top
            vtop = np.ones(len(top)) * -1
            vline[lbound:rbound] = vtop

            for i in range(2, len(boundary), 2):
                rbound = boundary[i] + 1
                lbound = boundary[i - 1] + 1
                gap = rbound - lbound
                if gap % 2 == 0:
                    gap = gap // 2

                    if gridres > 1e-9:
                        # use this statement to enforce nm scale DLVO @ boundaries
                        left = (np.arange(1, gap + 1) * gridres) - (gridres - 1e-9)
                    else:
                        left = np.arange(1, gap + 1) * gridres

                    right = left[::-1]
                    line[lbound:rbound] = np.append(left, right)

                    left = np.ones(gap)
                    right = np.ones(gap) * -1
                    vline[lbound:rbound] = np.append(left, right)

                else:
                    gap = gap // 2

                    if gridres > 1e-9:
                        left = (np.arange(1, gap + 2) * gridres) - (gridres - 1e-9)
                        right = (np.arange(1, gap + 1) * gridres) - (gridres - 1e-9)
                    else:
                        left = np.arange(1, gap + 2) * gridres
                        right = np.arange(1, gap + 1) * gridres

                    line[lbound:rbound] = np.append(left, right)

                    left = np.ones(gap + 1)
                    right = np.ones(gap) * -1
                    vline[lbound:rbound] = np.append(left, right)

            rbound = ylen
            lbound = boundary[-1] + 1
            gap = rbound - lbound

            if gridres > 1e-9:
                bottom = (np.arange(1, gap + 1) * gridres) - (gridres - 1e-9)
            else:
                bottom = np.arange(1, gap + 1) * gridres

            line[lbound:rbound] = bottom

            bottom = np.ones(gap)
            vline[lbound:rbound] = bottom
        else:
            pass
        return line, vline

def LBVArray(LBv, img):
    """
    Method to create a velocity array for use in the colloid simulation model

    Parameters:
    ----------
    :param np.ndarray LBv: lattice boltzmann velocity array
    :param np.ndarray img: boolean image array

    Returns:
    -------
    :return: np.ndarray of velocity
    """
    vel = np.zeros((len(LBv), len(LBv[0])))
    invert = np.invert(img.astype(bool))
    vel = np.array([LBv[i] * invert[i] for i in range(len(LBv))])
    print(vel.shape)
    return vel

def InterpV(LBv, gridsplit, img=False):
    """
    Interpolation method for the lattice Boltzmann velocity array

    Parameters:
    ----------
    :param np.ndarray LBv: lattice boltzmann velocity array with pore boundaries enforced
    :param float gridsplit: interpolation factor
    :param bool img: flag to indicate boolean image interpolation or velocity interpolation

    Returns:
    -------
    :return: (np.ndarray) interpolated velocity array
    """
    ylen = len(LBv)
    xlen = len(LBv[0])
    xindex = np.arange(0, xlen)
    yindex = np.arange(0, ylen)
    ifactor = 1. / gridsplit
    xnew = np.arange(0, xlen - 1 + (ifactor), ifactor)
    ynew = np.arange(0, ylen - 1 + (ifactor), ifactor)
    f = interpolate.interp2d(xindex, yindex, LBv, kind='linear')
    znew = f(xnew, ynew)

    if img:
        # correct pore boundaries from interpolation
        znew[znew >= ifactor * gridsplit / 2.] = 1  # 0] = 1# ifactor*gridsplit/2.] = 1
        znew[znew < ifactor * gridsplit / 2.] = 0  # 0] = 0# ifactor*gridsplit/2.] = 0

    return znew


import numpy as np
import time
import os
from atomistic_domains import Atom
from typing import List

class Lattice:
    """
    The class "lattice" defines methods which manipulate lattice structures.
    Lattice structures have the following attributes:
    :param type: Lattice type: e.g. fcc, hcp, bcc, etc.
    :param box: Unit cell box dimensions (including tilt)
    :param basis_atoms: List of atoms (of object class: atom) which make up the
                                                primitive unit cell (in relative coordinates)
    :param chem: List of chemical elements in alloy
    :param num_el: Number of different elements present
    :param num_atoms: number of atoms per each element (dictionary type)
    :param tot_atoms: Total number of atoms in unit cell
    """

    def __init__(
        self,
        x_vector: np.ndarray = np.zeros(shape=(1, 3)),
        y_vector: np.ndarray = np.zeros(shape=(1, 3)),
        z_vector: np.ndarray = np.zeros(shape=(1, 3)),
        basis_atoms: List[Atom] = None,
    ):
        self.x_vector = x_vector
        self.y_vector = y_vector
        self.z_vector = z_vector
        if not basis_atoms:
            basis_atoms = []
        else:
            self.basis_atoms = basis_atoms
        return

    def add_atom(self, new_atom):
        """
        Adds an atom to the list of basis atoms of the lattice. The new atom
        must be of the type: "atom",
        """
        self.basis_atoms.append(new_atom)
        self.tot_atoms += 1
        new = False
        for i in range(self.num_el):
            if new_atom.element == self.chem[i]:
                new = True
        if new == False:
            self.chem.append(new_atom.element)
            self.num_atoms[new_atom.element] = 0
        self.num_atoms[new_atom.element] += 1

        return

    def genfromposcar(self, filename):
        scale = np.genfromtxt(filename, skip_header=1, max_rows=1)
        box = np.genfromtxt(filename, delimiter="", skip_header=2, max_rows=3)
        self.box = box
        self.chem = np.genfromtxt(filename, delimiter="", skip_header=5, max_rows=1, dtype=str)
        self.num_el = len(self.chem)
        atom_numbers = np.genfromtxt(filename, delimiter="", skip_header=6, max_rows=1, dtype=int)
        for i in range(self.num_el):
            self.num_atoms[self.chem[i]] = atom_numbers[i]
        self.tot_atoms = sum(atom_numbers)
        ptype = np.genfromtxt(filename, delimiter=">", skip_header=7, max_rows=1, dtype=str)
        if ptype == "Selective Dynamics":
            eln = 0
            j = 0
            for i in range(self.tot_atoms):
                if j < self.num_atoms[self.chem[eln]]:
                    atomel = self.chem[eln]
                    j += 1
                else:
                    eln += 1
                    atomel = self.chem[eln]
                    j = 1
                acoord = np.genfromtxt(filename, delimiter="", skip_header=9 + i, max_rows=1)[
                    :, 0:3
                ]
                adyn = np.genfromtxt(
                    filename, delimiter="", skip_header=9 + i, max_rows=1, dtype=str
                )[:, 3:]
                newatom = atom(id=i + 1, element=atomel, pos=acoord, dyn=adyn)
                self.basis_atoms.append(newatom)
        else:
            eln = 0
            j = 0
            for i in range(self.tot_atoms):
                if j < self.num_atoms[self.chem[eln]]:
                    atomel = self.chem[eln]
                    j += 1
                else:
                    eln += 1
                    atomel = self.chem[eln]
                    j = 1
                acoord = np.genfromtxt(filename, delimiter="", skip_header=8 + i, max_rows=1)
                newatom = atom(id=i + 1, element=atomel, pos=acoord)
                self.basis_atoms.append(newatom)

        mag_header = int(self.tot_atoms + 9)
        self.mag_mom = np.genfromtxt(filename, delimiter="", skip_header=mag_header)
        if self.mag_mom.shape[0] != self.tot_atoms:
            print(
                "WARNING: Insufficient number of magnetic moments. \
					Magnetism ignored."
            )
            self.mag_mom = np.zeros(shape=(self.tot_atoms, 3))

        for i in range(self.tot_atoms):
            self.basis_atoms[i].mag_mom = self.mag_mom[i]

        self.tot_atoms = len(self.basis_atoms)
        print(self.tot_atoms)
        return

    def repeated_length(self, direction=np.zeros(shape=(1, 3))):
        """
        Calculates the repeatable euclidian distance length of the unit cell in
        the Miller Incex direction specified.
        :param direction: Crystallographic direction of axis of interest
        :return: repeatable length
        """
        lengths = np.zeros(shape=(3, 1))

        for i in range(3):
            if direction[i] == 0:
                lengths[i] = 1e6
            else:
                t = 0.5 * self.box[i, i] / float(direction[i])
                lengths[i] = 2 * sqrt(
                    (t * direction[0]) ** 2 + (t * direction[1]) ** 2 + (t * direction[2]) ** 2
                )
        rl = min(lengths)
        return rl

    def hexCij(c11, c12, c13, c33, c44):
        """
        Populates 6x6 elastic tensor for hexagonal crystals
        :param c11: c11 elastic constant
        :param c12: c12 elastic constant
        :param c13: c13 elastic constant
        :param c44: c44 elastic constant
        :return C: 6x6 elastic tensor for Voigt notation calculations
        :return S: 6x6 compliance tensor for Voigt notation calculations
        """
        C = np.zeros(shape=(6, 6))
        C[0, 0] = C[1, 1] = c11
        C[0, 1] = C[1, 0] = c12
        C[0, 2] = C[2, 0] = C[1, 2] = C[2, 1] = c13
        C[2, 2] = c33
        C[3, 3] = C[4, 4] = c44
        C[5, 5] = 0.5 * (c11 - c12)
        S = np.linalg.inv(C)
        return {"C": C, "S": S}


'''
REFER TO: https://github.com/ej0cl6/pytorch-adversarial-examples/blob/master/attackers.py
'''

# Copyright (c) 2018-present, Royal Bank of Canada.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
#

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import torch
import torch.nn as nn
import torch.optim as optim

from advertorch.utils import calc_l2distsq
from advertorch.utils import tanh_rescale
from advertorch.utils import torch_arctanh
from advertorch.utils import clamp
from advertorch.utils import to_one_hot
from advertorch.utils import replicate_input

from advertorch.attacks.base import Attack
from advertorch.attacks.base import LabelMixin
from advertorch.attacks.utils import is_successful

import numpy as np
import torch
import torch.nn.functional as F

class Attacker:
    def __init__(self, clip_max=0.5, clip_min=-0.5):
        self.clip_max = clip_max
        self.clip_min = clip_min

    def generate(self, model, x, y):
        pass

class FGSM(Attacker):
    """
    Fast Gradient Sign Method
    Ian J. Goodfellow, Jonathon Shlens, Christian Szegedy.
    Explaining and Harnessing Adversarial Examples.
    ICLR, 2015
    """

    def __init__(self, eps=0.15, clip_max=0.5, clip_min=-0.5):
        super(FGSM, self).__init__(clip_max, clip_min)
        self.eps = eps

    def generate(self, model, x, y):
        model.eval()
        nx = torch.unsqueeze(x, 0)
        ny = torch.unsqueeze(y, 0)
        nx.requires_grad_()
        out = model(nx)
        loss = F.cross_entropy(out, ny)
        loss.backward()
        x_adv = nx + self.eps * torch.sign(nx.grad.data)
        x_adv.clamp_(self.clip_min, self.clip_max)
        x_adv.squeeze_(0)

        return x_adv.detach()

class BIM(Attacker):
    """
    Basic Iterative Method
    Alexey Kurakin, Ian J. Goodfellow, Samy Bengio.
    Adversarial Examples in the Physical World.
    arXiv, 2016
    """

    def __init__(self, eps=0.15, eps_iter=0.01, n_iter=50, clip_max=0.5, clip_min=-0.5):
        super(BIM, self).__init__(clip_max, clip_min)
        self.eps = eps
        self.eps_iter = eps_iter
        self.n_iter = n_iter

    def generate(self, model, x, y):
        model.eval()
        nx = torch.unsqueeze(x, 0)
        ny = torch.unsqueeze(y, 0)
        nx.requires_grad_()
        eta = torch.zeros(nx.shape)

        for i in range(self.n_iter):
            out = model(nx + eta)
            loss = F.cross_entropy(out, ny)
            loss.backward()

            eta += self.eps_iter * torch.sign(nx.grad.data)
            eta.clamp_(-self.eps, self.eps)
            nx.grad.data.zero_()

        x_adv = nx + eta
        x_adv.clamp_(self.clip_min, self.clip_max)
        x_adv.squeeze_(0)

        return x_adv.detach()

class DeepFool(Attacker):
    """
    DeepFool
    Seyed-Mohsen Moosavi-Dezfooli, Alhussein Fawzi, Pascal Frossard
    DeepFool: A Simple and Accurate Method to Fool Deep Neural Networks.
    CVPR, 2016
    """

    def __init__(self, model,max_iter=50, clip_max=0.5, clip_min=-0.5,epsilon=0.00784):
        super(DeepFool, self).__init__(clip_max, clip_min)
        self.max_iter = max_iter
        self.model = model
        self.epsilon = epsilon

    def perturb(self,x,y):
        x = x.detach().clone()
        adv = torch.zeros_like(x)
        for i in range (x.size()[0]):
            adv[i] = self.generate(x[i],y[i])
        return adv.detach()

    def generate(self, x, y):
        model = self.model
        model.eval()
        nx = torch.unsqueeze(x.detach().clone(),0)
        nx.requires_grad_()
        eta = torch.zeros_like(nx) # instead of torch.zeors(nx.shape)

        out = model(nx + eta)
        n_class = out.shape[1]
        py = out.max(1)[1].item()
        ny = out.max(1)[1].item()

        i_iter = 0

        while py == ny and i_iter < self.max_iter  :
            out[0, py].backward(retain_graph=True)
            grad_np = nx.grad.data.clone()
            value_l = np.inf
            ri = None

            for i in range(n_class):
                if i == py:
                    continue

                nx.grad.data.zero_()
                out[0, i].backward(retain_graph=True)
                grad_i = nx.grad.data.clone()

                wi = grad_i - grad_np
                fi = out[0, i] - out[0, py]
                value_i = np.abs(fi.item()) / np.linalg.norm(wi.cpu().numpy().flatten())

                if value_i < value_l:
                    ri = value_i / np.linalg.norm(wi.cpu().numpy().flatten()) * wi

            eta += ri.clone()
            eta.clamp_(-self.epsilon,self.epsilon)

            nx.grad.data.zero_()
            out = model(nx + eta)
            py = out.max(1)[1].item()
            i_iter += 1

        x_adv = nx + eta

        x_adv.clamp_(self.clip_min, self.clip_max)
        x_adv.squeeze_(0)

        return x_adv.detach()

CARLINI_L2DIST_UPPER = 1e10
CARLINI_COEFF_UPPER = 1e10
INVALID_LABEL = -1
REPEAT_STEP = 10
ONE_MINUS_EPS = 0.999999
UPPER_CHECK = 1e9
PREV_LOSS_INIT = 1e6
TARGET_MULT = 10000.0
NUM_CHECKS = 10

class CarliniWagnerL2Attack(Attack, LabelMixin):
    """
    The Carlini and Wagner L2 Attack, https://arxiv.org/abs/1608.04644
    :param predict: forward pass function.
    :param num_classes: number of clasess.
    :param confidence: confidence of the adversarial examples.
    :param targeted: if the attack is targeted.
    :param learning_rate: the learning rate for the attack algorithm
    :param binary_search_steps: number of binary search times to find the
        optimum
    :param max_iterations: the maximum number of iterations
    :param abort_early: if set to true, abort early if getting stuck in local
        min
    :param initial_const: initial value of the constant c
    :param clip_min: mininum value per input dimension.
    :param clip_max: maximum value per input dimension.
    :param loss_fn: loss function
    """

    def __init__(self, predict, num_classes, epsilon=0.00784,confidence=0,
                 targeted=False, learning_rate=0.01,
                 binary_search_steps=9, max_iterations=10000,
                 abort_early=True, initial_const=1e-3,
                 clip_min=0., clip_max=1., loss_fn=None):
        """Carlini Wagner L2 Attack implementation in pytorch."""
        if loss_fn is not None:
            import warnings
            warnings.warn(
                "This Attack currently do not support a different loss"
                " function other than the default. Setting loss_fn manually"
                " is not effective."
            )

        loss_fn = None

        super(CarliniWagnerL2Attack, self).__init__(
            predict, loss_fn, clip_min, clip_max)

        self.learning_rate = learning_rate
        self.max_iterations = max_iterations
        self.binary_search_steps = binary_search_steps
        self.abort_early = abort_early
        self.confidence = confidence
        self.initial_const = initial_const
        self.num_classes = num_classes
        # The last iteration (if we run many steps) repeat the search once.
        self.repeat = binary_search_steps >= REPEAT_STEP
        self.targeted = targeted
        self.epsilon = epsilon

    def _loss_fn(self, output, y_onehot, l2distsq, const):
        # TODO: move this out of the class and make this the default loss_fn
        #   after having targeted tests implemented
        real = (y_onehot * output).sum(dim=1)

        # TODO: make loss modular, write a loss class
        other = ((1.0 - y_onehot) * output - (y_onehot * TARGET_MULT)
                 ).max(1)[0]
        # - (y_onehot * TARGET_MULT) is for the true label not to be selected

        if self.targeted:
            loss1 = clamp(other - real + self.confidence, min=0.)
        else:
            loss1 = clamp(real - other + self.confidence, min=0.)
        loss2 = (l2distsq).sum()
        loss1 = torch.sum(const * loss1)
        loss = loss1 + loss2
        return loss

    def _is_successful(self, output, label, is_logits):
        # determine success, see if confidence-adjusted logits give the right
        #   label

        if is_logits:
            output = output.detach().clone()
            if self.targeted:
                output[torch.arange(len(label)).long(),
                       label] -= self.confidence
            else:
                output[torch.arange(len(label)).long(),
                       label] += self.confidence
            pred = torch.argmax(output, dim=1)
        else:
            pred = output
            if pred == INVALID_LABEL:
                return pred.new_zeros(pred.shape).byte()

        return is_successful(pred, label, self.targeted)

    def _forward_and_update_delta(
            self, optimizer, x_atanh, delta, y_onehot, loss_coeffs):

        optimizer.zero_grad()
        adv = tanh_rescale(delta + x_atanh, self.clip_min, self.clip_max)
        transimgs_rescale = tanh_rescale(x_atanh, self.clip_min, self.clip_max)
        output = self.predict(adv)
        l2distsq = calc_l2distsq(adv, transimgs_rescale)
        loss = self._loss_fn(output, y_onehot, l2distsq, loss_coeffs)
        loss.backward()
        optimizer.step()

        return loss.item(), l2distsq.data, output.data, adv.data

    def _get_arctanh_x(self, x):
        result = clamp((x - self.clip_min) / (self.clip_max - self.clip_min),
                       min=0., max=1.) * 2 - 1
        return torch_arctanh(result * ONE_MINUS_EPS)

    def _update_if_smaller_dist_succeed(
            self, adv_img, labs, output, l2distsq, batch_size,
            cur_l2distsqs, cur_labels,
            final_l2distsqs, final_labels, final_advs):

        target_label = labs
        output_logits = output
        _, output_label = torch.max(output_logits, 1)

        mask = (l2distsq < cur_l2distsqs) & self._is_successful(
            output_logits, target_label, True)

        cur_l2distsqs[mask] = l2distsq[mask]  # redundant
        cur_labels[mask] = output_label[mask]

        mask = (l2distsq < final_l2distsqs) & self._is_successful(
            output_logits, target_label, True)
        final_l2distsqs[mask] = l2distsq[mask]
        final_labels[mask] = output_label[mask]
        final_advs[mask] = adv_img[mask]

    def _update_loss_coeffs(
            self, labs, cur_labels, batch_size, loss_coeffs,
            coeff_upper_bound, coeff_lower_bound):

        # TODO: remove for loop, not significant, since only called during each
        # binary search step
        for ii in range(batch_size):
            cur_labels[ii] = int(cur_labels[ii])
            if self._is_successful(cur_labels[ii], labs[ii], False):
                coeff_upper_bound[ii] = min(
                    coeff_upper_bound[ii], loss_coeffs[ii])

                if coeff_upper_bound[ii] < UPPER_CHECK:
                    loss_coeffs[ii] = (
                        coeff_lower_bound[ii] + coeff_upper_bound[ii]) / 2
            else:
                coeff_lower_bound[ii] = max(
                    coeff_lower_bound[ii], loss_coeffs[ii])
                if coeff_upper_bound[ii] < UPPER_CHECK:
                    loss_coeffs[ii] = (
                        coeff_lower_bound[ii] + coeff_upper_bound[ii]) / 2
                else:
                    loss_coeffs[ii] *= 10

    def perturb(self, x, y=None):
        x, y = self._verify_and_process_inputs(x, y)

        # Initialization
        if y is None:
            y = self._get_predicted_label(x)
        x = replicate_input(x)
        batch_size = len(x)
        coeff_lower_bound = x.new_zeros(batch_size)
        coeff_upper_bound = x.new_ones(batch_size) * CARLINI_COEFF_UPPER
        loss_coeffs = torch.ones_like(y).float() * self.initial_const
        final_l2distsqs = [CARLINI_L2DIST_UPPER] * batch_size
        final_labels = [INVALID_LABEL] * batch_size
        final_advs = x
        x_atanh = self._get_arctanh_x(x)
        y_onehot = to_one_hot(y, self.num_classes).float()

        final_l2distsqs = torch.FloatTensor(final_l2distsqs).to(x.device)
        final_labels = torch.LongTensor(final_labels).to(x.device)

        # Start binary search
        for outer_step in range(self.binary_search_steps):
            delta = nn.Parameter(torch.zeros_like(x))
            optimizer = optim.Adam([delta], lr=self.learning_rate)
            cur_l2distsqs = [CARLINI_L2DIST_UPPER] * batch_size
            cur_labels = [INVALID_LABEL] * batch_size
            cur_l2distsqs = torch.FloatTensor(cur_l2distsqs).to(x.device)
            cur_labels = torch.LongTensor(cur_labels).to(x.device)
            prevloss = PREV_LOSS_INIT

            if (self.repeat and outer_step == (self.binary_search_steps - 1)):
                loss_coeffs = coeff_upper_bound
            for ii in range(self.max_iterations):
                loss, l2distsq, output, adv_img = \
                    self._forward_and_update_delta(
                        optimizer, x_atanh, delta, y_onehot, loss_coeffs)
                if self.abort_early:
                    if ii % (self.max_iterations // NUM_CHECKS or 1) == 0:
                        if loss > prevloss * ONE_MINUS_EPS:
                            break
                        prevloss = loss

                self._update_if_smaller_dist_succeed(
                    adv_img, y, output, l2distsq, batch_size,
                    cur_l2distsqs, cur_labels,
                    final_l2distsqs, final_labels, final_advs)

            self._update_loss_coeffs(
                y, cur_labels, batch_size,
                loss_coeffs, coeff_upper_bound, coeff_lower_bound)

        #calculate perturbation
        eta = final_advs - x
        eta.clamp_(-self.epsilon,self.epsilon)
        final_advs = x + eta

        return final_advs