D-Net

package eu.dnetlib.data.search.solr;

import com.google.gson.Gson;

import eu.dnetlib.api.data.SearchServiceException;

import eu.dnetlib.data.search.transform.Transformer;

import eu.dnetlib.data.search.transform.TransformerException;

import eu.dnetlib.data.search.utils.solr.SolrResultSetOptionsUtil;

import eu.dnetlib.data.search.utils.solr.SolrResultsFormatter;

import eu.dnetlib.domain.EPR;

import eu.dnetlib.functionality.cql.CqlTranslatorImpl;

import gr.uoa.di.driver.enabling.resultset.ResultSet;

import org.apache.commons.lang.StringEscapeUtils;

import org.apache.log4j.Logger;

import org.apache.solr.client.solrj.SolrServerException;

import org.apache.solr.client.solrj.impl.CloudSolrServer;

import org.apache.solr.client.solrj.response.FacetField;

import org.apache.solr.client.solrj.response.QueryResponse;

import org.apache.solr.common.SolrDocumentList;

import org.apache.solr.common.params.SolrParams;

import org.apache.solr.common.util.NamedList;

import org.z3950.zing.cql.CQLParseException;

import javax.ws.rs.core.MediaType;

import java.io.IOException;

import java.io.OutputStream;

import java.util.ArrayList;

import java.util.HashMap;

import java.util.List;

import java.util.Map;

/**

 * Created by antleb on 2/4/14.

 */

public class SolrResultSet implements ResultSet<String> {

    private Logger logger = Logger.getLogger(getClass());

    private EPR epr = null;

    private CloudSolrServer solrClient = null;

    private NamedList<String> queryOpts = new NamedList<String>();

    long size = -1;

    public SolrResultSet(EPR epr, CloudSolrServer solrClient) throws IOException, CQLParseException {

        this.epr = epr;

        this.solrClient = solrClient;

        this.queryOpts = SolrResultSetOptionsUtil.extractQueryOptions(epr.getParameter("query"));

        String layout = epr.getParameter("layout");

        String mdformat = epr.getParameter("mdformat");

        String interpretation = epr.getParameter("interpretation");

        solrClient.setDefaultCollection(mdformat + "-" + layout + "-" + interpretation);

    }

    @Override

    public boolean isOpen() {

        return true;

    }

    @Override

    public boolean isAlive() {

        return true;

    }

    @Override

    public void close() {

        //solrClient.shutdown();

    }

    @Override

    public int size() {

        return (int) size;

    }

    @Override

    @Deprecated

    public List<String> getElements(int from, int to) {

        return get(from, to);

    }

    List<FacetField> facetFields = null;

    @Override

    @Deprecated

    public List<String> get(int from, int to) {

        List<String> res = new ArrayList<String>();

        QueryResponse rsp = null;

        HashMap<String, List<String>> map = new HashMap<String, List<String>>();

        logger.debug("from: " + from);

        logger.debug("to: " + to);

        queryOpts.add("start", from+1 + "");

        queryOpts.add("rows", to + 1+"");

        try {

            rsp = solrClient.query(SolrParams.toSolrParams(queryOpts));

            facetFields = rsp.getFacetFields();

            SolrDocumentList docs = rsp.getResults();

            if (facetFields!=null && !facetFields.isEmpty()) {

                for (int i = from - 1; i < to; i++) {

                    for (FacetField field : facetFields) {

                        if (field.getValueCount() > i) {

                            BrowseField bf = new BrowseField();

                            bf.setId(field.getValues().get(i).getName());

                            bf.setName(field.getValues().get(i).getName());

                            bf.setCount(field.getValues().get(i).getCount() + "");

                            if (map.get(field.getName()) == null) {

                                map.put(field.getName(), new ArrayList<String>());

                            }

                            map.get(field.getName()).add(new Gson().toJson(bf));

                        }

                    }

                }

                for (Map.Entry<String, List<String>> facetEntry : map.entrySet()) {

                    StringBuilder builder = new StringBuilder();

                    builder.append("\"" + facetEntry.getKey() + "\"" + " : ");

                    builder.append(facetEntry.getValue());

                    res.add(builder.toString());

                }

            }

            logger.debug("time: " + rsp.getElapsedTime());

            logger.debug("found: " + docs.getNumFound());

            logger.debug("docs: " + docs.size());

            for (int i = 0; i < docs.size(); i++) {

                String result = ((ArrayList<String>) docs.get(i).get("__result")).get(0);

                res.add(result);

            }

            return res;

        } catch (SolrServerException sse) {

            logger.error("Fail to get results from Solr. ", sse);

        }

        return null;

    }

    @Override

    public EPR getEpr() {

        return null;

    }

    public Map<String,List<String>> newGet(int from, int to, String format, Transformer transformer, Transformer oldRefineTransformer) {

        List<String> refineSolrResults = new ArrayList<String>();

        List<String> searchSolrResults = new ArrayList<String>();

        logger.debug("format: " + format);

        QueryResponse rsp = null;

        HashMap<String, List<String>> map = new HashMap<String, List<String>>();

        //logger.info("from: " + from*to);

        //logger.info("to: " + to);

        queryOpts.add("start", from*to + "");

        queryOpts.add("rows", to +"");

        try {

            long startTime = System.nanoTime();

            rsp = solrClient.query(SolrParams.toSolrParams(queryOpts));

            long estimatedTime = System.nanoTime() - startTime;

            logger.info("Solrj time " + estimatedTime/1000000 +  " milliseconds for query:" + queryOpts.get("q") +

                    " and facets " + queryOpts.getAll("facet.field") + " and fq " + queryOpts.getAll("fq") + " from: "

                    + from + " and size " + to);

            facetFields = rsp.getFacetFields();

            SolrDocumentList docs = rsp.getResults();

            this.size = docs.getNumFound();

            if (facetFields!=null && !facetFields.isEmpty()) {

                logger.debug("Checking " + (format != null && format.equals(MediaType.APPLICATION_JSON)));

                if (format != null && format.equals(MediaType.APPLICATION_JSON)) {

                    for (FacetField field : facetFields) {

                        map.put(field.getName(), new ArrayList<String>());

                        BrowseField bf = null;

                        for (int i = 0; i < field.getValueCount(); i++) {

                            bf = new BrowseField();

                            //bf.setId(org.apache.commons.lang3.StringEscapeUtils.escapeJson(field.getValues().get(i).getName()));

                            bf.setId(field.getValues().get(i).getName());

                            String[] facetedValues = field.getValues().get(i).getName().split("\\|\\|",2);

                            if (facetedValues.length > 1) {

                                //bf.setName(org.apache.commons.lang3.StringEscapeUtils.escapeJson(facetedValues[1]));

                                bf.setName(facetedValues[1]);

                            } else if (field.getValues().get(i).getName().split("_\\:\\:",2).length > 1) {

                                //bf.setName(org.apache.commons.lang3.StringEscapeUtils.escapeJson(field.getValues().get(i).getName().split("\\:\\:",2)[1]).replaceAll("\\:\\:", "\\|"));

                                bf.setName(field.getValues().get(i).getName().split("\\:\\:",2)[1].replaceAll("\\:\\:", "\\|"));

                            } else {

                                //bf.setName(org.apache.commons.lang3.StringEscapeUtils.escapeJson(field.getValues().get(i).getName()));

                                bf.setName(field.getValues().get(i).getName());

                            }

                            bf.setCount(field.getValues().get(i).getCount() + "");

                            map.get(field.getName()).add(new Gson().toJson(bf));

                        }

                    }

                    StringBuilder builder = null;

                    for (Map.Entry<String, List<String>> facetEntry : map.entrySet()) {

                        builder = new StringBuilder();

                        builder.append("\"" + facetEntry.getKey() + "\"" + " : ");

                        builder.append(facetEntry.getValue());

                        refineSolrResults.add(builder.toString());

                    }

                } else { //the old implementation & xml as default //TODO check compatibility

                    logger.debug("Creating old browse results.");

                    createXmlRefineFields(refineSolrResults, oldRefineTransformer);

                }

            }

            for (int i = 0; i < docs.size(); i++) {

                String result = ((ArrayList<String>) docs.get(i).get("__result")).get(0);

                try {

                    if (transformer != null) {

                        //logger.debug("1 >>>>>>" + result);

                        String xml = result.replaceAll("<em>","").replaceAll("</em>","");

                        result = transformer.transform(xml);

                        //logger.debug("2 >>>>>>" + result);

                    }

                } catch (TransformerException te) {

                    logger.warn("Error transforming " + result, te);

                    continue;

                }

                if (format != null && format.equals(MediaType.APPLICATION_JSON)) {

                    searchSolrResults.add(SolrResultsFormatter.xml2Json(result));

                } else { // default xml

                    searchSolrResults.add(result);

                }

            }

            Map<String,List<String>> response = new HashMap<String, List<String>>();

            logger.debug("refine results " + refineSolrResults);

            //logger.debug("search results " + searchSolrResults);

            response.put("refine",refineSolrResults);

            response.put("search", searchSolrResults);

            return response;

        } catch (SolrServerException sse) {

            logger.error("Error calling Solr.", sse);

        }

        return null;

    }

    public void cursorGet(Transformer transformer, OutputStream os) throws SolrServerException, SearchServiceException {

        queryOpts.add("start", "0");

        queryOpts.add("rows", "0");

        queryOpts.remove("rows");

        queryOpts.add("rows", "500");

        queryOpts.add("fl", "__result");

        queryOpts.add("shards.tolerant","true");

        queryOpts.add("cursorMark", "*");

        queryOpts.add("sort", "__indexrecordidentifier asc");

        String cursorMark = "*";

        String nextCursorMark = "";

        int curs = 0;

        QueryResponse resp = solrClient.query(SolrParams.toSolrParams(queryOpts));

        while (!cursorMark.equals(nextCursorMark)) {

            resp = solrClient.query(SolrParams.toSolrParams(queryOpts));

            cursorMark = nextCursorMark;

            nextCursorMark = resp.getNextCursorMark();

            for (int i = 0; i < resp.getResults().size(); i++) {

                if (transformer != null) {

                    String result = null;

                    try {

                        result = transformer.transform(((ArrayList<String>) resp.getResults().get(i).get("__result")).get(0));

                    } catch (TransformerException te) {

                        logger.warn("Error transforming " + result, te);

                        continue;

                    }

                    try {

                        os.write(result.getBytes());

                        os.flush();

                    } catch (IOException e) {

                        logger.error("Cursor get... ", e);

                        continue;

                    }

                }

            }

            queryOpts.remove("cursorMark");

            queryOpts.add("cursorMark", nextCursorMark);

            curs ++;

        }

        logger.debug("CURS " + curs);

    }

    //TODO get rid of this as soon as Joomla portal is out

    //Just copied and refactored the old one...

    @Deprecated

    private void createXmlRefineFields(List<String> res, Transformer oldRefineTransformer) {

        int max = -12;

        for (FacetField field:facetFields) {

            logger.debug("field " + field.getName() + " has count " + field.getValueCount());

            if (field.getValueCount() > max) {

                max = field.getValueCount();

            }

        }

        logger.debug("max " + max);

        for (int i = 0; i < max; i++) {

            StringBuilder sb = new StringBuilder();

            sb.append("<row>");

            for (FacetField field:facetFields) {

                if (field.getValueCount() > i) {

                    sb.append("<groupresult field=\"").append(field.getName()).append("\">");

                    sb.append("<count>").append(field.getValues().get(i).getCount()).append("</count>");

                    sb.append("<originalValue>").append(StringEscapeUtils.escapeXml(field.getValues().get(i).getName())).append("</originalValue>");

                    String[] facetValues = field.getValues().get(i).getName().split("\\|\\|");

                    if(facetValues.length > 1) {

                        sb.append("<value>").append(StringEscapeUtils.escapeXml(facetValues[1])).append("</value>");

                    } else {

                        sb.append("<value>").append(StringEscapeUtils.escapeXml(facetValues[0])).append("</value>");

                    }

                    sb.append("</groupresult>");

                }

            }

            sb.append("</row>");

            try {

                //logger.debug("row: " + sb.toString());

                //logger.debug("row2: " + oldRefineTransformer.transform(sb.toString()));

                //TODO remove

                res.add(oldRefineTransformer.transform(sb.toString()));

            } catch (TransformerException te) {

                logger.error("Cannot transform refine for: " + sb.toString(), te);

            }

        }

    }

   public static void main(String[] args) throws IOException, CQLParseException, SolrServerException {

        CloudSolrServer solrClient = new CloudSolrServer("beta.solr.openaire.eu:9983");

        solrClient.setDefaultCollection("TMF-index-openaire");

        NamedList<String> queryOpts = new NamedList<String>();

        //q=*:*&start=0&rows=10&cursorMark=*&sort=dateofcollection asc

        queryOpts.add("q", new CqlTranslatorImpl().getTranslatedQuery("oaftype=result").asLucene());

        queryOpts.add("start", "0");

        queryOpts.add("rows", "2");

        //queryOpts.add("fl", "__result");

        queryOpts.add("shards.tolerant","true");

       // queryOpts.add("cursorMark", "*");

        queryOpts.add("sort", "resultdateofacceptance desc");

        //queryOpts.add("q", new CqlTranslatorImpl().getTranslatedQuery("oaftype exact project").asLucene());

        NamedList<String> extraOpts = new NamedList<String>();

        QueryResponse resp = solrClient.query(SolrParams.toSolrParams(queryOpts));

        System.out.println("results " + resp.getResults().size());

        /*String cursorMark = "*";

        String nextCursorMark = "";

        int curs = 0;

        while (!cursorMark.equals(nextCursorMark)) {

            System.out.println("cursor " + cursorMark);

            System.out.println("next cursor " + nextCursorMark);

            cursorMark = nextCursorMark;

            for (int i = 0; i < resp.getResults().size(); i++) {

                String result = ((ArrayList<String>) resp.getResults().get(i).get("__result")).get(0);

                //System.out.println(result);

                resp = solrClient.query(SolrParams.toSolrParams(queryOpts));

            }

            nextCursorMark = resp.getNextCursorMark();

            queryOpts.add("cursorMark", nextCursorMark);

            System.out.println("CURS " + curs);

            curs ++;

        }

        //System.out.println((new CqlTranslatorImpl().getTranslatedQuery("objIdentifier = acm_________::0002c24f82c295e925a2bdf7bbf49bfc").asLucene()));

        //extraOpts.add("start", "1");

       // extraOpts.add("rows", "10");

       // extraOpts.addAll(queryOpts);

        //queryOpts.add("facet", "true");

        //TranslatedQuery translatedQuery = new CqlTranslatorImpl().getTranslatedQuery("oaftype=result sortBy resultdateofacceptance/sort.descending");

     //   queryOpts.add("q", "oaftype=project");

        //queryOpts.add("facet", "true");

        //queryOpts.add("facet.mincount", "1");

        //queryOpts.add("fq", "popularity");

//        queryOpts.put("fq", new CqlTranslatorImpl().getTranslatedQuery("").asLucene());

       // queryOpts.add("facet.field", "contextid");

       //  queryOpts.add("facet.field", "contextname");

       //  queryOpts.add("facet.mincount", "1");

       //  queryOpts.add("facet.threads", "10");

       // System.out.println(translatedQuery.getOptions().getSort().getMode());

       // System.out.println(translatedQuery.getOptions().getSort().getField());

        //queryOpts.add("sort", translatedQuery.getOptions().getSort().getField() + " " + translatedQuery.getOptions().getSort().getMode() );

/*        QueryResponse resp = null;

        synchronized (solrClient) {

            resp = solrClient.query(SolrParams.toSolrParams(extraOpts));

        }*/

//        System.out.println("time: " + resp.getElapsedTime());

    //System.out.println("results: " + resp.getResults());

/*      System.out.println(resp.getFacetField("contextname").getValueCount());

        for (FacetField.Count count:resp.getFacetField("contextname").getValues())

            System.out.println(count.getName() + " : " +  count.getCount());

        int max = -12;

        for (FacetField field:resp.getFacetFields()) {

            if (field.getValueCount() > max)

                max = field.getValueCount();

        }

        System.out.println("max: " + max);

*/

    }

//    @Override

//    public EPR getEpr() {

//        return epr;

//   }

}

class BrowseField {

    String name;

    public String getName() {

        return name;

    }

    public void setName(String name) {

        this.name = name;

    }

    public String getId() {

        return id;

    }

    public void setId(String id) {

        this.id = id;

    }

    public String getCount() {

        return count;

    }

    public void setCount(String count) {

        this.count = count;

    }

    String id;

    String count;

}