import inspect
import os
import torch
import torch.nn.functional as F
import numpy as np
import argparse
import decimal
from datetime import datetime
from forbiddenfruit import curse
from timeit import default_timer as timer


class Timer:
    def __init__(self, activity=None, units=1, shouldPrint=True, f=None):
        self.activity = activity
        self.units = units
        self.shouldPrint = shouldPrint
        self.f = f

    def __enter__(self):
        self.start = timer()
        return self

    def getUnitTime(self):
        return (self.end - self.start) / self.units

    def __str__(self):
        return "Avg time to " + self.activity + ": " + str(self.getUnitTime())

    def __exit__(self, *args):
        self.end = timer()
        if self.shouldPrint:
            printBoth(self, f=self.f)


def cudify(x):
    if use_cuda:
        return x.cuda(non_blocking=True)
    return x


def pyval(a, **kargs):
    return dten([a], **kargs)


def ifThenElse(cond, a, b):
    cond = cond.to_dtype()
    return cond * a + (1 - cond) * b


def ifThenElseL(cond, a, b):
    return cond * a + (1 - cond) * b


def product(it):
    if isinstance(it, int):
        return it
    product = 1
    for x in it:
        if x >= 0:
            product *= x
    return product


def getEi(batches, num_elem):
    return eye(num_elem).expand(batches, num_elem, num_elem).permute(1, 0, 2)


def one_hot(batch, d):
    bs = batch.size()[0]
    indexes = [list(range(bs)), batch]
    values = [1 for _ in range(bs)]
    return cudify(torch.sparse.FloatTensor(ltenCPU(indexes), ftenCPU(values), torch.Size([bs, d])))


def seye(n, m=None):
    if m is None:
        m = n
    mn = n if n < m else m
    indexes = [[i for i in range(mn)], [i for i in range(mn)]]
    values = [1 for i in range(mn)]
    return cudify(torch.sparse.ByteTensor(ltenCPU(indexes), dtenCPU(values), torch.Size([n, m])))


dtype = torch.float32
ftype = torch.float32
ltype = torch.int64
btype = torch.uint8

torch.set_default_dtype(dtype)

cpu = torch.device("cpu")

cuda_async = True

ftenCPU = lambda *args, **kargs: torch.tensor(*args, dtype=ftype, device=cpu, **kargs)
dtenCPU = lambda *args, **kargs: torch.tensor(*args, dtype=dtype, device=cpu, **kargs)
ltenCPU = lambda *args, **kargs: torch.tensor(*args, dtype=ltype, device=cpu, **kargs)
btenCPU = lambda *args, **kargs: torch.tensor(*args, dtype=btype, device=cpu, **kargs)

if torch.cuda.is_available() and not 'NOCUDA' in os.environ:
    print("using cuda")
    device = torch.device("cuda")
    ften = lambda *args, **kargs: torch.tensor(*args, dtype=ftype, device=device, **kargs).cuda(non_blocking=cuda_async)
    dten = lambda *args, **kargs: torch.tensor(*args, dtype=dtype, device=device, **kargs).cuda(non_blocking=cuda_async)
    lten = lambda *args, **kargs: torch.tensor(*args, dtype=ltype, device=device, **kargs).cuda(non_blocking=cuda_async)
    bten = lambda *args, **kargs: torch.tensor(*args, dtype=btype, device=device, **kargs).cuda(non_blocking=cuda_async)
    ones = lambda *args, **cargs: torch.ones(*args, **cargs).cuda(non_blocking=cuda_async)
    zeros = lambda *args, **cargs: torch.zeros(*args, **cargs).cuda(non_blocking=cuda_async)
    eye = lambda *args, **cargs: torch.eye(*args, **cargs).cuda(non_blocking=cuda_async)
    use_cuda = True
    print("set up cuda")
else:
    print("not using cuda")
    ften = ftenCPU
    dten = dtenCPU
    lten = ltenCPU
    bten = btenCPU
    ones = torch.ones
    zeros = torch.zeros
    eye = torch.eye
    use_cuda = False
    device = cpu


def smoothmax(x, alpha, dim=0):
    return x.mul(F.softmax(x * alpha, dim)).sum(dim + 1)


def str2bool(v):
    if v.lower() in ('yes', 'true', 't', 'y', '1'):
        return True
    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
        return False
    else:
        raise argparse.ArgumentTypeError('Boolean value expected.')


def flat(lst):
    lst_ = []
    for l in lst:
        lst_ += l
    return lst_


def printBoth(*st, f=None):
    print(*st)
    if not f is None:
        print(*st, file=f)


def hasMethod(cl, mt):
    return callable(getattr(cl, mt, None))


def getMethodNames(Foo):
    return [func for func in dir(Foo) if callable(getattr(Foo, func)) and not func.startswith("__")]


def getMethods(Foo):
    return [getattr(Foo, m) for m in getMethodNames(Foo)]


max_c_for_norm = 10000


def numel(arr):
    return product(arr.size())


def chunks(l, n):
    """Yield successive n-sized chunks from l."""
    for i in range(0, len(l), n):
        yield l[i:i + n]


def loadDataset(dataset, batch_size, type, test_slice=None):
    """
    load the dataset
    :param dataset: the name of the dataset, currently support SST2CharLevel and SST2WordLevel
    :param batch_size: the batch size of the dataset loader
    :param type: "train", "val", "test"
    :param test_slice: select a slice of the data
    :return: a dataset loader
    """
    assert type in ["train", "val", "test"]
    if type == "val":
        X, y = dataset.val_X, dataset.val_y
    elif type == "train":
        X, y = dataset.train_X, dataset.train_y
    else:
        X, y = dataset.test_X, dataset.test_y

    if test_slice is not None:
        X = X[test_slice]
        y = y[test_slice]

    x = torch.from_numpy(X)
    train_set = torch.utils.data.TensorDataset(x, torch.from_numpy(y))

    return torch.utils.data.DataLoader(
        train_set
        , batch_size=batch_size
        , shuffle=(type == "train"),
        **({'num_workers': 1, 'pin_memory': True} if use_cuda else {}))


def variable(Pt):
    class Point:
        def isSafe(self, target):
            pred = self.max(1, keepdim=True)[1]  # get the index of the max log-probability
            return pred.eq(target.data.view_as(pred))

        def isPoint(self):
            return True

        def labels(self):
            return [self[0].max(1)[1]]  # get the index of the max log-probability

        def softplus(self):
            return F.softplus(self)

        def elu(self):
            return F.elu(self)

        def selu(self):
            return F.selu(self)

        def sigm(self):
            return F.sigmoid(self)

        def conv3d(self, *args, **kargs):
            return F.conv3d(self, *args, **kargs)

        def conv2d(self, *args, **kargs):
            return F.conv2d(self, *args, **kargs)

        def conv1d(self, *args, **kargs):
            return F.conv1d(self, *args, **kargs)

        def conv_transpose3d(self, *args, **kargs):
            return F.conv_transpose3d(self, *args, **kargs)

        def conv_transpose2d(self, *args, **kargs):
            return F.conv_transpose2d(self, *args, **kargs)

        def conv_transpose1d(self, *args, **kargs):
            return F.conv_transpose1d(self, *args, **kargs)

        def max_pool2d(self, *args, **kargs):
            return F.max_pool2d(self, *args, **kargs)

        def avg_pool2d(self, *args, **kargs):
            return F.avg_pool2d(self, *args, **kargs)

        def adaptive_avg_pool2d(self, *args, **kargs):
            return F.adaptive_avg_pool2d(self, *args, **kargs)

        def cat(self, other, dim=0, **kargs):
            return torch.cat((self, other), dim=dim, **kargs)

        def addPar(self, a, b):
            return a + b

        def abstractApplyLeaf(self, foo, *args, **kargs):
            return self

        def diameter(self):
            return pyval(0)

        def to_dtype(self):
            return self.type(dtype=dtype, non_blocking=cuda_async)

        def loss(self, target, **kargs):
            if torch.__version__[0] == "0":
                return F.cross_entropy(self, target, reduce=False)
            else:
                return F.cross_entropy(self, target, reduction='none')

        def deep_loss(self, *args, **kargs):
            return 0

        def merge(self, *args, **kargs):
            return self

        def splitRelu(self, *args, **kargs):
            return self

        def lb(self):
            return self

        def vanillaTensorPart(self):
            return self

        def center(self):
            return self

        def ub(self):
            return self

        def cudify(self, cuda_async=True):
            return self.cuda(non_blocking=cuda_async) if use_cuda else self

    def log_softmax(self, *args, dim=1, **kargs):
        return F.log_softmax(self, *args, dim=dim, **kargs)

    if torch.__version__[0] == "0" and torch.__version__ != "0.4.1":
        Point.log_softmax = log_softmax

    def log_softmax(self, *args, dim=1, **kargs):
        return F.log_softmax(self, *args, dim=dim, **kargs)

    if torch.__version__[0] == "0" and torch.__version__ != "0.4.1":
        Point.log_softmax = log_softmax

    for nm in getMethodNames(Point):
        curse(Pt, nm, getattr(Point, nm))


variable(torch.autograd.Variable)
variable(torch.cuda.DoubleTensor)
variable(torch.DoubleTensor)
variable(torch.cuda.FloatTensor)
variable(torch.FloatTensor)
variable(torch.ByteTensor)
variable(torch.Tensor)


def default(dic, nm, d):
    if dic is not None and nm in dic:
        return dic[nm]
    return d


def softmaxBatchNP(x, epsilon, subtract=False):
    """Compute softmax values for each sets of scores in x."""
    x = x.astype(np.float64)
    ex = x / epsilon if epsilon is not None else x
    if subtract:
        ex -= ex.max(axis=1)[:, np.newaxis]
    e_x = np.exp(ex)
    sm = (e_x / e_x.sum(axis=1)[:, np.newaxis])
    am = np.argmax(x, axis=1)
    bads = np.logical_not(np.isfinite(sm.sum(axis=1)))

    if epsilon is None:
        sm[bads] = 0
        sm[bads, am[bads]] = 1
    else:
        epsilon *= (x.shape[1] - 1) / x.shape[1]
        sm[bads] = epsilon / (x.shape[1] - 1)
        sm[bads, am[bads]] = 1 - epsilon

    sm /= sm.sum(axis=1)[:, np.newaxis]
    return sm


def cadd(a, b):
    both = a.cat(b)
    a, b = both.split(a.size()[0])
    return a + b


def msum(a, b, l):
    if a is None:
        return b
    if b is None:
        return a
    return l(a, b)


class SubAct(argparse.Action):
    def __init__(self, sub_choices, *args, **kargs):
        super(SubAct, self).__init__(*args, nargs='+', **kargs)
        self.sub_choices = sub_choices
        self.sub_choices_names = None if sub_choices is None else getMethodNames(sub_choices)

    def __call__(self, parser, namespace, values, option_string=None):
        if self.sub_choices_names is not None and not values[0] in self.sub_choices_names:
            msg = 'invalid choice: %r (choose from %s)' % (values[0], self.sub_choices_names)
            raise argparse.ArgumentError(self, msg)

        prev = getattr(namespace, self.dest)
        setattr(namespace, self.dest, prev + [values])


def catLists(val):
    if isinstance(val, list):
        v = []
        for i in val:
            v += catLists(i)
        return v
    return [val]


def sumStr(val):
    s = ""
    for v in val:
        s += v
    return s


def catStrs(val):
    s = val[0]
    if len(val) > 1:
        s += "("
    for v in val[1:2]:
        s += v
    for v in val[2:]:
        s += ", " + v
    if len(val) > 1:
        s += ")"
    return s


def printNumpy(x):
    return "[" + sumStr([decimal.Decimal(float(v)).__format__("f") + ", " for v in x.data.cpu().numpy()])[:-2] + "]"


def printStrList(x):
    return "[" + sumStr(v + ", " for v in x)[:-2] + "]"


def printListsNumpy(val):
    if isinstance(val, list):
        return printStrList(printListsNumpy(v) for v in val)
    return printNumpy(val)


def parseValues(values, methods, *others):
    if len(values) == 1 and values[0]:
        x = eval(values[0], dict(pair for l in ([methods] + list(others)) for pair in l.__dict__.items()))

        return x() if inspect.isclass(x) else x
    args = []
    kargs = {}
    for arg in values[1:]:
        if '=' in arg:
            k = arg.split('=')[0]
            v = arg[len(k) + 1:]
            try:
                kargs[k] = eval(v)
            except:
                kargs[k] = v
        else:
            args += [eval(arg)]
    return getattr(methods, values[0])(*args, **kargs)


def preDomRes(outDom, target):  # TODO: make faster again by keeping sparse tensors sparse
    t = one_hot(target.long(), outDom.size()[1]).to_dense().to_dtype()
    tmat = t.unsqueeze(2).matmul(t.unsqueeze(1))

    tl = t.unsqueeze(2).expand(-1, -1, tmat.size()[1])

    inv_t = eye(tmat.size()[1]).expand(tmat.size()[0], -1, -1)
    inv_t = inv_t - tmat

    tl = tl.bmm(inv_t)

    fst = outDom.unsqueeze(1).matmul(tl).squeeze(1)
    snd = outDom.unsqueeze(1).matmul(inv_t).squeeze(1)

    return (fst - snd) + t


def mopen(shouldnt, *args, **kargs):
    if shouldnt:
        import contextlib
        return contextlib.suppress()
    return open(*args, **kargs)


def file_timestamp():
    return str(datetime.now()).replace(":", "").replace(" ", "")


def prepareDomainNameForFile(s):
    return s.replace(" ", "_").replace(",", "").replace("(", "_").replace(")", "_").replace("=", "_")


# delimited only
def callCC(foo):
    class RV(BaseException):
        def __init__(self, v):
            self.v = v

    def cc(x):
        raise RV(x)

    try:
        return foo(cc)
    except RV as rv:
        return rv.v
