from typing import Any

import numpy as np
import torch
import torch.nn as nn
from tqdm import tqdm

from .base_postprocessor import BasePostprocessor
from math import ceil
""" Code borrowed from https://github.com/snu-mllab/Neural-Relation-Graph
"""


def normalize(feat, nc=50000):
    with torch.no_grad():
        split = ceil(len(feat) / nc)
        for i in range(split):
            feat_ = feat[i * nc:(i + 1) * nc]
            feat[i * nc:(i + 1) *
                 nc] = feat_ / torch.sqrt((feat_**2).sum(-1) + 1e-10).reshape(-1, 1)

    return feat


def kernel(feat, feat_t, prob, prob_t, split=2):
    """Kernel function (assume feature is normalized)
    """
    size = ceil(len(feat_t) / split)
    rel_full = []
    for i in range(split):
        feat_t_ = feat_t[i * size:(i + 1) * size]
        prob_t_ = prob_t[i * size:(i + 1) * size]

        with torch.no_grad():
            dot = torch.matmul(feat, feat_t_.transpose(1, 0))
            dot = torch.clamp(dot, min=0.)

            sim = torch.matmul(prob, prob_t_.transpose(1, 0))
            rel = dot * sim

        rel_full.append(rel)

    rel_full = torch.cat(rel_full, dim=-1)
    return rel_full


def get_relation(feat, feat_t, prob, prob_t, pow=1, chunk=50, thres=0.03):
    """Get relation values (top-k and summation)
    
    Args:
        feat (torch.Tensor [N,D]): features of the source data
        feat_t (torch.Tensor [N',D]): features of the target data
        prob (torch.Tensor [N,C]): probabilty vectors of the source data
        prob_t (torch.Tensor [N',C]): probabilty vectors of the target data
        pow (int): Temperature of kernel function
        chunk (int): batch size of kernel calculation (trade off between memory and speed)
        thres (float): cut off value for small relation graph edges. Defaults to 0.03.

    Returns:
        graph: statistics of relation graph
    """

    n = feat.shape[0]
    n_chunk = ceil(n / chunk)

    score = []
    for i in range(n_chunk):
        feat_ = feat[i * chunk:(i + 1) * chunk]
        prob_ = prob[i * chunk:(i + 1) * chunk]

        rel = kernel(feat_, feat_t, prob_, prob_t)

        mask = (rel.abs() > thres)
        rel_mask = mask * rel
        edge_sum = (rel_mask.sign() * (rel_mask.abs()**pow)).sum(-1)

        score.append(edge_sum.cpu())

    score = torch.cat(score, dim=0)

    return score


class RelationPostprocessor(BasePostprocessor):
    def __init__(self, config):
        super(RelationPostprocessor, self).__init__(config)
        self.args = self.config.postprocessor.postprocessor_args
        self.pow = self.args.pow
        self.feature_log = None
        self.args_dict = self.config.postprocessor.postprocessor_sweep
        self.setup_flag = False

    def setup(self, net: nn.Module, id_loader_dict, ood_loader_dict):
        if not self.setup_flag:
            feature_log = []
            prob_log = []
            net.eval()
            with torch.no_grad():
                for batch in tqdm(id_loader_dict['train'], desc='Setup: ', position=0, leave=True):
                    data = batch['data'].cuda()
                    data = data.float()

                    logit, feature = net(data, return_feature=True)
                    prob = torch.softmax(logit, dim=1)
                    feature_log.append(normalize(feature))
                    prob_log.append(prob)

            self.feat_train = torch.cat(feature_log, axis=0)
            self.prob_train = torch.cat(prob_log, axis=0)

            self.setup_flag = True
        else:
            pass

    @torch.no_grad()
    def postprocess(self, net: nn.Module, data: Any):
        output, feature = net(data, return_feature=True)
        feature = normalize(feature)
        prob = torch.softmax(output, dim=1)

        score = get_relation(feature, self.feat_train, prob, self.prob_train, pow=self.pow)

        _, pred = torch.max(prob, dim=1)

        return pred, score

    def set_hyperparam(self, hyperparam: list):
        self.pow = hyperparam[0]

    def get_hyperparam(self):
        return self.pow