import disco
 
def estimate_map(e, params):
  z=dict([(elem,1) for elem in e[1].split(params.splitter)]).keys()
  x=[a for a in z if not a in params.ys]
  y=[a for a in z if a in params.ys]
 
  return [(b+params.splitter+a,1) for a in x for b in y] + [(a,1) for a in z] + [('',1)]
#[(b+params.splitter,1) for b in y] + [(params.splitter+a,1) for a in x] + [(params.splitter,1)]
 
 
def estimate_combiner(k, v, counts, done, params):
  if done: return counts.items()
  
  if not counts.has_key(k): counts[k]=v
  else: counts[k]+=v
 
 
def estimate_reduce(iter, out, params):
  counts={}
 
  for k,v in iter:
    v=int(v)
    if not counts.has_key(k): counts[k]=v
    else: counts[k]+=v
 
  for k,v in counts.iteritems(): out.add(k,repr(v))
 
 
def predict_map(e, params):
  ll=dict([(k,params.loglikelihoods[k]) for k in params.ys.keys()])
 
  for elem in e[1].split(params.splitter):
    if params.ys.has_key(elem): continue
    
    for y in params.ys:
      ll[y]+=params.loglikelihoods[y+params.splitter+elem]
 
  return [(e[0], k + ' ' + repr(ll[k])) for k in params.ys]
 
 
def estimate(input, ys, splitter=' ', host="disco://localhost", map_reader=disco.chain_reader):
  ys=dict([(id,1) for id in ys])
 
  results = disco.job(host, name = 'naive_bayes_estimate',
          input_files = input, 
          map_reader = map_reader, 
          fun_map = estimate_map,
          combiner = estimate_combiner,
          reduce = estimate_reduce,
          params = disco.Params(ys=ys,splitter=splitter),
          sort = False, clean = False)
 
  total=0
  items={}
  classes={}
  pairs={}
  for key,value in disco.result_iterator(results):
    l=key.split(splitter)
    value=int(value)
    if len(l)==1: 
      if l[0]=='': total=value
      elif ys.has_key(l[0]): classes[l[0]]=value
      else: items[l[0]]=value
    else:
      pairs[key]=value
 
#counts[key]=[[c,i], [not c, i], [c, not i], [not c, not i]]
  counts={}
  for i in items:
    for y in ys:
      key=y+splitter+i
      counts[key]=[0,0,0,0]
      if pairs.has_key(key): counts[key][0]=pairs[key]
      counts[key][1]=items[i]-counts[key][0]
      if not classes.has_key(y): counts[key][2]=0
      else: counts[key][2]=classes[y]-counts[key][0]
      counts[key][3]=total-sum(counts[key][:3])
 
      # add pseudocounts
      counts[key]=map(lambda x: x+1, counts[key])
  total+=4
 
  import math
  loglikelihoods={}
  for key,value in counts.iteritems():
    log_c=math.log(value[0]+value[2])
    l=key.split(splitter)
    if not loglikelihoods.has_key(l[0]): loglikelihoods[l[0]]=0.0
    loglikelihoods[l[0]]+=math.log(value[2])-log_c
    loglikelihoods[key]=math.log(value[0])-math.log(value[2])
 
  return loglikelihoods
 
 
def predict(input, loglikelihoods, ys, splitter, host="disco://localhost", map_reader=disco.chain_reader):
  ys=dict([(id,1) for id in ys])
  results = disco.job(host, name = 'naive_bayes_predict',
          input_files = input, 
          map_reader = map_reader, 
          fun_map = predict_map,
          params=disco.Params(loglikelihoods=loglikelihoods,ys=ys,splitter=splitter),
          sort = False, clean = False)
 
  return results