<commit_before>package owltools.solrj;

import java.io.IOException;
import java.net.MalformedURLException;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;

import org.apache.log4j.Logger;
import org.apache.solr.client.solrj.SolrServerException;
import org.apache.solr.common.SolrInputDocument;
import org.semanticweb.owlapi.model.OWLClass;
import org.semanticweb.owlapi.model.OWLObject;
import org.semanticweb.owlapi.model.OWLObjectProperty;
import org.semanticweb.owlapi.model.OWLPropertyExpression;

import owltools.gaf.Bioentity;
import owltools.gaf.ExtensionExpression;
import owltools.gaf.GafDocument;
import owltools.gaf.GeneAnnotation;
import owltools.gaf.WithInfo;
import owltools.graph.OWLGraphEdge;
import owltools.graph.OWLQuantifiedProperty;

import com.google.gson.*;

public class GafSolrDocumentLoader extends AbstractSolrLoader {

	private static Logger LOG = Logger.getLogger(GafSolrDocumentLoader.class);

	GafDocument gafDocument;
	int doc_limit_trigger = 1000; // the number of documents to add before pushing out to solr
	int current_doc_number;
	
	public GafSolrDocumentLoader(String url) throws MalformedURLException {
		super(url);
		current_doc_number = 0;
	}

	public GafDocument getGafDocument() {
		return gafDocument;
	}

	public void setGafDocument(GafDocument gafDocument) {
		this.gafDocument = gafDocument;
	}

	@Override
	public void load() throws SolrServerException, IOException {
		gafDocument.index();
		for (Bioentity e : gafDocument.getBioentities()) {
			add(e);
			current_doc_number++;
			if( current_doc_number % doc_limit_trigger == 0 ){
				LOG.info("Processed " + doc_limit_trigger + " bioentities at " + current_doc_number + " and committing...");
				incrementalAddAndCommit();
			}
		}
		LOG.info("Doing cleanup commit.");
		incrementalAddAndCommit(); // pick up anything that we didn't catch
		//LOG.info("Optimizing.");
		//server.optimize();
		LOG.info("Done.");
	}

//	private OWLObjectProperty getPartOfProperty() {
//		OWLObjectProperty p = graph.getOWLObjectPropertyByIdentifier("BFO:0000050");
//		return p;
//	}

	// Main wrapping for adding non-ontology documents to GOlr.
	// Also see OntologySolrLoader.
	private void add(Bioentity e) {

		String eid = e.getId();
		String esym = e.getSymbol();
		String edb = e.getDb();
		//LOG.info("Adding: " + eid + " " + esym);
		
		// We'll need this for serializing later.
		Gson gson = new Gson();
		
		SolrInputDocument bioentity_doc = new SolrInputDocument();
		
		// Bioentity document base.
		bioentity_doc.addField("document_category", "bioentity");
		bioentity_doc.addField("id", eid);
		bioentity_doc.addField("label", esym);
		bioentity_doc.addField("db", edb);
		bioentity_doc.addField("type", e.getTypeCls());
		String taxId = e.getNcbiTaxonId();
		bioentity_doc.addField("taxon", taxId);
		addLabelField(bioentity_doc, "taxon_label", taxId);

		Map<String,SolrInputDocument> evAggDocMap = new HashMap<String,SolrInputDocument>();
		
		// Annotation doc
		for (GeneAnnotation a : gafDocument.getGeneAnnotations(e.getId())) {
			SolrInputDocument annotation_doc = new SolrInputDocument();

			String clsId = a.getCls();
			String refId = a.getReferenceId();

			// Annotation document base.
			annotation_doc.addField("document_category", "annotation");
			annotation_doc.addField("id", eid + "_:_" + clsId); // TODO - make unique
			annotation_doc.addField("bioentity", eid);
			annotation_doc.addField("bioentity_label", esym);
			annotation_doc.addField("source", a.getAssignedBy());
			annotation_doc.addField("date", a.getLastUpdateDate());
			annotation_doc.addField("taxon", taxId);
			addLabelField(annotation_doc, "taxon_label", taxId);

			annotation_doc.addField("reference", refId);
			// TODO - ev. closure
			annotation_doc.addField("evidence_type", a.getEvidenceCls());
			//annotation_doc.addField("evidence_with", a.getWithExpression());

			// Drag in "with" (col 8).
			for (WithInfo wi : a.getWithInfos()) {
				annotation_doc.addField("evidence_with", wi.getWithXref());
			}

			annotation_doc.addField("annotation_class", clsId);
			addLabelField(annotation_doc, "annotation_class_label", clsId);

			///
			/// isa_partof_closure
			///
			
			OWLObject cls = graph.getOWLObjectByIdentifier(clsId);
			// TODO: This may be a bug workaround, or it may be the way things are.
			// getOWLObjectByIdentifier returns null on alt_ids, so skip them for now.
			if(cls != null ){
				//	System.err.println(clsId);
			
				// Add to annotation and bioentity isa_partof closures.
				List<String> idClosure = graph.getIsaPartofIDClosure(cls);
				List<String> labelClosure = graph.getIsaPartofLabelClosure(cls);
				annotation_doc.addField("isa_partof_closure", idClosure);
				annotation_doc.addField("isa_partof_closure_label", labelClosure);
				for( String tlabel : labelClosure){
					addFieldUnique(bioentity_doc, "isa_partof_closure_label", tlabel);
					addFieldUnique(bioentity_doc, "isa_partof_closure", tlabel);
				}
					
				// Compile closure maps to JSON.
				Map<String, String> isa_partof_map = graph.getIsaPartofClosureMap(cls);
				if( ! isa_partof_map.isEmpty() ){
					String jsonized_isa_partof_map = gson.toJson(isa_partof_map);
					annotation_doc.addField("isa_partof_closure_map", jsonized_isa_partof_map);
				}
	
				// Cycle through and pick up all the associated bits for the terms in the closure.
				for( String tid : idClosure ){
	
					String tlabel = isa_partof_map.get(tid);				
					//OWLObject c = graph.getOWLObjectByIdentifier(tid);
	
					// Annotation evidence aggregate base.
					String evAggId = eid + "_:ev:_" + clsId;
					SolrInputDocument ev_agg_doc;
					if (evAggDocMap.containsKey(evAggId)) {
						ev_agg_doc = evAggDocMap.get(evAggId);	
					} else {
						ev_agg_doc = new SolrInputDocument();
						evAggDocMap.put(evAggId, ev_agg_doc);<commit_msg>now properly outputing ids in closure for gp<commit_after>package owltools.solrj;

import java.io.IOException;
import java.net.MalformedURLException;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;

import org.apache.log4j.Logger;
import org.apache.solr.client.solrj.SolrServerException;
import org.apache.solr.common.SolrInputDocument;
import org.semanticweb.owlapi.model.OWLClass;
import org.semanticweb.owlapi.model.OWLObject;
import org.semanticweb.owlapi.model.OWLObjectProperty;
import org.semanticweb.owlapi.model.OWLPropertyExpression;

import owltools.gaf.Bioentity;
import owltools.gaf.ExtensionExpression;
import owltools.gaf.GafDocument;
import owltools.gaf.GeneAnnotation;
import owltools.gaf.WithInfo;
import owltools.graph.OWLGraphEdge;
import owltools.graph.OWLQuantifiedProperty;

import com.google.gson.*;

public class GafSolrDocumentLoader extends AbstractSolrLoader {

	private static Logger LOG = Logger.getLogger(GafSolrDocumentLoader.class);

	GafDocument gafDocument;
	int doc_limit_trigger = 1000; // the number of documents to add before pushing out to solr
	int current_doc_number;
	
	public GafSolrDocumentLoader(String url) throws MalformedURLException {
		super(url);
		current_doc_number = 0;
	}

	public GafDocument getGafDocument() {
		return gafDocument;
	}

	public void setGafDocument(GafDocument gafDocument) {
		this.gafDocument = gafDocument;
	}

	@Override
	public void load() throws SolrServerException, IOException {
		gafDocument.index();
		for (Bioentity e : gafDocument.getBioentities()) {
			add(e);
			current_doc_number++;
			if( current_doc_number % doc_limit_trigger == 0 ){
				LOG.info("Processed " + doc_limit_trigger + " bioentities at " + current_doc_number + " and committing...");
				incrementalAddAndCommit();
			}
		}
		LOG.info("Doing cleanup commit.");
		incrementalAddAndCommit(); // pick up anything that we didn't catch
		//LOG.info("Optimizing.");
		//server.optimize();
		LOG.info("Done.");
	}

//	private OWLObjectProperty getPartOfProperty() {
//		OWLObjectProperty p = graph.getOWLObjectPropertyByIdentifier("BFO:0000050");
//		return p;
//	}

	// Main wrapping for adding non-ontology documents to GOlr.
	// Also see OntologySolrLoader.
	private void add(Bioentity e) {

		String eid = e.getId();
		String esym = e.getSymbol();
		String edb = e.getDb();
		//LOG.info("Adding: " + eid + " " + esym);
		
		// We'll need this for serializing later.
		Gson gson = new Gson();
		
		SolrInputDocument bioentity_doc = new SolrInputDocument();
		
		// Bioentity document base.
		bioentity_doc.addField("document_category", "bioentity");
		bioentity_doc.addField("id", eid);
		bioentity_doc.addField("label", esym);
		bioentity_doc.addField("db", edb);
		bioentity_doc.addField("type", e.getTypeCls());
		String taxId = e.getNcbiTaxonId();
		bioentity_doc.addField("taxon", taxId);
		addLabelField(bioentity_doc, "taxon_label", taxId);

		Map<String,SolrInputDocument> evAggDocMap = new HashMap<String,SolrInputDocument>();
		
		// Annotation doc
		for (GeneAnnotation a : gafDocument.getGeneAnnotations(e.getId())) {
			SolrInputDocument annotation_doc = new SolrInputDocument();

			String clsId = a.getCls();
			String refId = a.getReferenceId();

			// Annotation document base.
			annotation_doc.addField("document_category", "annotation");
			annotation_doc.addField("id", eid + "_:_" + clsId); // TODO - make unique
			annotation_doc.addField("bioentity", eid);
			annotation_doc.addField("bioentity_label", esym);
			annotation_doc.addField("source", a.getAssignedBy());
			annotation_doc.addField("date", a.getLastUpdateDate());
			annotation_doc.addField("taxon", taxId);
			addLabelField(annotation_doc, "taxon_label", taxId);

			annotation_doc.addField("reference", refId);
			// TODO - ev. closure
			annotation_doc.addField("evidence_type", a.getEvidenceCls());
			//annotation_doc.addField("evidence_with", a.getWithExpression());

			// Drag in "with" (col 8).
			for (WithInfo wi : a.getWithInfos()) {
				annotation_doc.addField("evidence_with", wi.getWithXref());
			}

			annotation_doc.addField("annotation_class", clsId);
			addLabelField(annotation_doc, "annotation_class_label", clsId);

			///
			/// isa_partof_closure
			///
			
			OWLObject cls = graph.getOWLObjectByIdentifier(clsId);
			// TODO: This may be a bug workaround, or it may be the way things are.
			// getOWLObjectByIdentifier returns null on alt_ids, so skip them for now.
			if(cls != null ){
				//	System.err.println(clsId);
			
				// Add to annotation and bioentity isa_partof closures; label and id.
				List<String> idClosure = graph.getIsaPartofIDClosure(cls);
				List<String> labelClosure = graph.getIsaPartofLabelClosure(cls);
				annotation_doc.addField("isa_partof_closure", idClosure);
				annotation_doc.addField("isa_partof_closure_label", labelClosure);
				for( String tlabel : labelClosure){
					addFieldUnique(bioentity_doc, "isa_partof_closure_label", tlabel);
				}
				for( String tid : idClosure){
					addFieldUnique(bioentity_doc, "isa_partof_closure", tid);
				}
	
				// Compile closure maps to JSON.
				Map<String, String> isa_partof_map = graph.getIsaPartofClosureMap(cls);
				if( ! isa_partof_map.isEmpty() ){
					String jsonized_isa_partof_map = gson.toJson(isa_partof_map);
					annotation_doc.addField("isa_partof_closure_map", jsonized_isa_partof_map);
				}
	
				// Cycle through and pick up all the associated bits for the terms in the closure.
				for( String tid : idClosure ){
	
					String tlabel = isa_partof_map.get(tid);				
					//OWLObject c = graph.getOWLObjectByIdentifier(tid);
	
					// Annotation evidence aggregate base.
					String evAggId = eid + "_:ev:_" + clsId;
					SolrInputDocument ev_agg_doc;
					if (evAggDocMap.containsKey(evAggId)) {
						ev_agg_doc = evAggDocMap.get(evAggId);	
					} else {
						ev_agg_doc = new SolrInputDocument();
						evAggDocMap.put(evAggId, ev_agg_doc);
{'fraction_non_alphanumeric': 0.11698744769874476, 'fraction_numerical': 0.002510460251046025, 'mean_word_length': 4.4545869465997265, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2405338', 'n_tokens_mistral': 4552, 'n_tokens_neox': 3921, 'n_words': 803}