import numpy as np
import torch
import librosa

class DataLoader(torch.utils.data.Dataset):
    '''
    スペクトログラムに変換済みの音声データセットを読み込む。
    PyTorchでの学習用。
    可変長データ、ストレッチ済みデータ、ゼロパディングに対応。

    Parameters
    ----------
    person_list : list of int
        ファイル名のperson番号-1の番号リスト。
    voice_list : list of int
        ファイル名のvoice番号-1の番号リスト。
    batch_size : tuple of int
        (person_length, voice_length)のタプル。
        ただしdeform_type=='stretch'の場合は(1, 1)。
    nphonemes_path : str
        音素長ファイルのパス。
    dataset_path : str
        データセットのパス。
        例："%(deform_type)/%(person)s/%(voice).wav"
    deform_type : str
        変形の種類
    phonemes_length : int
        deform_type=='padding'の場合の整形後のサイズ。
    '''

    def __init__(self, person_list, voice_list, batch_size, dataset_path, deform_type, phonemes_length=None, mel=False, seed=None):
        self.person_list = person_list
        self.voice_list  = voice_list
        self.batch_size  = batch_size

        self.dataset_path    = dataset_path
        self.deform_type     = deform_type
        self.phonemes_length = phonemes_length
        self.mel_basis       = librosa.filters.mel(sr=24000, n_fft=1024) if mel else None

        self.person_nbatches  = len(person_list) // batch_size[0]
        self.phoneme_nbatches = len(voice_list) * 32 // batch_size[1]

        if seed is not None:
            self.reset_shuffle(seed)
        else:
            self.shuffle = np.arange(len(self))

    def __len__(self):
        return self.person_nbatches * self.phoneme_nbatches

    def reset_shuffle(self, seed=0):
        np.random.seed(seed)
        self.shuffle = np.random.permutation(len(self))

    def __getitem__(self, batch_idx):
        '''
        Returns
        -------
        data : [batch_length, phoenemes_length, nfft // 2]
        true : [batch_length]
        '''

        if batch_idx < 0 or batch_idx >= len(self):
            raise IndexError()

        batch_idx = self.shuffle[batch_idx]

        # インデックスをperson方向とphoneme方向に分解
        phoneme_batch_idx, person_batch_idx = divmod(batch_idx, self.person_nbatches)

        # 読み込むべきファイルのpersonの範囲を求める
        person_start_idx = person_batch_idx * self.batch_size[0]
        person_end_idx   = (person_batch_idx + 1) * self.batch_size[0]

        # 読み込むべきファイルのvoiceの範囲を求める
        voice_start_idx = phoneme_batch_idx * self.batch_size[1] // 32
        voice_end_idx   = ((phoneme_batch_idx + 1) * self.batch_size[1] - 1) // 32 + 1

        # 複数のvoiceを結合したデータから取り出すべき範囲を求める
        voice_start_phoneme = phoneme_batch_idx % 32 * self.batch_size[1]
        voice_end_phoneme   = voice_start_phoneme + self.batch_size[1]

        data = list()
        for person_idx in range(person_start_idx, person_end_idx):

            person_data = list()
            for voice_idx in range(voice_start_idx, voice_end_idx):
                specific = {
                    'person'     : self.person_list[person_idx] + 1,
                    'voice'      : self.voice_list[voice_idx] + 1,
                    'deform_type': self.deform_type,
                }
                pack = np.load(self.dataset_path % specific, allow_pickle=True)

                if self.deform_type == 'stretch':
                    sp = pack['sp']
                elif self.deform_type == 'padding':
                    sp = np.array([self._zero_padding(x[:self.phonemes_length], self.phonemes_length) for x in pack['sp']])
                elif self.deform_type == 'variable':
                    sp = np.array([x for x in pack['sp']])

                person_data.extend(sp)
            data.extend(person_data[voice_start_phoneme:voice_end_phoneme])
        data = np.array(data)

        if self.mel_basis is not None:
            data = np.dot(data, self.mel_basis.T)
        else:
            data = data[:, :, 1:]

        label  = np.concatenate([[person_idx] * self.batch_size[1] for person_idx in range(person_start_idx, person_end_idx)])

        data  = torch.from_numpy(data).float().to('cuda')
        label = torch.from_numpy(label).long().to('cuda')
        return data, label

    @staticmethod
    def _zero_padding(x, target_length):
        y_pad = target_length - len(x)
        return np.pad(x, ((0, y_pad), (0, 0)), mode='constant') if y_pad > 0 else x

def test():
    nphonemes_path = 'resource/jvs_ver1_nphonemes_%(condition)s.txt'
    dataset_path   = 'resource/jvs_ver1_phonemes/jvs%(person)03d/VOICEACTRESS100_%(voice)03d_%(deform_type)s.npz'

    complete_list  = np.arange(100)
    person_no_list = [8, 16, 17, 21, 23, 29, 35, 37, 42, 46, 47, 50, 54, 58, 59, 73, 88, 97]
    person_ok_list = list(filter(lambda x:x not in person_no_list, complete_list))
    voice_no_list  = [5, 18, 24, 40, 42, 44, 46, 55, 56, 59, 60, 61, 63, 65, 71, 73, 75, 81, 84, 85, 87, 93, 94, 98]
    voice_ok_list  = list(filter(lambda x:x not in voice_no_list, complete_list))

    loader = DataLoader(complete_list, voice_ok_list, (2, 32), nphonemes_path, dataset_path, 'padding', 32)
    for data in loader:
        print(data[0].shape, data[1].shape)


# python3
# Copyright 2021 InstaDeep Ltd. 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.

import collections
import math
import os
import queue
import time
from typing import Any

import numpy as np

try:
    os.environ["PYGAME_HIDE_SUPPORT_PROMPT"] = "hide"
    import pygame
    from pysc2.lib import colors, features, point, transform
    from pysc2.lib.renderer_human import _Surface
    from smac.env import StarCraft2Env
except ModuleNotFoundError:
    pass

def clamp(n: float, smallest: int, largest: int) -> float:
    return max(smallest, min(n, largest))

class Renderer:
    """
    This class helps render sc2 environments.
    Adapted from https://github.com/oxwhirl/smac/pull/71.
    """

    def __init__(self, env: "StarCraft2Env", mode: str):
        os.environ["PYGAME_HIDE_SUPPORT_PROMPT"] = "hide"

        self.env = env
        self.mode = mode
        self.obs = None
        self._window_scale = 0.75
        self.game_info = game_info = self.env._controller.game_info()
        self.static_data = self.env._controller.data()

        self._obs_queue: queue.Queue = queue.Queue()
        self._game_times: collections.deque = collections.deque(
            maxlen=100
        )  # Avg FPS over 100 frames.  # pytype: disable=wrong-keyword-args
        self._render_times: collections.deque = collections.deque(
            maxlen=100
        )  # pytype: disable=wrong-keyword-args
        self._last_time = time.time()
        self._last_game_loop = 0
        self._name_lengths: dict = {}

        self._map_size = point.Point.build(game_info.start_raw.map_size)
        self._playable = point.Rect(
            point.Point.build(game_info.start_raw.playable_area.p0),
            point.Point.build(game_info.start_raw.playable_area.p1),
        )

        window_size_px = point.Point(self.env.window_size[0], self.env.window_size[1])
        window_size_px = self._map_size.scale_max_size(
            window_size_px * self._window_scale
        ).ceil()
        self._scale = window_size_px.y // 32

        self.display = pygame.Surface(window_size_px)

        if mode == "human":
            self.display = pygame.display.set_mode(window_size_px, 0, 32)  # type: ignore  # noqa: E501
            pygame.display.init()

            pygame.display.set_caption("Starcraft Viewer")
        pygame.font.init()
        self._world_to_world_tl = transform.Linear(
            point.Point(1, -1), point.Point(0, self._map_size.y)
        )
        self._world_tl_to_screen = transform.Linear(scale=window_size_px / 32)
        self.screen_transform = transform.Chain(
            self._world_to_world_tl, self._world_tl_to_screen
        )

        surf_loc = point.Rect(point.origin, window_size_px)
        sub_surf = self.display.subsurface(pygame.Rect(surf_loc.tl, surf_loc.size))
        self._surf = _Surface(
            sub_surf, None, surf_loc, self.screen_transform, None, self.draw_screen
        )

        self._font_small = pygame.font.Font(None, int(self._scale * 0.5))
        self._font_large = pygame.font.Font(None, self._scale)

        self.upgrade_colors = [
            colors.black,  # unused...
            colors.white * 0.6,
            colors.white * 0.8,
            colors.white,
        ]

    def close(self) -> None:
        pygame.display.quit()
        pygame.quit()

    def _get_units(self) -> Any:
        if self.obs:
            for u in sorted(
                self.obs.observation.raw_data.units,
                key=lambda u: (u.pos.z, u.owner != 16, -u.radius, u.tag),
            ):
                yield u, point.Point.build(u.pos)

    def get_unit_name(self, surf: Any, name: str, radius: float) -> dict:
        """Get a length limited unit name for drawing units."""
        key = (name, radius)
        if key not in self._name_lengths:
            max_len = surf.world_to_surf.fwd_dist(radius * 1.6)
            for i in range(len(name)):
                if self._font_small.size(name[: i + 1])[0] > max_len:
                    self._name_lengths[key] = name[:i]
                    break
            else:
                self._name_lengths[key] = name
        return self._name_lengths[key]

    def render(self, mode: str) -> np.array:
        self.obs = self.env._obs
        self.score = self.env.reward
        self.step = self.env._episode_steps

        now = time.time()
        if self.obs:
            self._game_times.append(
                (
                    now - self._last_time,
                    max(
                        1,
                        self.obs.observation.game_loop - self.obs.observation.game_loop,
                    ),
                )
            )

        if mode == "human":
            pygame.event.pump()

        self._surf.draw(self._surf)

        observation = np.array(pygame.surfarray.pixels3d(self.display))

        if mode == "human":
            pygame.display.flip()

        self._last_time = now
        if self.obs:
            self._last_game_loop = self.obs.observation.game_loop
            # self._obs_queue.put(self.obs)
        return (
            np.transpose(observation, axes=(1, 0, 2)) if mode == "rgb_array" else None
        )

    def draw_base_map(self, surf: Any) -> None:
        """Draw the base map."""
        hmap_feature = features.SCREEN_FEATURES.height_map
        hmap = self.env.terrain_height * 255
        hmap = hmap.astype(np.uint8)
        if (
            self.env.map_name == "corridor"
            or self.env.map_name == "so_many_baneling"
            or self.env.map_name == "2s_vs_1sc"
        ):
            hmap = np.flip(hmap)
        else:
            hmap = np.rot90(hmap, axes=(1, 0))
        if not hmap.any():
            hmap = hmap + 100  # pylint: disable=g-no-augmented-assignment
        hmap_color = hmap_feature.color(hmap)
        out = hmap_color * 0.6

        surf.blit_np_array(out)

    def draw_units(self, surf: Any) -> None:  # noqa
        """Draw the units."""
        unit_dict = None  # Cache the units {tag: unit_proto} for orders.
        tau = 2 * math.pi
        for u, p in self._get_units():
            fraction_damage = clamp(
                (u.health_max - u.health) / (u.health_max or 1), 0, 1
            )
            surf.draw_circle(colors.PLAYER_ABSOLUTE_PALETTE[u.owner], p, u.radius)

            if fraction_damage > 0:
                surf.draw_circle(
                    colors.PLAYER_ABSOLUTE_PALETTE[u.owner] // 2,
                    p,
                    u.radius * fraction_damage,
                )
            surf.draw_circle(colors.black, p, u.radius, thickness=1)

            if self.static_data.unit_stats[u.unit_type].movement_speed > 0:
                surf.draw_arc(
                    colors.white,
                    p,
                    u.radius,
                    u.facing - 0.1,
                    u.facing + 0.1,
                    thickness=1,
                )

            def draw_arc_ratio(
                color: Any,
                world_loc: Any,
                radius: Any,
                start: Any,
                end: Any,
                thickness: Any = 1,
            ) -> None:
                surf.draw_arc(
                    color, world_loc, radius, start * tau, end * tau, thickness
                )

            if u.shield and u.shield_max:
                draw_arc_ratio(
                    colors.blue, p, u.radius - 0.05, 0, u.shield / u.shield_max
                )

            if u.energy and u.energy_max:
                draw_arc_ratio(
                    colors.purple * 0.9, p, u.radius - 0.1, 0, u.energy / u.energy_max
                )
            elif u.orders and 0 < u.orders[0].progress < 1:
                draw_arc_ratio(colors.cyan, p, u.radius - 0.15, 0, u.orders[0].progress)
            if u.buff_duration_remain and u.buff_duration_max:
                draw_arc_ratio(
                    colors.white,
                    p,
                    u.radius - 0.2,
                    0,
                    u.buff_duration_remain / u.buff_duration_max,
                )
            if u.attack_upgrade_level:
                draw_arc_ratio(
                    self.upgrade_colors[u.attack_upgrade_level],
                    p,
                    u.radius - 0.25,
                    0.18,
                    0.22,
                    thickness=3,
                )
            if u.armor_upgrade_level:
                draw_arc_ratio(
                    self.upgrade_colors[u.armor_upgrade_level],
                    p,
                    u.radius - 0.25,
                    0.23,
                    0.27,
                    thickness=3,
                )
            if u.shield_upgrade_level:
                draw_arc_ratio(
                    self.upgrade_colors[u.shield_upgrade_level],
                    p,
                    u.radius - 0.25,
                    0.28,
                    0.32,
                    thickness=3,
                )

            def write_small(loc: Any, s: Any) -> None:
                surf.write_world(self._font_small, colors.white, loc, str(s))

            name = self.get_unit_name(
                surf, self.static_data.units.get(u.unit_type, "<none>"), u.radius
            )

            if name:
                write_small(p, name)

            if self.obs:
                start_point = p
                for o in u.orders:
                    target_point = None
                    if o.HasField("target_unit_tag"):
                        if unit_dict is None:
                            unit_dict = {
                                t.tag: t for t in self.obs.observation.raw_data.units
                            }
                        target_unit = unit_dict.get(o.target_unit_tag)
                        if target_unit:
                            target_point = point.Point.build(target_unit.pos)
                    if target_point:
                        surf.draw_line(colors.cyan, start_point, target_point)
                        start_point = target_point
                    else:
                        break

    def draw_overlay(self, surf: Any) -> None:
        """Draw the overlay describing resources."""
        if self.obs:
            obs = self.obs.observation
            times, steps = zip(*self._game_times)
            sec = obs.game_loop // 22.4
            surf.write_screen(
                self._font_large,
                colors.green,
                (-0.2, 0.2),
                "Score: %s, Step: %s, %.1f/s, Time: %d:%02d"
                % (
                    self.score,
                    self.step,
                    sum(steps) / (sum(times) or 1),
                    sec // 60,
                    sec % 60,
                ),
                align="right",
            )
            surf.write_screen(
                self._font_large,
                colors.green * 0.8,
                (-0.2, 1.2),
                "APM: %d, EPM: %d, FPS: O:%.1f, R:%.1f"
                % (
                    obs.score.score_details.current_apm,
                    obs.score.score_details.current_effective_apm,
                    len(times) / (sum(times) or 1),
                    len(self._render_times) / (sum(self._render_times) or 1),
                ),
                align="right",
            )

    def draw_screen(self, surf: Any) -> None:
        """Draw the screen area."""
        self.draw_base_map(surf)
        self.draw_units(surf)
        self.draw_overlay(surf)


# @Shitong Zhu, SRA 2019
# This Foolbox module has been modified to reflect our "augmented"
# projection algorithm
# Please replace the "iterative_projected_gradient.py" file under
# foolbox Python library directory

from __future__ import division
import numpy as np
from abc import abstractmethod
import logging
import warnings
import os
from bs4 import BeautifulSoup

# for computing distance after each projection
from scipy.spatial import distance

from .base import Attack
from .base import call_decorator
from .. import distances
from ..utils import crossentropy
from .. import nprng

from ..distances import Distance
from ..distances import MSE

from ..map_back import compute_x_after_mapping_back

from ..perturb_html import featureMapbacks

from ..classify import predict
from ..classify import setup_clf, setup_one_hot
from ..classify import read_one_hot_feature_list

import matplotlib.pyplot as plt
import math

import json
from urllib.parse import urlsplit
from urllib.parse import urlparse

LABLE = {"AD": 1, "NONAD": 0}
NORM_MAP = {
    301: 59,
    302: 60,
    303: 61,
    304: 62,
    305: 63,
    306: 64,
    307: 65,
    309: 67,
    310: 68,
    311: 69,
}
HOME_DIR = os.getenv("HOME")

class IterativeProjectedGradientBaseAttack(Attack):
    """Base class for iterative (projected) gradient attacks.

    Concrete subclasses should implement __call__, _gradient
    and _clip_perturbation.

    TODO: add support for other loss-functions, e.g. the CW loss function,
    see https://github.com/MadryLab/mnist_challenge/blob/master/pgd_attack.py
    """

    def __init__(self, *args, **kwargs):
        def _reduce_url_to_domain(url):
            html_filename = url.split('/')[-1]
            html_filename = html_filename.split('_')[-1]
            html_filename = html_filename.strip('.source')
            html_filename = html_filename.strip('.html')
            return html_filename

        super(IterativeProjectedGradientBaseAttack, self).__init__(*args, **kwargs)

        self.BASE_CRAWLED_DIR = HOME_DIR + "/rendering_stream/html"
        self.BASE_HTML_DIR = HOME_DIR + "/rendering_stream/html"
        self.BASE_EVAL_HTML_DIR = HOME_DIR + "/rendering_stream/eval_html"
        self.all_html_filepaths = os.listdir(self.BASE_CRAWLED_DIR)
        self.BASE_MAPPING_DIR = HOME_DIR + "/rendering_stream/mappings"
        self.BASE_TIMELINE_DIR = HOME_DIR + "/rendering_stream/timeline"
        self.BASE_DATA_DIR = HOME_DIR + "/attack-adgraph-pipeline/data"
        self.BASE_MODEL_DIR = HOME_DIR + "/attack-adgraph-pipeline/model"
        self.original_dataset_fpath = self.BASE_DATA_DIR + '/dataset_1203.csv'
        self.final_domain_to_original_domain_mapping_fpath = HOME_DIR + "/map_local_list_unmod_new.csv"

        one_hot_feature_list = read_one_hot_feature_list(self.original_dataset_fpath)
        events = sorted(list(one_hot_feature_list["FEATURE_NODE_CATEGORY"]))
        tag_1 = sorted(list(one_hot_feature_list["FEATURE_FIRST_PARENT_TAG_NAME"]))
        tag_2 = sorted(
            list(one_hot_feature_list["FEATURE_FIRST_PARENT_SIBLING_TAG_NAME"]))
        tag_3 = sorted(list(one_hot_feature_list["FEATURE_SECOND_PARENT_TAG_NAME"]))
        tag_4 = sorted(
            list(one_hot_feature_list["FEATURE_SECOND_PARENT_SIBLING_TAG_NAME"]))
        setup_one_hot(events, tag_1, tag_2, tag_3, tag_4)

        self.final_domain_to_original_domain_mapping = {}
        self.original_domain_to_final_domain_mapping = {}

        with open(self.final_domain_to_original_domain_mapping_fpath, 'r') as fin:
            data = fin.readlines()
        for row in data:
            row = row.strip()
            original_domain, final_url = row.split(',', 1)
            final_domain = urlparse(final_url)[1]
            self.final_domain_to_original_domain_mapping[final_domain] = original_domain
            self.original_domain_to_final_domain_mapping[original_domain] = final_domain

    @abstractmethod
    def _gradient(self, a, x, class_, strict=True):
        raise NotImplementedError

    @abstractmethod
    def _clip_perturbation(self, a, noise, epsilon):
        raise NotImplementedError

    @abstractmethod
    def _check_distance(self, a):
        raise NotImplementedError

    def _compute_l2_distance(self, x1, x2):
        return distance.cdist(x1, x2, 'euclidean')

    def _compute_lp_distance(self, x1, x2):
        diff = np.array(x1) - np.array(x2)
        value = np.max(np.abs(diff)).astype(np.float64)
        return value

    def _generate_constrained_perturbation(self, perturbation,
                                           perturbable_idx_set,
                                           only_increase_idx_set,
                                           debug=False):

        for i in range(len(perturbation)):
            if i not in perturbable_idx_set:
                perturbation[i] = 0.0
            if i in only_increase_idx_set and perturbation[i] < 0.0:
                perturbation[i] = -perturbation[i]
        return perturbation

    def _unscale_feature(self, val, stats, is_float=False):
        [maxn, minn] = stats
        maxn, minn = float(maxn), float(minn)
        if is_float:
            return val * (maxn + minn) + minn
        else:
            return int(round(val * (maxn + minn) + minn))

    def _rescale_feature(self, val, stats):
        [maxn, minn] = stats
        maxn, minn = float(maxn), float(minn)
        if maxn == minn:
            return val
        return (val - minn) / (maxn - minn)

    def _calculate_diff(self, ori, per, stats):
        return self._unscale_feature(per, stats) - self._unscale_feature(ori, stats)

    def _reject_imperturbable_features(self, candidate,
                                       original, perturbable_idx_set,
                                       debug=False):

        assert len(candidate) == len(original), "[ERROR] Lengths of two input arrays not equal!"

        rejected_candidate = []

        for i in range(len(candidate)):
            if i in perturbable_idx_set:
                rejected_candidate.append(candidate[i])
            else:
                rejected_candidate.append(original[i])

        return np.array(rejected_candidate)

    def _reject_only_increase_features(self, candidate,
                                       original, only_increase_idx_set,
                                       debug=False):

        assert len(candidate) == len(original), "[ERROR] Lengths of two input arrays not equal!"

        rejected_candidate = []

        for i in range(len(candidate)):
            if i in only_increase_idx_set and candidate[i] < original[i]:
                rejected_candidate.append(original[i])
            else:
                rejected_candidate.append(candidate[i])

        return np.array(rejected_candidate)

    def _legalize_candidate(self, candidate,
                            original,
                            feature_types,
                            debug=False):
        if debug:
            print("[INFO] Entered candidate legalization method!")
            print(feature_types)
            input("Press Enter to continue...")
        adv_x = np.copy(candidate)
        processed_adv_x = np.copy(candidate)

        for i in range(len(adv_x)):
            adv_val = adv_x[i]
            ori_val = original[i]
            processed_adv_val = None
            if feature_types[i] == 'b':
                if abs(adv_val - 1.0) > abs(adv_val - 0.0):
                    processed_adv_val = 1
                else:
                    processed_adv_val = 0

            if processed_adv_val is not None:
                processed_adv_x[i] = processed_adv_val

        lookahead_cnt = 0
        for i in range(len(adv_x)):
            if lookahead_cnt - 1 > 0:
                lookahead_cnt -= 1
                continue
            if feature_types[i] == 'c':
                j = i
                while feature_types[j] == 'c' and j + 1 < len(adv_x):
                    j += 1
                categorical_interval_end = j
                maxn = -10000
                maxn_idx = i
                for j in range(i, categorical_interval_end):
                    if adv_x[j] > maxn:
                        maxn = adv_x[j]
                        maxn_idx = j
                for j in range(i, categorical_interval_end):
                    if j == maxn_idx:
                        processed_adv_val = 1
                    else:
                        processed_adv_val = 0
                    processed_adv_x[j] = processed_adv_val
                    lookahead_cnt += 1

        legalized_candidate = processed_adv_x

        return legalized_candidate

    def _get_mode_and_class(self, a):
        # determine if the attack is targeted or not
        target_class = a.target_class()
        targeted = target_class is not None

        if targeted:
            class_ = target_class
        else:
            class_ = a.original_class
        return targeted, class_

    def _run(self, a, binary_search,
             epsilon, stepsize, iterations,
             random_start, return_early,
             perturbable_idx_set, only_increase_idx_set,
             feature_defs, normalization_ratios,
             enforce_interval, request_id, model, browser_id,
             map_back_mode, feature_idx_map, logger):
        if not a.has_gradient():
            warnings.warn('applied gradient-based attack to model that'
                          ' does not provide gradients')
            return

        self._check_distance(a)

        targeted, class_ = self._get_mode_and_class(a)

        if binary_search:
            if isinstance(binary_search, bool):
                k = 20
            else:
                k = int(binary_search)
            return self._run_binary_search(
                a, epsilon, stepsize, iterations,
                random_start, targeted, class_, return_early, k=k,
                perturbable_idx_set=perturbable_idx_set,
                only_increase_idx_set=only_increase_idx_set,
                feature_defs=feature_defs,
                normalization_ratios=normalization_ratios,
                enforce_interval=enforce_interval,
                request_id=request_id)
        else:
            return self._run_one(
                a, epsilon, stepsize, iterations,
                random_start, targeted, class_, return_early,
                perturbable_idx_set=perturbable_idx_set,
                only_increase_idx_set=only_increase_idx_set,
                feature_defs=feature_defs,
                normalization_ratios=normalization_ratios,
                enforce_interval=enforce_interval,
                request_id=request_id, model=model, browser_id=browser_id,
                map_back_mode=map_back_mode,
                feature_idx_map=feature_idx_map, logger=logger)

    def _get_geometric_enforce_interval(self, base, iteations, i):
        progress = float(i + 1) / iteations
        if progress < 0.1:
            growth_ratio = 1
        elif progress >= 0.1 and progress < 0.5:
            growth_ratio = 3
        else:
            growth_ratio = 10
        grown = base * growth_ratio
        return grown

    def _run_binary_search(self, a, epsilon, stepsize, iterations,
                           random_start, targeted, class_, return_early, k,
                           perturbable_idx_set, only_increase_idx_set, feature_defs,
                           normalization_ratios, enforce_interval, request_id):
        factor = stepsize / epsilon

        def try_epsilon(epsilon):
            stepsize = factor * epsilon
            return self._run_one(
                a, epsilon, stepsize, iterations,
                random_start, targeted, class_, return_early,
                perturbable_idx_set, only_increase_idx_set, feature_defs,
                normalization_ratios, enforce_interval, request_id)

        for i in range(k):
            if try_epsilon(epsilon):
                logging.info('successful for eps = {}'.format(epsilon))
                break
            logging.info('not successful for eps = {}'.format(epsilon))
            epsilon = epsilon * 1.5
        else:
            logging.warning('exponential search failed')
            return

        bad = 0
        good = epsilon

        for i in range(k):
            epsilon = (good + bad) / 2
            if try_epsilon(epsilon):
                good = epsilon
                logging.info('successful for eps = {}'.format(epsilon))
            else:
                bad = epsilon
                logging.info('not successful for eps = {}'.format(epsilon))

    def _is_diff_zero(self, diff):
        for feature in diff:
            if abs(feature) > 0.00001:
                return False
        return True

    def _get_diff(self, x, original, perturbable_idx_set, normalization_ratios,
            debug=False):
        if debug:
            print(original[0], x[0])
            print(original[1], x[1])

        delta = {}
        for idx in list(perturbable_idx_set):
            if normalization_ratios[idx]['type'] == 'B':
                diff = float(x[idx]) - float(original[idx])
            else:
                diff = self._calculate_diff(original[idx], x[idx], normalization_ratios[idx]['val'])
            delta[idx] = diff

        return delta

    def _reverse_a_map(self, map):
        reversed = {}
        for feature_name, feature_id in map.items():
            reversed[feature_id] = feature_name
        return reversed

    def _read_curr_html(self, domain, read_modified):
        print("Reading HTML: %s" % domain)
        if read_modified:
            with open(self.BASE_CRAWLED_DIR + "/modified_" + domain + '.html', "r") as fin:
                curr_html = BeautifulSoup(fin, features="html.parser")
        else:
            with open(self.BASE_CRAWLED_DIR + "/" + domain + '.html', "r") as fin:
                curr_html = BeautifulSoup(fin, features="html.parser")

        return curr_html, domain + '.html'

    def _get_url_from_url_id(
        self,
        final_domain,
        target_url_id
    ):
        with open(self.BASE_MAPPING_DIR + '/' + self.final_domain_to_original_domain_mapping[final_domain] + '.csv') as fin:
            data = fin.readlines()
            for row in data:
                row = row.strip()
                url_id, url = row.split(',', 1)
                if url_id == target_url_id:
                    return url
        return None

    def _read_url_id_to_url_mapping(
        self,
        final_domain
    ):
        url_id_to_url_mapping = {}
        with open(self.BASE_MAPPING_DIR + '/' + self.final_domain_to_original_domain_mapping[final_domain] + '.csv') as fin:
            data = fin.readlines()
            for row in data:
                row = row.strip()
                url_id, url = row.split(',', 1)
                url_id_to_url_mapping[url_id] = url
        return url_id_to_url_mapping

    def _get_x_after_mapping_back(
        self,
        domain,
        final_domain,
        url_id,
        diff,
        browser_id,
        working_dir="~/AdGraphAPI/scripts",
        feature_idx_map=None,
        first_time=False
    ):
        reversed_feature_idx_map = self._reverse_a_map(feature_idx_map)
        html, html_fname = self._read_curr_html(domain, read_modified=False)

        if first_time:
            if not os.path.isfile(self.BASE_TIMELINE_DIR + '/' + domain + '.json'):
                cmd = "python3 ~/AdGraphAPI/scripts/load_page_adgraph.py --domain %s --id %s --final-domain %s --mode proxy" % (domain, browser_id, final_domain)
                os.system(cmd)
            cmd = "python ~/AdGraphAPI/scripts/rules_parser.py --target-dir ~/rendering_stream/timeline --domain %s" % domain
            os.system(cmd)
            cmd = "~/AdGraphAPI/adgraph ~/rendering_stream/ features/ mappings/ %s parsed_%s" % (domain, domain)
            os.system(cmd)
            self._url_id_to_url_mapping = self._read_url_id_to_url_mapping(final_domain)

        url = self._url_id_to_url_mapping[url_id]
        original_url = url

        at_least_one_diff_success = False
        new_html = None

        for feature_id, delta in diff.items():
            new_html, modified_url = featureMapbacks(
                name=reversed_feature_idx_map[feature_id],
                html=html,
                url=url,
                delta=delta,
                domain=final_domain
            )

            if new_html is None:
                continue
            else:
                html = new_html
                url = modified_url
                at_least_one_diff_success = True

        if not at_least_one_diff_success:
            print("[ERROR] No diff was successfully mapped back!")
            raise Exception

        # Write back to HTML file after circulating all outstanding perturbations in this iteration
        mapped_x, mapped_unnormalized_x = compute_x_after_mapping_back(
            domain,
            url_id,
            html,
            html_fname,
            working_dir=working_dir,
            browser_id=browser_id,
            final_domain=final_domain
        )

        return mapped_x, mapped_unnormalized_x, original_url, url

    def _deprocess_x(self, x, feature_types, verbal=False):
        FEATURE_TYPES = {'F', 'B', 'C', 'S', 'D', 'L', 'FF'}

        def deprocess_float_feature(val, ratio, ffloat=True):
            val = float(val)
            maxn = float(ratio[0])
            minn = float(ratio[1])
            if ffloat:
                return float(val * (maxn - minn) + minn)
            else:
                return round(val * (maxn - minn) + minn)

        def deprocess_nominal_feature(val, category_name):
            val = float(val)
            if val == 1.0:
                return category_name
            elif val == 0.0:
                return None
            else:
                print("[ERROR] WTF? val: %s %s" % (str(val), category_name))
                raise Exception

        def deprocess_shift_feature(val, offset):
            val = float(val)
            offset = float(offset)
            return int(math.ceil(val + offset))

        features = x
        deprocessed_features = []

        for j in range(len(feature_types)):
            assert feature_types[j]['type'] in FEATURE_TYPES, "[ERROR] Feature type not supported!"

            if feature_types[j]['type'] == 'F':
                ratio = feature_types[j]['val']
                deprocessed_features.append(
                    deprocess_float_feature(features[j], ratio, ffloat=False))
            elif feature_types[j]['type'] == 'FF':
                ratio = feature_types[j]['val']
                deprocessed_features.append(
                    deprocess_float_feature(features[j], ratio, ffloat=True))
            elif feature_types[j]['type'] == 'C':
                category_name = feature_types[j]['val']
                new_val = deprocess_nominal_feature(features[j], category_name)
                if new_val is not None:
                    deprocessed_features.append(new_val)
            elif feature_types[j]['type'] == 'S':
                offset = feature_types[j]['val']
                deprocessed_features.append(
                    deprocess_shift_feature(features[j], offset))
            elif feature_types[j]['type'] == 'B':
                val = features[j]
                deprocessed_features.append(int(float(val)))
            elif feature_types[j]['type'] == 'D':
                val = features[j]
                deprocessed_features.append(val)
            # label column
            elif feature_types[j]['type'] == 'L':
                label = features[j]
                deprocessed_features.append(label)
            else:
                print("???")

        return deprocessed_features

    def _compare_x(self, x1, x2, tags=["X1", "X2"], it=None, X3=None):
        assert len(x1) == len(x2), "[ERROR] Two inputs have different sizes!"
        if it is not None:
            print("Iter #%d" % it)
        for i in range(len(x1)):
            if x1[i] != x2[i]:
                if X3 is not None:
                    print("i:", i, tags[0], ":", x1[i], tags[1], ":", x2[i], "original:", X3[i])
                else:
                    print("i:", i, tags[0], ":", x1[i], tags[1], ":", x2[i])

    def _retrain_local_model(self, model, x, y):
        def __generate_mini_batch(x, y, whole_train_set):
            # import random
            train_set = whole_train_set.sample(n=50)
            train_y = train_set.pop('CLASS').to_numpy().tolist()
            train_x = train_set.to_numpy().tolist()
            for _ in range(50):
                train_x.append(x)
                train_y.append(y)
            return np.array(train_x), np.array(train_y)

        train_x, train_y = __generate_mini_batch(x, y, self._train_data_set)
        model.fit(
            x=train_x,
            y=train_y,
            batch_size=100,
            epochs=1
        )

    def _get_label(self, logits):
        if logits[1] > logits[0]:
            return "AD"
        else:
            return "NONAD"

    def _run_one(self, a, epsilon, stepsize, iterations,
                 random_start, targeted, class_, return_early, model, browser_id,
                 perturbable_idx_set=None, only_increase_idx_set=None,
                 feature_defs=None, normalization_ratios=None,
                 enforce_interval=1, request_id="URL_dummy", dynamic_interval=False,
                 debug=False, draw=False, map_back_mode=True, remote_model=True,
                 check_correctness=False, check_init_correctness=True, feature_idx_map=None, logger=None):
        if draw:
            x_axis, y_l2, y_lf = [], [], []

        domain, url_id = request_id.split(',')
        final_domain = urlparse(domain)[1]
        original_domain = self.final_domain_to_original_domain_mapping[final_domain]

        min_, max_ = a.bounds()
        s = max_ - min_

        original = a.unperturbed.copy()

        if random_start:
            # using uniform noise even if the perturbation clipping uses
            # a different norm because cleverhans does it the same way
            noise = nprng.uniform(
                -epsilon * s, epsilon * s, original.shape).astype(
                    original.dtype)
            x = original + self._clip_perturbation(a, noise, epsilon)
            strict = False  # because we don't enforce the bounds here
        else:
            x = original
            strict = True  # we don't care about the bounds because we are not attacking image clf

        success_cent = False

        if enforce_interval == iterations - 1:
            only_post_process = True
        else:
            only_post_process = False

        is_first_iter = True
        for i in range(iterations):

            if dynamic_interval:
                curr_interval = self._get_geometric_enforce_interval(enforce_interval, iterations, i)
            else:
                curr_interval = enforce_interval

            if i != 0 and (i % curr_interval == 0 or i == iterations - 1):
                should_enforce_policy = True
            else:
                should_enforce_policy = False

            gradient = self._gradient(a, x, class_, strict=strict)
            # non-strict only for the first call and
            # only if random_start is True
            if targeted:
                gradient = -gradient

            # untargeted: gradient ascent on cross-entropy to original class
            # targeted: gradient descent on cross-entropy to target class
            # (this is the actual attack step)
            x = x + stepsize * gradient

            if should_enforce_policy:
                x_before_projection = x.copy()

            # phase 1: reject disallowed perturbations by changing feature values
            # back to original
            if only_post_process:
                if should_enforce_policy and perturbable_idx_set is not None:
                    x = self._reject_imperturbable_features(
                        x,
                        original,
                        perturbable_idx_set)
            else:
                if perturbable_idx_set is not None:
                    x = self._reject_imperturbable_features(
                        x,
                        original,
                        perturbable_idx_set)

            # phase 1: reject decreasing perturbations by changing only-increase
            # feature values back to original
            if should_enforce_policy and only_increase_idx_set is not None:
                x = self._reject_only_increase_features(
                    x,
                    original,
                    only_increase_idx_set)

            # phase 2: change values back to allowed ranges
            if should_enforce_policy and feature_defs is not None:
                x = self._legalize_candidate(
                    x,
                    original,
                    feature_defs)

            if should_enforce_policy:
                l2_dist = self._compute_l2_distance([x_before_projection], [x])
                lf_dist = self._compute_lp_distance([x_before_projection], [x])
                if draw:
                    x_axis.append(i)
                    y_l2.append(l2_dist[0])
                    y_lf.append(lf_dist)
                if debug:
                    print("Step #%d, L2 distance: %f | LP distance: %f" % (i, l2_dist, lf_dist))

            x = original + self._clip_perturbation(a, x - original, original, epsilon, feature_defs)
            x = np.clip(x, min_, max_)

            if should_enforce_policy and map_back_mode:
                diff = self._get_diff(
                    x,
                    original,
                    perturbable_idx_set,
                    normalization_ratios,
                )
                print("Delta at iter #%d: %s" % (i, str(diff)))
                print("Domain: %s" % original_domain)
                print("URL ID: %s" % url_id)
                x_before_mapping_back = x.copy()
                try:
                    x_cent, unnorm_x_cent, original_url, url = self._get_x_after_mapping_back(
                        original_domain,
                        final_domain,
                        url_id,
                        diff,
                        browser_id=browser_id,
                        feature_idx_map=feature_idx_map,
                        first_time=is_first_iter
                    )
                except Exception as err:
                    print("Error occured mapping: %s" % err)
                    return False
                is_first_iter = False

                del unnorm_x_cent[-1]  # remove label

                for j in range(len(x_cent)):
                    if j in diff:
                        if j not in NORM_MAP:
                            continue
                        if diff[j] == 1.0:
                            x_cent[j] = 1.0
                            unnorm_x_cent[NORM_MAP[j]] = "1"
                        if diff[j] == -1.0:
                            x_cent[j] = 0.0
                            unnorm_x_cent[NORM_MAP[j]] = "0"

                print("unnorm_x_cent:", unnorm_x_cent)
                mapping_diff = self._get_diff(
                    x_cent,
                    original,
                    perturbable_idx_set,
                    normalization_ratios,
                )
                print("Delta between before and after mapping-back: %s" % mapping_diff)
                if self._is_diff_zero(mapping_diff):
                    print("[ERROR] Xs before and after mapping back did not change!")
                    return False

            if should_enforce_policy and remote_model and not map_back_mode:
                unnorm_x = self._deprocess_x(x, normalization_ratios)
                unnorm_unperturbed = self._deprocess_x(original, normalization_ratios)

            if i == 0 and check_init_correctness:
                unnorm_unperturbed = self._deprocess_x(original, normalization_ratios)
                prediction_unperturbed_local = self._get_label(model.predict(np.array([original]))[0])
                prediction_unperturbed_remote = predict(unnorm_unperturbed, self._remote_model)[0]
                if prediction_unperturbed_local == prediction_unperturbed_remote == "AD":
                    print("Sanity check passed!")
                else:
                    print("Local:", prediction_unperturbed_local)
                    print("Remote:", prediction_unperturbed_remote)
                    print("Sanity check failed!")
                    return False

            if should_enforce_policy:
                logits, is_adversarial = a.forward_one(x)
                if logging.getLogger().isEnabledFor(logging.DEBUG):
                    if targeted:
                        ce = crossentropy(a.original_class, logits)
                        logging.debug('crossentropy to {} is {}'.format(
                            a.original_class, ce))
                    ce = crossentropy(class_, logits)
                    logging.debug('crossentropy to {} is {}'.format(class_, ce))

                # Use X before mapping back as the starting point for next iteration
                # as there might be feature perturbations that we cannot incoperate in
                # feature space. That is, we should only "descend" by changing "perturbable"
                # features in a "legitimate" fashion
                if remote_model:
                    if map_back_mode:
                        prediction_remote_cent = predict(unnorm_x_cent, self._remote_model)[0]
                    else:
                        prediction_remote = predict(unnorm_x, self._remote_model)[0]
                    prediction_original = predict(unnorm_unperturbed, self._remote_model)[0]
                    if map_back_mode:
                        prediction_local_cent = self._get_label(model.predict(np.array([x_cent]))[0])
                    else:
                        prediction_local = self._get_label(model.predict(np.array([x]))[0])

                    if not only_post_process:
                        if map_back_mode:
                            retrain_cnt_cent = 0
                            print("Remote (cent): %s / local (cent): %s" %
                                  (prediction_remote_cent, prediction_local_cent))
                            while prediction_remote_cent != prediction_local_cent and retrain_cnt_cent < 10:
                                retrain_cnt_cent += 1
                                self._retrain_local_model(model, x_cent, LABLE[prediction_remote_cent])
                                prediction_local = self._get_label(model.predict(np.array([x_cent]))[0])
                                print("(cent) iter #%d, has retrained %d time(s)" % (i, retrain_cnt_cent))
                        else:
                            retrain_cnt = 0
                            print("Remote: %s / local: %s" % (prediction_remote, prediction_local))
                            while prediction_remote != prediction_local and retrain_cnt < 10:
                                retrain_cnt += 1
                                self._retrain_local_model(model, x, LABLE[prediction_remote])
                                prediction_local = self._get_label(model.predict(np.array([x]))[0])
                                print("iter #%d, has retrained %d time(s)" % (i, retrain_cnt))

                    if map_back_mode:
                        if prediction_original == "AD" and prediction_remote_cent == "NONAD":
                            success_cent = True

                        if success_cent:
                            msg = "SUCCESS, iter_%d, %s, %s, %s, %s, %s, %s" % (i, original_domain, final_domain, str(mapping_diff), url_id, original_url, url)
                            print(msg)
                            logger.info(msg)

                            cmd = "cp %s %s" % (
                                self.BASE_HTML_DIR + '/' + original_domain + '.html',
                                self.BASE_EVAL_HTML_DIR + '/original_' + original_domain + '.html'
                            )
                            os.system(cmd)

                            cmd = "cp %s %s" % (
                                self.BASE_HTML_DIR + '/modified_' + original_domain + '.html',
                                self.BASE_EVAL_HTML_DIR + '/' + original_domain + '_' + url_id + '.html'
                            )
                            os.system(cmd)
                            return True

                        x = x_before_mapping_back
                    else:
                        if prediction_original == "AD" and prediction_remote == "NONAD":
                            msg = "SUCCESS, iter_%d, %s, %s, %s" % (i, domain, url_id, str(diff))
                            print(msg)
                            logger.info(msg)
                            return True
        if draw:
            plt.plot(x_axis, y_l2, linewidth=3)
            plt.plot(x_axis, y_lf, linewidth=3)
            plt.show()
        msg = "FAIL, iter_%d, %s, %s, %s, %s, %s, %s" % (i, original_domain, final_domain, str(mapping_diff), url_id, original_url, url)
        print(msg)
        logger.info(msg)
        return False

class LinfinityGradientMixin(object):
    def _gradient(self, a, x, class_, strict=True):
        gradient = a.gradient_one(x, class_, strict=strict)
        gradient = np.sign(gradient)
        min_, max_ = a.bounds()
        gradient = (max_ - min_) * gradient
        return gradient

class L1GradientMixin(object):
    def _gradient(self, a, x, class_, strict=True):
        gradient = a.gradient_one(x, class_, strict=strict)
        # using mean to make range of epsilons comparable to Linf
        gradient = gradient / np.mean(np.abs(gradient))
        min_, max_ = a.bounds()
        gradient = (max_ - min_) * gradient
        return gradient

class L2GradientMixin(object):
    def _gradient(self, a, x, class_, strict=True):
        gradient = a.gradient_one(x, class_, strict=strict)
        # using mean to make range of epsilons comparable to Linf
        gradient = gradient / np.sqrt(np.mean(np.square(gradient)))
        min_, max_ = a.bounds()
        gradient = (max_ - min_) * gradient
        return gradient

class LinfinityClippingMixin(object):
    def _clip_perturbation(self, a, perturbation, original, epsilon, feature_defs):
        _LOCAL_EPSILON = 0.8  # hard-coded for now
        _LOCAL_FEATURE_IDX_SET = {0, 1}  # hard-coded for now as well

        original_perturbation = perturbation.copy()
        min_, max_ = a.bounds()
        s = max_ - min_
        clipped = np.clip(perturbation, -epsilon * s, epsilon * s)

        for i in range(len(clipped)):
            if feature_defs[i] in {"b", "c"}:
                clipped[i] = original_perturbation[i]
            if i in _LOCAL_FEATURE_IDX_SET:
                s_local = original[i]
                offset = min(epsilon * s, _LOCAL_EPSILON * s_local)  # use the smaller offset
                clipped[i] = np.clip([original_perturbation[i]], -offset, offset)[0]

        return clipped

class L1ClippingMixin(object):
    def _clip_perturbation(self, a, perturbation, epsilon):
        # using mean to make range of epsilons comparable to Linf
        norm = np.mean(np.abs(perturbation))
        norm = max(1e-12, norm)  # avoid divsion by zero
        min_, max_ = a.bounds()
        s = max_ - min_
        # clipping, i.e. only decreasing norm
        factor = min(1, epsilon * s / norm)
        return perturbation * factor

class L2ClippingMixin(object):
    def _clip_perturbation(self, a, perturbation, epsilon):
        # using mean to make range of epsilons comparable to Linf
        norm = np.sqrt(np.mean(np.square(perturbation)))
        norm = max(1e-12, norm)  # avoid divsion by zero
        min_, max_ = a.bounds()
        s = max_ - min_
        # clipping, i.e. only decreasing norm
        factor = min(1, epsilon * s / norm)
        # input(str(perturbation * factor))
        return perturbation * factor

class LinfinityDistanceCheckMixin(object):
    def _check_distance(self, a):
        if not isinstance(a.distance, distances.Linfinity):
            logging.warning('Running an attack that tries to minimize the'
                            ' Linfinity norm of the perturbation without'
                            ' specifying foolbox.distances.Linfinity as'
                            ' the distance metric might lead to suboptimal'
                            ' results.')

class L1DistanceCheckMixin(object):
    def _check_distance(self, a):
        if not isinstance(a.distance, distances.MAE):
            logging.warning('Running an attack that tries to minimize the'
                            ' L1 norm of the perturbation without'
                            ' specifying foolbox.distances.MAE as'
                            ' the distance metric might lead to suboptimal'
                            ' results.')

class L2DistanceCheckMixin(object):
    def _check_distance(self, a):
        if not isinstance(a.distance, distances.MSE):
            logging.warning('Running an attack that tries to minimize the'
                            ' L2 norm of the perturbation without'
                            ' specifying foolbox.distances.MSE as'
                            ' the distance metric might lead to suboptimal'
                            ' results.')

class LinfinityBasicIterativeAttack(
        LinfinityGradientMixin,
        LinfinityClippingMixin,
        LinfinityDistanceCheckMixin,
        IterativeProjectedGradientBaseAttack):

    """The Basic Iterative Method introduced in [1]_.

    This attack is also known as Projected Gradient
    Descent (PGD) (without random start) or FGMS^k.

    References
    ----------
    .. [1] Alexey Kurakin, Ian Goodfellow, Samy Bengio,
           "Adversarial examples in the physical world",
            https://arxiv.org/abs/1607.02533

    .. seealso:: :class:`ProjectedGradientDescentAttack`

    """

    @call_decorator
    def __call__(self, input_or_adv, label=None, unpack=True,
                 binary_search=True,
                 epsilon=0.3,
                 stepsize=0.05,
                 iterations=10,
                 random_start=False,
                 return_early=True):
        """Simple iterative gradient-based attack known as
        Basic Iterative Method, Projected Gradient Descent or FGSM^k.

        Parameters
        ----------
        input_or_adv : `numpy.ndarray` or :class:`Adversarial`
            The original, unperturbed input as a `numpy.ndarray` or
            an :class:`Adversarial` instance.
        label : int
            The reference label of the original input. Must be passed
            if `a` is a `numpy.ndarray`, must not be passed if `a` is
            an :class:`Adversarial` instance.
        unpack : bool
            If true, returns the adversarial input, otherwise returns
            the Adversarial object.
        binary_search : bool or int
            Whether to perform a binary search over epsilon and stepsize,
            keeping their ratio constant and using their values to start
            the search. If False, hyperparameters are not optimized.
            Can also be an integer, specifying the number of binary
            search steps (default 20).
        epsilon : float
            Limit on the perturbation size; if binary_search is True,
            this value is only for initialization and automatically
            adapted.
        stepsize : float
            Step size for gradient descent; if binary_search is True,
            this value is only for initialization and automatically
            adapted.
        iterations : int
            Number of iterations for each gradient descent run.
        random_start : bool
            Start the attack from a random point rather than from the
            original input.
        return_early : bool
            Whether an individual gradient descent run should stop as
            soon as an adversarial is found.
        """

        a = input_or_adv
        del input_or_adv
        del label
        del unpack

        assert epsilon > 0

        self._run(a, binary_search,
                  epsilon, stepsize, iterations,
                  random_start, return_early)

BasicIterativeMethod = LinfinityBasicIterativeAttack
BIM = BasicIterativeMethod

class L1BasicIterativeAttack(
        L1GradientMixin,
        L1ClippingMixin,
        L1DistanceCheckMixin,
        IterativeProjectedGradientBaseAttack):

    """Modified version of the Basic Iterative Method
    that minimizes the L1 distance.

    .. seealso:: :class:`LinfinityBasicIterativeAttack`

    """

    @call_decorator
    def __call__(self, input_or_adv, label=None, unpack=True,
                 binary_search=True,
                 epsilon=0.3,
                 stepsize=0.05,
                 iterations=10,
                 random_start=False,
                 return_early=True):
        """Simple iterative gradient-based attack known as
        Basic Iterative Method, Projected Gradient Descent or FGSM^k.

        Parameters
        ----------
        input_or_adv : `numpy.ndarray` or :class:`Adversarial`
            The original, unperturbed input as a `numpy.ndarray` or
            an :class:`Adversarial` instance.
        label : int
            The reference label of the original input. Must be passed
            if `a` is a `numpy.ndarray`, must not be passed if `a` is
            an :class:`Adversarial` instance.
        unpack : bool
            If true, returns the adversarial input, otherwise returns
            the Adversarial object.
        binary_search : bool or int
            Whether to perform a binary search over epsilon and stepsize,
            keeping their ratio constant and using their values to start
            the search. If False, hyperparameters are not optimized.
            Can also be an integer, specifying the number of binary
            search steps (default 20).
        epsilon : float
            Limit on the perturbation size; if binary_search is True,
            this value is only for initialization and automatically
            adapted.
        stepsize : float
            Step size for gradient descent; if binary_search is True,
            this value is only for initialization and automatically
            adapted.
        iterations : int
            Number of iterations for each gradient descent run.
        random_start : bool
            Start the attack from a random point rather than from the
            original input.
        return_early : bool
            Whether an individual gradient descent run should stop as
            soon as an adversarial is found.
        """

        a = input_or_adv
        del input_or_adv
        del label
        del unpack

        assert epsilon > 0

        self._run(a, binary_search,
                  epsilon, stepsize, iterations,
                  random_start, return_early)

class L2BasicIterativeAttack(
        L2GradientMixin,
        L2ClippingMixin,
        L2DistanceCheckMixin,
        IterativeProjectedGradientBaseAttack):

    """Modified version of the Basic Iterative Method
    that minimizes the L2 distance.

    .. seealso:: :class:`LinfinityBasicIterativeAttack`

    """

    def _read_dataset(self, fname):
        import pandas as pd
        from sklearn.utils import shuffle

        train_dataframe = pd.read_csv(fname)
        dataframe = shuffle(train_dataframe)

        return dataframe

    # overriden init method in ABC
    def _initialize(self):
        self._remote_model = setup_clf(
            self.BASE_MODEL_DIR + "/rf.pkl"
        )
        self._train_data_set = self._read_dataset(
            self.BASE_DATA_DIR + "/hand_preprocessed_trimmed_label_gt_augmented_train_set.csv"
        )

    @call_decorator
    def __call__(self, input_or_adv, label=None, unpack=True,
                 binary_search=True,
                 epsilon=0.3,
                 stepsize=0.01,
                 iterations=40,
                 random_start=False,
                 return_early=True,
                 perturbable_idx_set=None,
                 only_increase_idx_set=None,
                 feature_defs=None,
                 normalization_ratios=None,
                 enforce_interval=1,
                 request_id="URL_dummy",
                 model=None,
                 browser_id=None,
                 map_back_mode=False):
        """Simple iterative gradient-based attack known as
        Basic Iterative Method, Projected Gradient Descent or FGSM^k.

        Parameters
        ----------
        input_or_adv : `numpy.ndarray` or :class:`Adversarial`
            The original, unperturbed input as a `numpy.ndarray` or
            an :class:`Adversarial` instance.
        label : int
            The reference label of the original input. Must be passed
            if `a` is a `numpy.ndarray`, must not be passed if `a` is
            an :class:`Adversarial` instance.
        unpack : bool
            If true, returns the adversarial input, otherwise returns
            the Adversarial object.
        binary_search : bool or int
            Whether to perform a binary search over epsilon and stepsize,
            keeping their ratio constant and using their values to start
            the search. If False, hyperparameters are not optimized.
            Can also be an integer, specifying the number of binary
            search steps (default 20).
        epsilon : float
            Limit on the perturbation size; if binary_search is True,
            this value is only for initialization and automatically
            adapted.
        stepsize : float
            Step size for gradient descent; if binary_search is True,
            this value is only for initialization and automatically
            adapted.
        iterations : int
            Number of iterations for each gradient descent run.
        random_start : bool
            Start the attack from a random point rather than from the
            original input.
        return_early : bool
            Whether an individual gradient descent run should stop as
            soon as an adversarial is found.
        """
        a = input_or_adv
        del input_or_adv
        del label
        del unpack

        assert epsilon > 0

        print("Parameters:", epsilon, stepsize, iterations, enforce_interval, request_id)
        self._run(a, binary_search,
                  epsilon, stepsize, iterations,
                  random_start, return_early,
                  perturbable_idx_set, only_increase_idx_set,
                  feature_defs, normalization_ratios,
                  enforce_interval, request_id,
                  model, browser_id, map_back_mode)

class ProjectedGradientDescentAttack(
        LinfinityGradientMixin,
        LinfinityClippingMixin,
        LinfinityDistanceCheckMixin,
        IterativeProjectedGradientBaseAttack):

    """The Projected Gradient Descent Attack
    introduced in [1]_ without random start.

    When used without a random start, this attack
    is also known as Basic Iterative Method (BIM)
    or FGSM^k.

    References
    ----------
    .. [1] Aleksander Madry, Aleksandar Makelov, Ludwig Schmidt,
           Dimitris Tsipras, Adrian Vladu, "Towards Deep Learning
           Models Resistant to Adversarial Attacks",
           https://arxiv.org/abs/1706.06083

    .. seealso::

       :class:`LinfinityBasicIterativeAttack` and
       :class:`RandomStartProjectedGradientDescentAttack`

    """

    def _read_dataset(self, fname):
        import pandas as pd
        from sklearn.utils import shuffle

        train_dataframe = pd.read_csv(fname)
        dataframe = shuffle(train_dataframe)

        return dataframe

    # overriden init method in ABC
    def _initialize(self):
        self._remote_model = setup_clf("../model/rf.pkl")
        self.BASE_DATA_DIR = HOME_DIR + "/attack-adgraph-pipeline/data"
        self._train_data_set = self._read_dataset(
            self.BASE_DATA_DIR + "/hand_preprocessed_trimmed_label_gt_augmented_train_set.csv")

    @call_decorator
    def __call__(self, input_or_adv, label=None, unpack=True,
                 binary_search=True,
                 epsilon=0.3,
                 stepsize=0.01,
                 iterations=40,
                 random_start=False,
                 return_early=True,
                 perturbable_idx_set=None,
                 only_increase_idx_set=None,
                 feature_defs=None,
                 normalization_ratios=None,
                 enforce_interval=1,
                 request_id="URL_dummy",
                 model=None,
                 browser_id=None,
                 map_back_mode=False,
                 feature_idx_map=None,
                 logger=None):
        """Simple iterative gradient-based attack known as
        Basic Iterative Method, Projected Gradient Descent or FGSM^k.

        Parameters
        ----------
        input_or_adv : `numpy.ndarray` or :class:`Adversarial`
            The original, unperturbed input as a `numpy.ndarray` or
            an :class:`Adversarial` instance.
        label : int
            The reference label of the original input. Must be passed
            if `a` is a `numpy.ndarray`, must not be passed if `a` is
            an :class:`Adversarial` instance.
        unpack : bool
            If true, returns the adversarial input, otherwise returns
            the Adversarial object.
        binary_search : bool or int
            Whether to perform a binary search over epsilon and stepsize,
            keeping their ratio constant and using their values to start
            the search. If False, hyperparameters are not optimized.
            Can also be an integer, specifying the number of binary
            search steps (default 20).
        epsilon : float
            Limit on the perturbation size; if binary_search is True,
            this value is only for initialization and automatically
            adapted.
        stepsize : float
            Step size for gradient descent; if binary_search is True,
            this value is only for initialization and automatically
            adapted.
        iterations : int
            Number of iterations for each gradient descent run.
        random_start : bool
            Start the attack from a random point rather than from the
            original input.
        return_early : bool
            Whether an individual gradient descent run should stop as
            soon as an adversarial is found.
        """

        a = input_or_adv
        del input_or_adv
        del label
        del unpack

        assert epsilon > 0

        # if self._remote_model is None:
        #     self._remote_model = setup_clf("/home/shitong/Desktop/AdGraphAPI/scripts/model/rf.pkl")

        print("Parameters:", epsilon, stepsize, iterations, enforce_interval, request_id)
        self._run(a, binary_search,
                  epsilon, stepsize, iterations,
                  random_start, return_early,
                  perturbable_idx_set, only_increase_idx_set,
                  feature_defs, normalization_ratios,
                  enforce_interval, request_id,
                  model, browser_id, map_back_mode, feature_idx_map, logger)

ProjectedGradientDescent = ProjectedGradientDescentAttack
PGD = ProjectedGradientDescent

class RandomStartProjectedGradientDescentAttack(
        LinfinityGradientMixin,
        LinfinityClippingMixin,
        LinfinityDistanceCheckMixin,
        IterativeProjectedGradientBaseAttack):

    """The Projected Gradient Descent Attack
    introduced in [1]_ with random start.

    References
    ----------
    .. [1] Aleksander Madry, Aleksandar Makelov, Ludwig Schmidt,
           Dimitris Tsipras, Adrian Vladu, "Towards Deep Learning
           Models Resistant to Adversarial Attacks",
           https://arxiv.org/abs/1706.06083

    .. seealso:: :class:`ProjectedGradientDescentAttack`

    """

    @call_decorator
    def __call__(self, input_or_adv, label=None, unpack=True,
                 binary_search=True,
                 epsilon=0.3,
                 stepsize=0.01,
                 iterations=40,
                 random_start=True,
                 return_early=True):
        """Simple iterative gradient-based attack known as
        Basic Iterative Method, Projected Gradient Descent or FGSM^k.

        Parameters
        ----------
        input_or_adv : `numpy.ndarray` or :class:`Adversarial`
            The original, unperturbed input as a `numpy.ndarray` or
            an :class:`Adversarial` instance.
        label : int
            The reference label of the original input. Must be passed
            if `a` is a `numpy.ndarray`, must not be passed if `a` is
            an :class:`Adversarial` instance.
        unpack : bool
            If true, returns the adversarial input, otherwise returns
            the Adversarial object.
        binary_search : bool or int
            Whether to perform a binary search over epsilon and stepsize,
            keeping their ratio constant and using their values to start
            the search. If False, hyperparameters are not optimized.
            Can also be an integer, specifying the number of binary
            search steps (default 20).
        epsilon : float
            Limit on the perturbation size; if binary_search is True,
            this value is only for initialization and automatically
            adapted.
        stepsize : float
            Step size for gradient descent; if binary_search is True,
            this value is only for initialization and automatically
            adapted.
        iterations : int
            Number of iterations for each gradient descent run.
        random_start : bool
            Start the attack from a random point rather than from the
            original input.
        return_early : bool
            Whether an individual gradient descent run should stop as
            soon as an adversarial is found.
        """

        a = input_or_adv
        del input_or_adv
        del label
        del unpack

        assert epsilon > 0

        self._run(a, binary_search,
                  epsilon, stepsize, iterations,
                  random_start, return_early)

RandomProjectedGradientDescent = RandomStartProjectedGradientDescentAttack
RandomPGD = RandomProjectedGradientDescent

class MomentumIterativeAttack(
        LinfinityClippingMixin,
        LinfinityDistanceCheckMixin,
        IterativeProjectedGradientBaseAttack):

    """The Momentum Iterative Method attack
    introduced in [1]_. It's like the Basic
    Iterative Method or Projected Gradient
    Descent except that it uses momentum.

    References
    ----------
    .. [1] Yinpeng Dong, Fangzhou Liao, Tianyu Pang, Hang Su,
           Jun Zhu, Xiaolin Hu, Jianguo Li, "Boosting Adversarial
           Attacks with Momentum",
           https://arxiv.org/abs/1710.06081

    """

    def _gradient(self, a, x, class_, strict=True):
        # get current gradient
        gradient = a.gradient_one(x, class_, strict=strict)
        gradient = gradient / max(1e-12, np.mean(np.abs(gradient)))

        # combine with history of gradient as new history
        self._momentum_history = \
            self._decay_factor * self._momentum_history + gradient

        # use history
        gradient = self._momentum_history
        gradient = np.sign(gradient)
        min_, max_ = a.bounds()
        gradient = (max_ - min_) * gradient
        return gradient

    def _run_one(self, *args, **kwargs):
        # reset momentum history every time we restart
        # gradient descent
        self._momentum_history = 0
        return super(MomentumIterativeAttack, self)._run_one(*args, **kwargs)

    @call_decorator
    def __call__(self, input_or_adv, label=None, unpack=True,
                 binary_search=True,
                 epsilon=0.3,
                 stepsize=0.06,
                 iterations=10,
                 decay_factor=1.0,
                 random_start=False,
                 return_early=True):
        """Momentum-based iterative gradient attack known as
        Momentum Iterative Method.

        Parameters
        ----------
        input_or_adv : `numpy.ndarray` or :class:`Adversarial`
            The original, unperturbed input as a `numpy.ndarray` or
            an :class:`Adversarial` instance.
        label : int
            The reference label of the original input. Must be passed
            if `a` is a `numpy.ndarray`, must not be passed if `a` is
            an :class:`Adversarial` instance.
        unpack : bool
            If true, returns the adversarial input, otherwise returns
            the Adversarial object.
        binary_search : bool
            Whether to perform a binary search over epsilon and stepsize,
            keeping their ratio constant and using their values to start
            the search. If False, hyperparameters are not optimized.
            Can also be an integer, specifying the number of binary
            search steps (default 20).
        epsilon : float
            Limit on the perturbation size; if binary_search is True,
            this value is only for initialization and automatically
            adapted.
        stepsize : float
            Step size for gradient descent; if binary_search is True,
            this value is only for initialization and automatically
            adapted.
        iterations : int
            Number of iterations for each gradient descent run.
        decay_factor : float
            Decay factor used by the momentum term.
        random_start : bool
            Start the attack from a random point rather than from the
            original input.
        return_early : bool
            Whether an individual gradient descent run should stop as
            soon as an adversarial is found.
        """
        a = input_or_adv
        del input_or_adv
        del label
        del unpack

        assert epsilon > 0

        self._decay_factor = decay_factor

        self._run(a, binary_search,
                  epsilon, stepsize, iterations,
                  random_start, return_early)

MomentumIterativeMethod = MomentumIterativeAttack


import numpy as np
from robosuite.models.arenas import Arena
from robosuite.utils.mjcf_utils import xml_path_completion
from robosuite.utils.mjcf_utils import array_to_string, string_to_array
from collections import OrderedDict
from robosuite.models.objects import CylinderObject, PlateWithHoleObject, BoxObject, BallObject

import xml.etree.ElementTree as ET

import imageio

class HeightTableArena(Arena):
    """Workspace that contains an empty table."""

    def __init__(
        self, table_height_full_size=(0.8, 0.8, 0.8, 0.8), table_friction=(1, 0.005, 0.0001)
    ):
        """
        Args:
            table_full_size: full dimensions of the table height: x, y, height scale over the surface, box depth from the min
            friction: friction parameters of the table
        """
        super().__init__(xml_path_completion("arenas/height_table_arena.xml"))

        self.table_height_full_size = np.array(table_height_full_size)
        self.table_half_size = np.array([
                self.table_height_full_size[0],
                self.table_height_full_size[1],
                self.table_height_full_size[3]]) / 2

        self.table_friction = table_friction

        self.floor = self.worldbody.find("./geom[@name='floor']")
        self.table_body = self.worldbody.find("./geom[@name='table_hf_geom']")
        self.table_top_abs = self.worldbody.find("./geom[@name='table_hf_geom']")
        self.table_asset = self.asset.find("./hfield[@name='table_hf']")

        self.sensor_sites = {}

        self.configure_location()

    def configure_location(self):
        self.bottom_pos = np.array([0, 0, 0])
        self.floor.set("pos", array_to_string(self.bottom_pos))

        self.center_pos = self.bottom_pos + np.array([0, 0, 2*self.table_half_size[2]])
        self.table_body.set("pos", array_to_string(self.center_pos))
        self.table_asset.set("size", array_to_string(self.table_height_full_size))
        self.table_body.set("friction", array_to_string(self.table_friction))

        height_img = imageio.imread(xml_path_completion(self.table_asset.get("file")), as_gray=True)

        total_pixels = height_img.shape[0] * height_img.shape[1]

        #TODO: the number of sensors affects the performance. Do not increse over 500
        #800 sensors -> mujoco steps up to 0.25sec per policy step!
        num_sensors = min(500, total_pixels)

        sample_rate = np.floor(total_pixels / num_sensors)

        #Random Sampling:
        used_pairs = []
        sensor_counter = 0
        while sensor_counter < num_sensors:
            i = np.random.randint(0, high=height_img.shape[0])
            j = np.random.randint(0, high=height_img.shape[1])

            if (i,j) not in used_pairs: # We could add additional constraints like that sensors should be appart from each other

                self.mujoco_objects = OrderedDict()
                square = BallObject(size=[0.01],
                                   rgba=[0, 0, 1, 1],
                                   density=500,
                                   friction=0.05)
                square_name = 'contact_'+str(i)+"_"+str(j)
                collision_b = square.get_collision(name=square_name, site=True)

                x_pos = 0
                if height_img.shape[0]%2==0:
                    x_pos = -(np.floor(height_img.shape[0]/2) - 0.5 - i)*2*self.table_height_full_size[0]/height_img.shape[0]
                else:
                    x_pos = (i-np.floor(height_img.shape[0]/2))*2*self.table_height_full_size[0]/height_img.shape[0]

                y_pos = 0
                if height_img.shape[1]%2==0:
                    y_pos = -(np.floor(height_img.shape[1]/2)-0.5-j)*2*self.table_height_full_size[1]/height_img.shape[1]
                else:
                    y_pos = (j-np.floor(height_img.shape[1]/2))*2*self.table_height_full_size[1]/height_img.shape[1]

                #Compute mean of pixel value and neighbors -> better for low dim images
                acc = 0
                mean_ctr = 1
                pixel_value = height_img[i,j]

                # DO NOT DELETE! The following code helps for low dimensional pictures because the interpolation "buries" some sensors
                # if i+1< height_img.shape[0]:
                #     if height_img[i+1,j] > pixel_value:
                #         acc += height_img[i+1,j]
                #         mean_ctr += 1
                # if i-1<=0:
                #     if height_img[i-1,j] > pixel_value:
                #         acc += height_img[i-1,j]
                #         mean_ctr += 1
                # if j+1< height_img.shape[0]:
                #     if height_img[i,j+1] > pixel_value:
                #         acc += height_img[i,j+1]
                #         mean_ctr += 1
                # if j-1<=0:
                #     if height_img[i,j-1] > pixel_value:
                #         acc += height_img[i,j-1]
                #         mean_ctr += 1

                acc += height_img[i,j]
                z_pos = self.table_height_full_size[3] + self.table_height_full_size[2]*(acc/float(mean_ctr))/np.max(height_img)+0.005

                position_sensor = [y_pos, x_pos,z_pos]
                position_sensor[1] = -position_sensor[1]

                collision_b.find("site").set("pos", array_to_string(position_sensor))
                self.worldbody.append(collision_b.find("site"))
                self.sensor_sites[sensor_counter] = position_sensor

                # Add this pixel to the list of used pixels
                used_pairs += [(i,j)]

                # Add sensor linked to the sensor site
                ET.SubElement(self.sensor, "touch", attrib={"name" : square_name +"_sensor", "site" : square_name})# + "_site"})
                sensor_counter+=1


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

# ========================================================================
# the script is copied from
# https://github.com/DIFACQUIM/ECIF/blob/master/01_ECIF_Calculation.ipynb
# ========================================================================

import numpy as np
import pandas as pd
import os
from rdkit import Chem
from scipy.spatial.distance import cdist
from itertools import product
from rdkit.ML.Descriptors.MoleculeDescriptors import MolecularDescriptorCalculator
import argparse

##For ECIF
# Possible predefined protein atoms
ECIF_ProteinAtoms = ['C;4;1;3;0;0', 'C;4;2;1;1;1', 'C;4;2;2;0;0', 'C;4;2;2;0;1',
                     'C;4;3;0;0;0', 'C;4;3;0;1;1', 'C;4;3;1;0;0', 'C;4;3;1;0;1',
                     'C;5;3;0;0;0', 'C;6;3;0;0;0', 'N;3;1;2;0;0', 'N;3;2;0;1;1',
                     'N;3;2;1;0;0', 'N;3;2;1;1;1', 'N;3;3;0;0;1', 'N;4;1;2;0;0',
                     'N;4;1;3;0;0', 'N;4;2;1;0;0', 'O;2;1;0;0;0', 'O;2;1;1;0;0',
                     'S;2;1;1;0;0', 'S;2;2;0;0;0']

# Possible ligand atoms according to the PDBbind 2016 "refined set"
ECIF_LigandAtoms = ['Br;1;1;0;0;0', 'C;3;3;0;1;1', 'C;4;1;1;0;0', 'C;4;1;2;0;0',
                     'C;4;1;3;0;0', 'C;4;2;0;0;0', 'C;4;2;1;0;0', 'C;4;2;1;0;1',
                     'C;4;2;1;1;1', 'C;4;2;2;0;0', 'C;4;2;2;0;1', 'C;4;3;0;0;0',
                     'C;4;3;0;0;1', 'C;4;3;0;1;1', 'C;4;3;1;0;0', 'C;4;3;1;0;1',
                     'C;4;4;0;0;0', 'C;4;4;0;0;1', 'C;5;3;0;0;0', 'C;5;3;0;1;1',
                     'C;6;3;0;0;0', 'Cl;1;1;0;0;0', 'F;1;1;0;0;0', 'I;1;1;0;0;0',
                     'N;2;1;1;0;0',
                     'N;3;1;0;0;0', 'N;3;1;1;0;0', 'N;3;1;2;0;0', 'N;3;2;0;0;0',
                     'N;3;2;0;0;1', 'N;3;2;0;1;1', 'N;3;2;1;0;0', 'N;3;2;1;0;1',
                     'N;3;2;1;1;1', 'N;3;3;0;0;0', 'N;3;3;0;0;1', 'N;3;3;0;1;1',
                     'N;4;1;2;0;0', 'N;4;1;3;0;0', 'N;4;2;1;0;0', 'N;4;2;2;0;0',
                     'N;4;2;2;0;1', 'N;4;3;0;0;0', 'N;4;3;0;0;1', 'N;4;3;1;0;0',
                     'N;4;3;1;0;1', 'N;4;4;0;0;0', 'N;4;4;0;0;1', 'N;5;2;0;0;0',
                     'N;5;3;0;0;0', 'N;5;3;0;1;1', 'O;2;1;0;0;0', 'O;2;1;1;0;0',
                     'O;2;2;0;0;0', 'O;2;2;0;0;1', 'O;2;2;0;1;1', 'P;5;4;0;0;0',
                     'P;6;4;0;0;0', 'P;6;4;0;0;1', 'P;7;4;0;0;0', 'S;2;1;0;0;0',
                     'S;2;1;1;0;0', 'S;2;2;0;0;0', 'S;2;2;0;0;1', 'S;2;2;0;1;1',
                     'S;3;3;0;0;0', 'S;3;3;0;0;1', 'S;4;3;0;0;0', 'S;6;4;0;0;0',
                     'S;6;4;0;0;1', 'S;7;4;0;0;0']

PossibleECIF = [i[0]+"-"+i[1] for i in product(ECIF_ProteinAtoms, ECIF_LigandAtoms)]

##For ELEMENTS
ELEMENTS_ProteinAtoms = ["C","N","O", "S"]
ELEMENTS_LigandAtoms = ["Br", "C", "Cl", "F", "I", "N", "O", "P", "S"]
PossibleELEMENTS = [i[0]+"-"+i[1] for i in product(ELEMENTS_ProteinAtoms, ELEMENTS_LigandAtoms)]

##For ligand descriptors
LigandDescriptors = ['MaxEStateIndex', 'MinEStateIndex', 'MaxAbsEStateIndex', 'MinAbsEStateIndex',
                      'qed', 'MolWt', 'HeavyAtomMolWt', 'ExactMolWt', 'NumValenceElectrons',
                      'FpDensityMorgan1', 'FpDensityMorgan2', 'FpDensityMorgan3', 'BalabanJ',
                      'BertzCT', 'Chi0', 'Chi0n', 'Chi0v', 'Chi1', 'Chi1n', 'Chi1v', 'Chi2n',
                      'Chi2v', 'Chi3n', 'Chi3v', 'Chi4n', 'Chi4v', 'HallKierAlpha', 'Kappa1',
                      'Kappa2', 'Kappa3', 'LabuteASA', 'PEOE_VSA14', 'SMR_VSA1', 'SMR_VSA10',
                      'SMR_VSA2', 'SMR_VSA3', 'SMR_VSA4', 'SMR_VSA5', 'SMR_VSA6', 'SMR_VSA7',
                      'SMR_VSA9', 'SlogP_VSA1', 'SlogP_VSA10', 'SlogP_VSA11', 'SlogP_VSA12',
                      'SlogP_VSA2', 'SlogP_VSA3', 'SlogP_VSA4', 'SlogP_VSA5', 'SlogP_VSA6',
                      'SlogP_VSA7', 'SlogP_VSA8', 'TPSA', 'EState_VSA1', 'EState_VSA10',
                      'EState_VSA11', 'EState_VSA2', 'EState_VSA3', 'EState_VSA4', 'EState_VSA5',
                      'EState_VSA6', 'EState_VSA7', 'EState_VSA8', 'EState_VSA9', 'VSA_EState1',
                      'VSA_EState10', 'VSA_EState2', 'VSA_EState3', 'VSA_EState4', 'VSA_EState5',
                      'VSA_EState6', 'VSA_EState7', 'VSA_EState8', 'VSA_EState9', 'FractionCSP3',
                      'HeavyAtomCount', 'NHOHCount', 'NOCount', 'NumAliphaticCarbocycles',
                      'NumAliphaticHeterocycles', 'NumAliphaticRings', 'NumAromaticCarbocycles',
                      'NumAromaticHeterocycles', 'NumAromaticRings', 'NumHAcceptors', 'NumHDonors',
                      'NumHeteroatoms', 'NumRotatableBonds', 'NumSaturatedCarbocycles',
                      'NumSaturatedHeterocycles', 'NumSaturatedRings', 'RingCount', 'MolLogP',
                      'MolMR', 'fr_Al_COO', 'fr_Al_OH', 'fr_Al_OH_noTert', 'fr_ArN', 'fr_Ar_N',
                      'fr_Ar_NH', 'fr_Ar_OH', 'fr_COO', 'fr_COO2', 'fr_C_O', 'fr_C_O_noCOO',
                      'fr_C_S', 'fr_HOCCN', 'fr_Imine', 'fr_NH0', 'fr_NH1', 'fr_NH2', 'fr_N_O',
                      'fr_Ndealkylation1', 'fr_Ndealkylation2', 'fr_Nhpyrrole', 'fr_SH', 'fr_aldehyde',
                      'fr_alkyl_carbamate', 'fr_alkyl_halide', 'fr_allylic_oxid', 'fr_amide',
                      'fr_amidine', 'fr_aniline', 'fr_aryl_methyl', 'fr_azo', 'fr_barbitur',
                      'fr_benzene', 'fr_bicyclic', 'fr_dihydropyridine', 'fr_epoxide', 'fr_ester',
                      'fr_ether', 'fr_furan', 'fr_guanido', 'fr_halogen', 'fr_hdrzine', 'fr_hdrzone',
                      'fr_imidazole', 'fr_imide', 'fr_isocyan', 'fr_isothiocyan', 'fr_ketone',
                      'fr_ketone_Topliss', 'fr_lactam', 'fr_lactone', 'fr_methoxy', 'fr_morpholine',
                      'fr_nitrile', 'fr_nitro', 'fr_nitro_arom', 'fr_nitroso', 'fr_oxazole',
                      'fr_oxime', 'fr_para_hydroxylation', 'fr_phenol', 'fr_phenol_noOrthoHbond',
                      'fr_piperdine', 'fr_piperzine', 'fr_priamide', 'fr_pyridine', 'fr_quatN',
                      'fr_sulfide', 'fr_sulfonamd', 'fr_sulfone', 'fr_term_acetylene', 'fr_tetrazole',
                      'fr_thiazole', 'fr_thiocyan', 'fr_thiophene', 'fr_urea']

DescCalc = MolecularDescriptorCalculator(LigandDescriptors)

def GetAtomType(atom):
	# This function takes an atom in a molecule and returns its type as defined for ECIF
	'''An atom type from EFIC is defined as:
			Atom symbol;
			Explicit valence;
			Attached heavy atoms;
			Attached hydrogens;
			Aromaticity;
			Ring membership'''
	AtomType = [atom.GetSymbol(),
				str(atom.GetExplicitValence()),
				str(len([x.GetSymbol() for x in atom.GetNeighbors() if x.GetSymbol() != "H"])),
				str(len([x.GetSymbol() for x in atom.GetNeighbors() if x.GetSymbol() == "H"])),
				str(int(atom.GetIsAromatic())),
				str(int(atom.IsInRing())),
				]

	return(";".join(AtomType))

def LoadSDFasDF(SDF):
	# This function takes an SDF for a ligand as input and
	#	 returns it as a pandas DataFrame with its atom types labeled according to ECIF

	m = Chem.MolFromMolFile(SDF, sanitize=False)
	m.UpdatePropertyCache(strict=False)

	ECIF_atoms = []
	for atom in m.GetAtoms():
		if atom.GetSymbol() != "H": # Include only non-hydrogen atoms
			entry = [int(atom.GetIdx())]
			entry.append(GetAtomType(atom))
			pos = m.GetConformer().GetAtomPosition(atom.GetIdx())
			entry.append(float("{0:.4f}".format(pos.x)))
			entry.append(float("{0:.4f}".format(pos.y)))
			entry.append(float("{0:.4f}".format(pos.z)))
			ECIF_atoms.append(entry)

	df = pd.DataFrame(ECIF_atoms)
	df.columns = ["ATOM_INDEX", "ECIF_ATOM_TYPE","X","Y","Z"]
	if len(set(df["ECIF_ATOM_TYPE"]) - set(ECIF_LigandAtoms)) > 0:
		print("WARNING: Ligand contains unsupported atom types. Only supported atom-type pairs are counted.")
		#print(set(df["ECIF_ATOM_TYPE"]) - set(ECIF_LigandAtoms))
		#for x in set(ECIF_LigandAtoms):
		#	if x not in set(df["ECIF_ATOM_TYPE"]):
		#		print(x)
	return(df)

def LoadPDBasDF(PDB):
	# This function takes a PDB for a protein as input and
	#   returns it as a pandas DataFrame with its atom types labeled according to ECIF

	ECIF_atoms = []

	f = open(PDB)
	for i in f:
		if i[:4] == "ATOM":
			# Include only non-hydrogen atoms
			if (len(i[12:16].replace(" ","")) < 4 and i[12:16].replace(" ","")[0] != "H") or (len(i[12:16].replace(" ","")) == 4 and i[12:16].replace(" ","")[1] != "H" and i[12:16].replace(" ","")[0] != "H"):
				ECIF_atoms.append([int(i[6:11]),
						i[17:20]+"-"+i[12:16].replace(" ",""),
						float(i[30:38]),
						float(i[38:46]),
						float(i[46:54])
						])

	f.close()

	df = pd.DataFrame(ECIF_atoms, columns=["ATOM_INDEX","PDB_ATOM","X","Y","Z"])
	df = df.merge(Atom_Keys, left_on='PDB_ATOM', right_on='PDB_ATOM')[["ATOM_INDEX", "ECIF_ATOM_TYPE", "X", "Y", "Z"]].sort_values(by="ATOM_INDEX").reset_index(drop=True)
	if list(df["ECIF_ATOM_TYPE"].isna()).count(True) > 0:
		print("WARNING: Protein contains unsupported atom types. Only supported atom-type pairs are counted.")
	return(df)

def GetPLPairs(PDB_protein, SDF_ligand, distance_cutoff=6.0):
	# This function returns the protein-ligand atom-type pairs for a given distance cutoff

	# Load both structures as pandas DataFrames
	Target = LoadPDBasDF(PDB_protein)
	Ligand = LoadSDFasDF(SDF_ligand)

	# Take all atoms from the target within a cubic box around the ligand considering the "distance_cutoff criterion"
	for i in ["X","Y","Z"]:
		Target = Target[Target[i] < float(Ligand[i].max())+distance_cutoff]
		Target = Target[Target[i] > float(Ligand[i].min())-distance_cutoff]

	# Get all possible pairs
	Pairs = list(product(Target["ECIF_ATOM_TYPE"], Ligand["ECIF_ATOM_TYPE"]))
	Pairs = [x[0]+"-"+x[1] for x in Pairs]
	Pairs = pd.DataFrame(Pairs, columns=["ECIF_PAIR"])
	Distances = cdist(Target[["X","Y","Z"]], Ligand[["X","Y","Z"]], metric="euclidean")
	Distances = Distances.reshape(Distances.shape[0]*Distances.shape[1],1)
	Distances = pd.DataFrame(Distances, columns=["DISTANCE"])

	Pairs = pd.concat([Pairs,Distances], axis=1)
	Pairs = Pairs[Pairs["DISTANCE"] <= distance_cutoff].reset_index(drop=True)
	# Pairs from ELEMENTS could be easily obtained froms pairs from ECIF
	Pairs["ELEMENTS_PAIR"] = [x.split("-")[0].split(";")[0]+"-"+x.split("-")[1].split(";")[0] for x in Pairs["ECIF_PAIR"]]
	return Pairs

##Calculation of ECIF
def GetECIF(PDB_protein, SDF_ligand, distance_cutoff=6.0):
	# Main function for the calculation of ECIF
	Pairs = GetPLPairs(PDB_protein, SDF_ligand, distance_cutoff=distance_cutoff)
	ECIF = [list(Pairs["ECIF_PAIR"]).count(x) for x in PossibleECIF]
	return ECIF

##Calculation of ELEMENTS
def GetELEMENTS(PDB_protein, SDF_ligand, distance_cutoff=6.0):
	# Function for the calculation of ELEMENTS
	Pairs = GetPLPairs(PDB_protein, SDF_ligand, distance_cutoff=distance_cutoff)
	ELEMENTS = [list(Pairs["ELEMENTS_PAIR"]).count(x) for x in PossibleELEMENTS]
	return ELEMENTS

###Calculation of Ligand descriptors
def GetRDKitDescriptors(SDF):
	# Function for the calculation of ligand descriptors
	mol = Chem.MolFromMolFile(SDF, sanitize=False)
	mol.UpdatePropertyCache(strict=False)
	Chem.GetSymmSSSR(mol)
	return DescCalc.CalcDescriptors(mol)

parser = argparse.ArgumentParser(description='calculate the ECIF')
parser.add_argument('-p', '--protein', dest='protein', required=True,
					help='Input protein file (.pdb)')
parser.add_argument('-l', '--ligand', dest='ligand', required=True,
					help='Input ligand file (.sdf)')
parser.add_argument('--prot_name', dest='prot_name', type=str,
					help='The name of the protein.')
parser.add_argument('--lig_name', dest='lig_name', type=str,
					help='The name of the ligand.')
parser.add_argument('--pdb_atom_keys_file', dest='pdb_atom_keys_file', default='./PDB_Atom_Keys.csv',
					help='The file recording the atom information of the protein. (default: "./PDB_Atom_Keys.csv")')
parser.add_argument('-d', '--desc_type', dest='desc_type', default='all', choices=['all','ecif','elem','lig'],
					help='Which descritpors to generate.')

args = parser.parse_args()
if args.prot_name is None:
	args.prot_name = os.path.basename(args.protein).split('_')[0]
if args.lig_name is None:
	args.lig_name = os.path.basename(args.ligand).split('_')[0]

Atom_Keys = pd.read_csv(args.pdb_atom_keys_file, sep=",")
if args.desc_type == 'ecif':
	final = GetECIF(args.protein, args.ligand, distance_cutoff=6.0)
elif args.desc_type == 'elem':
	final = GetELEMENTS(args.protein, args.ligand, distance_cutoff=6.0)
elif args.desc_type == 'lig':
	final = list(GetRDKitDescriptors(args.ligand))
else:
	final = GetECIF(args.protein, args.ligand, distance_cutoff=6.0) + GetELEMENTS(args.protein, args.ligand, distance_cutoff=6.0) + list(GetRDKitDescriptors(args.ligand))

df_out = pd.DataFrame([args.prot_name, args.lig_name, final]).T
df_out.columns = ['pdb_id', 'lig_id', 'desc']
df_out.to_csv('ecif_%s.csv'%args.desc_type)