import torch
from scipy.stats import norm, binom_test
import numpy as np
from math import ceil
from statsmodels.stats.proportion import proportion_confint
import scipy.stats as stats
import math
import torch.nn.functional as F
from torch.nn import  Sigmoid, Softmax
from architectures import get_architecture, ReasonNN
from train_utils import setup_seed

class Smooth(object):
    """A smoothed classifier g """

    # to abstain, Smooth returns this int
    ABSTAIN = -1

    def __init__(self, dataset, base_classifier, num_classes, sigma, card = False, reason_model =""):
        """
        :param base_classifier: maps from [batch x channel x height x width] to [batch x num_classes]
        :param num_classes:
        :param sigma: the noise level hyperparameter
        """
        self.base_classifier = base_classifier
        self.num_classes = num_classes
        self.sigma = sigma
        self.card = card
        if card:
            self.reason_model = torch.nn.DataParallel(ReasonNN(dataset)).cuda()
            self.reason_model.load_state_dict(torch.load(reason_model + '/checkpoint.pth.tar')['state_dict'])
            self.reason_model.eval()

    def certify(self, x: torch.tensor, n0: int, n: int, alpha: float, batch_size: int) -> (int, float):
        """ Monte Carlo algorithm for certifying that g's prediction around x is constant within some L2 radius.
        With probability at least 1 - alpha, the class returned by this method will equal g(x), and g's prediction will
        robust within a L2 ball of radius R around x.

        :param x: the input [channel x height x width]
        :param n0: the number of Monte Carlo samples to use for selection
        :param n: the number of Monte Carlo samples to use for estimation
        :param alpha: the failure probability
        :param batch_size: batch size to use when evaluating the base classifier
        :return: (predicted class, certified radius)
                 in the case of abstention, the class will be ABSTAIN and the radius 0.
        """
        if self.card:
            counts_selection = self._card_sample_noise(x, n0, batch_size)
        else:
            self.base_classifier.eval()
            # draw samples of f(x+ epsilon)
            counts_selection = self._sample_noise(x, n0, batch_size)
        # use these samples to take a guess at the top class
        cAHat = counts_selection.argmax().item()
        # draw more samples of f(x + epsilon)
        if self.card:
            counts_estimation = self._card_sample_noise(x, n, batch_size)
        else:
            counts_estimation = self._sample_noise(x, n, batch_size)
        # use these samples to estimate a lower bound on pA
        nA = counts_estimation[cAHat].item()
        pABar = self._lower_confidence_bound(nA, n, alpha)
        if pABar < 0.5:
            return Smooth.ABSTAIN, 0.0
        else:
            radius = self.sigma * norm.ppf(pABar)
            return cAHat, radius
        
    def predict(self, x: torch.tensor, n: int, alpha: float, batch_size: int) -> int:
        """ Monte Carlo algorithm for evaluating the prediction of g at x.  With probability at least 1 - alpha, the
        class returned by this method will equal g(x).

        This function uses the hypothesis test described in https://arxiv.org/abs/1610.03944
        for identifying the top category of a multinomial distribution.

        :param x: the input [channel x height x width]
        :param n: the number of Monte Carlo samples to use
        :param alpha: the failure probability
        :param batch_size: batch size to use when evaluating the base classifier
        :return: the predicted class, or ABSTAIN
        """
        if self.card:
            counts = self._card_sample_noise(x, n, batch_size)
        else:
            self.base_classifier.eval()
            counts = self._sample_noise(x, n, batch_size)
        top2 = counts.argsort()[::-1][:2]
        count1 = counts[top2[0]]
        count2 = counts[top2[1]]
        if binom_test(count1, count1 + count2, p=0.5) > alpha:
            return Smooth.ABSTAIN
        else:
            return top2[0]

    def _sample_noise(self, x: torch.tensor, num: int, batch_size) -> np.ndarray:
        """ Sample the base classifier's prediction under noisy corruptions of the input x.

        :param x: the input [channel x width x height]
        :param num: number of samples to collect
        :param batch_size:
        :return: an ndarray[int] of length num_classes containing the per-class counts
        """
        with torch.no_grad():
            counts = np.zeros(self.num_classes, dtype=int)
            for _ in range(ceil(num / batch_size)):
                this_batch_size = min(batch_size, num)
                num -= this_batch_size

                batch = x.repeat((this_batch_size, 1, 1, 1))
                noise = torch.randn_like(batch, device='cuda') * self.sigma
                predictions = self.base_classifier(batch + noise).argmax(1)
                counts += self._count_arr(predictions.cpu().numpy(), self.num_classes)
            return counts

    def _card_sample_noise(self, x, num, batch_size):
        """ Sample the base classifier's prediction under noisy corruptions of the input x.

        :param x: the input [channel x width x height]
        :param num: number of samples to collect
        :param batch_size:
        :return: an ndarray[int] of length num_classes containing the per-class counts
        """
        all_path = [self.base_classifier]
        all_path.extend(['attribute_model/attribute_group_%d.pth.tar'%i for i in range(85)])
        all_path.extend(['hierarchy_model/hierarchy_group_%d.pth.tar'%i for i in range(50)])
        
        assemble = torch.FloatTensor()
        for model_id, path in enumerate(all_path):
            checkpoint = torch.load(path)
            cur_num = num
            if model_id == 0:
                model = get_architecture(checkpoint["arch"], "AWA", classes = 28)
                m = Softmax(dim=1)
            else:
                model = get_architecture(checkpoint["arch"], "AWA", classes = 1)
                m = Sigmoid()

            model.load_state_dict(checkpoint['state_dict'])
            model.eval()

            setup_seed(0)
            with torch.no_grad():
                confidence = torch.FloatTensor([])
                for _ in range(ceil(cur_num / batch_size)):
                    this_batch_size = min(batch_size, cur_num)
                    cur_num -= this_batch_size

                    batch = x.repeat((this_batch_size, 1, 1, 1))
                    noise = torch.randn_like(batch, device='cuda') * self.sigma
                    logits = model(batch + noise)
                    confidence = torch.cat((confidence,m(logits).detach().cpu()),dim=0)
            assemble = torch.cat((assemble,confidence),dim=1)
        with torch.no_grad():
            predictions = self.reason_model(assemble).argmax(1)
            counts = self._count_arr(predictions.cpu().numpy(), self.num_classes)
        return counts
    
    def _count_arr(self, arr: np.ndarray, length: int) -> np.ndarray:
        counts = np.zeros(length, dtype=int)
        for idx in arr:
            counts[idx] += 1
        return counts

    def _lower_confidence_bound(self, NA: int, N: int, alpha: float) -> float:
        """ Returns a (1 - alpha) lower confidence bound on a bernoulli proportion.

        This function uses the Clopper-Pearson method.

        :param NA: the number of "successes"
        :param N: the number of total draws
        :param alpha: the confidence level
        :return: a lower bound on the binomial proportion which holds true w.p at least (1 - alpha) over the samples
        """
        return proportion_confint(NA, N, alpha=2 * alpha, method="beta")[0]
