#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed Oct 28 18:51:07 2020

@author: tamilieberman
"""

'''
straincraker.py reassigns reads from a kraken output file and produces a new kraken output file

It uses a purity threshold to reassign taxa to the taxonomic level at which the percent
of kmers assigned at a node or its descendents is above the specified the specified treshold
Kmers above the taxonomic level in consideration are not included in this calculation
Requires a precompiled tree, generated from straincracker_loadtree.py and nodes.dmp

Copyright (C) 2020 Tami Lieberman, tami@mit.edu

Required Parameters:
   -i,--infile X....................kraken output file, gzipped
   -t,--treefile X.....................straincraker treefile
   -o, --output X......................modified kraken report file
   -p, --purity X......................purity threshold

Example: python straincracker.py --t treematrix.npy --i Meta-F_Extract-Single_Ch_krakSeq.txt.gz --o Meta-F_Extract-Single_Ch_krakSeq.txt.cracked --p .99

#steps to all of straincraker
#make treematrix, krakendatabase format
#kraken
#strainkrak
#convert
#bracken
'''

# %% Run the actual program

#################################################################################

import numpy as np
import gzip
import argparse
import os
#################################################################################

# %%
def main():
    #Parse arguments
    parser = argparse.ArgumentParser(description='Reassign reads using treefile from strainkraker_loadtree.py')
    parser.add_argument('--i', metavar='infile', nargs='?', required=True, help='gzipped output from kraken')
    parser.add_argument('--o', metavar='outfile', nargs='?', required=True, help='destination output file')
    parser.add_argument('--t', metavar='treefile', nargs='?', required=True,  help='source tree file')
    parser.add_argument('--p', metavar='purity', type=float, required=True, nargs='?', help='purity threshold')
    args = parser.parse_args()
    strainCrack(args.t,args.i,args.o,args.p)

#%%

def read_in_results(kstrings):

    kcounts=dict({0:0, 1:0})

    for i,s in enumerate(kstrings):
        x=s.split(':')
        if x[0] != '|':
            if int(x[0]) in kcounts:
                kcounts[int(x[0])]+=int(x[1])
            else:
                kcounts[int(x[0])]=int(x[1])

    kcounts.pop(0) #ignore kmers alinging to levels 0 or 1
    kcounts.pop(1) #ignore kmers alinging to levels 0 or 1

    t=np.array(list(kcounts.keys()))
    k=np.array(list(kcounts.values()))

    return t,k

# %%

def classify_read(taxa,kmers, treem, purity_thresh,tlevels):

    found_taxa=taxa.copy()
    kmer_values=kmers.copy()
    taxa_categorized=0
    subtree=treem[found_taxa,:];

    num_levels=max(np.where(subtree>0)[1])

    #descend tree
    #calculate kmers at each level
    for l in range(1,num_levels+1):

        taxa_at_this_level_or_below=np.logical_and(kmer_values>0,tlevels>=l)

        #first try to classify at this level
        if np.sum(taxa_at_this_level_or_below) == 1:
            taxa_categorized=found_taxa[taxa_at_this_level_or_below][0]
            break
        elif np.sum(taxa_at_this_level_or_below) < 1:
            #go with previous assignment
            break
        else:
         #need to coninue classifying

            #find all taxa at this level that have descendents w kmers
            level_i_classification_for_taxa_found=subtree[:,l]
            taxa_at_this_level_w_evidence=np.unique(level_i_classification_for_taxa_found[level_i_classification_for_taxa_found>0])

            #are there 0, 1, or more taxa at this level?
            if taxa_at_this_level_w_evidence.size == 1:
                #just 1 taxa, pure by definitaion, continue going down tree
                taxa_categorized = taxa_at_this_level_w_evidence[0]
            elif taxa_at_this_level_w_evidence.size < 1:
                #report prev taxa
                print('Warning: somehow got here')
                break
            else: # taxa_at_this_level_w_evidence.size > 1:

                #check which is best path to go down, or report prev
                kmers_at_level=np.zeros(np.shape(taxa_at_this_level_w_evidence),dtype=int)
                for i,t in enumerate(taxa_at_this_level_w_evidence):
                    self_and_children_taxa=np.equal(level_i_classification_for_taxa_found,t);
                    kmers_at_level[i]=sum(kmer_values[self_and_children_taxa])

                purity=max(kmers_at_level)/sum(kmers_at_level)
                if purity > purity_thresh:

                    taxa_categorized=taxa_at_this_level_w_evidence[kmers_at_level==max(kmers_at_level)]
                    taxa_categorized = taxa_categorized[0]

                    #set all counts on paths not below this to 0
                    to_delete=np.not_equal(level_i_classification_for_taxa_found,taxa_categorized)
                    kmer_values[to_delete]=0

                else:
                    #report prev taxa
                    break

    return taxa_categorized

# %%
def strainCrack(treefile,infile,outfile,purity_threshold):

    #i mport tree structure
    tree = np.load(treefile)

    # make an array that says what taxonmic level each taxa is at, helpful for later
    taxa_nums=np.array([range(tree.shape[0])])
    taxa_exists=np.sum(np.equal(tree,taxa_nums.T)>0,1)
    taxonomic_levels=np.zeros(taxa_exists.shape,dtype=int)
    taxonomic_levels[taxa_exists>0]=np.array((np.nonzero(np.equal(tree,taxa_nums.T))))[1]

    # read in input and output file

    f=gzip.open(infile,"rt")
    of=open(outfile,"w")
    l=f.readline()

    # test
    changed=0
    tested=0
    while len(l) > 0 and l != "\n":
        line=l.strip().split()

        if line[0]=='C':
            taxa, kmers =read_in_results(line[4:])
            if kmers.size > 1:
                tested+=1
                new_classification=classify_read(taxa,kmers,tree,purity_threshold, taxonomic_levels[taxa])
                if int(line[2]) != new_classification:
                    changed+=1
                    line[2]=str(new_classification)
            newline1='\t'.join(map(str, line[:4]))
            newline2=' '.join(map(str, line[4:]))
            of.write(newline1 + '\t' + newline2 + '\n')
        else:
            of.write(l)

        l=f.readline()

    of.close()
    f.close()

    return changed, tested

if __name__ == "__main__":
    main()


import numpy as np
import copy

from mahjong.ai.monte.game_state import GameState
from mahjong.rule.algo.judge_chain_maker import JudgeChainMaker
from mahjong.rule.algo.winchecker import WinChecker
from mahjong.rule.algo.hand_calculator import HandCalculator
from mahjong.ai.constant import constant
import logging
from mahjong.log.logger import logger
from mahjong.rule.util.tile_convert import TilesConv
from mahjong.rule.util.utils import check_ready_to_win, check_ready_to_touch

S0 = 10
S1 = 6
S2 = 6
K0 = 2
K1 = 1
K2 = 1
D0 = 3

ves_ai = JudgeChainMaker(1)

WEIGHT_FITST = 1
WEIGHT_SECOND = 0.1
WEIGHT_ZIMO = 1
WEIGHT_TOUCH_GANG = 0.7
WEIGHT_MELD3_GANG = 0.7
WEIGHT_DISCARD = 0.5
WEIGHT_FANGPAO = 1
WEIGHT_PASS_FANGPAO = 0.6
WEIGHT_PENG = 0.5
WEIGHT_BEAT = 0.4
WEIGHT_GANG = 0.6

class Attack:
    def __init__(self, s0=S0, s1=S1, s2=S2, k0=K0, k1=K1, k2=K2):
        self.s0 = s0
        self.s1 = s1
        self.s2 = s2
        self.k0 = k0
        self.k1 = k1
        self.k2 = k2

    @staticmethod
    def tile_remain_num(tile, game_state, index):
        s = game_state
        return 4 - s.hands[index][tile] - s.get_discards(tile)

    @staticmethod
    def get_first_level_ts(t):
        if t >= 0:
            p = t % 9
            if p * (p - 8):
                return [t + 1, t - 1]
            else:
                return [t + 1 - 1 * p // 4]
        return None

    @staticmethod
    def get_second_level_ts(t):
        if t >= 0:
            p = t % 9
            if p * (p - 8) * (p - 1) * (p - 7):
                return [t + 2, t - 2]
            else:
                if p < 2:
                    return [t + 2]
                else:
                    return [t - 2]
        return None

    @staticmethod
    def think_beat(node, index):
        expects = []
        for n in node.children:
            if n.reason == "fangpao":
                result = n.game_result
                score = result * WEIGHT_FANGPAO
                expects.append(score)
            elif n.gang_index >= 0:
                expect = Attack.think_expectation(n.game_state, index) * WEIGHT_GANG
                if expect < 0:
                    expect = 0
                expects.append(expect)
            elif n.peng_index >= 0:
                expect = Attack.think_expectation(n.game_state, index) * WEIGHT_PENG
                if expect < 0:
                    expect = 0
                expects.append(expect)
            else:
                expect = Attack.think_expectation(n.game_state, index) * WEIGHT_BEAT
                if expect < 0:
                    expect = 0
                expects.append(expect)

        expect_prob = np.array(Attack.expects2_probability(expects))

        rps = []
        for n in node.children:
            sum = 0
            for t in range(0, 18):
                sum += Attack.rp_t1(t, index, n.game_state)
            if n.reason == "fangpao":
                rps.append(sum * WEIGHT_FANGPAO)
            elif n.gang_index >= 0:
                rps.append(sum * WEIGHT_GANG)
            elif n.peng_index >= 0:
                rps.append(sum * WEIGHT_PENG)
            else:
                rps.append(sum * WEIGHT_BEAT)
        beat_rps_probability = np.array(Attack.weights2_probability(rps))
        final = expect_prob * 0.7 + beat_rps_probability * 0.3
        for i in range(0, len(node.children)):
            node.children[i].set_beat_probability(final[i])

    @staticmethod
    def think_after_touch(node):
        expects = []
        rps = []
        for n in node.children:
            if n.reason == "zimo":
                result = n.game_result
                score = result * WEIGHT_ZIMO
                expects.append(score)
                sum_rp = 0
                for t in range(0, 18):
                    sum_rp += Attack.rp_t1(t, node.get_turn(), n.game_state)
                rps.append(sum_rp * WEIGHT_ZIMO)
            elif n.touch_gang_index >= 0:
                expect = Attack.think_expectation(n.game_state, node.get_turn()) * WEIGHT_TOUCH_GANG
                if expect < 0:
                    expect = 0
                expects.append(expect)
                sum_rp = 0
                for t in range(0, 18):
                    sum_rp += Attack.rp_t1(t, node.get_turn(), n.game_state)
                rps.append(sum_rp * WEIGHT_TOUCH_GANG)
            elif n.melds3_gang_index >= 0:
                expect = Attack.think_expectation(n.game_state, node.get_turn()) * WEIGHT_MELD3_GANG
                if expect < 0:
                    expect = 0
                expects.append(expect)
                sum_rp = 0
                for t in range(0, 18):
                    sum_rp += Attack.rp_t1(t, node.get_turn(), n.game_state)
                rps.append(sum_rp * WEIGHT_MELD3_GANG)
            else:
                expect = Attack.think_expectation(n.game_state, node.get_turn()) * WEIGHT_DISCARD
                if expect < 0:
                    expect = 0
                expects.append(expect)
                sum_rp = 0
                for t in range(0, 18):
                    sum_rp += Attack.rp_t1(t, node.get_turn(), n.game_state)
                rps.append(sum_rp * WEIGHT_DISCARD)

        expect_prob = np.array(Attack.expects2_probability(expects))
        peng_rps_probability = np.array(Attack.weights2_probability(rps))
        final = expect_prob * 0.7 + peng_rps_probability * 0.3
        for i in range(0, len(node.children)):
            node.children[i].set_after_touch_probability(final[i])

    @staticmethod
    def think_expectation(game_state: GameState, index):
        """
        计算玩家index的手牌胡牌期望,要求手牌数目是%3 == 1
        :param game_state:
        :param index:
        :return:
        """
        check = check_ready_to_touch(game_state.hands[index])
        if not check:
            logger().error("hand length is not ready to touch")
            return -1
        hand = copy.deepcopy(game_state.hands[index])
        chains = ves_ai.calc_effective_cards(hand, 0)
        expect = 1
        logger().debug("Length of chains:{}".format(len(chains)))
        for chain in chains:
            # 计算每种胡牌链的期望
            expect_per_chain = 0.0
            pairs = chain.touchPlayPairs
            prob_pair = 1.0
            weight = WEIGHT_SECOND
            if len(pairs) == 0:
                weight = WEIGHT_FITST

            for pair in pairs:
                need = pair.touch_tile()
                remain = game_state.get_remain(need, index)
                prob_pair *= remain
                prob_pair /= constant.PLAYER_NUM
                if prob_pair > 1:
                    prob_pair = 1

            for dh in chain.drawHands:
                # 计算一个链种每种胡牌的期望
                chain.hand[dh] += 1
                meld3 = game_state.melds_3[index]
                meld4 = game_state.melds_4[index]
                cost = HandCalculator.estimate_max_score(chain.hand, dh, meld3, meld4)
                chain.hand[dh] -= 1
                remain_dh = game_state.get_remain(dh, index)
                prob_dh = remain_dh / constant.PLAYER_NUM
                if prob_dh > 1:
                    prob_dh = 1
                expect_per_chain += prob_pair * prob_dh * cost * weight
            expect += (expect_per_chain / len(chain.drawHands))
        return expect

    @staticmethod
    def expects2_probability(expects):
        """
        Calculate probability of card to be played according to expectation array.
        The array represents the expectation of score for hands playing the index's card.
        :param expects:
        :return:
        """
        tile_expects = copy.deepcopy(expects)
        sum1 = sum(tile_expects)
        for index in range(0, len(tile_expects)):
            if tile_expects[index] > 0:
                tile_expects[index] = tile_expects[index] / sum1
        return tile_expects

    @staticmethod
    def weights2_probability(tile_weights_in):
        """
        邻近权重越大概率越大
        :param tile_weights_in:输入邻近权重
        :return: 概率，归一化
        """
        tile_weights = copy.deepcopy(tile_weights_in)
        sum1 = sum(tile_weights)
        for index in range(0, len(tile_weights)):
            tile_weights[index] = tile_weights[index] / sum1
        return tile_weights

    @staticmethod
    def reverse_weights2_probability(tile_weights_in):
        """
        邻近权重越大概率越小，邻近权重越小概率越大
        :param tile_weights_in:输入邻近权重
        :return: 概率，归一化
        """
        tile_weights = copy.deepcopy(tile_weights_in)
        sum1 = sum(tile_weights)
        for index in range(0, len(tile_weights)):
            if tile_weights[index] > 0:
                tile_weights[index] = sum1 / tile_weights[index]
        sum2 = sum(tile_weights)
        for index in range(0, len(tile_weights)):
            if tile_weights[index] > 0:
                tile_weights[index] = tile_weights[index] / sum2
        return tile_weights

    @staticmethod
    def rp_t1(t, index, game_state):
        """
        计算邻近的牌的个数和权重，权重越高邻近的牌越多
        :param index: 计算玩家{index}的t牌邻近权重
        :param t:牌值
        :param game_state:牌局
        :return: 返回邻近牌的权重,邻近牌越多权重越高,如果手牌中没有该牌返回0
        """
        s = game_state
        p_t = 0
        sn_0 = s.hands[index][t]
        if sn_0:
            sn_1, sn_2, kn_0, kn_1, kn_2 = 0, 0, 0, 0, 0
            kn_0 = Attack.tile_remain_num(t, game_state, index)
            # 计算相邻的牌的个数
            for t1 in Attack.get_first_level_ts(t):
                sn_1 += s.hands[index][t1]
                kn_1 += Attack.tile_remain_num(t1, game_state, index)
            # 计算相隔一张牌的个数
            for t2 in Attack.get_second_level_ts(t):
                sn_2 += s.hands[index][t2]
                kn_2 += Attack.tile_remain_num(t2, game_state, index)
            p_t = sn_0 * S0 + (sn_1 * S1) + (sn_2 * S2) + kn_0 * K0 + (kn_1 * K1) + (
                    kn_2 * K2) + sn_0 * sn_0 * sn_0
        p_t /= TilesConv.tiles18_count(s.hands[index])
        return p_t


import binascii
import glob,os
import numpy as np
from keras.preprocessing.image import ImageDataGenerator, load_img, img_to_array
from keras.models import Sequential, load_model
import time
import matplotlib.pyplot as plt
model_path = './models/weights-improvement-10-0.92.hdf5'
model = load_model(model_path)
# model.load_weights(model_weights_path)
img_width, img_height = 150, 150

def predict(file):
  x = load_img(file, target_size=(img_width,img_height))
#   plt.imshow(x)
#   plt.show()
  x = img_to_array(x)
  x = x.reshape((1,) + x.shape)
  x /= 255
  array = model.predict(x)
  result = array[0,0]

  answer = np.argmax(result)
  if result < 0.5:
    print(file+" Predicted to be a cat")
  elif result > 0.5:
    print(file+" Predicted to be a dog")
  return answer

print("Enter path you want to scan:")
strDrive=input()
print()
print("File Name\t\t|\t\tMasqueraded?")
print("_____________________________________________________")
masqueraded=[]
i =0
for path in glob.glob(strDrive+"**/*", recursive=True):
    if os.path.isfile(path):
        name=os.path.basename(path)
        i += 1
        with open(path, 'rb') as f:
            content = f.read()
        var=(binascii.hexlify(content).decode("utf-8"))
        if name.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', )):
            print (f"{name}\t\t|\t\tNo")
        elif  var.startswith("ffd8ffdb"):
            print (f"{name}\t\t|\t\tYes  (JPG)")
            masqueraded.append(path)
#             predict(path)
        elif var.startswith("ffd8ffe000104a4649460001"):
            print (f"{name}\t\t|\t\tYes (JPG)")
            masqueraded.append(path)
#             predict(path)
        elif var.startswith("ffd8ffee"):
            print (f"{name}\t\t|\t\tYes (JPG)")
            masqueraded.append(path)
#             predict(path)
        elif var.startswith("89504e470d0a1a0a"):
            print (f"{name}\t\t|\t\tYes (PNG)")
            masqueraded.append(path)
#             predict(path)
        elif var.startswith("ffd8ffe1????457869660000"):
            print (f"{name}\t\t|\t\tYes  (JPG)")
            masqueraded.append(path)
#             predict(path)
        elif var.startswith("424D"):
            print (f"{name}\t\t|\t\tYes")
            masqueraded.append(path)
#             predict(path)
        elif var.startswith("474946383761"):
            print (f"{name}\t\t|\t\tYes (GIF)")
            masqueraded.append(path)
#             predict(path)
        else:
            print (f"{name}\t\t|\t\tNo")
print(f"{i} files scanned successfully.")
print()
print(f"{len(masqueraded)} masqueraded files scanned successfully.")
print("list of all masqueraded images:")
print()
for j in masqueraded:
    predict(j)


# Prediction interface for Cog ⚙
# Reference: https://github.com/replicate/cog/blob/main/docs/python.md

import os
os.environ['TRANSFORMERS_CACHE'] = '/src/cache'
os.environ['XDG_CACHE_HOME'] = '/src/cache'

from torch import nn
import numpy as np
import torch
import torch.nn.functional as nnf
import sys
from typing import Tuple, List, Union, Optional
from transformers import (
    GPT2Tokenizer,
    GPT2LMHeadModel,
    AdamW,
    get_linear_schedule_with_warmup,
)
import skimage.io as io
import PIL.Image

import cog
# import torch

N = type(None)
V = np.array
ARRAY = np.ndarray
ARRAYS = Union[Tuple[ARRAY, ...], List[ARRAY]]
VS = Union[Tuple[V, ...], List[V]]
VN = Union[V, N]
VNS = Union[VS, N]
T = torch.Tensor
TS = Union[Tuple[T, ...], List[T]]
TN = Optional[T]
TNS = Union[Tuple[TN, ...], List[TN]]
TSN = Optional[TS]
TA = Union[T, ARRAY]

WEIGHTS_PATHS = {
    "coco": "neuralhash.same.pt",
}

D = torch.device
CPU = torch.device("cpu")

class Predictor(cog.Predictor):
    def setup(self):
        """Load the model into memory to make running multiple predictions efficient"""
        self.device = torch.device("cuda")
        self.tokenizer = GPT2Tokenizer.from_pretrained("gpt2")

        self.models = {}
        self.prefix_length = 10
        for key, weights_path in WEIGHTS_PATHS.items():
            model = ClipCaptionModel(self.prefix_length)
            model.load_state_dict(torch.load(weights_path, map_location=CPU))
            model = model.eval()
            model = model.to(self.device)
            self.models[key] = model

    @cog.input("hash", type=str, help="Apple Neuralhash")
    @cog.input(
        "model",
        type=str,
        default="coco",
        help="Model to use",
    )
    @cog.input(
        "use_beam_search",
        type=bool,
        default=False,
        help="Whether to apply beam search to generate the output text",
    )
    def predict(self, hash, model, use_beam_search):
        """Run a single prediction on the model"""
        model = self.models[model]

        with torch.no_grad():
            binary = bin(int(hash, 16))[2:].zfill(96)
            print(hash, len(binary), binary)

            embedding = torch.zeros(1, 96, dtype=torch.float32)
            for i,b in enumerate(binary):
                if b == "1":
                    embedding[0][i] = 2
            embedding = torch.sub(embedding, 1)

            prefix = embedding.to(self.device, dtype=torch.float32)
            print("prefix", prefix.size(), prefix)
            # prefix_embed = model.clip_project(prefix)
            # prefix_embed = prefix_embed.reshape(1, self.prefix_length, -1)
            # print(prefix_embed)
        if use_beam_search:
            return generate_beam(model, self.tokenizer, embed=prefix_embed)[0]
        else:
            return generate2(model, self.tokenizer, prefix=prefix)

class MLP(nn.Module):

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.model(x)

    def __init__(self, sizes: Tuple[int, ...], bias=True, act=nn.Tanh):
        super(MLP, self).__init__()
        layers = []
        for i in range(len(sizes) - 1):
            layers.append(nn.Linear(sizes[i], sizes[i + 1], bias=bias))
            if i < len(sizes) - 2:
                layers.append(act())
        self.model = nn.Sequential(*layers)

class ClipCaptionModel(nn.Module):

    def get_dummy_token(self, batch_size: int, device: torch.device) -> torch.Tensor:
        return torch.zeros(batch_size, self.prefix_length, dtype=torch.int64, device=device)

    def forward(self, tokens: torch.Tensor, prefix: torch.Tensor, mask: Optional[torch.Tensor] = None,
                labels: Optional[torch.Tensor] = None):
        embedding_text = self.gpt.transformer.wte(tokens)
        prefix_projections = self.clip_project(prefix).view(-1, self.prefix_length, self.gpt_embedding_size)
        embedding_cat = torch.cat((prefix_projections, embedding_text), dim=1)
        if labels is not None:
            dummy_token = self.get_dummy_token(tokens.shape[0], tokens.device)
            labels = torch.cat((dummy_token, tokens), dim=1)
        out = self.gpt(inputs_embeds=embedding_cat, labels=labels, attention_mask=mask)
        return out

    def __init__(self, prefix_length: int, prefix_size: int = 96):
        super(ClipCaptionModel, self).__init__()
        self.prefix_length = prefix_length
        self.gpt = GPT2LMHeadModel.from_pretrained('gpt2')
        self.gpt_embedding_size = self.gpt.transformer.wte.weight.shape[1]
        self.clip_project = MLP((prefix_size, (self.gpt_embedding_size * prefix_length) // 2,
                                 self.gpt_embedding_size * prefix_length))

class ClipCaptionPrefix(ClipCaptionModel):

    def parameters(self, recurse: bool = True):
        return self.clip_project.parameters()

    def train(self, mode: bool = True):
        super(ClipCaptionPrefix, self).train(mode)
        self.gpt.eval()
        return self

def generate2(
        model,
        tokenizer,
        tokens=None,
        prompt=None,
        prefix=None,
        entry_count=1,
        entry_length=67,  # maximum number of words
        top_p=0.8,
        temperature=1.,
        stop_token: str = '\n',
):
    model.eval()
    generated_num = 0
    generated_list = []
    stop_token_index = tokenizer.encode(stop_token)[0]
    device = next(model.parameters()).device

    with torch.no_grad():
        for entry_idx in range(entry_count):
            generated = torch.tensor([[31383]]).to(device)

            for i in range(entry_length):
                outputs = model(generated, prefix)
                logits = outputs.logits
                logits = logits[:, -1, :]

                next_token = torch.argmax(logits, -1).unsqueeze(0)
                # next_token_embed = model.gpt.transformer.wte(next_token)
                if tokens is None:
                    tokens = next_token
                else:
                    tokens = torch.cat((tokens, next_token), dim=1)
                # print(generated.shape, next_token.shape, generated, next_token)
                generated = torch.cat((generated, next_token), dim=1)
                # print(generated)
                if stop_token_index == next_token.item():
                    break

            output_list = list(tokens.squeeze().cpu().numpy())
            output_text = tokenizer.decode(output_list)
            generated_list.append(output_text)

    return generated_list[0]


import torch
import os
# dir_list=os.listdir("./total/")
# print(dir_list)
import math
import numpy as np
import glob
import argparse

# import matplotlib as mpl
from matplotlib.legend_handler import HandlerLine2D, HandlerTuple
import matplotlib
matplotlib.use('TkAgg')
import matplotlib.pyplot as plt

mode='ImageNet'
save_suffix=''
parser = argparse.ArgumentParser(description='Runs GADA')
parser.add_argument('-c', '--config', type=str, default='_3DDFA_V2/configs/mb1_120x120.yml')
parser.add_argument('-f', '--img_fp', type=str, default='_3DDFA_V2/examples/inputs/trump_hillary.jpg')
parser.add_argument('-m', '--mode', type=str, default='cpu', help='gpu or cpu mode')
parser.add_argument('-o', '--opt', type=str, default='pncc',
                    choices=['2d_sparse', '2d_dense', '3d', 'depth', 'pncc', 'uv_tex', 'pose', 'ply', 'obj'])

parser.add_argument('--onnx', action='store_true', default=True)
parser.add_argument('--model', type=int, default=2, help='index of configurations')
parser.add_argument('--dict_model', type=int, default=3, help='index of configurations')

parser.add_argument('--attack', type=str,  default='HSJA', help='attack method')
parser.add_argument('--defense', type=str, default='gaussian', help='defense method')

parser.add_argument('--dataset', type=str, default='LFW', help='dataset')
parser.add_argument('--result_dir', type=str, default='results', help='directory for saving results')
parser.add_argument('--sampled_image_dir', type=str, default='save', help='directory to cache sampled images')

parser.add_argument('--batch_size', type=int, default=1, help='number of image samples')

parser.add_argument('--max_num_queries', type=int, default=10000, help='number of image samples')
parser.add_argument('--log_interval', type=int, default=1000, help='number of image samples')
parser.add_argument('--num_runs', type=int, default=500, help='number of image samples')

parser.add_argument('--seed', type=int, default=1234, help='number of image samples')
parser.add_argument('--attack_batch_size', type=int, default=100, help='number of image samples')

parser.add_argument('--resume', type=bool, default=False, help='resume attack')
parser.add_argument('--targeted', type=bool, default=True, help='perform targeted attack')
parser.add_argument('--epsilon', type=float, default=0.2, help='step size per iteration')

parser.add_argument('--save_suffix', type=str, default='', help='suffix appended to save file')

args = parser.parse_args()
def get_pth_file_all(directory_path,ext=".pth"):
    return glob.glob(os.path.abspath(directory_path)+"/*"+ext)
def mean(l, logarithm=False):
    if logarithm:
        l_lg = [np.log(x) for x in l]
        return np.exp(sum(l_lg) / len(l_lg))
    else:
        return sum(l) / len(l)

if __name__=='__main__':

    file_list = get_pth_file_all('./')  # full path
    name_list = [os.path.split(x)[1] for x in file_list]
    print(name_list)

    def load_and_plot_linf(file_name, start_img_idx=0, end_img_idx=500, start_iter=100, end_iter=10000, threshold=0.05,
                      logarithm=False):
        log_data = torch.load(file_name)
        total_num_queries = log_data['total_num_queries']
        total_log_l_2 = log_data['total_log_l_2']
        total_log_l_inf = log_data['total_log_l_inf']
        total_log_is_adversarial = log_data['total_log_is_adversarial']
        total_log_dist = log_data['total_log_dist']
        total_log_last_adv_imgs = log_data['total_log_last_adv_imgs']
        total_log_best_l_2_adv_imgs = log_data['total_log_best_l_2_adv_imgs']
        total_log_best_l_inf_adv_imgs = log_data['total_log_best_l_inf_adv_imgs']
        total_log_best_l_2 = log_data['total_log_best_l_2']
        total_log_best_l_inf = log_data['total_log_best_l_inf']

        end_idx = log_data['upper_idx'] // args.batch_size
        print(file_name, end_idx)
        end_idx = min(end_idx, end_img_idx)
        best_log_l_inf = torch.zeros_like(total_log_l_inf[start_img_idx:end_idx])
        best_val = torch.ones(end_idx - start_img_idx) * 5e1
        for i in range(end_iter):
            best_val = torch.min(best_val, 5e1 * (
                        1 - total_log_is_adversarial[start_img_idx:end_idx, i]) + total_log_is_adversarial[
                                                                                start_img_idx:end_idx,
                                                                                i] * total_log_l_inf[
                                                                                start_img_idx:end_idx, i])
            best_log_l_inf[:, i] = best_val
        avg_X = torch.arange(start_iter, end_iter)
        if logarithm == True:
            avg_Y = torch.exp(torch.mean(torch.log(best_log_l_inf[:, start_iter:end_iter]), 0))
        else:
            avg_Y = torch.mean(best_log_l_inf[:, start_iter:end_iter], 0)
        target_idx = [1000, 2000, 5000, 10000]

        norms = torch.mean(best_log_l_inf, 0)[torch.LongTensor(target_idx) - 1].numpy()

        thresholds = [0.05, 0.03]
        Q = np.zeros(len(thresholds))
        for i, t in enumerate(thresholds):
            best_log_l_inf_t = best_log_l_inf.clone()
            best_log_l_inf_t[best_log_l_inf <= t] = 0
            best_log_l_inf_t[best_log_l_inf > t] = 1
            Q[i] = ((torch.sum(best_log_l_inf_t) / (end_idx - start_img_idx)).item())
        return avg_X, avg_Y, norms, Q

    def load_and_plot(file_name,start_img_idx=0,end_img_idx=500,start_iter=100,end_iter=10000, threshold=4,logarithm=True):
        log_data=torch.load(file_name)
        total_num_queries = log_data['total_num_queries']
        total_log_l_2 = log_data['total_log_l_2']
        total_log_l_inf = log_data['total_log_l_inf']
        total_log_is_adversarial = log_data['total_log_is_adversarial']
        total_log_dist = log_data['total_log_dist']
        total_log_last_adv_imgs = log_data['total_log_last_adv_imgs']
        total_log_best_l_2_adv_imgs = log_data['total_log_best_l_2_adv_imgs']
        total_log_best_l_inf_adv_imgs = log_data['total_log_best_l_inf_adv_imgs']
        total_log_best_l_2 = log_data['total_log_best_l_2']
        total_log_best_l_inf = log_data['total_log_best_l_inf']

        end_idx = log_data['upper_idx'] // args.batch_size
        print(file_name,end_idx)
        end_idx=min(end_idx,end_img_idx)
        best_log_l2=torch.zeros_like(total_log_l_2[start_img_idx:end_idx])
        best_val=torch.ones(end_idx-start_img_idx)*5e1
        for i in range(end_iter):
            best_val=torch.min(best_val,5e1*(1-total_log_is_adversarial[start_img_idx:end_idx,i])+total_log_is_adversarial[start_img_idx:end_idx,i]*total_log_l_2[start_img_idx:end_idx,i])
            best_log_l2[:,i]=best_val
        avg_X=torch.arange(start_iter,end_iter)
        if logarithm==True:
            avg_Y=torch.exp(torch.mean(torch.log(best_log_l2[:,start_iter:end_iter]),0))
        else:
            avg_Y=torch.mean(best_log_l2[:,start_iter:end_iter],0)
        target_idx=[1000,2000,5000,10000]

        norms=torch.mean(best_log_l2,0)[torch.LongTensor(target_idx)-1].numpy()

        thresholds=[4,2]
        Q=np.zeros(len(thresholds))
        for i,t in enumerate(thresholds):
            best_log_l2_t=best_log_l2.clone()
            best_log_l2_t[best_log_l2<=t]=0
            best_log_l2_t[best_log_l2>t]=1
            Q[i]=((torch.sum(best_log_l2_t)/(end_idx-start_img_idx)).item())
        return avg_X,avg_Y,norms,Q

    def load_and_plot_SD(file_name, start_img_idx=0, end_img_idx=500, start_iter=100, end_iter=10000, threshold=4,
                      logarithm=True):
        log_data = torch.load(file_name)
        total_num_queries = log_data['total_num_queries']
        total_log_l_2 = log_data['total_log_l_2']
        total_log_l_inf = log_data['total_log_l_inf']
        total_log_is_adversarial = log_data['total_log_is_adversarial']
        total_log_dist = log_data['total_log_dist']
        total_log_last_adv_imgs = log_data['total_log_last_adv_imgs']
        total_log_best_l_2_adv_imgs = log_data['total_log_best_l_2_adv_imgs']
        total_log_best_l_inf_adv_imgs = log_data['total_log_best_l_inf_adv_imgs']
        total_log_k_avg_dist = log_data['total_log_k_avg_dist']
        total_log_best_l_inf = log_data['total_log_best_l_inf']

        end_idx = log_data['upper_idx'] // args.batch_size
        print(file_name, end_idx)
        end_idx = min(end_idx, end_img_idx)
        best_log_l2 = torch.zeros_like(total_log_l_2[start_img_idx:end_idx])
        best_val = torch.ones(end_idx - start_img_idx) * 5e1
        for i in range(end_iter):
            best_val = torch.min(best_val, 5e1 * (
                        1 - total_log_is_adversarial[start_img_idx:end_idx, i]) + total_log_is_adversarial[
                                                                                start_img_idx:end_idx,
                                                                                i] * total_log_l_2[
                                                                                start_img_idx:end_idx, i])
            best_log_l2[:, i] = best_val
        avg_X = torch.arange(start_iter, end_iter)
        avg_Y = best_log_l2[10, start_iter:end_iter]

        target_idx = [2000, 5000, 10000]

        norms = torch.mean(best_log_l2, 0)[torch.LongTensor(target_idx) - 1].numpy()

        total_log_k_avg_dist_t=total_log_k_avg_dist.clone()
        total_log_k_avg_dist_t[total_log_k_avg_dist < 0.002] = 1
        total_log_k_avg_dist_t[total_log_k_avg_dist >= 0.002] = 0

        avg_k_avg_dist=total_log_k_avg_dist[10,start_iter:end_iter]
        Q=torch.mean(torch.sum(total_log_k_avg_dist_t,1)).item()

        return avg_X, avg_Y, avg_k_avg_dist,norms, Q

    def load_and_plot_Dist(file_name, start_img_idx=0, end_img_idx=500, start_iter=100, end_iter=10000, threshold=4,
                      logarithm=True):
        log_data = torch.load(file_name)
        total_num_queries = log_data['total_num_queries']
        total_log_l_2 = log_data['total_log_l_2']
        total_log_l_inf = log_data['total_log_l_inf']
        total_log_is_adversarial = log_data['total_log_is_adversarial']
        total_log_dist = log_data['total_log_dist']
        total_log_last_adv_imgs = log_data['total_log_last_adv_imgs']
        total_log_best_l_2_adv_imgs = log_data['total_log_best_l_2_adv_imgs']
        total_log_best_l_inf_adv_imgs = log_data['total_log_best_l_inf_adv_imgs']
        total_log_best_l_inf = log_data['total_log_best_l_inf']

        end_idx=log_data['upper_idx']
        end_idx = min(end_idx, end_img_idx)
        print(file_name, end_idx)
        best_log_l2 = torch.zeros_like(total_log_l_2[start_img_idx:end_idx,:end_iter])
        best_val = torch.ones(end_idx - start_img_idx) * 5e1
        for i in range(end_iter):
            best_val = torch.min(best_val, 5e1 * (
                        1 - total_log_is_adversarial[start_img_idx:end_idx, i]) + total_log_is_adversarial[
                                                                                start_img_idx:end_idx,
                                                                                i] * total_log_l_2[
                                                                                start_img_idx:end_idx, i])
            best_log_l2[:, i] = best_val

        return best_log_l2
    def load_and_plot_IMG(file_name, start_img_idx=0, end_img_idx=500, start_iter=100, end_iter=10000, idx=0, title=False,first=0):
        log_data = torch.load(file_name)
        total_num_queries = log_data['total_num_queries']
        total_log_l_2 = log_data['total_log_l_2']
        total_log_l_inf = log_data['total_log_l_inf']
        total_log_is_adversarial = log_data['total_log_is_adversarial']
        total_log_dist = log_data['total_log_dist']
        total_log_last_adv_imgs = log_data['total_log_last_adv_imgs']
        total_log_best_l_2_adv_imgs = log_data['total_log_best_l_2_adv_imgs']
        total_log_best_l_inf_adv_imgs = log_data['total_log_best_l_inf_adv_imgs']
        total_log_best_l_inf = log_data['total_log_best_l_inf']

        end_idx = log_data['upper_idx'] // args.batch_size
        print(file_name, end_idx)
        end_idx = min(end_idx, end_img_idx)
        best_log_l2 = torch.zeros_like(total_log_l_2[idx])
        best_val = torch.ones(1) * 5e1
        for i in range(end_iter):
            best_val = torch.min(best_val, 5e1 * (
                        1 - total_log_is_adversarial[idx, i]) + total_log_is_adversarial[
                                                                                idx,
                                                                                i] * total_log_l_2[
                                                                                idx, i])
            best_log_l2[i] = best_val

        left, width = .25, .5
        bottom, height = .25, .5
        right = left + width
        top = bottom + height
        #fig, ax = plt.subplots(figsize=(10, 1), tight_layout=True)
        im_idx=[0,1,2,5,9]

        fig, axs = plt.subplots(nrows=1, ncols=len(im_idx), gridspec_kw={'wspace': 0, 'hspace': 0},
                                   squeeze=True, tight_layout=True)
        for i,k in enumerate(im_idx):
            if first==0:
                axs[i].text(0.5 * (left + right), 1+0.1, '%dK'%(k+1),
                        horizontalalignment='center',
                        verticalalignment='center',
                        fontsize=11, color='black',
                        transform=axs[i].transAxes)
            batch = total_log_best_l_2_adv_imgs[idx, k]
            x_adv_i = np.transpose(batch.numpy(), [1, 2, 0])
            #axs[i].set_title('%.2f'%best_log_l2[(k+1)*1000-1].item())
            axs[i].text(0.5 * (left + right), -0.1, '%.2f'%best_log_l2[(k+1)*1000-1].item(),
                    horizontalalignment='center',
                    verticalalignment='center',
                    fontsize=11, color='black',
                    transform=axs[i].transAxes)

            axs[i].axis("off")
            axs[i].imshow(x_adv_i)
        # build a rectangle in axes coords

        axs[0].text(-0.1, 0.5 * (bottom + top), title,
                horizontalalignment='right',
                verticalalignment='center',
                rotation='vertical',
                    fontsize=11, color='black',
                transform=axs[0].transAxes)

    if False:
        img_indices=[190,258,298,430,455]
        for img_idx in img_indices:
            dt='CPLFW'
            targeted='1'
            attacks = ['SO','EA']
            model = 'IR_50'
            defense = 'gaussian'
            colors = ['#1f77b4',
                      '#ff7f0e',
                      '#2ca02c',
                      '#d62728',
                      '#9467bd',
                      '#8c564b',
                      '#e377c2',
                      '#7f7f7f',
                      '#bcbd22',
                      '#17becf',
                      '#1a55FF']
            # colors=['tab:blue','tab:red','g','tab:purple','tab:red','tab:cyan','tab:brown','tab:pink']

            PLOT_INFO = []
            for i, at in enumerate(attacks):
                fname = dt + '_500_' + targeted + '_' + at + '_' + model + '_' + defense + '_.pth'
                if os.path.exists(fname):
                    PLOT_INFO.append((fname, colors[i], at, [0, 500, 1000, 10000], ''))
            fig, ax = plt.subplots(figsize=(6, 4), tight_layout=True)
            avg_X = []
            avg_y = []
            plot_info = []
            labels = []
            for i, logs in enumerate(PLOT_INFO):
                at = logs[2]
                load_and_plot_IMG(logs[0], start_img_idx=logs[3][0], end_img_idx=logs[3][1],
                                  start_iter=logs[3][2], end_iter=logs[3][3], idx=img_idx, title=at, first=i)
                plt.savefig(dt + '_500_' + targeted + '_' + model + '_IMG_' + str(img_idx) + '_'+ at  + '.pdf', dpi=300, bbox_inches='tight')
                plt.close()

    #
    PLOT_INFO = [
    ('CPLFW_500_0_EA_IR_50_gaussian_.pth', 'm', 'EA',[0,500,1000,10000],''),
    ('CPLFW_500_0_EAGD_IR_50_gaussian_.pth', 'r', 'EAGD',[0,500,1000,10000],''),
    ]
    fig, ax = plt.subplots(figsize=(6, 4), tight_layout=True)
    avg_X = []
    avg_y = []
    labels = []
    for i, logs in enumerate(PLOT_INFO):
        l2_norm = load_and_plot_Dist(logs[0], start_img_idx=logs[3][0], end_img_idx=logs[3][1],
                                             start_iter=logs[3][2], end_iter=logs[3][3], logarithm=False)

        if i==1:
            print(torch.topk(l2_norm_o[:,-1]-l2_norm[:,-1],30,0))
            #print(torch.sum((l2_norm_o[:,-1]-l2_norm[:,-1])>0))
        l2_norm_o=l2_norm.clone()

    ##############################################
    if False:
        print('SD Graph')
        datasets = ['LFW']  # 'LFW',
        attacks = ['EA','EAR','EAG','EAGR']  #'SO', 'HSJA','GD',
        model = 'IR_50'
        defense = 'SD'
        colors = ['#1f77b4',
                  '#ff7f0e',
                  '#2ca02c',
                  '#d62728',
                  '#9467bd',
                  '#8c564b',
                  '#e377c2',
                  '#7f7f7f',
                  '#bcbd22',
                  '#17becf',
                  '#1a55FF']
        PLOT_INFO = []
        for i, at in enumerate(attacks):
            fname = 'LFW_100_0_' + at + '_' + model + '_' + defense + '_.pth'
            if os.path.exists(fname):
                PLOT_INFO.append((fname, colors[i], at, [0, 100, 500, 10000], ''))
        fig, ax1 = plt.subplots(figsize=(6, 4), tight_layout=True)

        ax1.set_xlabel('# Queries', fontsize=12)
        ax1.set_ylabel('$\ell_2$ norm of perturbation', fontsize=12)
        ax1.tick_params(axis='y')
        ax2 = ax1.twinx()  # instantiate a second axes that shares the same x-axis
        ax2.set_ylabel('k-NN dist in stateful detection', fontsize=12)
        ax2.tick_params(axis='y')
        avg_X = []
        avg_y = []
        plot_info = []
        labels = []
        legends=[]
        legends_labels=[]
        for i, logs in enumerate(PLOT_INFO):
            x, y, k_avg_dist,norms,  detections = load_and_plot_SD(logs[0], start_img_idx=logs[3][0], end_img_idx=logs[3][1],
                                                 start_iter=logs[3][2], end_iter=logs[3][3], logarithm=False)
            x = x[::2].numpy()
            y = y[::2].numpy()
            k_avg_dist = k_avg_dist[::2]
            for n_item in norms:
                print("%.2f " % (n_item), end='')
            print("%.1f" %detections, end='')
            print()
            l=ax1.plot(x, y, logs[1], label=logs[2])
            r=ax2.plot(x, k_avg_dist, logs[1],label=logs[2], alpha=0.5) #linestyle='dashed',
            legends.append((l[0],r[0]))
            legends_labels.append(logs[2])
        threshold=0.002*np.ones_like(k_avg_dist)
        r = ax2.plot(x, threshold, 'k', label='Detection threshold', alpha=1)  # linestyle='dashed',
        legends.append((r[0]))
        legends_labels.append('Detection threshold')
        plt.legend(handles=legends,labels=legends_labels,numpoints=1, handler_map={tuple: HandlerTuple(ndivide=None)}) #
        #ax2.set_ylim(0,1.5)
        ax2.set_yscale('log')
        #legend = ax.legend(loc='upper center', shadow=True, fontsize='x-large')
        #legend.get_frame().set_facecolor('white')
        #ax2.ylim(0,0.1)
        plt.xticks([1000, 2000, 5000, 10000], ['1K', '2K', '5K', '10K'])
        plt.savefig('SD_Graph.pdf', dpi=300, bbox_inches='tight')
        plt.show()
    ##############################################
    if False:
        datasets = ['LFW']  # 'LFW',
        attacks = ['EAGDR', 'EAGDO', 'EAGD']  #'EV', 'EVD', 'EVG','EV', 'EVD', 'EVG','EAGD'
        model = 'IR_50'
        defense = 'gaussian'
        colors = ['#1f77b4',
                  '#ff7f0e',
                  '#2ca02c',
                  '#d62728',
                  '#9467bd',
                  '#8c564b',
                  '#e377c2',
                  '#7f7f7f',
                  '#bcbd22',
                  '#17becf',
                  '#1a55FF']
        # colors=['tab:blue','tab:red','g','tab:purple','tab:red','tab:cyan','tab:brown','tab:pink']
        for targeted in ['0', '1']:
            for dt in datasets:
                PLOT_INFO = []
                for i, at in enumerate(attacks):
                    fname = dt + '_500_' + targeted + '_' + at + '_' + model + '_' + defense + '_.pth'
                    at=at.replace('EAGDO','EAGD1')
                    if os.path.exists(fname):
                        PLOT_INFO.append((fname, colors[i], at, [0, 500, 1000,3000], ''))
                fig, ax = plt.subplots(figsize=(6, 4), tight_layout=True)

                avg_X = []
                avg_y = []
                plot_info = []
                labels = []
                for i, logs in enumerate(PLOT_INFO):
                    x, y, norms, queries = load_and_plot(logs[0], start_img_idx=logs[3][0], end_img_idx=logs[3][1],
                                                         start_iter=logs[3][2], end_iter=logs[3][3], logarithm=False)
                    x = x[::2].numpy()
                    y = y[::2].numpy()
                    for n_item in norms:
                        print("%.2f " % (n_item), end='')
                    for n_item in queries:
                        print("%d " % (n_item), end='')
                    print()
                    plt.plot(x, y, logs[1], label=logs[2])
                plt.legend()
                plt.title('Dodging attacks on the LFW dataset', fontsize=12)
                #plt.xticks([1000, 2000, 5000, 10000], ['1K', '2K', '5K', '10K'])
                plt.xticks([1000, 2000, 3000], ['1K', '2K', '3K'])
                ax.set_xlabel('# Queries', fontsize=12)
                #ax.set_yscale('log')
                ax.set_ylabel('$\ell_2$ norm of perturbation', fontsize=12)
                plt.savefig('Dodging_LFW.pdf', dpi=300, bbox_inches='tight')
                plt.show()

    # Main Table Results
    # LFW CPLFW
    # Impersonation Attack
    # Dogging Attack
    if True:
        datasets=['LFW','CPLFW'] #'LFW',

        attacks=['SO','HSJA','EA','EAD','EAG','EAGD']
        model='IR_101'
        defense='gaussian'
        colors = ['#1f77b4',
                  '#ff7f0e',
                  '#2ca02c',
                  '#d62728',
                  '#9467bd',
                  '#8c564b',
                  '#e377c2',
                  '#7f7f7f',
                  '#bcbd22',
                  '#17becf',
                  '#1a55FF']
        # colors=['tab:blue','tab:red','g','tab:purple','tab:red','tab:cyan','tab:brown','tab:pink']
        for targeted in ['0','1']:#
            for dt in datasets:
                if targeted == '1':
                    if dt=='LFW':
                        attacks = ['SO', 'HSJA', 'EA', 'EAG']  # 'GD'
                    else:
                        attacks = ['SO', 'HSJA', 'EA', 'EAG']  # 'GD'
                else:
                    attacks = ['SO', 'HSJA', 'EA', 'EAD', 'EAG', 'EAGD']  # 'GD'

                PLOT_INFO=[]
                for i, at in enumerate(attacks):
                    fname=dt+'_500_'+targeted+'_'+at+'_'+model+'_'+defense+'_.pth'
                    if os.path.exists(fname):
                        PLOT_INFO.append((fname,colors[i],at,[0,500,1000,10000],''))
                fig, ax = plt.subplots(figsize=(6,4), tight_layout=True)
                avg_X = []
                avg_y = []
                plot_info = []
                labels = []
                for i, logs in enumerate(PLOT_INFO):
                    x, y ,norms,queries= load_and_plot(logs[0], start_img_idx=logs[3][0], end_img_idx=logs[3][1],
                                              start_iter=logs[3][2], end_iter=logs[3][3], logarithm=False)
                    x = x[::10].numpy()
                    y = y[::10].numpy()
                    for n_item in norms:
                        print(" %.2f &"%(n_item),end='')
                    for n_item in queries:
                        print(" %d &"%( n_item), end='')
                    print()
                    at=logs[2]
                    plt.plot(x, y, logs[1], label=at)
                plt.legend()
                #plt.grid(alpha=0.3)
                plt.xticks([1000, 2000, 5000, 10000], ['1K', '2K', '5K', '10K'])
                if targeted=='1':
                    plt.title('Impersonation attacks on the '+dt+' dataset')
                else:
                    plt.title('Dodging attacks on the '+dt+' dataset')

                ax.set_xlabel('# Queries', fontsize=12)
                ax.set_ylabel('$\ell_2$ norm of perturbation', fontsize=12)
                plt.savefig(dt+'_'+targeted+'_'+model+'.pdf', dpi=300, bbox_inches='tight')
                plt.show()

    if False:
        datasets = ['LFW', 'CPLFW']  # 'LFW',
        attacks = ['SFA', 'SFAD', 'SFAG', 'SFAGD']  #
        model = 'IR_50'
        defense = 'gaussian'
        colors = ['#1f77b4',
                  '#ff7f0e',
                  '#2ca02c',
                  '#d62728',
                  '#9467bd',
                  '#8c564b',
                  '#e377c2',
                  '#7f7f7f',
                  '#bcbd22',
                  '#17becf',
                  '#1a55FF']
        for targeted in ['0', '1']:
            for dt in datasets:
                PLOT_INFO = []
                for i, at in enumerate(attacks):
                    fname = dt + '_500_' + targeted + '_' + at + '_' + model + '_' + defense + '_.pth'
                    if os.path.exists(fname):
                        PLOT_INFO.append((fname, colors[i], at, [0, 500, 1000, 10000], ''))
                fig, ax1 = plt.subplots(figsize=(6, 4), tight_layout=True)

                ax1.set_xlabel('# Queries', fontsize=12)
                ax1.set_ylabel('$\ell_\infty$ norm of perturbation', fontsize=12)
                ax1.tick_params(axis='y')
                ax2 = ax1.twinx()  # instantiate a second axes that shares the same x-axis
                ax2.set_ylabel('$\ell_2$ norm of perturbation', fontsize=12)
                ax2.tick_params(axis='y')
                avg_X = []
                avg_y = []
                plot_info = []
                labels = []
                legends = []
                legends_labels = []

                for i, logs in enumerate(PLOT_INFO):
                    x_2, y_2, norms, queries = load_and_plot(logs[0], start_img_idx=logs[3][0], end_img_idx=logs[3][1],
                                                             start_iter=logs[3][2], end_iter=logs[3][3],
                                                             logarithm=False)

                    x, y, norms, queries = load_and_plot_linf(logs[0], start_img_idx=logs[3][0], end_img_idx=logs[3][1],
                                                              start_iter=logs[3][2], end_iter=logs[3][3],
                                                              logarithm=False)
                    x = x[::10].numpy()
                    y = y[::10].numpy()
                    for n_item in norms:
                        print("%.4f " % (n_item), end='')
                    for n_item in queries:
                        print("%d " % (n_item), end='')
                    print()
                    l = ax1.plot(x, y, logs[1], label=logs[2])
                    r = ax2.plot(x_2, y_2, logs[1], label=logs[2], linestyle='dashed',alpha=0.6)  # linestyle='dashed',
                    legends.append((l[0], r[0]))
                    legends_labels.append(logs[2])

                plt.legend(handles=legends, labels=legends_labels, numpoints=1, handlelength=3,
                           handler_map={tuple: HandlerTuple(ndivide=None)})  #

                plt.xticks([1000, 2000, 5000, 10000], ['1K', '2K', '5K', '10K'])
                plt.savefig('L_infty_Graph.pdf', dpi=300, bbox_inches='tight')
                plt.show()


from __future__ import print_function, division
import torch

torch.set_printoptions(profile="full")
import matplotlib.pyplot as plt
import argparse, os
import numpy as np
import shutil
from torch.utils.data import DataLoader
from torchvision import transforms

from models.CDCNs_u import Conv2d_cd, CDCN_u

from Load_OULUNPUcrop_valtest import Spoofing_valtest, Normaliztion_valtest, ToTensor_valtest

import torch.optim as optim

from utils import performances

# Dataset root

val_image_dir = '/export2/home/wht/oulu_img_crop/dev_file_flod/'
test_image_dir = '/export2/home/wht/oulu_img_crop/test_file_flod/'

val_map_dir = '/export2/home/wht/oulu_img_crop/dev_depth_flod/'
test_map_dir = '/export2/home/wht/oulu_img_crop/test_depth_flod/'

val_list = '/export2/home/wht/oulu_img_crop/protocols/Protocol_1/Dev.txt'
test_list = '/export2/home/wht/oulu_img_crop/protocols/Protocol_1/Test.txt'

# main function
def test():
    # GPU  & log file  -->   if use DataParallel, please comment this command
    os.environ["CUDA_VISIBLE_DEVICES"] = '0, 1, 2, 3'

    isExists = os.path.exists(args.log)
    if not isExists:
        os.makedirs(args.log)
    log_file = open(args.log + '/' + args.log + '_log_P1.txt', 'w')

    log_file.write('Oulu-NPU, P1:\n ')
    log_file.flush()

    print('test!\n')
    log_file.write('test!\n')
    log_file.flush()

    model = CDCN_u(basic_conv=Conv2d_cd, theta=0.7)
    # model = ResNet18_u()

    model = model.cuda()
    model = torch.nn.DataParallel(model)
    model.load_state_dict(torch.load('./DUM/checkpoint/CDCN_U_P1.pkl', map_location='cuda:0'))

    print(model)

    optimizer = optim.Adam(model.parameters(), lr=0.001, weight_decay=0.00005)

    for epoch in range(args.epochs):

        model.eval()

        with torch.no_grad():
            ###########################################
            '''                val             '''
            ###########################################
            # val for threshold
            val_data = Spoofing_valtest(val_list, val_image_dir, val_map_dir,
                                        transform=transforms.Compose([Normaliztion_valtest(), ToTensor_valtest()]))
            dataloader_val = DataLoader(val_data, batch_size=1, shuffle=False, num_workers=4)

            map_score_list = []

            for i, sample_batched in enumerate(dataloader_val):
                # get the inputs
                inputs, spoof_label = sample_batched['image_x'].cuda(), sample_batched['spoofing_label'].cuda()
                val_maps = sample_batched['val_map_x'].cuda()  # binary map from PRNet

                optimizer.zero_grad()

                # pdb.set_trace()
                map_score = 0.0
                for frame_t in range(inputs.shape[1]):
                    mu, logvar, map_x, x_concat, x_Block1, x_Block2, x_Block3, x_input = model(
                        inputs[:, frame_t, :, :, :])
                    score_norm = torch.sum(mu) / torch.sum(val_maps[:, frame_t, :, :])
                    map_score += score_norm

                map_score = map_score / inputs.shape[1]

                map_score_list.append('{} {}\n'.format(map_score, spoof_label[0][0]))
                # pdb.set_trace()
            map_score_val_filename = args.log + '/' + args.log + '_map_score_val.txt'
            with open(map_score_val_filename, 'w') as file:
                file.writelines(map_score_list)

            ###########################################
            '''                test             '''
            ##########################################
            # test for ACC
            test_data = Spoofing_valtest(test_list, test_image_dir, test_map_dir,
                                         transform=transforms.Compose([Normaliztion_valtest(), ToTensor_valtest()]))
            dataloader_test = DataLoader(test_data, batch_size=1, shuffle=False, num_workers=4)

            map_score_list = []

            for i, sample_batched in enumerate(dataloader_test):
                # get the inputs
                inputs, spoof_label = sample_batched['image_x'].cuda(), sample_batched['spoofing_label'].cuda()
                test_maps = sample_batched['val_map_x'].cuda()

                optimizer.zero_grad()

                # pdb.set_trace()
                map_score = 0.0
                for frame_t in range(inputs.shape[1]):
                    mu, logvar, map_x, x_concat, x_Block1, x_Block2, x_Block3, x_input = model(
                        inputs[:, frame_t, :, :, :])

                    score_norm = torch.sum(mu) / torch.sum(test_maps[:, frame_t, :, :])
                    map_score += score_norm

                map_score = map_score / inputs.shape[1]

                map_score_list.append('{} {}\n'.format(map_score, spoof_label[0][0]))

            map_score_test_filename = args.log + '/' + args.log + '_map_score_test.txt'
            with open(map_score_test_filename, 'w') as file:
                file.writelines(map_score_list)

            #############################################################
            #       performance measurement both val and test
            #############################################################
            val_threshold, test_threshold, val_ACC, val_ACER, test_ACC, test_APCER, test_BPCER, test_ACER, test_ACER_test_threshold = performances(
                map_score_val_filename, map_score_test_filename)

            print('epoch:%d, Val:  val_threshold= %.4f, val_ACC= %.4f, val_ACER= %.4f' % (
                epoch + 1, val_threshold, val_ACC, val_ACER))
            log_file.write('\n epoch:%d, Val:  val_threshold= %.4f, val_ACC= %.4f, val_ACER= %.4f \n' % (
                epoch + 1, val_threshold, val_ACC, val_ACER))

            print('epoch:%d, Test:  ACC= %.4f, APCER= %.4f, BPCER= %.4f, ACER= %.4f' % (
                epoch + 1, test_ACC, test_APCER, test_BPCER, test_ACER))
            log_file.write('epoch:%d, Test:  ACC= %.4f, APCER= %.4f, BPCER= %.4f, ACER= %.4f \n' % (
                epoch + 1, test_ACC, test_APCER, test_BPCER, test_ACER))
            log_file.flush()

    print('Finished Training')
    log_file.close()

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="save quality using landmarkpose model")
    parser.add_argument('--gpus', type=str, default='0,1,2,3', help='the gpu id used for predict')
    parser.add_argument('--lr', type=float, default=0.001, help='initial learning rate')
    parser.add_argument('--batchsize', type=int, default=32, help='initial batchsize')
    parser.add_argument('--step_size', type=int, default=500, help='how many epochs lr decays once')  # 500
    parser.add_argument('--gamma', type=float, default=0.5, help='gamma of optim.lr_scheduler.StepLR, decay of lr')
    parser.add_argument('--kl_lambda', type=float, default=0.001, help='')
    parser.add_argument('--echo_batches', type=int, default=50, help='how many batches display once')  # 50
    parser.add_argument('--epochs', type=int, default=1, help='total training epochs')
    parser.add_argument('--log', type=str, default="CDCN_U_P1_test", help='log and save model name')
    parser.add_argument('--finetune', action='store_true', default=False, help='whether finetune other models')
    parser.add_argument('--test', action='store_true', default=True, help='')

    args = parser.parse_args()
    test()