#!/usr/bin/env python

import os

import sys

import shutil

import random

import tempfile

import Queue

from org.apache.pig.scripting import Pig

from org.codehaus.jackson.map import ObjectMapper

EPS = 10e-6 # maximum distance between consective weights for convergence

pig_script = sys.argv[1] # pig script to run iteratively

data_dir = sys.argv[2] # directory where intermediate weights will be written

features = sys.argv[3] # location, inside data_dir, where the data to fit exists

num_features = sys.argv[4] # number of features

#

# Cleanup data dir

#

cmd = "rmr %s/weight-*" % data_dir

Pig.fs(cmd)

#

# Initialize weights

#

w0_fields = []

weights = []

for i in xrange(int(num_features)):

weights.append(str(random.random()))

w0_fields.append({"name":"w%s" % i,"type":25,"schema":None}) # See Pig's DataType.java

path = tempfile.mkdtemp()

w0 = open("%s/part-r-00000" % path, 'w')

w0.write("\t".join(weights)+"\n")

w0.close()

#

# Create schema for weights, place under weight-0 dir

#

w0_schema = {"fields":w0_fields,"version":0,"sortKeys":[],"sortKeyOrders":[]}

w0_schema_file = open("%s/.pig_schema" % path, 'w')

ObjectMapper().writeValue(w0_schema_file, w0_schema);

w0_schema_file.close()

#

# Copy initial weights to fs

#

copyFromLocal = "copyFromLocal %s %s/%s" % (path, data_dir, "weight-0")

Pig.fs(copyFromLocal)

#

# Iterate until converged

#

features = "%s/%s" % (data_dir,features)

script = Pig.compileFromFile(pig_script)

weight_queue = Queue.Queue(25) # for moving average

avg_weight = [0.0 for i in xrange(int(num_features))]

converged = False

prev = 0

weight_dir = tempfile.mkdtemp()

while not converged:

input_weights = "%s/weight-%s" % (data_dir,prev)

output_weights = "%s/weight-%s" % (data_dir,prev+1)

bound = script.bind({'input_weights':input_weights,'output_weights':output_weights,'data':features})

bound.runSingle()

#

# Copy schema for weights to each output

#

copyOutputSchema = "cp %s/.pig_schema %s/.pig_schema" % (input_weights, output_weights)

Pig.fs(copyOutputSchema)

#

# The first few iterations the weights bounce all over the place

#

if (prev > 1):

copyToLocalPrev = "copyToLocal %s/part-r-00000 %s/weight-%s" % (input_weights, weight_dir, prev)

copyToLocalNext = "copyToLocal %s/part-r-00000 %s/weight-%s" % (output_weights, weight_dir, prev+1)

Pig.fs(copyToLocalPrev)

Pig.fs(copyToLocalNext)

localPrev = "%s/weight-%s" % (weight_dir, prev)

localNext = "%s/weight-%s" % (weight_dir, prev+1)

x1 = open(localPrev,'r').readlines()[0]

x2 = open(localNext,'r').readlines()[0]

x1 = [float(x.strip()) for x in x1.split("\t")]

x2 = [float(x.strip()) for x in x2.split("\t")]

weight_queue.put(x1)

avg_weight = [x[1] + (x[0] - x[1])/(prev-1.0) for x in zip(x1,avg_weight)]

#

# Make sure to collect enough weights into the average before

# checking for convergence

#

if (prev > 25):

first_weight = weight_queue.get()

avg_weight = [(x[0] - x[1]/25.0 + x[2]/25.0) for x in zip(avg_weight, first_weight, x1)]

#

# Compute distance from weight centroid to new weight

#

d = sum([(pair[0] - pair[1])**2 for pair in zip(x2,avg_weight)])

converged = (d < EPS)

os.remove(localPrev)

os.remove(localNext)

prev += 1

#

# Cleanup

#

shutil.rmtree(path)

shutil.rmtree(weight_dir)