from typing import Union, Tuple
import eagerpy as ep

from ..models import Model

from ..criteria import Misclassification

from ..devutils import atleast_kd, flatten

from .base import MinimizationAttack
from .base import get_criterion
from .base import T


class NewtonFoolAttack(MinimizationAttack):
    """NewtonFool Attack introduced in [1]_

    References
    ----------
    .. [1] Uyeong Jang et al., "Objective Metrics and Gradient Descent
           Algorithms for Adversarial Examples in Machine Learning",
           https://dl.acm.org/citation.cfm?id=3134635
    """

    def __init__(self, steps: int = 100, stepsize: float = 0.01):
        self.steps = steps
        self.stepsize = stepsize

    def __call__(
        self, model: Model, inputs: T, criterion: Union[Misclassification, T]
    ) -> T:

        x, restore_type = ep.astensor_(inputs)
        criterion_ = get_criterion(criterion)
        del inputs, criterion

        N = len(x)
        rows = range(N)

        if isinstance(criterion_, Misclassification):
            classes = criterion_.labels
        else:
            raise ValueError("unsupported criterion")

        if classes.shape != (N,):
            raise ValueError(
                f"expected labels to have shape ({N},), got {classes.shape}"
            )

        min_, max_ = model.bounds

        x_l2_norm = flatten(x.square()).sum(1)

        def loss_fun(x: ep.Tensor) -> Tuple[ep.Tensor, Tuple[ep.Tensor, ep.Tensor]]:
            # TODO: this is wrong!
            logits = model(x)
            scores = ep.softmax(logits)
            pred = scores.argmax(-1)
            loss = scores.sum()
            return loss, (scores, pred)

        for i in range(self.steps):
            # (1) get the scores and gradients
            _, (scores, pred), gradients = ep.value_aux_and_grad(loss_fun, x)

            pred_scores = scores[rows, classes]

            num_classes = pred.shape[-1]

            # (2) calculate gradient norm
            gradients_l2_norm = flatten(gradients.square()).sum(1)

            # (3) calculate delta
            a = self.stepsize * x_l2_norm * gradients_l2_norm
            b = pred_scores - 1.0 / num_classes

            delta = ep.minimum(a, b)

            # (4) stop the attack if an adversarial example has been found
            # this is not described in the paper but otherwise once the prob. drops
            # below chance level the likelihood is not decreased but increased
            is_adversarial = (pred != classes).float32()
            delta *= is_adversarial

            # (5) calculate & apply current perturbation
            x -= (
                atleast_kd(delta / gradients_l2_norm.square(), gradients.ndim)
                * gradients
            )
            x = ep.clip(x, min_, max_)

        return restore_type(x)

@@ -956,32 +956,8 @@
 n(x)
-%0A        rows = range(N)
 %0A%0A  
@@ -1445,43 +1445,8 @@
 %5D%5D:%0A
-            # TODO: this is wrong!%0A
     
@@ -1527,16 +1527,23 @@
     pred
+_scores
  = score
@@ -1547,19 +1547,27 @@
 ores
-.argmax(-1)
+%5Brange(N), classes%5D
 %0A   
@@ -1578,24 +1578,29 @@
      loss = 
+pred_
 scores.sum()
@@ -1638,16 +1638,23 @@
 es, pred
+_scores
 )%0A%0A     
@@ -1759,16 +1759,23 @@
 es, pred
+_scores
 ), gradi
@@ -1825,39 +1825,32 @@
             pred
-_scores
  = scores%5Brows, 
@@ -1846,24 +1846,19 @@
 ores
-%5Brows, classes%5D%0A
+.argmax(-1)
 %0A   
@@ -1880,20 +1880,22 @@
 asses = 
-pred
+scores
 .shape%5B-
@@ -2440,24 +2440,28 @@
          is_
+not_
 adversarial 
@@ -2468,17 +2468,17 @@
 = (pred 
-!
+=
 = classe
@@ -2514,16 +2514,20 @@
 a *= is_
+not_
 adversar
@@ -2604,30 +2604,11 @@
     
-x -= (%0A               
+a =
  atl
@@ -2679,19 +2679,22 @@
         
-   
+x -= a
  * gradi
@@ -2698,29 +2698,16 @@
 adients%0A
-            )
 %0A       

Fix loss function in newtonfool (#10)