import numpy as np
from collections import Counter
from sklearn import tree
from sklearn.datasets import load_iris
from copy import deepcopy

def emperical_entropy(y_train,train_no):
    y_dict=Counter(y_train)
    prob={key:y_dict[key]/train_no for key in y_dict.keys()}
    H=0
    for key in prob.keys():
        H-=prob[key]*np.log2(prob[key])
    return H

def emperical_conditional_entropy(y_train,A,train_no):
    A_value=list(set(A))
    A_no=len(A_value)
    A_dict=Counter(A)
    H_A=0
    for i in range(A_no):
        H_A+=A_dict[A_value[i]]/train_no*emperical_entropy(y_train[A==A_value[i]],A_dict[A_value[i]])
    return H_A
        
def find_feature_no(X_train,y_train,epsilon,method):
#    print(1)
    train_no,feature_no=np.shape(X_train)
    H=emperical_entropy(y_train,train_no)
    max_feature_no=0
    maximum=0
    
    for i in range(feature_no):
        H_A=emperical_conditional_entropy(y_train,X_train[:,i],train_no)
        if method=='ID3':
            g_A=H-H_A
        if method=='C4.5':
            g_A=(H-H_A)/H
#        print(g_A)
        
#        print(i,H,H_A,g_A)
        if g_A>maximum:
            maximum=g_A
            max_feature_no=i
    if maximum<epsilon:
        max_feature_no=0;length=0;set_A=[];leaf_node=True
#        print(maximum)
        
    else:
        set_A=list(set(X_train[:,max_feature_no]))
        length=len(set_A)
        leaf_node=False
#        print(max_feature_no,set_A)
    return max_feature_no,length,set_A,leaf_node
        
def generate(X_train,y_train,father_no,epsilon,method):
#    print(2)
    train_no,feature_no=np.shape(X_train)
    
    y_dict=Counter(y_train)
    C=max(y_dict,key=y_dict.get)
    if len(y_dict)==1 or train_no==0:
        max_feature_no=0;length=0;set_A=[];leaf_node=True
#        print('a')
    else:
        max_feature_no,length,set_A,leaf_node=find_feature_no(X_train,y_train,epsilon,method)
#        print('b')

    return (X_train,y_train,father_no,C,max_feature_no,length,set_A,leaf_node)

def pruning(node,leaf):
    
    
    
    for key in leaf:
        layer_no,node_no=key
#        print(layer_no,node_no)
        
    i=0
    while True:
        if node_no+1-node[(layer_no-1,i)][5]>0:
#            print("check:",node[(layer_no-1,i)][5])
            node_no=node_no-node[(layer_no-1,i)][5]
            i+=1
            
        else:
            node_no=i
            break
        
    X_train,y_train,father_no,C,max_feature_no,length,set_A,leaf_node=node[(layer_no-1,node_no)]
    node[(layer_no-1,node_no)] = X_train,y_train,father_no,C,max_feature_no,length,set_A,True
    print(layer_no-1,node_no)

        
    father_node_list=[(layer_no-1,node_no)]
    
    father_node_list,node=delete(father_node_list,node)
        
    maximum=0
    for key in node:
        if key[0]>maximum:
            maximum=key[0]
        layer_no=maximum

    return layer_no, node

def delete(father_node_list,node):

    if len(father_node_list)==0:      
        return [],node

    else:    
        temp_list=[]
        temp_node=deepcopy(node)
        for key,value in temp_node.items():
            if value[2] in father_node_list:

                temp_list.append(key)
                node.pop(key)
                
        node_list=deepcopy(temp_list)
        return delete(node_list,node)

def Loss_function(nobe,alpha):
    leaf_dict=find_leaf(nobe)
    L=0
    for key,value in leaf_dict.items():
        train_no=len(value[1])
        L+=train_no*emperical_entropy(value[1],train_no)
    return L+alpha*len(leaf_dict)
        

def find_leaf(tree):
    leaf={}
#    print(tree)
    for key,value in tree.items():
        if value[7]==True:
            leaf[key]=value
    return leaf

class DecisionTree():
    def __init__(self,X_train,y_train,X_test,epsilon=10**-3,alpha=1,criterion='entropy',method='ID3'):
        self.X_train=X_train
        self.y_train=y_train
        self.X_test=X_test
        train_no,feature_no=np.shape(self.X_train)
        test_no,_=np.shape(self.X_test)
        self.train_no=train_no
        self.feature_no=feature_no
        self.test_no=test_no
        self.epsilon=epsilon
        self.alpha=alpha
        self.criterion=criterion
        self.method=method

    def train(self):
        self.node={}
        self.layer_no=0
        node_no=1

        self.node[(0,0)]=(self.X_train,self.y_train,(np.nan,np.nan))
    

        while True:
            temp_no=0
            for i in range(node_no):
                self.node[(self.layer_no,i)]=generate(self.node[(self.layer_no,i)][0],self.node[(self.layer_no,i)][1],self.node[self.layer_no,i][2],self.epsilon,self.method)
                X_train,y_train,father_no,C,max_feature_no,length,set_A,leaf_node=self.node[(self.layer_no,i)]
#                print(length,max_feature_no)
                if leaf_node==False:
                    for j in range(length):
                        X_train_son=np.delete(X_train[X_train[:,max_feature_no]==set_A[j]],max_feature_no,1)
                        y_train_son=y_train[X_train[:,max_feature_no]==set_A[j]]
                        if y_train_son.size==0:
                            y_train_son=y_train
#                            X_train_son=X_train

                        self.node[(self.layer_no+1,temp_no+j)]=(X_train_son,y_train_son,(self.layer_no,i))
                    temp_no+=length           

            if temp_no==0:
                break
#            print('layer_no:',self.layer_no)
            node_no=temp_no
            self.layer_no+=1
#        print(self.node.keys())


   
        minimum=Loss_function(self.node,self.alpha)
        print(minimum)
        flag=True
        while flag:
            leaf_dict=find_leaf(self.node)        
            flag=False
            leaf={}
            print(leaf_dict.keys())
            for key in leaf_dict:
                print(key)
                if key[0]==0:
                    flag=True
                    break
                leaf[key]=leaf_dict[key]
                temp_node=deepcopy(self.node) 
                layer_no,temp_node=pruning(temp_node, leaf)
                L=Loss_function(temp_node,self.alpha)
                print(L)
                leaf={}
                if L<minimum:
                    minimum=L
                    self.node=deepcopy(temp_node)
                    self.layer_no=layer_no
                    flag=True
                    break
        for key,value in self.node.items():
            if value[7]==True:
               print(key,value[2],value[3],value[6],value[7])
  
    def predict(self):
        y_test=np.zeros((self.test_no))


        for row_no in range(self.test_no):
            i=0
            node_no=0
            X_test=self.X_test[row_no,:]
            while i<=self.layer_no:
                X_train,y_train,father_no,C,max_feature_no,length,set_A,leaf_node=self.node[(i,node_no)]
#                print(i,node_no)
                if leaf_node==True:
                    y_test[row_no]=C
#                    print(y_test)
                    break
             
                index=set_A.index(X_test[max_feature_no])
                X_test=np.delete(X_test,max_feature_no)
                temp_no=0
                if node_no>=1:
                    for j in range(node_no):
                        temp_no+=self.node[(i,j)][5]
                    node_no=temp_no+index
                else:
                    node_no=index

                i+=1
        return y_test
                
if __name__=='__main__':
#    X_train=np.array([[0,0,0,0],[0,0,0,1],[0,1,0,1],[0,1,1,0],[0,0,0,0],[1,0,0,0],[1,0,0,1],[1,1,1,1],[1,0,1,2],[1,0,1,2],[2,0,1,2],[2,0,1,1],[2,1,0,1],[2,1,0,2],[2,0,0,0]])
#    y_train=np.array([0,0,1,1,0,0,0,1,1,1,1,1,1,1,0])
#    X_test=X_train
    iris=load_iris()
    X_train=iris.data[:,2:]
    y_train=iris.target
    X_train=X_train.astype(int)
    X_test=X_train

#    print(X_train)
#    print(y_train)
    
    model=DecisionTree(X_train,y_train,X_test,method='ID3')
    model.train()
    y_test=model.predict()
    print('Decision Tree', y_test)
    print(sum(y_test==y_train))
    model = tree.DecisionTreeClassifier(criterion='entropy')
    model.fit(X_train, y_train)
    y_test=model.predict(X_test)
    print('sklearn', y_test)
    print(sum(y_test==y_train))

            
            
            
            
            
            
            
            
            
            
            
            
            