source: gs3-extensions/mars-src/trunk/src/java/org/greenstone/gsdl3/service/GS2WekaDBSearch.java@ 36859

/*
 *    GS2WekaDBSearch.java
 *    Copyright (C) 2011 New Zealand Digital Library, http://www.nzdl.org
 *
 *    This program is free software; you can redistribute it and/or modify
 *   the Free Software Foundation; either version 2 of the License, or
 *    (at your option) any later version.
 *
 *    This program is distributed in the hope that it will be useful,
 *    but WITHOUT ANY WARRANTY; without even the implied warranty of
 *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *    GNU General Public License for more details.
 *
 *    You should have received a copy of the GNU General Public License
 *    along with this program; if not, write to the Free Software
 *    Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
 */
package org.greenstone.gsdl3.service;


// Greenstone classes
import org.greenstone.gsdl3.util.*;

// XML classes
import org.w3c.dom.Document;
import org.w3c.dom.Element;
import org.w3c.dom.NodeList;

// java
import java.util.Vector;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;
import java.util.Iterator;
import java.io.File;

import org.apache.log4j.*;

// To fit in with the class hierarchy that has been developed to data
// around music content based retrieval, the class name shifts from
// the external facing 'Recommender' suffix to using 'Query' and/or
// 'Search'

public class GS2WekaDBSearch extends AbstractGS2AudioSearch {

    static Logger logger = Logger.getLogger (org.greenstone.gsdl3.service.GS2WekaDBSearch.class.getName ());

    protected static final String OFFSET_PARAM  = "offset";
    protected static final String LENGTH_PARAM  = "length";
    protected static final String RADIUS_PARAM  = "radius";
    protected static final String MAXDOCS_PARAM = "maxDocs";
    protected static final String AROUSAL_PARAM  = "arousal";
    protected static final String VALENCE_PARAM  = "valence";

    //protected static final String WEKA_MODEL_DEFAULT_DIRECTORY = "weka-model";
    //protected static final String KNN_MODEL_FILENAME = "av-and-features-knn.ser";

    protected static final String WEKA_DB_DEFAULT_DIRECTORY = "wekaDB";
    protected static final String KNN_FEATURES_FILENAME = "av-features.csv";

    
    protected WekaDBWrapper wekadb_src = null;

    public GS2WekaDBSearch() {

    if(this.wekadb_src == null) {
        logger.info("Initializing WekaDBWrapper");
        this.wekadb_src = new WekaDBWrapper();
        }
    }
    
    /** do the actual query */
    protected Element processAudioQuery (Element request) {

    // As the MG version needs to be java-synchronized (this inspiration for this class)
    // And since it is not known how concurrent (thread-safe) Weka can be ...
    // => Play it safe for now and restrict access to 'wekadb_src' using synchronized also
        synchronized(this.wekadb_src) {
        // Create a new (empty) result message ('doc' is in ServiceRack.java)
        Document result_doc = XMLConverter.newDOM();
        Element result = result_doc.createElement (GSXML.RESPONSE_ELEM);

        // Rather than QUERY_SERVICE use "TextQuery" 
        // => makes the result looks the same as a text query
        result.setAttribute (GSXML.FROM_ATT, "TextQuery"); 
        result.setAttribute (GSXML.TYPE_ATT, GSXML.REQUEST_TYPE_PROCESS);
        
        // Get the parameters of the request
        Element param_list = (Element) GSXML.getChildByTagName (request, GSXML.PARAM_ELEM+GSXML.LIST_MODIFIER);
        if (param_list == null) {
        logger.error ("Weka AudioQuery request had no paramList.");
        return result;  // Return the empty result
        }
        
        // Process the request parameters
        HashMap params = GSXML.extractParams (param_list, false);
        
        // Make sure a query has been specified
        String query = (String) params.get (QUERY_PARAM);
        if (query == null || query.equals ("")) {
        return result;  // Return the empty result
        }

        // If an index hasn't been specified, use the default
        String index = (String) params.get (INDEX_PARAM);
        if (index == null) {
        index = WEKA_DB_DEFAULT_DIRECTORY; 
        }
        
        // The location of the Weka db index
        String toplevel_index_dir = GSFile.collectionIndexDir (this.site_home, this.cluster_name);
        String weka_db_index_dir = toplevel_index_dir + File.separatorChar + index;
        String assoc_index_dir = toplevel_index_dir + File.separatorChar + "assoc"; // ****

        // set the Weka DB query parameters to the values the user has specified
        setStandardQueryParams (params); // ****
        
        this.wekadb_src.runQuery(weka_db_index_dir, KNN_FEATURES_FILENAME, assoc_index_dir, query);
        Vector docs = this.wekadb_src.getQueryResult();

        if (docs.isEmpty()) {
        // something has gone wrong
        GSXML.addError (result, "Couldn't query the Weka DB", GSXML.ERROR_TYPE_SYSTEM);
        return result;
        }
        long totalDocs = docs.size();
        
        // Get the docnums out, and convert to HASH ids
        if (docs.size () == 0) {
        logger.error ("No results found...\n");
        }
        
        // Create a metadata list to store information about the query results
        Element metadata_list = result_doc.createElement (GSXML.METADATA_ELEM+GSXML.LIST_MODIFIER);
        result.appendChild (metadata_list);
        
        // Add a metadata element specifying the number of matching documents
        // because the total number is just the number returned, use numDocsReturned, not numDocsMatched
        GSXML.addMetadata (metadata_list, "numDocsReturned", ""+totalDocs);
        // add a metadata item to specify what actual query was done - eg if stuff was stripped out etc. and then we can use the query later, cos we don't know which parameter was the query
        GSXML.addMetadata (metadata_list, "query", query);
        
        if (docs.size () > 0) {
        // Create a document list to store the matching documents, and add them
        Element document_list = result_doc.createElement (GSXML.DOC_NODE_ELEM+GSXML.LIST_MODIFIER);
        result.appendChild (document_list);
        for (int d = 0; d < docs.size (); d++) {
            WekaDBDocInfo wdb_doc = (WekaDBDocInfo) docs.elementAt(d);

            String doc_id  = wdb_doc.getDocID();
            double rank    = wdb_doc.getTopRank();
            String offsets = wdb_doc.getOffsetList();

            Element doc_node = createDocNode (result_doc, doc_id, Double.toString (rank));
            doc_node.setAttribute("frameOffset", offsets);

            document_list.appendChild (doc_node);
        }
        }
        
        // Create an empty term list as a place holder for the term information
        Element term_list = result_doc.createElement (GSXML.TERM_ELEM+GSXML.LIST_MODIFIER);
        result.appendChild (term_list);
        
        return result;
    }//end of synchronized
    }
    
    // should probably use a list rather than map
    protected boolean setStandardQueryParams(HashMap params)
    {
    Set entries = params.entrySet();
    Iterator i = entries.iterator();
    while (i.hasNext()) {
        Map.Entry m = (Map.Entry)i.next();
        String name = (String)m.getKey();
        String value = (String)m.getValue();

        if (name.equals(OFFSET_PARAM)) {
        int offset = Integer.parseInt(value);
        this.wekadb_src.setOffset(offset);
        }
        else if (name.equals(LENGTH_PARAM)) {
        int length = Integer.parseInt(value);
        this.wekadb_src.setLength(length);
        }
        else if (name.equals(RADIUS_PARAM)) {
        double radius = Double.parseDouble(value);
        this.wekadb_src.setRadius(radius);
        }
        else if (name.equals(MAXDOCS_PARAM)) {
        int docs = Integer.parseInt(value);
        this.wekadb_src.setMaxDocs(docs);
        }
        else if (name.equals(AROUSAL_PARAM)) {
        double arousal = Double.parseDouble(value);
        this.wekadb_src.setArousal(arousal);
        }       
        else if (name.equals(VALENCE_PARAM)) {
        double valence = Double.parseDouble(value);
        this.wekadb_src.setValence(valence);        
        } // ignore any others
    }
    return true;
    }
    
    
}