import torch
import math

import a3t.diffai.helpers as h


class Const():
    def __init__(self, c):
        self.c = c if c is None else float(c)

    def getVal(self, c=None, **kargs):
        return self.c if self.c is not None else c

    def __str__(self):
        return str(self.c)

    def initConst(x):
        return x if isinstance(x, Const) else Const(x)


class Lin(Const):
    def __init__(self, start, end, steps, initial=0, quant=False):
        self.start = float(start)
        self.end = float(end)
        self.steps = float(steps)
        self.initial = float(initial)
        self.quant = quant

    def getVal(self, time=0, **kargs):
        if self.quant:
            time = math.floor(time)
        return (self.end - self.start) * max(0, min(1, float(time - self.initial) / self.steps)) + self.start

    def __str__(self):
        return "Lin(%s,%s,%s,%s, quant=%s)".format(str(self.start), str(self.end), str(self.steps), str(self.initial),
                                                   str(self.quant))


class Until(Const):
    def __init__(self, thresh, a, b):
        self.a = Const.initConst(a)
        self.b = Const.initConst(b)
        self.thresh = thresh

    def getVal(self, *args, time=0, **kargs):
        return self.a.getVal(*args, time=time, **kargs) if time < self.thresh else self.b.getVal(*args,
                                                                                                 time=time - self.thresh,
                                                                                                 **kargs)

    def __str__(self):
        return "Until(%s, %s, %s)" % (str(self.thresh), str(self.a), str(self.b))


class Scale(Const):  # use with mix when aw = 1, and 0 <= c < 1
    def __init__(self, c):
        self.c = Const.initConst(c)

    def getVal(self, *args, **kargs):
        c = self.c.getVal(*args, **kargs)
        if c == 0:
            return 0
        assert c >= 0
        assert c < 1
        return c / (1 - c)

    def __str__(self):
        return "Scale(%s)" % str(self.c)


def MixLin(*args, **kargs):
    return Scale(Lin(*args, **kargs))


class Normal(Const):
    def __init__(self, c):
        self.c = Const.initConst(c)

    def getVal(self, *args, shape=[1], **kargs):
        c = self.c.getVal(*args, shape=shape, **kargs)
        return torch.randn(shape, device=h.device).abs() * c

    def __str__(self):
        return "Normal(%s)" % str(self.c)


class Clip(Const):
    def __init__(self, c, l, u):
        self.c = Const.initConst(c)
        self.l = Const.initConst(l)
        self.u = Const.initConst(u)

    def getVal(self, *args, **kargs):
        c = self.c.getVal(*args, **kargs)
        l = self.l.getVal(*args, **kargs)
        u = self.u.getVal(*args, **kargs)
        if isinstance(c, float):
            return min(max(c, l), u)
        else:
            return c.clamp(l, u)

    def __str__(self):
        return "Clip(%s, %s, %s)" % (str(self.c), str(self.l), str(self.u))


class Fun(Const):
    def __init__(self, foo):
        self.foo = foo

    def getVal(self, *args, **kargs):
        return self.foo(*args, **kargs)

    def __str__(self):
        return "Fun(...)"


class Complement(Const):  # use with mix when aw = 1, and 0 <= c < 1
    def __init__(self, c):
        self.c = Const.initConst(c)

    def getVal(self, *args, **kargs):
        c = self.c.getVal(*args, **kargs)
        assert c >= 0
        assert c <= 1
        return 1 - c

    def __str__(self):
        return "Complement(%s)" % str(self.c)


class TrainInfo:
    out_y = None
    adv = False
    cur_ratio = 0
    total_batches_seen = 0
    abs_perturb = []
    victim_model = None
