import numpy as np
import itertools

def set_trace():
    from IPython.core.debugger import Pdb
    import sys
    Pdb(color_scheme='Linux').set_trace(sys._getframe().f_back)

# CART decision tree algorithm
# Starter code borrowed from Programming Collective Intelligence (Segaran)

# Helper method for converting label dictionary to label
def convertToLabel(dictionary):
    best_count = None
    for key, count in dictionary.items():
        if best_count is None:
            label = key
            best_count = count
        elif count > best_count:
            best_count = count
            label = key
    return label

# Create a new class that represents each node in the tree
class decisionnode:
  def __init__(self,col=-1,value=None,results=None,tb=None,fb=None):
    self.col = col # column index of the criteria to be tested
    self.value = value # value that the column must match to get true result
    self.results = results # dictionary of results for branch
    self.tb = tb # decision node if the result is true
    self.fb = fb # decision node if the result is false

# Create function that divides rows into two sets
# based on data in specific column
def divideset(rows, column, value):
    # Input value can be nominal or numeric
    # Returns two lists of rows:
    # (1) containing rows where data in specified column matches value;
    # (2) containing rows where it does not
    """
    if isinstance(value, (int, float)):
      split_function = lambda row:row[column] >= value
    else:
      split_function = lambda row:row[column] == value


    set1 = [row for row in rows if split_function(row)]
    set2 = [row for row in rows if not split_function(row)]
    """

    if isinstance(value, (int, float)):
        split_function = lambda row:row >= value
    else:
        split_function = lambda row:row == value

    # Divide the rows into two sets and return them
    pos_idx = split_function(rows[:,column])
    set1 = rows[pos_idx,:]
    set2 = rows[~pos_idx,:]

    return (set1, set2)


# Create function to get counts of each result in a set
def uniquecounts(rows):
    # Returns dictionary of how many times possible outcomes appear
    """
    results={}
    for row in rows:
      # The result is the last column
      r=row[len(row)-1]
      if r not in results: results[r]=0
      results[r]+=1
    """
    (keys, vals) = np.unique(rows[:,-1], return_counts=True)
    return dict(zip(keys, vals))


# Create function for Gini impurity
def giniimpurity(rows):
      # Returns expected error rate
    total = float(len(rows))
    counts = uniquecounts(rows)
    """
    imp=0

    for k1 in counts:
      p1=float(counts[k1])/total
      for k2 in counts:
        if k1==k2: continue
        p2 = float(counts[k2])/total
        imp += p1*p2
    """

    imp = 0.0
    for k1, k2 in itertools.combinations(counts.keys(), 2):
        p1 = float(counts[k1]) / total
        p2 = float(counts[k2]) / total
        imp += p1*p2

    imp *= 2

    return imp


# Entropy is the sum of p(x)log(p(x)) across all the different possible results
def entropy(rows):
    """
    from math import log
    log2=lambda x:log(x)/log(2)
    results=uniquecounts(rows)
    # Now calculate the entropy
    ent=0.0
    for r in results.keys():
      p=float(results[r])/len(rows)
      ent=ent-p*log2(p)
    """
    results = uniquecounts(rows)
    ent = reduce(lambda ent, p : ent - p * np.log2(p), [0.0] + \
                 map(lambda x: float(x) / len(rows), results.values()))
    return ent


def printtree(tree,indent=''):
   # Is this a leaf node?
   if tree.results!=None:
      print str(tree.results)
   else:
      # Print the criteria
      print str(tree.col)+':'+str(tree.value)+'? '

      # Print the branches
      print indent+'T->',
      printtree(tree.tb,indent+'  ')
      print indent+'F->',
      printtree(tree.fb,indent+'  ')


def getwidth(tree):
  if tree.tb==None and tree.fb==None: return 1
  return getwidth(tree.tb)+getwidth(tree.fb)

def getdepth(tree):
  if tree.tb==None and tree.fb==None: return 0
  return max(getdepth(tree.tb),getdepth(tree.fb))+1


from PIL import Image,ImageDraw

def drawtree(tree,png='tree.png'):
  w=getwidth(tree)*100
  h=getdepth(tree)*100+120

  img=Image.new('RGB',(w,h),(255,255,255))
  draw=ImageDraw.Draw(img)

  drawnode(draw,tree,w/2,20)
  img.save(png,'PNG')

def drawnode(draw,tree,x,y):
  if tree.results==None:
    # Get the width of each branch
    w1=getwidth(tree.fb)*100
    w2=getwidth(tree.tb)*100

    # Determine the total space required by this node
    left=x-(w1+w2)/2
    right=x+(w1+w2)/2

    # Draw the condition string
    draw.text((x-20,y-10),str(tree.col)+':'+str(tree.value),(0,0,0))

    # Draw links to the branches
    draw.line((x,y,left+w1/2,y+100),fill=(255,0,0))
    draw.line((x,y,right-w2/2,y+100),fill=(255,0,0))

    # Draw the branch nodes
    drawnode(draw,tree.fb,left+w1/2,y+100)
    drawnode(draw,tree.tb,right-w2/2,y+100)
  else:
    txt=' \n'.join(['%s:%d'%v for v in tree.results.items()])
    draw.text((x-20,y),txt,(0,0,0))


def classify(observation, tree):
  if tree.results != None:
      return tree.results
  else:
      v = observation[tree.col]
      branch = None
      if isinstance(v, (float, int)):
          if v >= tree.value:
              branch = tree.tb
          else:
              branch = tree.fb
      else:
          if v == tree.value:
              branch = tree.tb
          else:
              branch = tree.fb
      return classify(observation, branch)

def prune(tree, mingain):
  # If the branches aren't leaves, then prune them
  if tree.tb.results == None:
      prune(tree.tb, mingain)
  if tree.fb.results == None:
      prune(tree.fb,mingain)

  # If both the subbranches are now leaves, see if they should merged
  if tree.tb.results!=None and tree.fb.results!=None:
    # Build a combined dataset
    tb,fb = [],[]
    for v,c in tree.tb.results.items():
      tb += [[v]]*c
    for v,c in tree.fb.results.items():
      fb += [[v]]*c

    # Test the reduction in entropy
    delta = entropy(np.array(tb+fb)) - \
           (entropy(np.array(tb)) + entropy(np.array(fb))/2)

    if delta < mingain:
      # Merge the branches
      tree.tb, tree.fb = None, None
      tree.results = uniquecounts(np.array(tb+fb))

def mdclassify(observation,tree):
  if tree.results!=None:
    return tree.results
  else:
    v=observation[tree.col]
    if v==None:
      tr,fr=mdclassify(observation,tree.tb),mdclassify(observation,tree.fb)
      tcount=sum(tr.values())
      fcount=sum(fr.values())
      tw=float(tcount)/(tcount+fcount)
      fw=float(fcount)/(tcount+fcount)
      result={}
      for k,v in tr.items(): result[k]=v*tw
      for k,v in fr.items(): result[k]=v*fw
      return result
    else:
      if isinstance(v,int) or isinstance(v,float):
        if v>=tree.value: branch=tree.tb
        else: branch=tree.fb
      else:
        if v==tree.value: branch=tree.tb
        else: branch=tree.fb
      return mdclassify(observation,branch)

def variance(rows):
  if len(rows)==0: return 0
  data=[float(row[len(row)-1]) for row in rows]
  mean=sum(data)/len(data)
  variance=sum([(d-mean)**2 for d in data])/len(data)
  return variance

def buildtree(rows, scoref=entropy, min_gain=0.1):
  if len(rows) == 0:
      return decisionnode()

  current_score = scoref(rows)

  # Set up some variables to track the best criteria
  best_gain = 0.0
  best_criteria = None
  best_sets = None

  column_count = len(rows[0]) - 1
  for col in range(0, column_count):
    # Generate the list of unique values in this column
    """
    column_values={}
    for row in rows:
        column_values[row[col]]=1
    """

    # Now try dividing the rows up for each value in this column
    for value in np.unique(rows[:,col]):
      (set1, set2) = divideset(rows, col, value)

      # Information gain
      p = float(len(set1)) / len(rows)
      gain = current_score - p*scoref(set1) - (1-p)*scoref(set2)
      if all((gain > best_gain,
              len(set1) > 0,
              len(set2) > 0)):
          best_gain = gain
          best_criteria = (col, value)
          best_sets = (set1, set2)

  # Create the sub branches
  if best_gain > 0:
      trueBranch = buildtree(best_sets[0])
      falseBranch = buildtree(best_sets[1])
      return decisionnode(col = best_criteria[0],
                          value = best_criteria[1],
                          tb = trueBranch,
                          fb = falseBranch)
  else:
      return decisionnode(results = uniquecounts(rows))