import disco
 
def estimate_map(e, params):
  x=map(float,e[1].split(' '))
  y=x[params.y_id]
  del x[params.y_id]
 
  return [(e[0],(x,y))]
 
 
def estimate_combiner(k, v, vals, done, params):
  if vals=={}: 
    vals['x']=[0.0]*len(v[0])
    vals['x2']=[0.0]*len(v[0])
    vals['xy']=[0.0]*len(v[0])
    vals['y']=0.0
    vals['c']=0
 
  if done:
    return [(k, ' '. join(map(repr,vals['x'] + vals['x2'] + vals['xy'] + [ vals['y'], vals['c'] ])))]
 
  for i in range(len(v)):
    vals['x'][i]+=v[0][i]
    vals['x2'][i]+=v[0][i]*v[0][i]
    vals['xy'][i]+=v[0][i]*v[1]
  vals['y']+=v[1]
  vals['c']+=1
 
 
 
def predict_map(e, params):
  x=map(float,e[1].split(' '))
  return [(e[0],' '.join(map(repr,[params[i][0]+params[i][1]*x[i] for i in range(len(params))])))]
 
 
def estimate(input, y_id, host="disco://localhost", map_reader=disco.chain_reader):
  results = disco.job(host, name = 'naive_linear_regression_estimate',
          input_files = input, 
          map_reader = map_reader, 
          fun_map = estimate_map,
          combiner = estimate_combiner, 
          params=disco.Params(y_id=y_id),
          sort = False, clean = False)
 
  c=0
  y=0.0
  l=None
  x=None
  x2=None
  xy=None
 
  for key,value in disco.result_iterator(results):
    v=map(float,value.split(' '))
    
    if l==None:
      l=(len(v)-2)/3
      x=[0.0]*l
      x2=[0.0]*l
      xy=[0.0]*l
 
    c+=v[-1]
    y+=v[-2]
    for i in range(l):
      x[i]+=v[i]
      x2[i]+=v[l+i]
      xy[i]+=v[2*l+i]
 
  b = [ (c*xy[i] - x[i]*y)/(c*x2[i]+x[i]*x[i]) for i in range(l) ]
  a = [ (y-b[i]*x[i])/c for i in range(l) ]
 
  return zip(*(a,b))
 
 
def predict(input, model, host="disco://localhost", map_reader=disco.chain_reader):
  results = disco.job(host, name = 'naive_linear_regression_predict',
          input_files = input, 
          map_reader = map_reader, 
          fun_map = predict_map,
          params=model,
          sort = False, clean = False)
 
  return results