import numpy as np

def EvaluateFeature(Word, WordStats, Nc, num_classes):
    # A : number of times w and c co-occur
    # B : number of times w occurs without c
    # C : number of times c occurs without w
    # D : number of times neither c nor w occur
    # N : total number of documents
    
    # Returns an evaluation score

    N = np.sum(Nc)
    scores = []
    total_pres = np.sum(WordStats[Word][1])
    for c in range(num_classes):
        A = WordStats[Word][1][c]
        B = total_pres - A
        C = Nc[c] - A
        D = N - (A + B + C)
        scores.append(X2(A, B, C, D, N))
    
    # The paper doesn’t specify how to aggregate X2 across classes.
    # We use the max instead of the mean to capture the strongest class association.
    return np.max(scores)

def X2(A, B, C, D, N):
    # calculates the chi square statistics of a feature
    numerator = N * (A*D - C*B)**2
    denominator = max((A+C) * (B+D) * (A+B) * (C+D), 1e-12)
    if numerator == 0:
        return 0
    else:
        return numerator/denominator


class ACC_Q:
    '''A queue maintaining the accuracy over last n and 2n instances'''
    def __init__(self, N=50):
        self.N = N

        self.accuracies = []
        self.Acc_over_N = 0
        self.Acc_over_2N = 0
    
    def Push(self, val):
        self.accuracies.append(val)
        if len(self.accuracies) > 2*self.N:
            self.accuracies.pop(0)
        
        self.Acc_over_N = sum(self.accuracies[-self.N:])
        self.Acc_over_2N = sum(self.accuracies[-2 * self.N:])


class LEARNER:
    def __init__(self, num_classes, Document=None, DocClass=None, Features=[], N=50, data_name = None):
        
        self.num_classes = num_classes
        self.Nc = np.zeros(num_classes, dtype=float) # Number of documents corresponding to each class

        self.Vocabulary = Features # All words seen by the model till time t

        self.WordStats = {Word : np.zeros((2, num_classes)) for Word in self.Vocabulary} 
        # key: word (feature)
        # value: (2 x num_classes) matrix 
            # (0, c) -> number of times word is absent in class c document
            # (1, c) -> number of times word is present in class c document
        
        self.Pc = self.Nc / sum(self.Nc) if sum(self.Nc) != 0 \
                else 1.0/num_classes * np.ones(num_classes, dtype=np.float64) # Probability of occurance of each class
        self.P = dict() # Dictionary to store probability of each word (feature) w.r.t each class
                        # key : word | value : [prob. wrt class 0, prob. wrt class 1] (2d vector)

        self.acc = ACC_Q(N) # Queue of accuracy over last n instances
        self.Accuracy = 0
        self.age = 0
        self.probation = 0
        self.data_name = data_name

        if (Document is not None and DocClass is not None):
            self.Accuracy = 1 # If pre-trained, accuracy considered to be 100% till that instant
            self.Update(Document, DocClass)

    def score(self, Document):
        scores = self.Pc.copy()
        if len(self.P) != 0:
            for Word in Document:
                if Word in self.Vocabulary:
                    div = 1.0
                    if self.data_name in ['crowdsense_c3', 'crowdsense_c5']:
                        div = 1000.0
                    # For crowdsense data the scores shoots up to infinity. To curtail that we divide the input values by 1000.
                    scores *= self.P[Word]*(Document[Word]/div)

        return scores
    
    # Performs prediction of Learner and updates accuracy
    def use(self, Document, DocClass = None):
        scores = self.score(Document)
        logits = scores/sum(scores) if sum(scores)!=0 else scores
        pred_class =  int(np.argmax(logits))

        if DocClass is not None:
            self.acc.Push(int(pred_class == DocClass))
            self.Accuracy = self.acc.Acc_over_N
            
        return pred_class, logits.tolist()
    
    def UpdateWordStats(self, Document, DocClass):
        # Update the WordStat of every word(feature) in vocabulary of the learner
        for Word in self.Vocabulary:
            if Word in Document:
                self.WordStats[Word][1][DocClass] += 1
            else:
                self.WordStats[Word][0][DocClass] += 1
    
    def UpdateProbabilities(self):
        # Probability of Document belonging to class nc
        self.Pc = self.Nc/sum(self.Nc)  if sum(self.Nc) != 0 else self.Nc

        # Conditional probability of each word
        for Word in self.Vocabulary:
            self.P[Word] = (self.WordStats[Word][1] + 1) / (self.Nc + 2)
            # The addition in the numerator and denominator are for Laplace smoothing to avoid division by zero.
    
    def Update(self, Document, DocClass):
        # Wrapper to perform updates and call other update functions
        self.age += 1

        self.Nc[DocClass] += 1
        self.UpdateWordStats(Document, DocClass)
        self.UpdateProbabilities()    


class FAE:

    def __init__(self, num_classes, m=5, p=3, f=0.15, r=10, N=50, M=250, data_name = None, Document=None, DocClass=None):
        
        # Initialize parameters
        self.num_classes = num_classes
        self.m = m # (maturity) Number of instances needed before a learner’s classifications are used by the ensemble
        self.p = p # (probation time) is the number of times in a row a learner is allowed to be under the threshold before being removed
        self.f = f # (feature change threshold) is the threshold placed on the amount of change between the youngest learner’s set of features (yfs) and the top M features (mfs);
        self.r = r # (growth rate) is the number of instances between when the last learner was added and when the ensemble’s accuracy is checked for the addition of a new learner
        self.N = N # Number of instances over which to compute an accuracy measure;
        self.M = M # Number of features selected by the feature selection algorithm for a newly created learner

        self.FeatureDict = dict()   # key : word
                                    # value : evaluation score
        
        self.FeatureList = list() # list of words (features) sorted according to evaluation score
        
        self.TopMFeatures = list() # list of top M words (features) with highest evaluation score
        
        self.Vocabulary = set() # All words seen by the model till time t
        
        self.WordStats = dict()     # key: word (feature)
                                    # value: (2 x num_classes) matrix 
                                        # (0, c) -> number of times word is absent in class c document
                                        # (1, c) -> number of times word is present in class c document
        
        self.Nc = np.zeros(num_classes, dtype=float) # Number of documents corresponding to each class


        # Initialize other required attributed
        self.acc = ACC_Q(self.N)
        self.probation = np.array([])
        self.age = np.array([])
        self.data_name = data_name

        # Initialize the first learner
        if (Document is not None) and (DocClass is not None):
            self.InitialTraining(Document, DocClass)
            self.learners = [LEARNER(num_classes, Document=Document, DocClass=DocClass, Features=self.TopMFeatures,
                                      N=self.N, data_name = self.data_name)]
        else:
            self.learners = [LEARNER(num_classes, Features=self.TopMFeatures, N=self.N, data_name = self.data_name)]
        
        self.threshold = self.get_threshold()
    
    def UpdateWordStats(self, Document, DocClass):
        # Update the WordStat of every word(feature) in vocabulary of the learner
        for Word in self.Vocabulary:
            if Word in Document:
                self.WordStats[Word][1][DocClass] += 1
            else:
                self.WordStats[Word][0][DocClass] += 1
    
    def InitialTraining(self, Document, DocClass):
        self.Nc[DocClass] += 1

        # Create Initial Vocabulary
        self.Vocabulary = set(Document)

        # Create Initial WordStats and update
        self.WordStats = {Word: np.zeros((2, self.num_classes)) for Word in self.Vocabulary}
        self.UpdateWordStats(Document, DocClass)

        # Evaluate and calculate X2 statistic of a word
        for Word in self.Vocabulary:
            Evaluation = EvaluateFeature(Word, self.WordStats, self.Nc, self.num_classes)
            self.FeatureDict[Word] = Evaluation
        
        # Sort features in decreasing order of their evaluation value
        self.FeatureList = [Word for Word, Evaluation in sorted(self.FeatureDict.items(), key=lambda x:x[1], reverse=True)]        
        self.TopMFeatures = self.FeatureList[:self.M]

    def get_threshold(self):
        # Consider only mature learners (age >= m)
        mature_accuracies = [l.Accuracy for l in self.learners if l.age >= self.m]
        
        # If no mature learners yet, fallback to all learners (early phase)
        if len(mature_accuracies) == 0:
            mature_accuracies = [l.Accuracy for l in self.learners]

        th = (max(mature_accuracies) + min(mature_accuracies)) / 2.0
        return th
    
    def predict(self, Document, DocClass = None):
        
        age = np.array([learner.age for learner in self.learners])
        probation = np.array([learner.probation for learner in self.learners])

        mask = ((age >= self.m) * (probation == 0)).astype(int)
        if sum(mask) == 0:
            weights = age
        else:
            weights = np.array([l.Accuracy if m else 0 for m, l in zip(mask, self.learners)])

        scores = np.zeros(self.num_classes, dtype=float)

        for weight, learner in zip(weights, self.learners):
            pred, logit = learner.use(Document, DocClass)

            # Weighted vote of learners
            scores[pred] += weight

        scores = np.array(scores)/sum(scores) if sum(scores) !=0 else scores
        y_pred = int(np.argmax(scores))
        y_logit = scores

        if self.num_classes == 2:
            return y_pred, y_logit[1]

        return y_pred, y_logit.tolist()
   
    def Update(self, Document, DocClass):
        # Update Threshold
        self.threshold = self.get_threshold()
        
        # Update Probations
        for learner in self.learners:
            if learner.age >= self.m:
                if learner.Accuracy < self.threshold:
                    learner.probation += 1
                else:
                    learner.probation = 0
        
        # Remove learners whose probation is more than p
        self.learners = [l for l in self.learners if l.probation < self.p]
        
        # Create new learner on current instance
        self.learners.append(LEARNER(self.num_classes, Document=Document, DocClass=DocClass, Features=self.TopMFeatures, data_name=self.data_name))

    def UpdateFeatureSet(self, Document, DocClass=None):
        if DocClass is not None:
            self.Nc[DocClass] += 1
        # Update Vocabulary and Initiate WordStats for new words.
        for Word in set(Document):
            if Word not in self.Vocabulary:
                self.Vocabulary.add(Word)
                self.WordStats[Word] = np.zeros((2, self.num_classes))
        
        # Evaluate and calculate X2 statistic of a word
        for Word in self.Vocabulary:
            Evaluation = EvaluateFeature(Word, self.WordStats, self.Nc, self.num_classes)
            self.FeatureDict[Word] = Evaluation
        
        # Sort features in decreasing order of their evaluation value
        self.FeatureList = [Word for Word, Evaluation in sorted(self.FeatureDict.items(), key=lambda x:x[1], reverse=True)]        
        self.TopMFeatures = self.FeatureList[:self.M]

    def partial_fit(self, Document, DocClass):
        DocClass = int(DocClass)

        # Update Ensemble Accuracy
        y_pred, y_logit = self.predict(Document, DocClass)
        ensemble_accuracy = int(y_pred == DocClass)
        self.acc.Push(ensemble_accuracy)

        # Update Feature Set
        self.UpdateFeatureSet(Document, DocClass)

        # Train all learners on current instance
        for learner in self.learners:
            learner.Update(Document, DocClass)
        
        # Calculate change in features
        youngest_learner = self.learners[-1]
        
        nfs = set(self.TopMFeatures)
        yfs = set(youngest_learner.Vocabulary)

        delta = len(nfs.symmetric_difference(yfs)) / self.M
        
        if (delta>self.f) or ( (youngest_learner.age>self.r) and (self.acc.Acc_over_N < self.acc.Acc_over_2N) ):
           
            self.Update(Document, DocClass)
        
        return y_pred, y_logit
    
    def learner_count(self):
        return len(self.learners)