#Play some ai's against each other in a game
import sys
import imp
import risktools
import time
import os
import random
import json
import math

class DNode():
    def __init__(self, parent, data, data_focus):
        self.parent = parent
        #print 'Creating D Tree Node with data_focus: ', data_focus
        self.depth = 1
        self.used_features = []
        if parent is not None:
            self.depth = parent.depth + 1
            self.used_features = parent.used_features[:]
            
        self.positive_child = None
        self.negative_child = None
        
        self.data = data
        self.data_focus = data_focus
        if data_focus is not None and len(data_focus) > 0:
            self.compute_pos_prob()
        self.split_feature = None
    
    def compute_pos_prob(self):
        positives = 0
        for f in self.data_focus:
            d = self.data[f]
            if d.label == 1:
                positives += 1
        self.positive_probability = float(positives) / float(len(self.data_focus))
        self.entropy = self.compute_entropy(self.positive_probability)
    
    def compute_entropy(self,p):
        if p == 0 or p == 1:
            return 0
        return -(p*math.log(p,2) + (1-p)*math.log(1-p,2))
    
    def save_node(self, savefile):
        #Save out this node's information, then recursively for its children
        if self.split_feature is None:
            savefile.write(json.dumps(self.positive_probability) + '\n')
        else:
            savefile.write(json.dumps(self.positive_probability) + '|' + json.dumps(self.split_feature) + '\n')
            self.positive_child.save_node(savefile)
            self.negative_child.save_node(savefile)
    
    def load_node(self, loadfile):
        loadline = loadfile.readline()
        splitline = loadline.split('|')
        if len(splitline) == 1:
            self.positive_probability = json.loads(splitline[0])
        else:
            self.positive_probability = json.loads(splitline[0])    
            self.split_feature = json.loads(splitline[1])
            self.positive_child = DNode(self,None,None)
            self.positive_child.load_node(loadfile)
            self.negative_child = DNode(self,None,None)
            self.negative_child.load_node(loadfile)
    
    def print_node(self):
        for d in range(self.depth):
            print '-', 
        if self.split_feature is not None:
            if self.data:   
                print self.data[0].feature_names[self.split_feature], '<', self.split_feature, '> (', self.positive_probability, ' , ', 1 - self.positive_probability, ' ) , ', len(self.data_focus), ' examples   reached here'
            else:
                print '<', self.split_feature, '> (', self.positive_probability, ' , ', 1 - self.positive_probability, ' )'
            self.positive_child.print_node()
            self.negative_child.print_node()
        else:
            if self.data:
                print '(', self.positive_probability, ' , ', 1 - self.positive_probability, ' ) , ', len(self.data_focus), ' examples reached here'
            else:
                print '(', self.positive_probability, ' , ', 1 - self.positive_probability, ' )'
    
    def determine_info_gain(self, feature_index):
        """Determine the information gain for the given feature"""
        positive_count = 0 
        positive_positive_count = 0
        negative_count = 0
        negative_positive_count = 0
        
        for f in self.data_focus:
            d = self.data[f]
            if d.features[feature_index] == 1:
                positive_count += 1
                if d.label == 1:
                    positive_positive_count += 1
            else:
                negative_count += 1
                if d.label == 1:
                    negative_positive_count += 1
        
        positive_fraction = float(positive_count)/float(len(self.data_focus))
        negative_fraction = 1 - positive_fraction
        
        #Compute positive remaining entropy
        positive_remainder = 0
        if positive_count > 0:
            positive_remainder = positive_fraction * self.compute_entropy(float(positive_positive_count) / float(positive_count))
        
        #Compute negative remaining entropy
        negative_remainder = 0
        if negative_count > 0:
            negative_remainder = negative_fraction * self.compute_entropy(float(negative_positive_count) / float(negative_count))
        
        return self.entropy - positive_remainder - negative_remainder
    
    def classify(self, instance):
        if self.split_feature == None:
            return self.positive_probability
        else:
            if instance[self.split_feature] == 1:
                return self.positive_child.classify(instance)
            else:
                return self.negative_child.classify(instance)
         
    def determine_split(self, depth_limit):
        #Just return if tree is deep enough or if all examples are labelled correctly
        if self.depth == depth_limit or self.positive_probability == 1 or self.positive_probability == 0:
            return
          
        best_feature = None
        best_feature_information_gain = 0
            
        for i in range(len(self.data[0].features)):
            if i not in self.used_features:
                
                current_gain = self.determine_info_gain(i)
                
                if best_feature is None or current_gain > best_feature_information_gain:
                    best_feature = i
                    best_feature_information_gain = current_gain
            
        #Split on best_feature
        self.split_feature = best_feature
        
    
        positive_focus = []
        negative_focus = []
        
        self.used_features.append(best_feature)
        
        for f in self.data_focus:
            d = self.data[f]
            if d.features[best_feature] == 1:
                positive_focus.append(f)
            else:
                negative_focus.append(f)
            
        
        self.positive_child = DNode(self,self.data,positive_focus)
        
        self.negative_child = DNode(self,self.data,negative_focus)
    
        self.positive_child.determine_split(depth_limit)
        self.negative_child.determine_split(depth_limit)
    
class DTree():
    def __init__(self, data):
        data_focus = None
        if data is not None:
            data_focus = range(len(data))
        self.root = DNode(None,data,data_focus)
    
    def learn_tree(self, depth_limit):
        self.root.determine_split(depth_limit)
    
    def get_prob_of_win(self, instance):
        return self.root.classify(instance)
    
    def print_tree(self):
        self.root.print_node()
    
    def save_tree(self, savename):
        savefile = open(savename, 'w')
        self.root.save_node(savefile)
        savefile.close()
        
    def load_tree(self, loadname):
        loadfile = open(loadname, 'r')
        self.root.load_node(loadfile)
        loadfile.close()
    
def loadDTree(loadfilename):
    tree = DTree(None)
    tree.load_tree(loadfilename)
    return tree
    
class DData():
    def __init__(self, features, label, feature_names):
        self.features = features
        self.label = label
        self.feature_names = feature_names

def print_usage():
    print 'USAGE: python learn_d_tree.py data_filename depth'
       
def read_data(datafile):
    #First read feature_names
    newline = datafile.readline()
    feature_names = json.loads(newline)
    
    dataover = False
    
    dataset = []
    
    while not dataover:
        newline = datafile.readline()
        splitline = newline.split('|')
        if not newline:
            dataover = True
        else:
            features = json.loads(splitline[0])
            label = json.loads(splitline[1])

            datum = DData(features, label, feature_names)
            dataset.append(datum)
    return dataset
    
if __name__ == "__main__":
    #Get ais from command line arguments
    if len(sys.argv) != 3:
        print_usage()
        sys.exit()
    
    #get data file name
    datafilename = sys.argv[1]
    
    #get depth of tree
    depth = int(sys.argv[2])
    
    #Open the logfile
    datafile = open(datafilename, 'r')    
    
    data = read_data(datafile)
    
    datafile.close()
    
    tree = DTree(data)
    tree.learn_tree(depth)
    tree.print_tree()
    print 'Learned tree. Saving to file'
    
    savefilename = 'decision_trees\\' + datafilename[9:-4] + '_' + str(depth) + '.dtree'
    tree.save_tree(savefilename)
    print 'Saved file to ', savefilename
    
    print 'Loading tree.'
    tree2 = loadDTree(savefilename)
    tree2.print_tree()