//  BSD License (http://www.galagosearch.org/license)
package org.lemurproject.galago.core.eval;

import org.lemurproject.galago.utility.WrappedMap;

import java.io.*;
import java.util.Map;
import java.util.TreeMap;

/**
 * This class store a relevance judgments for a set of queries
 *  - the judgments for each query are stored in a QueryJudgments object
 * 
 * Relevance is represented by an integer
 *  - positive indicates that the document is relevant
 *  - zero indicates that the document is not relevant
 *  - negative indicates that the document is detrimental to results (irrelevant)
 * 
 * @author sjh, jfoley
 */
public class QuerySetJudgments extends WrappedMap<String, QueryJudgments> {


  /**
   * Loads a TREC judgments file.
   *
   * @param filename The filename of the judgments file to load.
   * @param makeBinary whether to force the judgments into a binary form
   * @param makePositive whether to force the judgments into a positive form.
   * @throws java.io.IOException file issues
   */
  public QuerySetJudgments(String filename, boolean makeBinary, boolean makePositive) throws IOException {
    this(loadJudgments(filename, makeBinary, makePositive));
  }
  
  /**
   * Creates a galago evaluation object from your home-grown judgments.
   *
   * @param data something like a Map<String,Map<String,Int>>
   */
  public QuerySetJudgments(Map<String, QueryJudgments> data) {
    super(data);
  }

  public QueryJudgments get(String query) {
    return wrapped.get(query);
  }

  // LOADING FUNCTIONS
  /**
   * Loads a TREC judgments file.
   *
   * @param filename The filename of the judgments file to load.
   * @param makeBinary whether to force the judgments into a binary form
   * @param makePositive whether to force the judgments into a positive form.
   * @return Maps from query numbers to lists of judgments for each query.
   * @throws java.io.IOException file issues
   */
  public static Map<String, QueryJudgments> loadJudgments(String filename, boolean makeBinary, boolean makePositive) throws IOException {
    BufferedReader in = null;
    Map<String, QueryJudgments> judgments = new TreeMap<String,QueryJudgments>();

    try {
      in = new BufferedReader(new FileReader(filename));

      while (true) {
        String line = in.readLine();
        if (line == null) {
          break;
        }
        
        String[] cols = line.split("\\s+");
        assert(cols.length == 4);

        String queryNumber = cols[0];
        String unused = cols[1];
        String docno = cols[2];
        String judgment = cols[3];

        // ensure the query is stored
        if (!judgments.containsKey(queryNumber)) {
          judgments.put(queryNumber, new QueryJudgments(queryNumber));
        }

        // add this judgment to the query
        int j = Integer.parseInt(judgment);
        if (makeBinary) {
          j = (j > 0) ? 1 : 0;
        } else if (makePositive) {
          j = (j > 0) ? j : 0;
        }
        judgments.get(queryNumber).add(docno, j);
      }
    } finally {
      if (in != null) {
        in.close();
      }
    }

    if(judgments.isEmpty()) {
      throw new RuntimeException("Error: no judgments found in file: "+filename);
    }
    return judgments;
  }

  public void saveJudgments(File fp) throws FileNotFoundException {
    PrintWriter out = null;

    try {
      out = new PrintWriter(fp);

      // for each qid
      for (String qid : this.keySet()) {
        QueryJudgments qj = this.get(qid);

        // for each document, rel pair
        for (Map.Entry<String, Integer> judgement : qj.entrySet()) {
          String doc = judgement.getKey();
          int rel = judgement.getValue();
          out.println(String.format("%s 0 %s %d", qid, doc, rel));
        }
      }
    } finally {
      if(out != null) out.close();
    }
  }
}

@@ -91,16 +91,70 @@
 .eval;%0A%0A
+import org.lemurproject.galago.utility.StreamCreator;%0A
 import o
@@ -2106,38 +2106,8 @@
 n %7B%0A
-    BufferedReader in = null;%0A
     
@@ -2162,102 +2162,137 @@
 Map%3C
-String,QueryJudgments%3E();%0A%0A    try %7B%0A      in = new BufferedReader(new FileReader(filename));%0A
+%3E();%0A%0A    try (BufferedReader in = new BufferedReader(new InputStreamReader(StreamCreator.openInputStream(filename), %22UTF-8%22))) %7B
 %0A   
@@ -2303,15 +2303,20 @@
 hile
- 
 (true)
+%0A     
  %7B%0A 
@@ -2406,24 +2406,16 @@
       %7D%0A
-        
 %0A       
@@ -2465,16 +2465,17 @@
   assert
+ 
 (cols.le
@@ -3083,76 +3083,8 @@
    %7D
- finally %7B%0A      if (in != null) %7B%0A        in.close();%0A      %7D%0A    %7D
 %0A%0A  

Support compressed qrels in loadJudgments()