from utils.momentum import momentum
from utils.adam import adam
from utils.manip import proj_with_clip, proj_l2_ball, proj_linf_ball
import numpy as np
from fool.deepfool import deepfool
import torch
import sys
from tqdm import tqdm
sys.path.append('../')


class UniversalPert():
    def __init__(self, audioLoader, pipeline, desired_fooling_rate, max_iter, eps=4, p=2, clip=None, device='cpu', verbose=False, learning_rate=1.0, track_rate=True, deepfool_max_iter=50, _lambda=3.0, optimizer="adam"):
        self.audioLoader = audioLoader
        self.pipeline = pipeline
        self.desired_fooling_rate = desired_fooling_rate
        self.max_iter = max_iter
        self.eps = eps
        self.p = p
        self.clip = clip
        self.device = device
        self.verbose = verbose
        self.learning_rate = learning_rate
        self._lambda = _lambda
        self.fool = deepfool(device=self.device, max_iter=deepfool_max_iter)
        self.track_rate = track_rate
        self.optimizer = optimizer
        self.univ_pert = self.pipeline.transformWrapper.init_pert()
        if self.optimizer == "momentum":
            self.optimizer = momentum(self.univ_pert, lr=self.learning_rate)
        elif self.optimizer == "adam":
            self.optimizer = adam(self.univ_pert, lr=self.learning_rate)
        else:
            raise ValueError("Optimizer not supported")

    def compute_confusion_matrix(self, univ_pert, target=None):
        self.pipeline.model.eval()
        confusion_matrix = np.zeros((self.num_classes, self.num_classes))
        for waveforms, labels in self.audioLoader.fixed_test_loader:
            waveforms = waveforms.to(self.device)
            labels = labels.to(self.device)
            orig_labels = torch.argmax(
                self.pipeline._forward(waveforms), dim=1)
            univ_pert_labels = torch.argmax(
                self.pipeline.forward(waveforms, univ_pert), dim=1)
            for i in range(len(labels)):
                confusion_matrix[orig_labels[i], univ_pert_labels[i]] += 1
        return confusion_matrix

    def compute_fooling_rate(self, univ_pert, pbar, target=None):
        self.pipeline.model.eval()
        fool_count = 0
        ex_count = 0
        for waveforms, labels in self.audioLoader.fixed_test_loader:
            pbar.set_description('Computing fooling rate {}/{}'.format(
                ex_count, len(self.audioLoader.fixed_test_loader.dataset)))
            waveforms = waveforms.to(self.device)
            labels = labels.to(self.device)
            orig_labels = torch.argmax(
                self.pipeline._forward(waveforms), dim=1)
            univ_pert_labels = torch.argmax(
                self.pipeline.forward(waveforms, univ_pert), dim=1)
            if target is None:
                fool_count += torch.sum(univ_pert_labels != orig_labels).item()
            else:
                fool_count += torch.sum(univ_pert_labels == target).item()
            ex_count += len(labels)
        return fool_count / ex_count

    def project(self, x):
        if self.p == 2:
            if self.clip:
                return proj_with_clip(x, self.eps, self.clip)
            else:
                return proj_l2_ball(x, self.eps)
        elif self.p == np.inf:
            return proj_linf_ball(x, self.eps)

    def update(self, delta_pert):
        self.univ_pert = self.optimizer.update(delta_pert)
        self.univ_pert = self.pipeline.transformWrapper.clamp_pert(
            self.univ_pert)

    def update_pert(self, n_iters=None, target=None):
        itr = 0
        fooling_rate = 0
        fooling_rates = []  # keep track of fooling rates at each iter
        if n_iters is None:
            n_iters = self.max_iter
        with tqdm(total=n_iters) as pbar:
            while itr < n_iters and fooling_rate < self.desired_fooling_rate:
                itr += 1
                pbar.set_description('Getting batch from dataset')
                # Get new batch from dataset
                waveforms, _ = next(iter(self.audioLoader.train_loader))
                waveforms = waveforms.to(self.device)

                pbar.set_description(
                    'Getting output for the unperturbed batch')

                # Get the model output for the batch
                orig_labels = torch.argmax(
                    self.pipeline._forward(waveforms), dim=1)

                pbar.set_description('Getting output for the perturbed batch')

                # Get the model output for the perturbed batch
                univ_pert_labels = torch.argmax(
                    self.pipeline.forward(waveforms, self.univ_pert), dim=1)

                pbar.set_description('Checking perturbation results')

                # Get the data that was not perturbed enough
                if target is None:  # if target is not provided, use the original labels
                    # get the waveforms that were not fooled
                    unfooled_wavs = waveforms[univ_pert_labels == orig_labels]
                    # get the labels corresponding to the waveforms that were not fooled
                    unfooled_labels = orig_labels[univ_pert_labels == orig_labels]
                else:  # if target is provided, use the target labels
                    unfooled_wavs = waveforms[univ_pert_labels != target]
                    unfooled_labels = orig_labels[univ_pert_labels != target]
                if len(unfooled_wavs) == 0:  # all wavs were fooled
                    pbar.set_description("All wavs were fooled!")
                    break
                pbar.set_description('Initializing new perturbations')

                # Compute new perturbations for the data that was not perturbed enough
                perturbed_waveforms = torch.zeros(unfooled_wavs.shape).to(
                    self.device)  # init perturbed waveforms, to be filled later
                # build shape of perturbations
                perts_shape = [unfooled_wavs.shape[0]]
                for dim in self.pipeline.transformWrapper.pert_shape:
                    perts_shape.append(dim)
                # init perturbations, to be filled later
                perts = torch.zeros(tuple(perts_shape)).to(self.device)
                unfooled_wavs = unfooled_wavs.to(self.device)

                pbar.set_description('Perturbing unfooled wavs')

                # Add the current perturbation to the waveforms
                pert_in_wav_domain = self.pipeline.transformWrapper.preprocess(
                    self.univ_pert)
                unfooled_wavs += pert_in_wav_domain

                # for each waveform, compute the necessary perturbation to fool the model
                for i in range(len(unfooled_wavs)):
                    pbar.set_description(
                        'Computing perturbations for wav {}/{}'.format(i+1, len(unfooled_wavs)))
                    _, pert_wav, pert = self.fool._perturb(
                        self.pipeline, unfooled_wavs[i:i+1].to(self.device), target)
                    # track the perturbed waveforms
                    perturbed_waveforms[i] = pert_wav
                    # track the perturbations, in the perturbation space
                    perts[i] = pert * self._lambda
                # extract perturbations
                wav_perturbations = perturbed_waveforms - unfooled_wavs

                # Get labels of the perturbed waveforms
                pert_labels = torch.argmax(
                    self.pipeline._forward(perturbed_waveforms), dim=1)

                # Get perturbations that fool the classifier
                fooling_perts = perts[pert_labels != unfooled_labels]
                # operate on the perturbations that fool the classifier
                if len(fooling_perts) > 0:
                    # Compute the mean of the fooling perts
                    delta_pert = torch.mean(
                        fooling_perts, dim=0).unsqueeze(0).to(self.device)
                    # Add perturbation to univ pert
                    self.update(delta_pert)
                elif self.verbose:
                    print(
                        "Iteration: {}, [SKIPPED] deepfool fooled no new wavs".format(itr))
                # Compute fooling rate
                if self.track_rate:
                    fooling_rate = self.compute_fooling_rate(
                        self.univ_pert, pbar, target)
                pbar.update()
                fooling_rates += [fooling_rate]
                if self.verbose:
                    print("Iteration: {}, fooling rate: {}".format(
                        itr, fooling_rate))
            pbar.close()
        if self.verbose:
            print("Finished fooling rate: {}".format(fooling_rate))
            print("============================")
        return self.univ_pert, fooling_rate, fooling_rates
