from uuid import uuid4
from typing import Sequence, Generator
from multiprocessing import cpu_count
from gensim.models.doc2vec import Doc2Vec, TaggedDocument
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.feature_selection import mutual_info_classif
from sklearn.metrics import roc_auc_score
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
import numpy as np
from antivirals.schema import Molecules
from antivirals.parser import trivial_generator

class _WrapGenerator:
    """
    Creates iterators out of replayable generators. Needed for gensim.
    """

    def __init__(self, func):
        self.func = func
        self.generator = func()

    def __iter__(self):
        self.generator = self.func()
        return self

    def __next__(self):
        res = next(self.generator)
        if not res:
            raise StopIteration
        else:
            return res

class Hyperparameters:
    """
    Hyperparameters for all chemistry models.
    """
    # Language hyperparams
    max_ngram = 2
    vector_algo = 'doc2vec'

    # Doc2Vec hyperparmas
    vec_dims = 460
    vec_window = 4
    max_vocab = 35000
    doc_epochs = 110
    alpha = 0.05

    # LDA hyperparams
    topic_epochs = 5
    topic_iterations = 20
    topics = 16
    decay = 0.5

    # Toxicity hyperparams
    estimators = 400
    min_samples_split = 2
    min_samples_leaf = 12
    tree_criterion = 'entropy'

    @staticmethod
    def from_dict(values):
        hp = Hyperparameters()
        hp.__dict__.update(values)
        return hp

class Chemistry:
    """
    A master model that encodes all aspects of chemistry.
    """
    hyperparams: Hyperparameters
    uuid: str

    def __init__(
            self,
            hyperparams: Hyperparameters = Hyperparameters()):
        self.hyperparams = hyperparams
        self.language = Language(self.hyperparams)
        self.toxicity = Toxicity(self.hyperparams, self.language)
        self.uuid = str(uuid4())

    def from_molecules(self, mols: Molecules):
        tox_data = mols.get_mols_with_passfail_labels()
        X = tox_data.index
        y = tox_data.astype('int')
        self.language.fit(mols.get_all_mols(), X, y)
        self.toxicity = Toxicity(self.hyperparams, self.language)
        self.toxicity.build(X, y)
        print(f"Trained new chemistry model: {self.uuid}")

class Toxicity:
    """
    Implments the toxicity model ontop of the latent vectors of the chemical language model.
    """
    auc: Sequence[float]

    def __init__(self, hyperparams: Hyperparameters, language_model):
        self.hyperparams = hyperparams
        self.language = language_model

    def fit(self, X: Sequence[str], Y: np.ndarray):
        self.classif = RandomForestClassifier(
            bootstrap=False,
            criterion=self.hyperparams.tree_criterion,
            max_features=0.25,
            min_samples_leaf=self.hyperparams.min_samples_leaf,
            min_samples_split=self.hyperparams.min_samples_split,
            n_estimators=self.hyperparams.estimators
        )

        self.classif.fit(self.language.to_vecs(X), Y)

    def build(self, X: Sequence[str] = None, Y: np.ndarray = None):
        Xt, Xv, Yt, Yv = train_test_split(X, Y, test_size=0.2, random_state=18)
        self.fit(Xt, Yt)
        self.auc = self.audit(Xv, Yv)

    def audit(self, X: Sequence[str], Y: np.ndarray):
        Yhats = self.classif.predict_proba(self.language.to_vecs(X))
        if 'to_numpy' in dir(Y):
            Y = Y.to_numpy()
        res = []
        for i, Yhat in enumerate(Yhats):
            res.append(roc_auc_score(Y[:, i], Yhat[:, 1]))
        return res

class Language:
    """
    A chemical language model that creates semantic latent vectors from chemicals,
    based on the mutual information between subtokens of a chemical discriptor and
    a surrigate prediction set encoding chemistry semantics.
    """

    def __init__(self, hyperparams: Hyperparameters):
        self.hyperparams = hyperparams
        if hyperparams.vector_algo != 'doc2vec' and hyperparams.vector_algo != 'lda':
            raise RuntimeError(
                f"Unsupported algorithm: {hyperparams.vector_algo}")
        self.hyperparams = hyperparams

    def _smiles_to_advanced_lang(self, smiles_seq: Generator
                                 [str, None, None],
                                 training: bool = False) -> Generator[str, None,
                                                                      None]:
        for i, sent in enumerate(smiles_seq):
            sent: Sequence[str] = self._analyzer(sent)
            res: Sequence[str] = []
            for token in sent:
                if token in self.vocab:
                    res.append(token.replace(' ', 'A'))
            if training:
                yield TaggedDocument(words=res, tags=[i])
            else:
                yield res

    def _make_iterator(
            self, smiles_seq: Sequence[str],
            training: bool = False) -> _WrapGenerator:
        return _WrapGenerator(
            lambda: self._smiles_to_advanced_lang(
                trivial_generator(smiles_seq),
                training))

    def make_generator(self, X):
        return self._smiles_to_advanced_lang(trivial_generator(X))

    def to_vecs(self, X: Sequence[str]) -> np.ndarray:
        if self.hyperparams.vector_algo == 'lda':
            bows = [self.dictionary.doc2bow(i) for i in self.make_generator(X)]
            latent_vecs, _ = self.topic_model.inference(bows)
            return latent_vecs
        elif self.hyperparams.vector_algo == 'doc2vec':
            # Preallocate memory for performance
            latent_vecs = np.empty(
                (len(X), self.document_model.vector_size))

            for i, sent in enumerate(self.make_generator(X)):
                latent_vecs[i] = self.document_model.infer_vector(sent)

            return latent_vecs

    def _fit_language(
            self, X_unmapped: Sequence[str],
            X: Sequence[str],
            Y: np.ndarray):
        cv = CountVectorizer(
            max_df=0.95, min_df=2, lowercase=False,
            ngram_range=(1, self.hyperparams.max_ngram),
            max_features=(self.hyperparams.max_vocab * 18),
            token_pattern='[a-zA-Z0-9$&+,:;=?@_/~#\\[\\]|<>.^*()%!-]+')

        X_vec = cv.fit_transform(trivial_generator(X))

        local_vocab = set()
        for feat in Y.columns:
            res = zip(cv.get_feature_names(),
                      mutual_info_classif(
                          X_vec, Y[feat], discrete_features=True)
                      )
            local_vocab.update(res)
        self.vocab = {
            i[0]
            for i in sorted(
                local_vocab, key=lambda i: i[1],
                reverse=True)[: self.hyperparams.max_vocab]}

        self._analyzer = cv.build_analyzer()

    def _fit_document_model(
            self, X_unmapped: Sequence[str],
            X: Sequence[str],
            Y: np.ndarray):
        generator = list(self._make_iterator(X_unmapped, training=True))

        document_model = Doc2Vec(
            vector_size=self.hyperparams.vec_dims,
            alpha=self.hyperparams.alpha,
            workers=max(1, cpu_count() - 2),
            window=self.hyperparams.vec_window)
        document_model.build_vocab(generator)
        document_model.train(generator, total_examples=len(
            X_unmapped), epochs=self.hyperparams.doc_epochs)

        self.document_model = document_model

    def _fit_topic_model(
            self, X_unmapped: Sequence[str],
            X: Sequence[str],
            Y: np.ndarray):
        from gensim.models.ldamulticore import LdaMulticore
        from gensim.corpora.dictionary import Dictionary
        iterator = list(self.make_generator(X_unmapped))
        bow = Dictionary(iterator)

        docs = [bow.doc2bow(i) for i in iterator]

        self.topic_model = LdaMulticore(
            docs,
            id2word=bow, num_topics=self.hyperparams.topics,
            iterations=self.hyperparams.topic_iterations,
            passes=self.hyperparams.topic_epochs, random_state=18,
            decay=self.hyperparams.decay)
        self.dictionary = bow

    def fit(self, X_unmapped: Sequence[str], X: Sequence[str], Y: np.ndarray):
        self._fit_language(X_unmapped, X, Y)
        if self.hyperparams.vector_algo == 'lda':
            self._fit_topic_model(X_unmapped, X, Y)
        elif self.hyperparams.vector_algo == 'doc2vec':
            self._fit_document_model(X_unmapped, X, Y)


import torch
import torch.nn.functional as F
from torch import nn
import numpy as np
from torch.optim.lr_scheduler import CosineAnnealingLR

class BiC(nn.Module):
    def __init__(self, lr, scheduling, lr_decay_factor, weight_decay, batch_size, epochs):
        super(BiC, self).__init__()
        self.beta = torch.nn.Parameter(torch.ones(1))  #.cuda()
        self.gamma = torch.nn.Parameter(torch.zeros(1))  #.cuda()
        self.lr = lr
        self.scheduling = scheduling
        self.lr_decay_factor = lr_decay_factor
        self.weight_decay = weight_decay
        self.class_specific = False
        self.batch_size = batch_size
        self.epochs = epochs
        self.bic_flag = False

    def reset(self, lr=None, scheduling=None, lr_decay_factor=None, weight_decay=None, n_classes=-1):
        with torch.no_grad():
            if lr is None:
                lr = self.lr
            if scheduling is None:
                scheduling = self.scheduling
            if lr_decay_factor is None:
                lr_decay_factor = self.lr_decay_factor
            if weight_decay is None:
                weight_decay = self.weight_decay
            if self.class_specific:
                assert n_classes != -1
                self.beta = torch.nn.Parameter(torch.ones(n_classes).cuda())
                self.gamma = torch.nn.Parameter(torch.zeros(n_classes).cuda())
            else:
                self.beta = torch.nn.Parameter(torch.ones(1).cuda())
                self.gamma = torch.nn.Parameter(torch.zeros(1).cuda())
            self.optimizer = torch.optim.SGD([self.beta, self.gamma], lr=lr, momentum=0.9, weight_decay=weight_decay)
            # self.scheduler = CosineAnnealingLR(self.optimizer, 10)
            self.scheduler = torch.optim.lr_scheduler.MultiStepLR(self.optimizer, scheduling, gamma=lr_decay_factor)

    def extract_preds_and_targets(self, model, loader):
        preds, targets = [], []
        with torch.no_grad():
            for (x, y) in loader:
                preds.append(model(x.cuda())['logit'])
                targets.append(y.cuda())
        return torch.cat((preds)), torch.cat((targets))

    def update(self, logger, task_size, model, loader, loss_criterion=None):
        if task_size == 0:
            logger.info("no new task for BiC!")
            return
        if loss_criterion is None:
            loss_criterion = F.cross_entropy

        self.bic_flag = True
        logger.info("Begin BiC ...")
        model.eval()

        for epoch in range(self.epochs):
            preds_, targets_ = self.extract_preds_and_targets(model, loader)
            order = np.arange(preds_.shape[0])
            np.random.shuffle(order)

            preds, targets = preds_.clone(), targets_.clone()
            preds, targets = preds[order], targets[order]
            _loss = 0.0
            _correct = 0
            _count = 0
            for start in range(0, preds.shape[0], self.batch_size):
                if start + self.batch_size < preds.shape[0]:
                    out = preds[start:start + self.batch_size, :].clone()
                    lbls = targets[start:start + self.batch_size]
                else:
                    out = preds[start:, :].clone()
                    lbls = targets[start:]
                if self.class_specific is False:
                    out1 = out[:, :-task_size].clone()
                    out2 = out[:, -task_size:].clone()
                    outputs = torch.cat((out1, out2 * self.beta + self.gamma), 1)
                else:
                    outputs = out * self.beta + self.gamma
                loss = loss_criterion(outputs, lbls)
                self.optimizer.zero_grad()
                loss.backward()
                self.optimizer.step()
                _, pred = outputs.max(1)
                _correct += (pred == lbls).sum()
                _count += lbls.size(0)
                _loss += loss.item() * outputs.shape[0]
            logger.info("epoch {} loss {:4f} acc {:4f}".format(epoch, _loss / preds.shape[0], _correct / _count))

            self.scheduler.step()
        logger.info("beta {:.4f} gamma {:.4f}".format(self.beta.cpu().item(), self.gamma.cpu().item()))

    @torch.no_grad()
    def post_process(self, preds, task_size):
        if self.class_specific is False:
            if task_size != 0:
                preds[:, -task_size:] = preds[:, -task_size:] * self.beta + self.gamma
        else:
            preds = preds * self.beta + self.gamma
        return preds

class WA(object):
    def __init__(self):
        self.gamma = None

    @torch.no_grad()
    def update(self, classifier, task_size):
        old_weight_norm = torch.norm(classifier.weight[:-task_size], p=2, dim=1)
        new_weight_norm = torch.norm(classifier.weight[-task_size:], p=2, dim=1)
        self.gamma = old_weight_norm.mean() / new_weight_norm.mean()
        print(self.gamma.cpu().item())

    @torch.no_grad()
    def post_process(self, logits, task_size):
        logits[:, -task_size:] = logits[:, -task_size:] * self.gamma
        return logits


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Aug  6 14:58:35 2019

@author: bressler
"""

import SBCcode as sbc
from os import listdir
from os.path import isfile,join
from collections import Counter
import numpy as np
import matplotlib.pyplot as plt
import scipy
from gaincalc import get_gain
from PICOcode.REFPROP.SeitzModel import SeitzModel
import runlistscatalogue as rlc

OKLONGEVENTS = ['20170701_6-81', '20170627_3-55', '20170628_5-29', '20170623_7-38',
                '20170623_7-76', '20170624_0-35', '20170624_0-72', '20170801_2-39',
                '20170717_0-67', '20170626_5-6', '20170626_5-20', '20170626_9-2',
                '20170626_9-12', '20170626_9-77', '20170626_9-90', '20170627_0-31',
                '20170627_0-73', '20170627_0-74', '20170702_4-19', '20170702_4-28',
                '20170703_0-59', '20170703_6-37', '20170704_1-97']

def getRate(xyzf, runs, LTpreq,PT):
    LT = []
    totbub = 0
    bubperpset = []
    doublesperpset = []
    triplesperpset = []
    setpoints = []
    elt = 0
    feedbackTransducer = int(PT[-1])-1
    temperatures = []
    counts = []

    event_livetimes = []
    event_counter = 0
    for run in runs:
        runpath = '/bluearc/storage/SBC-17-data/'+run+'/'
        events = [evnt for evnt in listdir(runpath) if not isfile(join(runpath,evnt))]
        Nevents = len(events)
        print("run "+run+": %d events"%Nevents)
        indices = [i for i,x in enumerate(xyzf["runid"]) if str(x[0])+"_"+str(x[1]) == run]
        runposreco = {"ev":[xyzf["ev"][indices]],"x":[xyzf["bubX"][indices]],
                      "y":[xyzf["bubY"][indices]],"z":[xyzf["bubZ"][indices]]}

        runrawpath = '/bluearc/storage/SBC-17-data/'+run
        runreconpath = "/pnfs/coupp/persistent/grid_output/SBC-17/output/%s/"%run
        historyfilename = runreconpath+"HistoryAnalysis_%s.bin"%run
        history = sbc.DataHandling.ReadBinary.ReadBlock(historyfilename)

        getbubfile = "/coupp/data/home/coupp/HumanGetBub_output_SBC-17/HumanGetBub_%s.bin"%run
        c = sbc.DataHandling.ReadBinary.ReadBlock(getbubfile)
        #print(c.keys())
        evn = c["ev"]
        #print(evn)
        #print(c["nbubimage"])
        count = Counter(evn)
        #print(count)
        nbubs = np.zeros(Nevents)
        #print(len(nbubs))
        for c in count:
            #print(c)
            N = count[c]
            #print(N)
            counts.append(N)
            if N%2 != 0 and N != 1:
                print("lines isn't even")
                print(c)
                print(N)
            elif N%2 == 0:
                n = N/2
                nbubs[c]=n
        e0 = sbc.DataHandling.GetSBCEvent.GetEvent(runrawpath,0,"slowDAQ","event")
        T25 = np.mean(e0["slowDAQ"]["T1"])
        seitz25 = SeitzModel(25,T25,'xenon')
        Q25 = seitz25.Q
        #print(nbubs)
        for eventn in range(Nevents):
            try:
                nobub = np.isnan(runposreco["z"][0][int(eventn)])
                n = nbubs[int(eventn)]
                #print("n=%d"%n)

                if nobub and n==1:
                    print("nobub and n=1: " + run +'-'+str(eventn))
                """
                elif (not nobub) and n==1:
                    print("(not nobub) and n = 1: " + run + '-'+str(eventn))
                elif n != 1 and nobub:
                    print("nobub and n != 1: "+run+'-'+str(eventn))
                elif n !=1 and not nobub:
                    print("not nobub and n != 1: "+run+'-'+str(eventn))
                """
                #print(n==1)
                e = sbc.DataHandling.GetSBCEvent.GetEvent(runrawpath,eventn,"slowDAQ","event")
                lt_event = e["event"]["livetime"]
                if(lt_event>800):
                    if run+'-'+str(eventn) not in OKLONGEVENTS:
                        print("check long event, excluding now: "+run+"-"+str(eventn))
                        continue
                event_livetimes.append(lt_event)
                event_counter += 1
                if lt_event < 20:
                    #print("short event excluded")
                    continue
                T = np.mean(e["slowDAQ"]["T1"])
                #print(T)
                if T>-1000:
                    temperatures.append(T)
                else:
                    print("Not including unpysical temperature %f C"%T)
                pset = e["event"]["Pset"]
                #seitz = SeitzModel(pset,T,'xenon')
                #Q = seitz.Q
                elt += e["event"]["livetime"]
                if pset not in setpoints:
                    setpoints.append(pset)
                    LT.append(0)
                    bubperpset.append(0)
                    doublesperpset.append(0)
                    triplesperpset.append(0)
                ind = setpoints.index(pset)

                indices_back = 30
                trig_pressure = e["slowDAQ"][PT][list(e["slowDAQ"]["TriggerOut"]).index(1.0)-indices_back]

                #if not nobub:
                    #print("adding bubble at P = %f"%trig_pressure)
                if n > 0:
                    totbub += 1
                if (np.abs(trig_pressure-pset)<LTpreq/2):
                    if n == 1: bubperpset[ind]+=1
                    elif n==2: doublesperpset[ind] += 1
                    elif n==3: triplesperpset[ind] += 1

                edges = history["PressureEdge"][feedbackTransducer]
                times = history["PressureBins"][eventn][feedbackTransducer][:]
                centers = [(edges[i]+edges[i+1])/2 for i in range(len(edges)-1)]
                #print(edges)
                for i in range(len(centers)):
                    if abs(centers[i]-pset)<LTpreq/2:
                        LT[ind]+=times[i]

            except Exception as x:
                print(x)
                break
    totlt = sum(LT)
    r = totbub/totlt
    rerr = np.sqrt(totbub)/totlt

    rperpset = np.divide(bubperpset,LT)
    r_err_perpset = np.divide(np.sqrt(bubperpset),LT)

    for i in range(len(rperpset)):
        if LT[i] > 50:
            seitz = SeitzModel(setpoints[i],np.mean(temperatures),'xenon')
            #print(setpoints[i])
            #print(min(temperatures))
            #print(temperatures)
            #print(seitz)
            print("@@@@@@@@@@@@@@@@@@@@@@@@@@")
            print("Pressure bin: "+str(setpoints[i]))
            print("Live Time: "+str(LT[i]))
            print("Single Bubbles: "+str(bubperpset[i]))
            print("Singles Rate: " +str(rperpset[i])+"+/-"+str(r_err_perpset[i]))
            print("Double Bubbles: %d"%doublesperpset[i])
            print("Doubles Rate: %.2e"%(doublesperpset[i]/LT[i]))
            print("Triple Bubbles: %d"%triplesperpset[i])
            print("Triples Rate: %.2e"%(triplesperpset[i]/LT[i]))
            print("Seitz Threshold: %f keV"%seitz.Q)
            print("Temperature: %.2fC"%np.mean(temperatures))
            print("\n")
    if 25.0 in setpoints:
        print("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")
        print("25 psia bin: ")
        print(str(LT[setpoints.index(25.0)])+"s")
        print(str(bubperpset[setpoints.index(25.0)])+" bubbles")
        print(str(Q25) + " keV")
        print(str(T25) + " Celsius")
        print("rate: "+str(rperpset[setpoints.index(25.0)])+"+/-"+str(r_err_perpset[setpoints.index(25.0)])+" Hz")
        print("\n")
        print("overall rate: "+str(r)+"+/-"+str(rerr)+"Hz")
        print("total events: %d"%event_counter)
        print("total bubbles: "+str(totbub))
        print("fraction of events with bubbles: %f"%(totbub/event_counter))
        print("total live time: "+str(totlt))
        print("total live time from events: " +str(elt))
        print("mean live time: %f sec"%np.mean(event_livetimes))
        print("trigger rate: %f Hz"%(1/float(np.mean(event_livetimes))))
        print("bubble rate from bub fraction  and LT: %f Hz"%(float((totbub/event_counter))/float(np.mean(event_livetimes))))
        print("\n")
    plt.figure()
    plt.scatter(np.arange(event_counter),event_livetimes, facecolor='none', edgecolor='b')
    plt.grid()
    plt.xlabel('event count')
    plt.ylabel('event livetime (s)')
    plt.show()
    plt.figure()
    plt.hist(event_livetimes,int(np.ceil(np.sqrt(len(event_livetimes)))))
    plt.show()
    return [totbub,totlt,rperpset,r_err_perpset,r,rerr,setpoints]

def main():
    allxyzfname = "/pnfs/coupp/persistent/grid_output/SBC-17/output/SimpleXYZ_all.bin"
    xyzf = sbc.DataHandling.ReadBinary.ReadBlock(allxyzfname)

    [totbub,totlt,rperpset,r_err_perpset,r,rerr,setpoints] = getRate(xyzf,rlc.bgJuly3and4,1,"PT6")
    print(totbub)
    print(totlt)
    print(r)
    print(rerr)

if __name__ == "__main__":
    main()


"""
Launcher for experiments with FOCAL

"""
import os
import pathlib
import numpy as np
import click
import json
import torch
import multiprocessing as mp
from itertools import product

from rlkit.envs import ENVS
from rlkit.envs.wrappers import NormalizedBoxEnv
from rlkit.torch.sac.policies import TanhGaussianPolicy
from rlkit.torch.networks import FlattenMlp, MlpEncoder, RecurrentEncoder
from rlkit.torch.sac.sac import FOCALSoftActorCritic
from rlkit.torch.sac.agent import PEARLAgent
from rlkit.launchers.launcher_util import setup_logger
import rlkit.torch.pytorch_util as ptu
from configs.default import default_config
from numpy.random import default_rng

rng = default_rng()

def global_seed(seed=0):
    torch.manual_seed(seed)
    np.random.seed(seed)

def experiment(variant, seed=None):

    # create multi-task environment and sample tasks, normalize obs if provided with 'normalizer.npz'
    if 'normalizer.npz' in os.listdir(variant['algo_params']['data_dir']):
        obs_absmax = np.load(os.path.join(variant['algo_params']['data_dir'], 'normalizer.npz'))['abs_max']
        env = NormalizedBoxEnv(ENVS[variant['env_name']](**variant['env_params']), obs_absmax=obs_absmax)
    else:
        env = NormalizedBoxEnv(ENVS[variant['env_name']](**variant['env_params']))

    if seed is not None:
        global_seed(seed)
        env.seed(seed)

    tasks = env.get_all_task_idx()
    obs_dim = int(np.prod(env.observation_space.shape))
    action_dim = int(np.prod(env.action_space.shape))
    reward_dim = 1

    # instantiate networks
    latent_dim = variant['latent_size']
    context_encoder_input_dim = 2 * obs_dim + action_dim + reward_dim if variant['algo_params']['use_next_obs_in_context'] else obs_dim + action_dim + reward_dim
    context_encoder_output_dim = latent_dim * 2 if variant['algo_params']['use_information_bottleneck'] else latent_dim
    net_size = variant['net_size']
    recurrent = variant['algo_params']['recurrent']
    encoder_model = RecurrentEncoder if recurrent else MlpEncoder

    context_encoder = encoder_model(
        hidden_sizes=[200, 200, 200],
        input_size=context_encoder_input_dim,
        output_size=context_encoder_output_dim,
        output_activation=torch.tanh,
    )
    qf1 = FlattenMlp(
        hidden_sizes=[net_size, net_size, net_size],
        input_size=obs_dim + action_dim + latent_dim,
        output_size=1,
    )
    qf2 = FlattenMlp(
        hidden_sizes=[net_size, net_size, net_size],
        input_size=obs_dim + action_dim + latent_dim,
        output_size=1,
    )
    vf = FlattenMlp(
        hidden_sizes=[net_size, net_size, net_size],
        input_size=obs_dim + latent_dim,
        output_size=1,
    )

    policy = TanhGaussianPolicy(
        hidden_sizes=[net_size, net_size, net_size],
        obs_dim=obs_dim + latent_dim,
        latent_dim=latent_dim,
        action_dim=action_dim,
    )

    agent = PEARLAgent(
        latent_dim,
        context_encoder,
        policy,
        **variant['algo_params']
    )
    if variant['algo_type'] == 'FOCAL':
        # critic network for divergence in dual form (see BRAC paper https://arxiv.org/abs/1911.11361)
        c = FlattenMlp(
            hidden_sizes=[net_size, net_size, net_size],
            input_size=obs_dim + action_dim + latent_dim,
            output_size=1
        )
        if 'randomize_tasks' in variant.keys() and variant['randomize_tasks']:
            rng = default_rng()
            train_tasks = rng.choice(len(tasks), size=variant['n_train_tasks'], replace=False)
            eval_tasks = set(range(len(tasks))).difference(train_tasks)
            if 'goal_radius' in variant['env_params']:
                algorithm = FOCALSoftActorCritic(
                    env=env,
                    train_tasks=train_tasks,
                    eval_tasks=eval_tasks,
                    nets=[agent, qf1, qf2, vf, c],
                    latent_dim=latent_dim,
                    goal_radius=variant['env_params']['goal_radius'],
                    **variant['algo_params']
                )
            else:
                algorithm = FOCALSoftActorCritic(
                    env=env,
                    train_tasks=list(tasks[:variant['n_train_tasks']]),
                    eval_tasks=list(tasks[-variant['n_eval_tasks']:]),
                    nets=[agent, qf1, qf2, vf, c],
                    latent_dim=latent_dim,
                    **variant['algo_params']
                )
        else:
            if 'goal_radius' in variant['env_params']:
                algorithm = FOCALSoftActorCritic(
                    env=env,
                    train_tasks=list(tasks[:variant['n_train_tasks']]),
                    eval_tasks=list(tasks[-variant['n_eval_tasks']:]),
                    nets=[agent, qf1, qf2, vf, c],
                    latent_dim=latent_dim,
                    goal_radius=variant['env_params']['goal_radius'],
                    **variant['algo_params']
                )
            else:
                algorithm = FOCALSoftActorCritic(
                    env=env,
                    train_tasks=list(tasks[:variant['n_train_tasks']]),
                    eval_tasks=list(tasks[-variant['n_eval_tasks']:]),
                    nets=[agent, qf1, qf2, vf, c],
                    latent_dim=latent_dim,
                    **variant['algo_params']
                )
    else:
        NotImplemented

    # optional GPU mode
    ptu.set_gpu_mode(variant['util_params']['use_gpu'], variant['util_params']['gpu_id'])
    if ptu.gpu_enabled():
        algorithm.to()

    # debugging triggers a lot of printing and logs to a debug directory
    DEBUG = variant['util_params']['debug']
    os.environ['DEBUG'] = str(int(DEBUG))

    # create logging directory
    # TODO support Docker
    exp_id = 'debug' if DEBUG else None
    experiment_log_dir = setup_logger(
        variant['env_name'],
        variant=variant,
        exp_id=exp_id,
        base_log_dir=variant['util_params']['base_log_dir'],
        seed=seed,
        snapshot_mode="all"
    )

    # optionally save eval trajectories as pkl files
    if variant['algo_params']['dump_eval_paths']:
        pickle_dir = experiment_log_dir + '/eval_trajectories'
        pathlib.Path(pickle_dir).mkdir(parents=True, exist_ok=True)

    # run the algorithm
    algorithm.train()

def deep_update_dict(fr, to):
    ''' update dict of dicts with new values '''
    # assume dicts have same keys
    for k, v in fr.items():
        if type(v) is dict:
            deep_update_dict(v, to[k])
        else:
            to[k] = v
    return to

@click.command()
@click.argument('config', default=None)
@click.option('--gpu', default=0)
def main(config, gpu):

    variant = default_config
    if config:
        with open(os.path.join(config)) as f:
            exp_params = json.load(f)
        variant = deep_update_dict(exp_params, variant)
    variant['util_params']['gpu_id'] = gpu

    # multi-processing
    p = mp.Pool(mp.cpu_count())
    if len(variant['seed_list']) > 0:
        p.starmap(experiment, product([variant], variant['seed_list']))
    else:
        experiment(variant)

if __name__ == "__main__":
    main()


import numpy as np
import torch
from torch import nn
from torch.optim import SGD
import torch.nn.functional as F
from inclearn.tools.metrics import ClassErrorMeter, AverageValueMeter

def finetune_last_layer(
    logger,
    network,
    loader,
    n_class,
    nepoch=30,
    lr=0.1,
    scheduling=[15, 35],
    lr_decay=0.1,
    weight_decay=5e-4,
    loss_type="ce",
    temperature=5.0,
    test_loader=None,
):
    network.eval()

    for conv in network.module.convnets:
        for p in conv.parameters():
            p.requires_grad = False

    #if hasattr(network.module, "convnets"):
    #    for net in network.module.convnets:
    #        net.eval()
    #else:
    #    network.module.convnet.eval()
    optim = SGD(network.module.classifier.parameters(), lr=lr, momentum=0.9, weight_decay=weight_decay)
    scheduler = torch.optim.lr_scheduler.MultiStepLR(optim, scheduling, gamma=lr_decay)

    if loss_type == "ce":
        criterion = nn.CrossEntropyLoss()
    else:
        criterion = nn.BCEWithLogitsLoss()

    logger.info("Begin finetuning last layer")

    for i in range(nepoch):
        total_loss = 0.0
        total_correct = 0.0
        total_count = 0
        # print(f"dataset loader length {len(loader.dataset)}")
        for inputs, targets in loader:
            inputs, targets = inputs.cuda(), targets.cuda()
            if loss_type == "bce":
                targets = to_onehot(targets, n_class)
            outputs = network(inputs)['logit']
            _, preds = outputs.max(1)
            optim.zero_grad()
            loss = criterion(outputs / temperature, targets)
            loss.backward()
            optim.step()
            total_loss += loss * inputs.size(0)
            total_correct += (preds == targets).sum()
            total_count += inputs.size(0)

        if test_loader is not None:
            test_correct = 0.0
            test_count = 0.0
            with torch.no_grad():
                for inputs, targets in test_loader:
                    outputs = network(inputs.cuda())['logit']
                    _, preds = outputs.max(1)
                    test_correct += (preds.cpu() == targets).sum().item()
                    test_count += inputs.size(0)

        scheduler.step()
        if test_loader is not None:
            logger.info(
                "Epoch %d finetuning loss %.3f acc %.3f Eval %.3f" %
                (i, total_loss.item() / total_count, total_correct.item() / total_count, test_correct / test_count))
        else:
            logger.info("Epoch %d finetuning loss %.3f acc %.3f" %
                        (i, total_loss.item() / total_count, total_correct.item() / total_count))
    return network

def extract_features(model, loader):
    targets, features = [], []
    model.eval()
    with torch.no_grad():
        for _inputs, _targets in loader:
            _inputs = _inputs.cuda()
            _targets = _targets.numpy()
            _features = model(_inputs)['feature'].detach().cpu().numpy()
            features.append(_features)
            targets.append(_targets)

    return np.concatenate(features), np.concatenate(targets)

def calc_class_mean(network, loader, class_idx, metric):
    EPSILON = 1e-8
    features, targets = extract_features(network, loader)
    # norm_feats = features/(np.linalg.norm(features, axis=1)[:,np.newaxis]+EPSILON)
    # examplar_mean = norm_feats.mean(axis=0)
    examplar_mean = features.mean(axis=0)
    if metric == "cosine" or metric == "weight":
        examplar_mean /= (np.linalg.norm(examplar_mean) + EPSILON)
    return examplar_mean

def update_classes_mean(network, inc_dataset, n_classes, task_size, share_memory=None, metric="cosine", EPSILON=1e-8):
    loader = inc_dataset._get_loader(inc_dataset.data_inc,
                                     inc_dataset.targets_inc,
                                     shuffle=False,
                                     share_memory=share_memory,
                                     mode="test")
    class_means = np.zeros((n_classes, network.module.features_dim))
    count = np.zeros(n_classes)
    network.eval()
    with torch.no_grad():
        for x, y in loader:
            feat = network(x.cuda())['feature']
            for lbl in torch.unique(y):
                class_means[lbl] += feat[y == lbl].sum(0).cpu().numpy()
                count[lbl] += feat[y == lbl].shape[0]
        for i in range(n_classes):
            class_means[i] /= count[i]
            if metric == "cosine" or metric == "weight":
                class_means[i] /= (np.linalg.norm(class_means) + EPSILON)
    return class_means


from __future__ import absolute_import
from __future__ import print_function, division
import sys
import time
import datetime
import argparse
import os
import numpy as np
import os.path as osp
import math
from random import sample
from scipy import io

import torchvision
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
import torch.backends.cudnn as cudnn

import models
from models.PCB import PCB_test
# from ReID_attr import get_target_withattr # Need Attribute file
from opts import get_opts, Imagenet_mean, Imagenet_stddev
from GD import Generator, MS_Discriminator, Pat_Discriminator, GANLoss, weights_init
from advloss import DeepSupervision, adv_CrossEntropyLoss, adv_CrossEntropyLabelSmooth, adv_TripletLoss
from util import data_manager
from util.dataset_loader import ImageDataset
from util.utils import fliplr, Logger, save_checkpoint, visualize_ranked_results
from util.eval_metrics import make_results
from util.samplers import RandomIdentitySampler, AttrPool

# Training settings
parser = argparse.ArgumentParser(description='adversarial attack')
parser.add_argument('--root', type=str, default='data', help="root path to data directory")
parser.add_argument('--targetmodel', type=str, default='aligned', choices=models.get_names())
parser.add_argument('--dataset', type=str, default='market1501', choices=data_manager.get_names())
# PATH
parser.add_argument('--G_resume_dir', type=str, default='', metavar='path to resume G')
parser.add_argument('--pre_dir', type=str, default='models', help='path to be attacked model')
parser.add_argument('--attr_dir', type=str, default='', help='path to attribute file')
parser.add_argument('--save_dir', type=str, default='logs', help='path to save model')
parser.add_argument('--vis_dir', type=str, default='vis', help='path to save visualization result')
parser.add_argument('--ablation', type=str, default='', help='for ablation study')
# var
parser.add_argument('--mode', type=str, default='train', help='train/test')
parser.add_argument('--D', type=str, default='MSGAN', help='Type of discriminator: PatchGAN or Multi-stage GAN')
parser.add_argument('--normalization', type=str, default='bn', help='bn or in')
parser.add_argument('--loss', type=str, default='xent_htri', choices=['cent', 'xent', 'htri', 'xent_htri'])
parser.add_argument('--ak_type', type=int, default=-1, help='-1 if non-targeted, 1 if attribute attack')
parser.add_argument('--attr_key', type=str, default='upwhite', help='[attribute, value]')
parser.add_argument('--attr_value', type=int, default=2, help='[attribute, value]')
parser.add_argument('--mag_in', type=float, default=16.0, help='l_inf magnitude of perturbation')
parser.add_argument('--temperature', type=float, default=-1, help="tau in paper")
parser.add_argument('--usegumbel', action='store_true', default=False, help='whether to use gumbel softmax')
parser.add_argument('--use_SSIM', type=int, default=2, help="0: None, 1: SSIM, 2: MS-SSIM ")
# Base
parser.add_argument('--train_batch', default=32, type=int,help="train batch size")
parser.add_argument('--test_batch', default=32, type=int, help="test batch size")
parser.add_argument('--epoch', type=int, default=50, help='number of epochs to train for')

parser.add_argument('--margin', type=float, default=0.3, help="margin for triplet loss")
parser.add_argument('--num_ker', type=int, default=32, help='generator filters in first conv layer')
parser.add_argument('--lr', type=float, default=0.0002, help='Learning Rate. Default=0.002')
parser.add_argument('--beta1', type=float, default=0.5, help='beta1 for adam. default=0.5')
parser.add_argument('--seed', type=int, default=123, help='random seed to use. Default=123')
parser.add_argument('--print_freq', type=int, default=20, help="print frequency")
parser.add_argument('--eval_freq', type=int, default=1, help="eval frequency")
parser.add_argument('--usevis', action='store_true', default=False, help='whether to save vis')

args = parser.parse_args()
is_training = args.mode == 'train'
attr_list = [args.attr_key, args.attr_value]
attr_matrix = None
if args.attr_dir:
  assert args.dataset in ['dukemtmcreid', 'market1501']
  attr_matrix = io.loadmat(args.attr_dir)
  args.ablation = osp.join('attr', args.attr_key + '=' + str(args.attr_value))

pre_dir = osp.join(args.pre_dir, args.targetmodel, args.dataset+'.pth.tar')
save_dir = osp.join(args.save_dir, args.targetmodel, args.dataset, args.ablation)
vis_dir = osp.join(args.vis_dir, args.targetmodel, args.dataset, args.ablation)

def main(opt):
  if not osp.exists(save_dir): os.makedirs(save_dir)
  if not osp.exists(vis_dir): os.makedirs(vis_dir)

  use_gpu = torch.cuda.is_available()
  pin_memory = True if use_gpu else False

  if args.mode == 'train':
    sys.stdout = Logger(osp.join(save_dir, 'log_train.txt'))
  else:
    sys.stdout = Logger(osp.join(save_dir, 'log_test.txt'))
  print("==========\nArgs:{}\n==========".format(args))

  if use_gpu:
    print("GPU mode")
    cudnn.benchmark = True
    torch.cuda.manual_seed(args.seed)
  else:
    print("CPU mode")

  ### Setup dataset loader ###
  print("Initializing dataset {}".format(args.dataset))
  dataset = data_manager.init_img_dataset(root=args.root, name=args.dataset, split_id=opt['split_id'], cuhk03_labeled=opt['cuhk03_labeled'], cuhk03_classic_split=opt['cuhk03_classic_split'])
  if args.ak_type < 0:  # non-targeted
    trainloader = DataLoader(ImageDataset(dataset.train, transform=opt['transform_train']), sampler=RandomIdentitySampler(dataset.train, num_instances=opt['num_instances']), batch_size=args.train_batch, num_workers=opt['workers'], pin_memory=pin_memory, drop_last=True)
  elif args.ak_type > 0:  # attribute attack
    trainloader = DataLoader(ImageDataset(dataset.train, transform=opt['transform_train']), sampler=AttrPool(dataset.train, args.dataset, attr_matrix, attr_list, sample_num=16), batch_size=args.train_batch, num_workers=opt['workers'], pin_memory=pin_memory, drop_last=True)
  queryloader = DataLoader(ImageDataset(dataset.query, transform=opt['transform_test']), batch_size=args.test_batch, shuffle=False, num_workers=opt['workers'], pin_memory=pin_memory, drop_last=False)
  galleryloader = DataLoader(ImageDataset(dataset.gallery, transform=opt['transform_test']), batch_size=args.test_batch, shuffle=False, num_workers=opt['workers'], pin_memory=pin_memory, drop_last=False)

  ### Prepare criterion ###
  if args.ak_type<0:
    clf_criterion = adv_CrossEntropyLabelSmooth(num_classes=dataset.num_train_pids, use_gpu=use_gpu) if args.loss in ['xent', 'xent_htri'] else adv_CrossEntropyLoss(use_gpu=use_gpu)
  else:
    clf_criterion = nn.MultiLabelSoftMarginLoss()
  metric_criterion = adv_TripletLoss(margin=args.margin, ak_type=args.ak_type)
  criterionGAN = GANLoss()

  ### Prepare pretrained model ###
  target_net = models.init_model(name=args.targetmodel, pre_dir=pre_dir, num_classes=dataset.num_train_pids)
  check_freezen(target_net, need_modified=True, after_modified=False)

  ### Prepare main net ###
  G = Generator(3, 3, args.num_ker, norm=args.normalization).apply(weights_init)
  if args.D == 'PatchGAN':
    D = Pat_Discriminator(input_nc=6, norm=args.normalization).apply(weights_init)
  elif args.D == 'MSGAN':
    D = MS_Discriminator(input_nc=6, norm=args.normalization, temperature=args.temperature, use_gumbel=args.usegumbel).apply(weights_init)
  check_freezen(G, need_modified=True, after_modified=True)
  check_freezen(D, need_modified=True, after_modified=True)
  print("Model size: {:.5f}M".format((sum(g.numel() for g in G.parameters())+sum(d.numel() for d in D.parameters()))/1000000.0))
  # setup optimizer
  optimizer_G = optim.Adam(G.parameters(), lr=args.lr, betas=(args.beta1, 0.999))
  optimizer_D = optim.Adam(D.parameters(), lr=args.lr, betas=(args.beta1, 0.999))

  if use_gpu:
    test_target_net = nn.DataParallel(target_net).cuda() if not args.targetmodel == 'pcb' else nn.DataParallel(PCB_test(target_net)).cuda()
    target_net = nn.DataParallel(target_net).cuda()
    G = nn.DataParallel(G).cuda()
    D = nn.DataParallel(D).cuda()

  if args.mode == 'test':
    epoch = 'test'
    test(G, D, test_target_net, dataset, queryloader, galleryloader, epoch, use_gpu, is_test=True)
    return 0

  # Ready
  start_time = time.time()
  train_time = 0
  worst_mAP, worst_rank1, worst_rank5, worst_rank10, worst_epoch = np.inf, np.inf, np.inf, np.inf, 0
  best_hit, best_epoch = -np.inf, 0
  print("==> Start training")

  for epoch in range(1,args.epoch+1):
    start_train_time = time.time()
    train(epoch, G, D, target_net, criterionGAN, clf_criterion, metric_criterion, optimizer_G, optimizer_D, trainloader, use_gpu)
    train_time += round(time.time() - start_train_time)

    if epoch % args.eval_freq == 0:
      print("==> Eval at epoch {}".format(epoch))
      if args.ak_type < 0:
        cmc, mAP = test(G, D, test_target_net, dataset, queryloader, galleryloader, epoch, use_gpu, is_test=False)
        is_worst = cmc[0]<=worst_rank1 and cmc[1]<=worst_rank5 and cmc[2]<=worst_rank10 and mAP<=worst_mAP
        if is_worst:
          worst_mAP, worst_rank1, worst_epoch = mAP, cmc[0], epoch
        print("==> Worst_epoch is {}, Worst mAP {:.1%}, Worst rank-1 {:.1%}".format(worst_epoch, worst_mAP, worst_rank1))
        save_checkpoint(G.state_dict(), is_worst, 'G', osp.join(save_dir, 'G_ep' + str(epoch) + '.pth.tar'))
        save_checkpoint(D.state_dict(), is_worst, 'D', osp.join(save_dir, 'D_ep' + str(epoch) + '.pth.tar'))

      else:
        all_hits = test(G, D, target_net, dataset, queryloader, galleryloader, epoch, use_gpu, is_test=False)
        is_best = all_hits[0]>=best_hit
        if is_best:
          best_hit, best_epoch = all_hits[0], epoch
        print("==> Best_epoch is {}, Best rank-1 {:.1%}".format(best_epoch, best_hit))
        save_checkpoint(G.state_dict(), is_best, 'G', osp.join(save_dir, 'G_ep' + str(epoch) + '.pth.tar'))
        save_checkpoint(D.state_dict(), is_best, 'D', osp.join(save_dir, 'D_ep' + str(epoch) + '.pth.tar'))

  elapsed = round(time.time() - start_time)
  elapsed = str(datetime.timedelta(seconds=elapsed))
  train_time = str(datetime.timedelta(seconds=train_time))
  print("Finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}.".format(elapsed, train_time))

def train(epoch, G, D, target_net, criterionGAN, clf_criterion, metric_criterion, optimizer_G, optimizer_D, trainloader, use_gpu):
  G.train()
  D.train()
  global is_training
  is_training = True

  for batch_idx, (imgs, pids, _, pids_raw) in enumerate(trainloader):
    if use_gpu:
      imgs, pids, pids_raw = imgs.cuda(), pids.cuda(), pids_raw.cuda()

    new_imgs, mask = perturb(imgs, G, D, train_or_test='train')
    new_imgs = new_imgs.cuda()
    mask = mask.cuda()
    # Fake Detection and Loss
    pred_fake_pool, _ = D(torch.cat((imgs, new_imgs.detach()), 1))
    loss_D_fake = criterionGAN(pred_fake_pool, False)

    # Real Detection and Loss
    num = args.train_batch//2
    pred_real, _ = D(torch.cat((imgs[0:num,:,:,:], imgs[num:,:,:,:].detach()), 1))
    loss_D_real = criterionGAN(pred_real, True)

    # GAN loss (Fake Passability Loss)
    pred_fake, _ = D(torch.cat((imgs, new_imgs), 1))
    loss_G_GAN = criterionGAN(pred_fake, True)

    # Re-ID advloss
    ls = target_net(new_imgs, is_training)
    if len(ls) == 1: new_outputs = ls[0]
    if len(ls) == 2: new_outputs, new_features = ls
    if len(ls) == 3: new_outputs, new_features, new_local_features = ls
    xent_loss, global_loss, loss_G_ssim = 0, 0, 0
    targets = None

    if args.loss in ['cent', 'xent', 'xent_htri']:
      if args.ak_type < 0:
        xent_loss = DeepSupervision(clf_criterion, new_outputs, pids) if isinstance(new_features, (tuple, list)) else clf_criterion(new_outputs, pids)

      elif args.ak_type > 0:
        targets = get_target_withattr(attr_matrix, args.dataset, attr_list, pids, pids_raw).float().cuda()
        xent_loss = 0#DeepSupervision(clf_criterion, new_outputs, targets) if isinstance(new_features, (tuple, list)) else clf_criterion(new_outputs, targets)

    if args.loss in ['htri', 'xent_htri']:
      assert len(ls) >= 2
      global_loss = DeepSupervision(metric_criterion, new_features, pids, targets) if isinstance(new_features, (tuple, list)) else metric_criterion(new_features, pids, targets)

    loss_G_ReID = (xent_loss+ global_loss)*opt['ReID_factor']

    # # SSIM loss
    if not args.use_SSIM == 0:
      from util.ms_ssim import msssim, ssim
      loss_func = msssim if args.use_SSIM == 2 else ssim
      loss_G_ssim = (1-loss_func(imgs, new_imgs))*0.1

    ############## Forward ###############
    loss_D = (loss_D_fake + loss_D_real)/2
    loss_G = loss_G_GAN + loss_G_ReID + loss_G_ssim
    ############## Backward #############
    # update generator weights
    optimizer_G.zero_grad()
    # loss_G.backward(retain_graph=True)
    loss_G.backward()
    optimizer_G.step()
    # update discriminator weights
    optimizer_D.zero_grad()
    loss_D.backward()
    optimizer_D.step()
    if (batch_idx+1) % args.print_freq == 0:
      print("===> Epoch[{}]({}/{}) loss_D: {:.4f} loss_G_GAN: {:.4f} loss_G_ReID: {:.4f} loss_G_SSIM: {:.4f}".format(epoch, batch_idx, len(trainloader), loss_D.item(), loss_G_GAN.item(), loss_G_ReID.item(), loss_G_ssim))

def test(G, D, target_net, dataset, queryloader, galleryloader, epoch, use_gpu, is_test=False, ranks=[1, 5, 10, 20]):
  global is_training
  is_training = False
  if args.mode == 'test' and args.G_resume_dir:
    G_resume_dir, D_resume_dir = args.G_resume_dir, args.G_resume_dir.replace('G', 'D')
    G_checkpoint, D_checkpoint = torch.load(G_resume_dir), torch.load(D_resume_dir)
    G_state_dict = G_checkpoint['state_dict'] if isinstance(G_checkpoint, dict) and 'state_dict' in G_checkpoint else G_checkpoint
    D_state_dict = D_checkpoint['state_dict'] if isinstance(D_checkpoint, dict) and 'state_dict' in D_checkpoint else D_checkpoint

    G.load_state_dict(G_state_dict)
    D.load_state_dict(D_state_dict)
    print("Sucessfully, loading {} and {}".format(G_resume_dir, D_resume_dir))

  with torch.no_grad():
    qf, lqf, new_qf, new_lqf, q_pids, q_camids = extract_and_perturb(queryloader, G, D, target_net, use_gpu, query_or_gallery='query', is_test=is_test, epoch=epoch)
    gf, lgf, g_pids, g_camids = extract_and_perturb(galleryloader, G, D, target_net, use_gpu, query_or_gallery='gallery', is_test=is_test, epoch=epoch)

    if args.ak_type > 0:
      distmat, hits, ignore_list = make_results(new_qf, gf, new_lqf, lgf, q_pids, g_pids, q_camids, g_camids, args.targetmodel, args.ak_type, attr_matrix, args.dataset, attr_list)
      print("Hits rate, Rank-{}: {:.1%}, Rank-{}: {:.1%}, Rank-{}: {:.1%}, Rank-{}: {:.1%}".format(ranks[0], hits[ranks[0]-1], ranks[1], hits[ranks[1]-1], ranks[2], hits[ranks[2]-1], ranks[3], hits[ranks[3]-1]))
      if not is_test:
        return hits

    else:
      if is_test:
        distmat, cmc, mAP = make_results(qf, gf, lqf, lgf, q_pids, g_pids, q_camids, g_camids, args.targetmodel, args.ak_type)
        new_distmat, new_cmc, new_mAP = make_results(new_qf, gf, new_lqf, lgf, q_pids, g_pids, q_camids, g_camids, args.targetmodel, args.ak_type)
        print("Results ----------")
        print("Before, mAP: {:.1%}, Rank-{}: {:.1%}, Rank-{}: {:.1%}, Rank-{}: {:.1%}, Rank-{}: {:.1%}".format(mAP, ranks[0], cmc[ranks[0]-1], ranks[1], cmc[ranks[1]-1], ranks[2], cmc[ranks[2]-1], ranks[3], cmc[ranks[3]-1]))
        print("After , mAP: {:.1%}, Rank-{}: {:.1%}, Rank-{}: {:.1%}, Rank-{}: {:.1%}, Rank-{}: {:.1%}".format(new_mAP, ranks[0], new_cmc[ranks[0]-1], ranks[1], new_cmc[ranks[1]-1], ranks[2], new_cmc[ranks[2]-1], ranks[3], new_cmc[ranks[3]-1]))
        if args.usevis:
          visualize_ranked_results(distmat, dataset, save_dir=osp.join(vis_dir, 'origin_results'), topk=20)
        if args.usevis:
          visualize_ranked_results(new_distmat, dataset, save_dir=osp.join(vis_dir, 'polluted_results'), topk=20)
      else:
        _, new_cmc, new_mAP = make_results(new_qf, gf, new_lqf, lgf, q_pids, g_pids, q_camids, g_camids, args.targetmodel, args.ak_type)
        print("mAP: {:.1%}, Rank-{}: {:.1%}, Rank-{}: {:.1%}, Rank-{}: {:.1%}, Rank-{}: {:.1%}".format(new_mAP, ranks[0], new_cmc[ranks[0]-1], ranks[1], new_cmc[ranks[1]-1], ranks[2], new_cmc[ranks[2]-1], ranks[3], new_cmc[ranks[3]-1]))
        return new_cmc, new_mAP

def extract_and_perturb(loader, G, D, target_net, use_gpu, query_or_gallery, is_test, epoch):
  f, lf, new_f, new_lf, l_pids, l_camids = [], [], [], [], [], []
  ave_mask, num = 0, 0
  for batch_idx, (imgs, pids, camids, pids_raw) in enumerate(loader):
    if use_gpu:
      imgs = imgs.cuda()
    ls = extract(imgs, target_net)
    if len(ls) == 1: features = ls[0]
    if len(ls) == 2:
      features, local_features = ls
      lf.append(local_features.detach().data.cpu())

    f.append(features.detach().data.cpu())
    l_pids.extend(pids)
    l_camids.extend(camids)

    if query_or_gallery == 'query':
      G.eval()
      D.eval()
      new_imgs, delta, mask = perturb(imgs, G, D, train_or_test='test')
      ave_mask += torch.sum(mask.detach()).cpu().numpy()
      num += imgs.size(0)

      ls = extract(new_imgs, target_net)
      if len(ls) == 1: new_features = ls[0]
      if len(ls) == 2:
        new_features, new_local_features = ls
        new_lf.append(new_local_features.detach().data.cpu())
      new_f.append(new_features.detach().data.cpu())

      ls = [imgs, new_imgs, delta, mask]
      if is_test:
        save_img(ls, pids, camids, epoch, batch_idx)

  f = torch.cat(f, 0)
  if not lf == []: lf = torch.cat(lf, 0)
  l_pids, l_camids = np.asarray(l_pids), np.asarray(l_camids)

  print("Extracted features for {} set, obtained {}-by-{} matrix".format(query_or_gallery, f.size(0), f.size(1)))
  if query_or_gallery == 'gallery':
    return [f, lf, l_pids, l_camids]
  elif query_or_gallery == 'query':
    new_f = torch.cat(new_f, 0)
    if not new_lf == []:
      new_lf = torch.cat(new_lf, 0)
    return [f, lf, new_f, new_lf, l_pids, l_camids]

def extract(imgs, target_net):
  if args.targetmodel in ['pcb', 'lsro']:
    ls = [target_net(imgs, is_training)[0] + target_net(fliplr(imgs), is_training)[0]]
  else:
    ls = target_net(imgs, is_training)
  for i in range(len(ls)): ls[i] = ls[i].data.cpu()
  return ls

def perturb(imgs, G, D, train_or_test='test'):
  n,c,h,w = imgs.size()
  delta = G(imgs)
  delta = L_norm(delta, train_or_test)
  new_imgs = torch.add(imgs.cuda(), delta[0:imgs.size(0)].cuda())

  _, mask = D(torch.cat((imgs, new_imgs.detach()), 1))
  delta = delta * mask
  new_imgs = torch.add(imgs.cuda(), delta[0:imgs.size(0)].cuda())

  for c in range(3):
    new_imgs.data[:,c,:,:] = new_imgs.data[:,c,:,:].clamp(new_imgs.data[:,c,:,:].min(), new_imgs.data[:,c,:,:].max()) # do clamping per channel
  if train_or_test == 'train':
    return new_imgs, mask
  elif train_or_test == 'test':
    return new_imgs, delta, mask

def L_norm(delta, mode='train'):
  delta.data += 1
  delta.data *= 0.5

  for c in range(3):
    delta.data[:,c,:,:] = (delta.data[:,c,:,:] - Imagenet_mean[c]) / Imagenet_stddev[c]

  bs = args.train_batch if (mode == 'train') else args.test_batch
  for i in range(bs):
    # do per channel l_inf normalization
    for ci in range(3):
      try:
        l_inf_channel = delta[i,ci,:,:].data.abs().max()
        # l_inf_channel = torch.norm(delta[i,ci,:,:]).data
        mag_in_scaled_c = args.mag_in/(255.0*Imagenet_stddev[ci])
        delta[i,ci,:,:].data *= np.minimum(1.0, mag_in_scaled_c / l_inf_channel.cpu()).float().cuda()
      except IndexError:
        break
  return delta

def save_img(ls, pids, camids, epoch, batch_idx):
  image, new_image, delta, mask = ls
  # undo normalize image color channels
  delta_tmp = torch.zeros(delta.size())
  for c in range(3):
    image.data[:,c,:,:] = (image.data[:,c,:,:] * Imagenet_stddev[c]) + Imagenet_mean[c]
    new_image.data[:,c,:,:] = (new_image.data[:,c,:,:] * Imagenet_stddev[c]) + Imagenet_mean[c]
    delta_tmp.data[:,c,:,:] = (delta.data[:,c,:,:] * Imagenet_stddev[c]) + Imagenet_mean[c]

  if args.usevis:
    torchvision.utils.save_image(image.data, osp.join(vis_dir, 'original_epoch{}_batch{}.png'.format(epoch, batch_idx)))
    torchvision.utils.save_image(new_image.data, osp.join(vis_dir, 'polluted_epoch{}_batch{}.png'.format(epoch, batch_idx)))
    torchvision.utils.save_image(delta_tmp.data, osp.join(vis_dir, 'delta_epoch{}_batch{}.png'.format(epoch, batch_idx)))
    torchvision.utils.save_image(mask.data*255, osp.join(vis_dir, 'mask_epoch{}_batch{}.png'.format(epoch, batch_idx)))

def check_freezen(net, need_modified=False, after_modified=None):
  # print(net)
  cc = 0
  for child in net.children():
    for param in child.parameters():
      if need_modified: param.requires_grad = after_modified
      # if param.requires_grad: print('child', cc , 'was active')
      # else: print('child', cc , 'was forzen')
    cc += 1

if __name__ == '__main__':
  opt = get_opts(args.targetmodel)
  main(opt)


import numpy as np
import pandas as pd
from Data import Data
import copy
import sys
sys.path.append("..")
import load_data
import feature_extraction
import feature_conversational_extraction
import feature_recommendation_task_extraction
import helper

class ConstructFeature(object):

    @staticmethod
    def constructFeature():

        # Prepared Feature Vector From Extracted Conversations
        # (1) Utterance Features (from the perspective of Seeker)
        # (2) Response Features (from the perspective of Recommender)
        # (3) Conversational Features
        # (4) Recommendation Task-Specific Features

        Data.feature_vector = []
        fitted_vectorizer = None
        entity_label_set = None
        POS_tag_set = None
        first_level_intent_dict = None
        whole_utterance_text = Data.whole_utterance_dataframe[['text']].values.reshape(-1).tolist()
        if Data.content_features == 1:
            fitted_vectorizer = feature_extraction.initialize_fitted_vectorizer (Data.whole_utterance_dataframe)
            entity_label_set = feature_extraction.get_occurred_entity_label_set(whole_utterance_text)
        if Data.discourse_features == 1:
            POS_tag_set = feature_extraction.get_occurred_POS_tag_set(whole_utterance_text)
        if Data.task_specific_features == 1:
            first_level_intent_dict = feature_recommendation_task_extraction.get_first_level_intent_dict(Data.higher_level_intent_dict_file)

        for each_conversation_info in Data.conversation_previous_N_turns:
            # Step 1: Extract Seekers' Utterance and Recommender' Responses

            Data.utterance_dataframe = load_data.load_utterance_data_from_extracted_conv(each_conversation_info)
            Data.user_utterance_dataframe  = Data.utterance_dataframe [Data.utterance_dataframe['role'] == 'seeker']
            Data.recommender_response_dataframe  = Data.utterance_dataframe [Data.utterance_dataframe['role'] == 'recommender']

            # print(Data.utterance_dataframe)
            # print(Data.user_utterance_dataframe)
            # print(Data.recommender_response_dataframe)

            # Step 2: Extract Utterance/Response Features [sentence-level features (1,2)]
            user_utterance_feature_vector_avg = []
            recommender_response_feature_vector_avg = []

            user_utterance_feature_vector_list = []
            recommender_response_feature_vector_list = []

            if len(Data.user_utterance_dataframe) > 0:
                if Data.content_features == 1 or Data.discourse_features == 1 or Data.sentiment_features == 1:

                    user_utterance_feature_vector_list = feature_extraction.extract_features(Data.user_utterance_dataframe, fitted_vectorizer, entity_label_set, POS_tag_set,  Data.content_features, Data.discourse_features,Data.sentiment_features, Data.pos_file, Data.neg_file)
                    user_utterance_feature_vector_avg = np.mean( np.array(user_utterance_feature_vector_list), axis=0)

                    assert (len(user_utterance_feature_vector_avg) == len(user_utterance_feature_vector_list[0]))

            if len(Data.recommender_response_dataframe) > 0:
                if Data.content_features == 1 or Data.discourse_features == 1 or Data.sentiment_features == 1:
                    recommender_response_feature_vector_list = feature_extraction.extract_features(Data.recommender_response_dataframe, fitted_vectorizer,  entity_label_set, POS_tag_set, Data.content_features, Data.discourse_features,Data.sentiment_features, Data.pos_file, Data.neg_file)
                    recommender_response_feature_vector_avg = np.mean( np.array(recommender_response_feature_vector_list), axis=0)
                    assert (len(recommender_response_feature_vector_avg) == len(recommender_response_feature_vector_list[0]))

            # Step 3: Extract Conversational label Features [Conversation-level features (1,2)]
            user_intent_feature_vector = []
            recommender_action_feature_vector = []

            if len(Data.user_utterance_dataframe) > 0:

                if Data.conversational_features == 1:
                    user_intent_feature_vector = feature_conversational_extraction.extract_multi_turn_label_features(Data.user_utterance_dataframe, Data.intent_label_indext_dict)

            if len(Data.recommender_response_dataframe) > 0:
                if Data.conversational_features == 1:
                    recommender_action_feature_vector = feature_conversational_extraction.extract_multi_turn_label_features(Data.recommender_response_dataframe, Data.action_label_indext_dict)

            # Step 4: Extract Recommendation Interaction Task-Specific Features [Conversation-level features (1,2)]
            task_specific_features = []
            if len(Data.user_utterance_dataframe) > 0 and len(Data.recommender_response_dataframe) > 0 and Data.task_specific_features == 1:

                task_specific_features = feature_recommendation_task_extraction.extract_task_specific_features(Data.user_utterance_dataframe, Data.recommender_response_dataframe, first_level_intent_dict)

            integrated_feature_vector = np.concatenate((user_utterance_feature_vector_avg , recommender_response_feature_vector_avg, user_intent_feature_vector,recommender_action_feature_vector, task_specific_features), axis=None)
            # print(len(integrated_feature_vector))
            assert( len(integrated_feature_vector) == len(user_utterance_feature_vector_avg) + len(recommender_response_feature_vector_avg) + len(user_intent_feature_vector) + len(recommender_action_feature_vector) + len(task_specific_features) )
            Data.feature_vector.append(integrated_feature_vector)

        Data.feature_vector = np.array(Data.feature_vector)

        # save features into csv file
        feature_csv_file = "feature_file/pre_turns_%d_content_%d_discourse_%d_sentiment_%d_conversation_%d_task_%d_feat.csv" % (Data.num_previous_turns, Data.content_features, Data.discourse_features,Data.sentiment_features, Data.conversational_features, Data.task_specific_features)
        np.savetxt(feature_csv_file, Data.feature_vector , delimiter= ",", fmt='%1.4e')
        # exit()
        #------------------------------------------------
        helper.print_current_time()
        print("Extract Features For Each extracted Conversations (N utterances) and Construct Feature Vector")

        print("[ %d conversations with the size of features vector (%d) ]" % (len (Data.feature_vector) , len (Data.feature_vector[0]) ))
        #------------------------------------------------

        # save features into csv file
        feature_csv_file = "feature_file/pre_turns_%d_content_%d_discourse_%d_sentiment_%d_conversation_%d_task_%d_feat.csv" % (Data.num_previous_turns, Data.content_features, Data.discourse_features,Data.sentiment_features, Data.conversational_features, Data.task_specific_features)
        np.savetxt(feature_csv_file, Data.feature_vector , delimiter= ",", fmt='%1.4e')
        # exit()


#!/usr/bin/env python3
"""Play some random bleeps.

This example shows the feature of playing a buffer at a given absolute time.
Since all play_buffer() invocations are (deliberately) done at once, this puts
some strain on the "action queue".
The "qsize" has to be increased in order to handle this.

This example also shows that NumPy arrays can be used, as long as they are
C-contiguous and use the 'float32' data type.

"""

from __future__ import division  # Only needed for Python 2.x
import numpy as np
import rtmixer
import sounddevice as sd

seed = 99

device = None
blocksize = 0
latency = 'low'
samplerate = 44100

bleeps = 300
qsize = 512  # Must be a power of 2

attack = 0.005
release = 0.1
pitch_min = 40
pitch_max = 80
duration_min = 0.2
duration_max = 0.6
amplitude_min = 0.05
amplitude_max = 0.15
start_min = 0
start_max = 10
channels = None

if duration_min < max(attack, release):
    raise ValueError('minimum duration is too short')

fade_in = np.linspace(0, 1, num=int(samplerate * attack))
fade_out = np.linspace(1, 0, num=int(samplerate * release))

r = np.random.RandomState(seed)

bleeplist = []

if channels is None:
    channels = sd.default.channels['output']
    if channels is None:
        channels = sd.query_devices(device, 'output')['max_output_channels']

for _ in range(bleeps):
    duration = r.uniform(duration_min, duration_max)
    amplitude = r.uniform(amplitude_min, amplitude_max)
    pitch = r.uniform(pitch_min, pitch_max)
    # Convert MIDI pitch (https://en.wikipedia.org/wiki/MIDI_Tuning_Standard)
    frequency = 2 ** ((pitch - 69) / 12) * 440
    t = np.arange(int(samplerate * duration)) / samplerate
    bleep = amplitude * np.sin(2 * np.pi * frequency * t, dtype='float32')
    bleep[:len(fade_in)] *= fade_in
    bleep[-len(fade_out):] *= fade_out

    # Note: Arrays must be 32-bit float and C contiguous!
    assert bleep.dtype == 'float32'
    assert bleep.flags.c_contiguous
    bleeplist.append(bleep)

actionlist = []
with rtmixer.Mixer(device=device, channels=channels, blocksize=blocksize,
                   samplerate=samplerate, latency=latency, qsize=qsize) as m:
    start_time = m.time
    actionlist = [
        m.play_buffer(bleep,
                      channels=[r.randint(channels) + 1],
                      start=start_time + r.uniform(start_min, start_max),
                      allow_belated=False)
        for bleep in bleeplist
    ]
    while m.actions:
        sd.sleep(100)
print('{0} buffer underflows in {1} processed audio blocks'.format(
    m.stats.output_underflows, m.stats.blocks))

belated = 0
min_delay = np.inf
max_delay = -np.inf
for action in actionlist:
    assert action.type == rtmixer.PLAY_BUFFER
    # NB: action.allow_belated might have been invalidated
    assert action.requested_time != 0
    if not action.actual_time:
        belated += 1
        assert action.done_frames == 0
        continue
    assert action.done_frames == action.total_frames
    delay = action.actual_time - action.requested_time
    if delay > max_delay:
        max_delay = delay
    if delay < min_delay:
        min_delay = delay

print('total number of bleeps:', len(actionlist))
print('belated bleeps (not played):', belated)
print('maxiumum delay (in usec):', max_delay * 1000 * 1000)
print('maxiumum negative delay: ', -min_delay * 1000 * 1000)
print('half sampling period:    ', 0.5 * 1000 * 1000 / samplerate)


from __future__ import absolute_import
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable

from model.utils.config import cfg
from .proposal_layer import _ProposalLayer
from .anchor_target_layer import _AnchorTargetLayer
from model.utils.net_utils import _smooth_l1_loss

import numpy as np
import math
import pdb
import time

class _RPN(nn.Module):
    """ region proposal network """
    def __init__(self, din):
        super(_RPN, self).__init__()

        self.din = din  # get depth of input feature map, e.g., 512
        self.anchor_scales = cfg.ANCHOR_SCALES  # 是[4,8,16,32]
        self.anchor_ratios = cfg.ANCHOR_RATIOS  # 是[0.5,1,2]
        self.feat_stride = cfg.FEAT_STRIDE[0]   # 是16,vgg下采样倍数

        # define the convrelu layers processing input feature map
        # 输入,输出 -> 512, 卷积核 -> 3*3, 步长 -> 1, 补周边 -> 1,需要偏置
        self.RPN_Conv = nn.Conv2d(self.din, 512, 3, 1, 1, bias=True)

        # define bg/fg classifcation score layer
        # 进行fg/bg判断, 4 * 3 * 2 作为分类卷积输出的维度
        self.nc_score_out = len(self.anchor_scales) * len(self.anchor_ratios) * 2 # 2(bg/fg) * 9 (anchors)??
        self.RPN_cls_score = nn.Conv2d(512, self.nc_score_out, 1, 1, 0)

        # define anchor box offset prediction layer
        # 对bbox进行分类, 维度4 * 3 * 4 = 48
        self.nc_bbox_out = len(self.anchor_scales) * len(self.anchor_ratios) * 4 # 4(coords) * 9 (anchors)
        self.RPN_bbox_pred = nn.Conv2d(512, self.nc_bbox_out, 1, 1, 0)

        # define proposal layer
        # 实例化 proposallayer
        self.RPN_proposal = _ProposalLayer(self.feat_stride, self.anchor_scales, self.anchor_ratios)

        # define anchor target layer
        self.RPN_anchor_target = _AnchorTargetLayer(self.feat_stride, self.anchor_scales, self.anchor_ratios)

        self.rpn_loss_cls = 0
        self.rpn_loss_box = 0

    @staticmethod
    def reshape(x, d):
        input_shape = x.size()
        x = x.view(
            input_shape[0],
            int(d),
            int(float(input_shape[1] * input_shape[2]) / float(d)),
            input_shape[3]
        )
        return x

    def forward(self, base_feat, im_info, gt_boxes, num_boxes):
        '''
        输入:
            base_feat   ->  feature_map     ->  size(1,512,H,W)
            im_info     ->  image W,H,ratio ->  size(1,3)
            gt_boxes    ->  GroundTruthBox  ->  size(1,数量,5)
            num_boxes   ->  目标框的数量      ->  size(1)
        '''
        # base_feat的维度0,是1,base_feat是四维,但只有一张图
        # base_feat.size(), !! 例如 [1, 512, H, W] (一张图, 512维, 大小H*W), 大小视输入像素确定
        batch_size = base_feat.size(0)

        # return feature map after convrelu layer
        # 进行RPN卷积后relu,  rpn_conv1 -> [1, 512, H, W]
        rpn_conv1 = F.relu(self.RPN_Conv(base_feat), inplace=True)

        # 第一路 ,进行(fg/bg)预测
        # get rpn classification score
        # 经过卷积, 得到分类得分,只分前景背景(fg/bg)??, rpn_cls_score -> [1, 24, H, W]
        rpn_cls_score = self.RPN_cls_score(rpn_conv1)

        # 改边前景背景分类得分形状 [1, 24, H, W] -> [1, 2, 12*H, W]
        rpn_cls_score_reshape = self.reshape(rpn_cls_score, 2)
        # 对rpn_cls_score_reshape的第一个维度进行softmax, 相当于分成两类
        rpn_cls_prob_reshape = F.softmax(rpn_cls_score_reshape, 1)
        # softmax后再把形状变回来nc_score_out = 3*4*2 = 12 形状:[1, 2, 12*H, W] -> [1, 24, H, W]
        rpn_cls_prob = self.reshape(rpn_cls_prob_reshape, self.nc_score_out)

        # 第二路 ,进行bbox的预测
        # get rpn offsets to the anchor boxes
        # 进行一层卷积, rpn_cls_score -> [1, 48, H, W]
        rpn_bbox_pred = self.RPN_bbox_pred(rpn_conv1)

        # proposal layer
        cfg_key = 'TRAIN' if self.training else 'TEST'

        # 进行roi的预测
        '''
        形成两个支路,(1)预测fg/bg (2)进行Anchor的预测并根据fg/bg进行非极大值(nms)删选
        输入:
            是个tuple(),元素如下
            rpn_cls_prob    ->  torch.size([1, 24, H, W]), 每层的像素是fg/bg的概率
            rpn_bbox_pred   ->  预测bbox的feature_map(大小:[1, 48, H, W])
            im_info         ->  图片的大小H, W, ratio
            cfg_key         ->  'TRAIN' or 'TEST'
        输出:
            rois            ->  size([1, num_proposal, 5])
            rois是anchor经过fg/bg预测 + nms 筛选过后的proposal, num_proposal<=2000(目标域是300), 最后一维[第一个元素恒定为0,x1,y1,x2,y2]
        '''
        rois = self.RPN_proposal((rpn_cls_prob.data, rpn_bbox_pred.data,
                                 im_info, cfg_key))

        self.rpn_loss_cls = 0
        self.rpn_loss_box = 0

        # generating training labels and build the rpn loss
        # 源域 需要进行, 目标域 跳过
        if self.training:
            # 如果 gt_boxes 不存在就警告
            assert gt_boxes is not None
            '''
            输入:
                rpn_cls_score   :FeatureMap经过一层分类的卷积后  ->  [1, 24, H, W]
                gt_boxes   ->  [1, 数量, 5]
                im_info    ->  [1,3]
                num_boxes  ->  [1]
            输出:
                返回output列表[标签, bbox坐标变化, in权重, out权重]
                    labels                  大小 -> (1, 1, 12*H, W)   : 标签(0:背景, 1:前景, -1:屏蔽)
                    bbox_targets            大小 -> (1, 12*4, H, W)   : 每个anchor与IOU最大的GTbbox的位移和缩放参数
                    bbox_inside_weights     大小 -> (1, 12*4, H, W)   : in权重
                    bbox_outside_weights    大小 -> (1, 12*4, H, W)   : out权重
            '''
            rpn_data = self.RPN_anchor_target((rpn_cls_score.data, gt_boxes, im_info, num_boxes))

            # compute classification loss 计算分类损失
            # rpn_cls_score前景背景分类得分     形状[1, 2, 12*H, W] -> [1, 12*H, W, 2] -> [1, 12*H*W, 2]
            rpn_cls_score = rpn_cls_score_reshape.permute(0, 2, 3, 1).contiguous().view(batch_size, -1, 2)
            # 取出rpn_data的label 变化形状为 -> (1, 12*H*W)
            rpn_label = rpn_data[0].view(batch_size, -1)

            # 找到rpn_label中的非屏蔽项(≠-1)的索引号index 最后是一维的数组
            rpn_keep = Variable(rpn_label.view(-1).ne(-1).nonzero().view(-1))
            # 前背景分类得分rpn_cls_score 先变形状为(12*H*W, 2) 对第0维度 寻找需要匹配的anchor
            rpn_cls_score = torch.index_select(rpn_cls_score.view(-1,2), 0, rpn_keep)
            # 前景背景label 先变形状为(12*H*W) 对第0维度 寻找需要匹配的anchor
            rpn_label = torch.index_select(rpn_label.view(-1), 0, rpn_keep.data)
            # 变成变量(新版本中就不用了)
            rpn_label = Variable(rpn_label.long())
            # 计算交叉损失(分类)
            self.rpn_loss_cls = F.cross_entropy(rpn_cls_score, rpn_label)
            # 计算前景的个数
            fg_cnt = torch.sum(rpn_label.data.ne(0))

            # 取出[bbox坐标变化, in权重, out权重]
            rpn_bbox_targets, rpn_bbox_inside_weights, rpn_bbox_outside_weights = rpn_data[1:]

            # compute bbox regression loss 都变成变量(新版本中就不用了)
            rpn_bbox_inside_weights = Variable(rpn_bbox_inside_weights)
            rpn_bbox_outside_weights = Variable(rpn_bbox_outside_weights)
            rpn_bbox_targets = Variable(rpn_bbox_targets)

            # 计算L1loss(回归)
            self.rpn_loss_box = _smooth_l1_loss(rpn_bbox_pred, rpn_bbox_targets, rpn_bbox_inside_weights,
                                                            rpn_bbox_outside_weights, sigma=3, dim=[1,2,3])

        # 返回roi,分类和回归的损失
        return rois, self.rpn_loss_cls, self.rpn_loss_box


from modulator import module_audio
import numpy as np
from parameters import CHANNELS, CHUNK, NP_TYPE, PA_TYPE, RATE
import pyaudio as pa
import scipy.signal as signal
import time

p = pa.PyAudio()

def filter_audio(data):
    """
    Performs digital filter audio filtering.
    :param data: input sound data, numpy array
    :return: filtered sound data, numpy array
    """
    # elliptic filter with those parameters proved
    # experimentally to be the best
    b, a = signal.iirdesign(0.2, 0.3, 4, 60, ftype="ellip")
    data = signal.filtfilt(b, a, data).astype(NP_TYPE)
    return data

# function for manipulating input data
def callback(data, frame_count, time_info, flag):
    """
    Applies sound effects to input and outputs it in real time.
    :param data: input sound data, byte stream
    :param frame_count: unused, see PyAudio documentation
    :param time_info: unused, see PyAudio documentation
    :param flag: unused, see PyAudio documentation
    :return: tuple of: audio data after applying sound effects (numpy array),
    PyAudio constant to continue recording
    """
    data = np.frombuffer(data, dtype=NP_TYPE)
    data = filter_audio(data)
    data[abs(data) < 10] = 0

    data = module_audio(data, RATE, NP_TYPE)

    data[abs(data) < 10] = 0

    return data, pa.paContinue

def playback(callback_function):
    """
    Records audio in endless loop and plays it back with callback function.
    """
    stream = p.open(format=PA_TYPE,
                    rate=RATE,
                    channels=CHANNELS,
                    input=True,
                    output=True,
                    frames_per_buffer=CHUNK,
                    stream_callback=callback_function)

    stream.start_stream()

    # main input loop
    while True:
        time.sleep(1)

    stream.stop_stream()
    stream.close()
    p.terminate()

if __name__ == "__main__":
    playback(callback)


# -*- coding: utf-8 -*-
"""
Created on Wed Oct 24 06:05:57 2018

@author: Mikko Pohja, modified by Tommi Honkanen
"""

import os
import sys
import cv2
import dlib
import numpy as np
import time
import threading
import glob

from skimage import io

face_recognition = dlib.face_recognition_model_v1('./models/dlib_face_recognition_resnet_model_v1.dat')
face_detector = dlib.get_frontal_face_detector()
shape_predictor = dlib.shape_predictor('./models/shape_predictor_68_face_landmarks.dat')

enrolled_faces = {}

# Semaphor for identifying thread
identifying = False
successes = 0

def face_to_vector(image, face):
  return (
    np
    .array(face_recognition.compute_face_descriptor(image, face))
    .astype(float)
  )

def faces_from_image(image):
  UPSAMPLING_FACTOR = 0
  faces = [
    (face.height() * face.width(), shape_predictor(image, face))
    for face in face_detector(image, UPSAMPLING_FACTOR)
  ]
  return [face for _, face in sorted(faces, reverse=True)]

def image_from_file(path):
  return io.imread(path)

def identify(image):
  # Get all faces
  faces = faces_from_image(image)
  # Pick largest face
  face = faces[0] if faces else None

  # Calculate face descriptor
  descriptor = face_recognition.compute_face_descriptor(image, face)
  face_vector = np.array(descriptor).astype(float)

  # THIS is probably hazardous as ordering may not be always the same?
  enroll_identifiers = np.array(list(enrolled_faces.keys()))
  enroll_matrix = np.array(list(enrolled_faces.values()))

  # Calculate differences between the face and all enrolled faces
  differences = np.subtract(np.array(enroll_matrix), face_vector)
  distances = np.linalg.norm(differences, axis=1)
  # and pick the closest one
  closest_index = np.argmin(distances)

  return enroll_identifiers[closest_index], distances[closest_index]

def handle_frame(origFrame, cb, name):
  #global identifying
  global successes
  try:
    frame = cv2.resize(origFrame, None, fx=0.5, fy=0.5, interpolation=cv2.INTER_CUBIC)

    start = time.time()
    identifier, distance = identify(frame)
    if (distance < 0.6):
       cb(identifier, distance, time.time() - start)
       if identifier == name:
           successes +=1
    else:
       cb('-', distance, time.time() - start)
    sys.stdout.flush()

  except Exception as e:
    exc = e
    cb(None, 0, time.time() - start)
    # print(e)

  #identifying = False

def webcam(cb):
  global identifying

  video_capture = cv2.VideoCapture(0)

  while True:
    video_capture.grab()
    if (not identifying):
      ret, frame = video_capture.retrieve()

      if (ret == False):
        print('No frame')
        break
      identifying = True

      thread = threading.Thread(target=handle_frame, args=(frame, cb))
      thread.daemon=True
      thread.start()

  # When everything is done, release the capture
  video_capture.release()

def load_enrolled_faces():
  global enrolled_faces
  enrolled_faces = np.load('faces.npy').item()

def enroll_face(image, name):
  try:
    # find face
    faces = faces_from_image(image)
    # Pick largest face
    face = faces[0] if faces else None
    # face to vector
    face_vector = face_to_vector(image, face)
    # save to enrolled faces list
    enrolled_faces[name] = face_vector
    # save npy file
  except Exception as e:
    print('Unable to enroll', name)
    print(e)

def logger(identifier, distance, duration):
  if (identifier == '-'):
    print('Unknown person, operation took', duration, 'seconds.')
  elif (identifier == None):
    print('No face recognised, operation took',duration, 'seconds.')
  else:
    print('The person in question is', identifier, 'with calculated distance of', distance, ',operation took', duration, 'seconds.')

def enrollImages():
  print('Enrolling all images')
  for filename in glob.glob('faces/*.jpg'):
    print(filename)
    name = filename.split('/')[0].split('.')[0]
    name = name.split('\\')[1]
    image = image_from_file(filename)
    enroll_face(image, name)
    print('Enrolled', name)
  np.save('faces.npy', enrolled_faces)
  print('Saved faces.npy')

def loadTestImages():
  print('Enrolling all test images')
  for filename in glob.glob('test_faces/*.jpg'):
    print(filename)
    name = filename.split('/')[1].split('.')[0]
    image = image_from_file(filename)
    enroll_face(image, name)
    print('Enrolled', name)
  np.save('test_faces.npy', enrolled_faces)
  print('Saved test_faces.npy')

#The main script that runs the image enrolling and testing
#To reset the faces, delete faces.npy from the script's folder
def main():

    if (not os.path.isfile('faces.npy')):
        enrollImages()
    else:
        print('Loading enrolled faces from faces.npy')
        load_enrolled_faces()

    tries = 0
    for filename in glob.glob('test_faces/*.jpg'):

        tries += 1
        print('Testing', filename)
        name = filename.split('/')[0].split('.')[0]
        name = name.split('\\')[1]

        image = image_from_file(filename)
        handle_frame(image, logger, name)

    print('The accuracy of the test was', successes/tries * 100, 'percent.')
main()


# encoding:utf-8
import sys

sys.path.append("..")
import numpy as np
from collections import defaultdict
from random import randint
from random import shuffle,choice
from math import log
import gensim.models.word2vec as w2v
from prettyprinter import cpprint

from mf import MF
from reader.trust import TrustGetter
from utility.matrix import SimMatrix
from utility.similarity import cosine
from utility import util

class CUNE(MF):
    """
    docstring for CUNE

    Zhang et al. Collaborative User Network Embedding for Social Recommender Systems. SDM
    """

    def __init__(self):
        super(CUNE, self).__init__()
        self.config.lambdaP = 0.01
        self.config.lambdaQ = 0.01
        self.config.alpha = 0.01
        self.config.isEarlyStopping = True
        self.tg = TrustGetter()
        self.config.walkCount = 30
        self.config.walkLength = 20
        self.config.walkDim = 20
        self.config.winSize = 5
        self.config.topK = 50

    def init_model(self, k):
        super(CUNE, self).init_model(k)
        self.user_sim = SimMatrix()
        self.generate_cu_net()
        self.deep_walk()
        self.compute_social_sim()

    def generate_cu_net(self):
        print('Building collaborative user network...')
        itemNet = {}
        for item in self.rg.trainSet_i:
            if len(self.rg.trainSet_i[item])>1:
                itemNet[item] = self.rg.trainSet_i[item]

        filteredRatings = defaultdict(list)
        for item in itemNet:
            for user in itemNet[item]:
                if itemNet[item][user] > 0:
                    filteredRatings[user].append(item)

        self.CUNet = defaultdict(list)

        for user1 in filteredRatings:
            s1 = set(filteredRatings[user1])
            for user2 in filteredRatings:
                if user1 != user2:
                    s2 = set(filteredRatings[user2])
                    weight = len(s1.intersection(s2))
                    if weight > 0:
                        self.CUNet[user1]+=[user2] # * weight

        # cpprint(self.CUNet)
        pass
    def deep_walk(self):
        print('Generating random deep walks...')
        self.walks = []
        self.visited = defaultdict(dict)
        for user in self.CUNet:
            for t in range(self.config.walkCount):
                path = [str(user)]
                lastNode = user
                for i in range(1,self.config.walkLength):
                    nextNode = choice(self.CUNet[lastNode])
                    count=0
                    while(nextNode in self.visited[lastNode]):
                        nextNode = choice(self.CUNet[lastNode])
                        #break infinite loop
                        count+=1
                        if count==self.config.walkLength: # 10
                            break
                    path.append(str(nextNode))
                    self.visited[user][nextNode] = 1
                    lastNode = nextNode
                self.walks.append(path)
        # shuffle(self.walks)
        # cpprint(self.walks)
        print('Generating user embedding...')
        self.model = w2v.Word2Vec(self.walks, size=self.config.walkDim, window=5, min_count=0, iter=3)
        print('User embedding generated.')
        pass

    def compute_social_sim(self):
        print('Constructing similarity matrix...')
        # self.W = np.zeros((self.rg.get_train_size()[0], self.config.walkDim))
        self.topKSim = defaultdict(dict)
        i = 0
        for user1 in self.CUNet:
            sims = {}
            for user2 in self.CUNet:
                if user1 != user2:
                    wu1 = self.model[str(user1)]
                    wu2 = self.model[str(user2)]
                    sims[user2]=cosine(wu1,wu2) #若为空咋整
            self.topKSim[user1] = sorted(sims.items(), key=lambda d: d[1], reverse=True)[:self.config.topK]
            i += 1
            if i % 200 == 0:
                print('progress:', i, '/', len(self.CUNet))
        # print(self.topKSim)
        #构建被关注列表
        print('Constructing desimilarity matrix...')
        self.topKSimBy = defaultdict(dict)
        for user in self.topKSim:
            users=self.topKSim[user]
            for user2 in users:
                self.topKSimBy[user2[0]][user] = user2[1]
        print('Similarity matrix finished.')

    def train_model(self, k):
        super(CUNE, self).train_model(k)
        iteration = 0
        while iteration < self.config.maxIter:
            self.loss = 0
            for index, line in enumerate(self.rg.trainSet()):
                user, item, rating = line
                u = self.rg.user[user]
                i = self.rg.item[item]
                error = rating - self.predict(user, item)
                self.loss += 0.5 * error ** 2
                p, q = self.P[u], self.Q[i]

                social_term_p, social_term_loss = np.zeros((self.config.factor)), 0.0
                followees = self.topKSim[user] #self.tg.get_followees(user) #self.topKSim[user]
                # print(followees)
                for followee in followees:
                    if self.rg.containsUser(followee[0]):
                        # s = self.user_sim[user][followee]
                        uf = self.P[self.rg.user[followee[0]]]
                        social_term_p += followee[1]* (p - uf)
                        social_term_loss += followee[1]* ((p - uf).dot(p - uf))

                social_term_m = np.zeros((self.config.factor))
                followers = self.topKSimBy[user]
                followers = sorted(followers.items(), key=lambda d: d[1], reverse=True)[:self.config.topK]
                for follower in followers:
                    if self.rg.containsUser(follower[0]):
                        ug = self.P[self.rg.user[follower[0]]]
                        social_term_m += follower[1]*(p - ug)

                # update latent vectors
                self.P[u] += self.config.lr * (
                        error * q - self.config.alpha * (social_term_p + social_term_m) - self.config.lambdaP * p)
                self.Q[i] += self.config.lr * (error * p - self.config.lambdaQ * q)

                self.loss += 0.5 * self.config.alpha * social_term_loss

            self.loss += 0.5 * self.config.lambdaP * (self.P * self.P).sum() + 0.5 * self.config.lambdaQ * (
                    self.Q * self.Q).sum()

            iteration += 1
            if self.isConverged(iteration):
                break

if __name__ == '__main__':
    # srg = CUNE()
    # srg.train_model(0)
    # coldrmse = srg.predict_model_cold_users()
    # print('cold start user rmse is :' + str(coldrmse))
    # srg.show_rmse()

    rmses = []
    maes = []
    cunemf = CUNE()
    # cunemf.init_model(0)
    # cunemf.generate_cu_net()
    # cunemf.deep_walk()
    # print(bmf.rg.trainSet_u[1])
    cunemf.config.k_fold_num = 5
    for i in range(cunemf.config.k_fold_num):
        print('the %dth cross validation training' % i)
        cunemf.train_model(i)
        rmse, mae = cunemf.predict_model()
        rmses.append(rmse)
        maes.append(mae)
    rmse_avg = sum(rmses) / cunemf.config.k_fold_num
    mae_avg = sum(maes) / cunemf.config.k_fold_num
    print("the rmses are %s" % rmses)
    print("the maes are %s" % maes)
    print("the average of rmses is %s " % rmse_avg)
    print("the average of maes is %s " % mae_avg)


from kivy.app import App
from kivy.uix.widget import Widget
from kivy.graphics import Color, Line
from kivy.uix.floatlayout import FloatLayout
from kivy.uix.boxlayout import BoxLayout
from math import cos, sin, pi
from kivy.clock import Clock
from kivy.lang import Builder
from kivy.properties import NumericProperty, ListProperty, BooleanProperty, StringProperty
from kivy.uix.popup import Popup
from kivy.core.window import Window
import datetime
import pyaudio
import numpy as np
import time
import keras.models as kmodels
import pprint
import threading
import copy

from kivy.config import Config

Config.set('graphics', 'width', '320')
Config.set('graphics', 'height', '200')

Builder.load_file('clock.kv')

class ClockApp(FloatLayout):
    pass

class CustomPopup(Popup):

    def __init__(self, dismiss_callback, dismiss_check):
        super().__init__()
        self.d_callback = dismiss_callback
        self.d_check = dismiss_check

    def on_dismiss(self):
        if not self.d_check():
            self.d_callback()
        pass

class Ticks(Widget):
    notification_stack = ListProperty()
    internal_dismiss = BooleanProperty()
    detection_str = StringProperty()

    def __init__(self, **kwargs):
        super(Ticks, self).__init__(**kwargs)
        self.bind(pos=self.update_clock)
        self.bind(size=self.update_clock)
        self.internal_dismiss = False

    def update_clock(self, *args):
        self.canvas.clear()
        with self.canvas:
            time = datetime.datetime.now()
            Color(0.1, 0.3, 0.1)
            Line(points=[self.center_x, self.center_y, self.center_x + 0.85 * self.r * sin(pi / 30 * time.second),
                         self.center_y + 0.8 * self.r * cos(pi / 30 * time.second)], width=1)
            Color(0.2, 0.5, 0.2)
            Line(points=[self.center_x, self.center_y, self.center_x + 0.7 * self.r * sin(pi / 30 * time.minute),
                         self.center_y + 0.7 * self.r * cos(pi / 30 * time.minute)], width=2, cap="round")
            Color(0.4, 0.7, 0.4)
            th = time.hour * 60 + time.minute
            Line(points=[self.center_x, self.center_y, self.center_x + 0.5 * self.r * sin(pi / 360 * th),
                         self.center_y + 0.5 * self.r * cos(pi / 360 * th)], width=2, cap="round")

    def dismiss_check(self):
        return self.internal_dismiss

    def on_popup_dismiss(self):
        print("dismissed")
        p = self.notification_stack.pop()
        del p
        if len(self.notification_stack) > 0:
            self.notification_stack[-1].open()
        pass

    def detect_check(self, *args):
        global str_detect
        if str_detect != "":
            self.detection_str = str_detect
            if len(self.notification_stack) > 0:
                if self.notification_stack[-1].detect != self.detection_str:
                    p = CustomPopup(self.on_popup_dismiss, self.dismiss_check)
                    p.detect = self.detection_str
                    self.notification_stack.append(p)
                    p.title = "{1} {0}".format(len(self.notification_stack), p.detect)
                    if len(self.notification_stack) > 1:
                        self.internal_dismiss = True
                        self.notification_stack[-2].dismiss()
                        self.internal_dismiss = False
                    p.open()
            else:
                p = CustomPopup(self.on_popup_dismiss, self.dismiss_check)
                p.detect = self.detection_str
                self.notification_stack.append(p)
                p.title = "{1} {0}".format(len(self.notification_stack), p.detect)
                if len(self.notification_stack) > 1:
                    self.internal_dismiss = True
                    self.notification_stack[-2].dismiss()
                    self.internal_dismiss = False
                p.open()
        str_detect = ""

class MyClockApp(App):
    def build(self):
        clock = ClockApp()
        Clock.schedule_interval(clock.ticks.update_clock, 1.0 / 20)
        Clock.schedule_interval(clock.ticks.detect_check, 1.0)
        return clock

class Audio_t(threading.Thread):
    def __init__(self, name: object):
        threading.Thread.__init__(self)
        self.name = name
        self.last = time.time()

    def run(self):
        global flag
        global n
        global stream
        global player
        global c
        while True:
            data = np.fromstring(stream.read(CHUNK), dtype=np.int16)
            c.acquire()
            n[0] = np.append(n[0][CHUNK:], data)
            c.notify_all()
            c.release()

class Predict_t(threading.Thread):
    def __init__(self, name: object):
        threading.Thread.__init__(self)
        self.name = name
        self.buff = np.array([])
        self.cl = ['door_wood_knock', 'clock_alarm', 'car_horn', 'siren', 'ambient']
        self.prev_det = ""

    def run(self):
        global flag
        global n
        global mx
        global c
        global player
        global str_detect
        while True:
            c.acquire()
            self.buff = copy.deepcopy(n)
            c.notify_all()
            c.release()
            if self.buff.max() <= 11000:
                time.sleep(1)
                continue
            self.buff = np.diff(self.buff)
            self.buff = np.append(self.buff, [0.0])
            self.buff = (self.buff - self.buff.mean()) / self.buff.max()
            self.buff = self.buff.reshape(1, self.buff.shape[0])
            activations = mx.predict(self.buff)
            u = activations[0]
            scale_factor = np.ptp(u)
            b = (u - np.min(u)) / np.ptp(u)
            if max(b) != b[-1]:
                detect, val = self.cl[np.argmax(b)], np.argmax(b)
                print(detect, b, scale_factor)
                if scale_factor >= 0.9 and detect == self.prev_det:
                    str_detect = detect
                    print("detection reported")
                self.prev_det = detect
            else:
                self.prev_det = ""

if __name__ == '__main__':
    p = pyaudio.PyAudio()
    c = threading.Condition()
    CHUNK = 2**14
    CHANNELS = 1
    RATE = 9600
    LEN = 1000
    flag = 0
    str_detect = ""
    mx = kmodels.load_model("saved_wavenet_classifier_diffmax_4.h5")
    stream = p.open(format=pyaudio.paInt16, channels=1, rate=RATE, input=True, frames_per_buffer=CHUNK)
    stream.start_stream()

    np.set_printoptions(precision=2, suppress=True)

    print("force memory initialization")
    n = np.zeros(19200, dtype=np.int16)
    n = np.reshape(n, (1, n.shape[0],))
    mx.predict(n)
    t1 = Audio_t("audio_t")
    t2 = Predict_t("Predict_t")

    t1.start()
    t2.start()
    MyClockApp().run()

    while True:
        pass

    stream.stop_stream()
    stream.close()
    p.terminate()