package org.thunlp.learning.lda;

import java.io.DataInput;
import java.io.DataOutput;
import java.io.IOException;

import org.apache.hadoop.io.Writable;

public class DocumentWritable implements Writable {
  public int [] words = null;
  public int [] topics = null;
  private int numWords = 0;
  private byte [] buffer = new byte[10240];
  
  public int getNumWords() {
    return numWords;
  }
  
  public void setNumWords( int n ) {
    if ( words == null || words.length < n ) {
      words = new int[n];
      topics = new int[n];
    }
    numWords = n;
  }

  public void readFields(DataInput input) throws IOException {
    int size = input.readInt();
    if (buffer.length < size) {
      buffer = new byte[size + 1024];
    }
    
    input.readFully(buffer, 0, size);
    setNumWords(size / 4 / 2);
    for (int i = 0; i < numWords; i++) {
      words[i] = fourBytesToInt(buffer, i * 2 * 4);
      topics[i] = fourBytesToInt(buffer, (i * 2 + 1) * 4);
    }
  }

  public void write(DataOutput output) throws IOException {
    int size = numWords * 2 * 4;
      if (buffer.length < size) {
        buffer = new byte[size + 1024];
      }
    for (int i = 0; i < numWords; i++) {
      intToFourBytes(buffer, i * 2 * 4, words[i]);
      intToFourBytes(buffer, (i * 2 + 1) * 4, topics[i]);
    }
    output.writeInt(size);
    output.write(buffer, 0, size);
  }
  
  public static int fourBytesToInt(byte [] b, int offset) {
    int i = 
      (b[offset] << 24) + 
      ((b[offset + 1] & 0xFF) << 16) +
      ((b[offset + 2] & 0xFF) << 8) +
      (b[offset + 3] & 0xFF);
    return i;
  }

  public static void intToFourBytes(byte [] b, int offset, int i) {
    b[offset] = (byte) (i >>> 24);
    b[offset + 1] = (byte) (i >>> 16);
    b[offset + 2] = (byte) (i >>> 8);
    b[offset + 3] = (byte) (i);
  }
  
  public String toString() {
    StringBuilder sb = new StringBuilder();
    for ( int i = 0 ; i < numWords ; i++ ) {
      sb.append(i > 0 ? " " : "");
      sb.append(words[i]);
      sb.append(":");
      sb.append(topics[i]);
    }
    return sb.toString();
  }

}

@@ -268,26 +268,86 @@
 %0A  p
-rivate int numWord
+ublic int %5B%5D labels = null;%0A  private int numWords = 0;%0A  private int numLabel
 s = 
@@ -615,24 +615,220 @@
 s = n;%0A  %7D%0A%0A
+  public int getNumLabels() %7B%0A    return numLabels;%0A  %7D%0A%0A  public void setNumLabels(int n) %7B%0A    if(labels == null %7C%7C labels.length %3C n) %7B%0A      labels = new int%5Bn%5D;%0A    %7D%0A    numLabels = n;%0A  %7D%0A%0A
   public voi
@@ -882,25 +882,33 @@
 n %7B%0A    int 
-s
+wordByteS
 ize = input.
@@ -914,24 +914,110 @@
 .readInt();%0A
+    int labelByteSize = input.readInt();%0A    int size = wordByteSize + labelByteSize;%0A
     if (buff
@@ -1139,17 +1139,25 @@
 umWords(
-s
+wordByteS
 ize / 4 
@@ -1319,16 +1319,166 @@
 );%0A    %7D
+%0A%0A    setNumLabels(labelByteSize / 4);%0A    for (int i = 0; i %3C numLabels; i++) %7B%0A      labels%5Bi%5D = fourBytesToInt(buffer, wordByteSize + i * 4);%0A    %7D
 %0A  %7D%0A%0A  
@@ -1539,25 +1539,33 @@
 n %7B%0A    int 
-s
+wordByteS
 ize = numWor
@@ -1576,16 +1576,99 @@
  2 * 4;%0A
+    int labelByteSize = numWords * 4;%0A    int size = wordByteSize + labelByteSize;%0A
       if
@@ -1905,25 +1905,179 @@
     
-output.writeInt(s
+for(int i = 0; i %3C numLabels; i++) %7B%0A      intToFourBytes(buffer, wordByteSize + i * 2, labels%5Bi%5D);%0A    %7D%0A    output.writeInt(wordByteSize);%0A    output.writeInt(labelByteS
 ize)
@@ -2634,24 +2634,176 @@
 gBuilder();%0A
+    sb.append(%22%5B%22);%0A    for(int i = 0; i %3C numLabels; i++) %7B%0A      sb.append(i %3E 0 ? %22 %22 : %22%22);%0A      sb.append(labels%5Bi%5D);%0A    %7D%0A    sb.append(%22%5D%5Ct%22);%0A
     for ( in

Add labels to DocumentWritable