import torch as th
from utils.meters import AverageMeter


def spatial_mean(x):
    return x.mean(tuple(range(2, x.dim()))) if x.dim() > 2 else x


def dim_reducer(x):
    if x.dim() <= 3:
        return x

    return x.view(x.shape[0], x.shape[1], -1)


# k will ignore k-1 most extreme values
def spatial_edges(x,k=1,is_max=True,):
    if x.dim() < 3:
        return x
    x_ = dim_reducer(x)

    ret = x_.topk(k, -1, is_max)[0]
    if is_max:
        return ret[:,:,0]

    return ret[:,:,k-1]


def spatial_min(x,k=1):
    return spatial_edges(x,k,False)


def spatial_max(x,k=1):
    return spatial_edges(x,k,True)


def spatial_margin(x,k=1):
    return spatial_max(x,k)-spatial_min(x,k)


def spatial_min_max(x,k=1):
    return th.stack([spatial_max(x,k),spatial_min(x,k)],1).view(x.shape[0],-1)


def spatial_mean_max(x,k=1):
    return th.stack([spatial_max(x,k),spatial_mean(x)],1).view(x.shape[0],-1)


def spatial_min_mean_max(x, k=1):
    return th.stack([spatial_max(x, k), spatial_mean(x), spatial_min(x, k)], 1).view(x.shape[0], -1)


def spatial_l2(x):
    if x.dim() < 3:
        return x
    return th.norm(dim_reducer(x), dim=-1)


def max_local_response(x, grid=3):
    if x.dim() < 3:
        return x
    if x.dim() < 4:
        return th.nn.functional.adaptive_max_pool1d(x, grid).view(x.shape[0], x.shape[1], -1).log_softmax(-1).sum(
            -1).neg_()
    return th.nn.functional.adaptive_max_pool2d(x, grid).sum((-1, -2)).view(x.shape[0], x.shape[1], -1).log_softmax(
        -1).sum(-1).neg_()


def sum_outer_product(x):
    if x.dim() < 3:
        return x
    # r, c = th.tril_indices(x.shape[1], x.shape[1]).split(1)
    x_ = x.view(x.shape[0], x.shape[1], -1)
    x_ = x_.bmm(x_.transpose(2, 1)).log_softmax(-1).sum(-1).neg_()
    # [:, :, r, c].view(x.shape[0],-1)
    # x_=x_.sum(2)
    return x_


def max_outer_product_flat(x):
    if x.dim() < 3:
        return x
    r, c = th.tril_indices(x.shape[1], x.shape[1]).split(1)
    x_ = x.view(x.shape[0], x.shape[1], -1).max(-1)[0].unsqueeze(2)
    x_ = x_.bmm(x_.transpose(2, 1))[:, r, c].view(x.shape[0], -1)
    #
    # x_=x_.sum(2)
    return x_


def G_p(ob, p=10):
    temp = ob ** p
    temp = temp.reshape(temp.shape[0], temp.shape[1], -1)
    temp = ((th.matmul(temp, temp.transpose(dim0=2, dim1=1)))).sum(dim=2)
    temp = (temp.sign() * th.abs(temp) ** (1 / p)).reshape(temp.shape[0], -1)
    return temp


class MahalanobisDistance():
    def __init__(self,mean,inv_cov):
        self.mean = mean
        self.inv_cov = inv_cov
        self.use_mean_device = False

    def __call__(self, x):
        if x.device != self.mean.device:
            # if not hasattr(self, 'use_mean_device') or self.use_mean_device != _USE_PERCENTILE_DEVICE:
            #     self.use_mean_device = _USE_PERCENTILE_DEVICE
            if self.use_mean_device:
                x = x.to(self.mean.device)
            else:
                self.mean = self.mean.to(x.device)
                self.inv_cov = self.inv_cov.to(x.device)
        x_c = x - self.mean
        return (x_c.matmul(self.inv_cov).matmul(x_c.t())).diag().unsqueeze(1).sqrt()


class SumL1ChannelsDiff():
    def __init__(self, mean):
        self.mean = mean
        self.use_mean_device = False

    def __call__(self, x):
        if x.device != self.mean.device:
            # if not hasattr(self, 'use_mean_device') or self.use_mean_device != _USE_PERCENTILE_DEVICE:
            #     self.use_mean_device = _USE_PERCENTILE_DEVICE
            if self.use_mean_device:
                x = x.to(self.mean.device)
            else:
                self.mean = self.mean.to(x.device)

        # return x.sub(self.mean).abs_().mean(1).unsqueeze(1)
        return x.abs().mean(1).unsqueeze(1)


# clac Simes per batch element (samples x variables)
def calc_simes(pval):
    pval, _ = th.sort(pval, 1)
    view_shape = [-1, pval.shape[1]] + [1] * (pval.dim() - 2)
    rank = th.arange(1, pval.shape[1] + 1, device=pval.device).view(view_shape)
    simes_pval, _ = th.min(pval.shape[1] * pval / rank, 1)
    return simes_pval.clamp(0,1).unsqueeze(1)


def calc_cond_fisher(pval, thresh=1):
    pval[pval > thresh] = 1
    return -2 * pval.log().sum(1).unsqueeze(1)


# rescaled fisher test
def calc_mean_fisher(pval):
    return -2 * pval.log().mean(1).unsqueeze(1)



class StatefulReductionFactory():
    def __init__(self, stateful_reduction, **reduction_kwargs):
        self.reduction_ctor = stateful_reduction
        self.reduction_kwargs = reduction_kwargs
        self.reset()

    def reset(self):
        self.all_records = {}
        self.curr_class = 0
        self.curr_input = 0
        self.set_class_input(0, 0)

    def set_class(self, c, create_new=True):
        self.curr_class = c
        self.set_class_input(self.curr_class, self.curr_input, create_new)

    def set_input(self, i, create_new=True):
        self.curr_input = i
        self.set_class_input(self.curr_class, self.curr_input, create_new)

    def set_class_input(self, c, i, create_new=True):
        self.class_input_specifier = (c, i)
        if create_new and self.class_input_specifier not in self.all_records:
            self.all_records[self.class_input_specifier] = {}
        self.record = self.all_records[self.class_input_specifier]

    def __call__(self, trace_name, create_new=True):
        if create_new and trace_name not in self.record:
            fn_ob = self.reduction_ctor(trace_name, **self.reduction_kwargs)
            self.record[trace_name] = fn_ob
        return self.record[trace_name]


class StatefulReduction():
    def __init__(self, trace_name):
        self.trace_name = trace_name

    def __call__(self, x, measuring=False):
        return self.specialized_reduction(self.shared_reduction(x), measuring)

    def specialized_reduction(self, observations_list, measuring=False):
        pass

    def shared_reduction(self, x):
        pass


class PdistanceGram(StatefulReduction):
    def __init__(self, trace_name, precentile=0.01, powers=range(1, 11)):
        super().__init__(trace_name)
        self.precentile = precentile
        self.powers = powers
        self._quantile_high = [AverageMeter() for p in powers]
        self._quantile_low = [AverageMeter() for p in powers]

    def shared_reduction(self, x):
        observations = []
        for i in self.powers:
            observations.append(G_p(x, i))
        return observations

    def specialized_reduction(self, observations_list, measuring=False):
        l = 0
        for e, i in enumerate(self.powers):
            observations = observations_list[e]
            if measuring:
                tmp = observations.topk(int(self.precentile * observations.shape[0]), 0, True)[0]
                p_95 = tmp[-1, :]
                tmp = observations.topk(int(self.precentile * observations.shape[0]), 0, False)[0]
                p_5 = tmp[-1, :]
                self._quantile_high[e].update(p_95, observations.shape[0])
                self._quantile_low[e].update(p_5, observations.shape[0])
            if observations.device != self._quantile_high[e].mean.device:
                self._quantile_high = self._quantile_high[e].mean.to(observations.device)
                self._quantile_low = self._quantile_low[e].mean.to(observations.device)
            l += (observations - self._quantile_high[e].mean).clamp(min=0) / self._quantile_high[e].mean.abs()
            l += (self._quantile_low[e].mean - observations).clamp(min=0) / self._quantile_low[e].mean.abs()
        return l


class Pdistance(StatefulReduction):
    def __init__(self, trace_name, precentile=0.05, reduction_fn=spatial_max):
        super().__init__(trace_name)
        self.precentile = precentile
        self.reduction_fn = reduction_fn
        self._quantile_high = AverageMeter()
        self._quantile_low = AverageMeter()

    def shared_reduction(self, x):
        observations = self.reduction_fn(x)
        return observations

    def specialized_reduction(self, observations, measuring=False):
        if measuring:
            tmp = observations.topk(int(self.precentile * observations.shape[0]), 0, True)[0]
            p_95 = tmp[-1, :]
            tmp = observations.topk(int(self.precentile * observations.shape[0]), 0, False)[0]
            p_5 = tmp[-1, :]
            self._quantile_high.update(p_95, observations.shape[0])
            self._quantile_low.update(p_5, observations.shape[0])
        if observations.device != self._quantile_high.mean.device:
            self._quantile_high = self._quantile_high.mean.to(observations.device)
            self._quantile_low = self._quantile_low.mean.to(observations.device)
        l = (observations - self._quantile_high.mean).clamp(min=0) / self._quantile_high.mean.abs()
        l += (self._quantile_low.mean - observations).clamp(min=0) / self._quantile_low.mean.abs()
        return l

    def __call__(self, x, measuring=False):
        return self.specialized_reduction(self.shared_reduction(x), measuring)

        # dev = (F.relu(mins[L][p] - g_p) / torch.abs(mins[L][p] + 1e-6)).sum(dim=1, keepdim=True)
        # dev += (F.relu(g_p - maxs[L][p]) / torch.abs(maxs[L][p] + 1e-6)).sum(dim=1, keepdim=True)


_DEFAULT_SPATIAL_REDUCTIONS = {
    'spatial-mean': spatial_mean,
    # 'gram_p1':partial(G_p,p=1),
    # 'gram_p2':partial(G_p,p=2),
    # 'max_deviations_0.95': StatefulReductionFactory(Pdistance,**dict(precentile=0.05)),
    # 'max_deviations_0.99': StatefulReductionFactory(Pdistance,**dict(precentile=0.01)),
    # 'gram_div_0.95_p1':StatefulReductionFactory(PdistanceGram,**dict(powers=range(1,2),precentile=0.05)),
    # 'gram_div_0.95_p2':StatefulReductionFactory(PdistanceGram,**dict(powers=range(1,3),precentile=0.05)),
    # 'gram_div_0.95_p10':StatefulReductionFactory(PdistanceGram,**dict(powers=range(10,11),precentile=0.05)),
    # 'gram_div_0.99_p1_p10':StatefulReductionFactory(PdistanceGram,**dict(precentile=0.01)),
    'spatial-max': spatial_max,
    # 'max_outer_product_flat':max_outer_product_flat,
    # 'max_local_response' : max_local_response,
    # 'sum_outer_product': sum_outer_product,
    # 'spatial-min': spatial_min,
    # 'spatial-min_max': spatial_min_max,
    # 'spatial-mean_max': spatial_mean_max,
    # 'spatial-min_mean_max': spatial_min_mean_max,
    # 'spatial-l2':spatial_l2,
    # 'spatial-l2':spatial_l2
}
