jpayne@68: package sketch;
jpayne@68: 
jpayne@68: import java.util.ArrayList;
jpayne@68: import java.util.concurrent.atomic.AtomicInteger;
jpayne@68: import java.util.concurrent.atomic.AtomicLong;
jpayne@68: 
jpayne@68: import kmer.AbstractKmerTable;
jpayne@68: import kmer.HashBuffer;
jpayne@68: import kmer.KmerTableSet;
jpayne@68: import shared.Shared;
jpayne@68: import shared.Timer;
jpayne@68: import shared.Tools;
jpayne@68: import structures.AbstractBitSet;
jpayne@68: import structures.IntHashMap;
jpayne@68: import structures.IntHashSetList;
jpayne@68: import structures.IntList;
jpayne@68: 
jpayne@68: public class SketchIndex extends SketchObject {
jpayne@68: 	
jpayne@68: 	public SketchIndex(ArrayList<Sketch> refs){
jpayne@68: 		refSketches=refs;
jpayne@68: 		tables=new KmerTableSet(new String[] {"ways="+WAYS, "tabletype="+AbstractKmerTable.ARRAYHF, "prealloc="+(prealloc>0 ? ""+prealloc : "f")}, 
jpayne@68: 				20+(defaultParams.trackCounts() ? 4 : 0)+18);//An extra 18 bytes per kmer because a lot of them occur multiple times
jpayne@68: 		tables.allocateTables();
jpayne@68: 		tableArray=tables.tables();
jpayne@68: 	}
jpayne@68: 	
jpayne@68: 	public void load(){
jpayne@68: 		spawnIndexThreads();
jpayne@68: 		if(useWhitelist){
jpayne@68: 			assert(!Whitelist.exists());
jpayne@68: 			Whitelist.initialize(tableArray);
jpayne@68: 		}
jpayne@68: 	}
jpayne@68: 	
jpayne@68: 	/** Spawn index threads */
jpayne@68: 	private void spawnIndexThreads(){
jpayne@68: 		
jpayne@68: 		//Do anything necessary prior to processing
jpayne@68: 		
jpayne@68: 		//Determine how many threads may be used
jpayne@68: 		final int threads=Shared.threads();
jpayne@68: 		ArrayList<IndexThread> alht=new ArrayList<IndexThread>(threads);
jpayne@68: 		AtomicInteger ai=new AtomicInteger(0);
jpayne@68: 		AtomicLong totalKeys=new AtomicLong(0);
jpayne@68: 		AtomicLong uniqueKeys=new AtomicLong(0);
jpayne@68: 		for(int i=0; i<threads; i++){
jpayne@68: 			alht.add(new IndexThread(ai, totalKeys, uniqueKeys));
jpayne@68: 		}
jpayne@68: 		
jpayne@68: 		//Start the threads
jpayne@68: 		for(IndexThread pt : alht){
jpayne@68: 			pt.start();
jpayne@68: 		}
jpayne@68: 		
jpayne@68: 		//Wait for completion of all threads
jpayne@68: 		boolean success=true;
jpayne@68: 		long codesProcessed=0;
jpayne@68: 		for(IndexThread pt : alht){
jpayne@68: 			
jpayne@68: 			//Wait until this thread has terminated
jpayne@68: 			while(pt.getState()!=Thread.State.TERMINATED){
jpayne@68: 				try {
jpayne@68: 					//Attempt a join operation
jpayne@68: 					pt.join();
jpayne@68: 					synchronized(pt){
jpayne@68: 						codesProcessed+=pt.codesProcessedT;
jpayne@68: 					}
jpayne@68: 				} catch (InterruptedException e) {
jpayne@68: 					//Potentially handle this, if it is expected to occur
jpayne@68: 					e.printStackTrace();
jpayne@68: 				}
jpayne@68: 			}
jpayne@68: 			success&=pt.success;
jpayne@68: 		}
jpayne@68: 		
jpayne@68: 		//Track whether any threads failed
jpayne@68: 		if(!success){errorState=true;}
jpayne@68: 		
jpayne@68: 		System.err.println("Indexed "+uniqueKeys+" unique and "+totalKeys+" total hashcodes."); //For some reason codesProcessed is nondeterministic.
jpayne@68: 		
jpayne@68: 		//Do anything necessary after processing
jpayne@68: //		System.gc();
jpayne@68: 	}
jpayne@68: 	
jpayne@68: 	/*--------------------------------------------------------------*/
jpayne@68: 	
jpayne@68: 	public SketchResults getSketches(Sketch a, DisplayParams params){
jpayne@68: 		if(useIntMap){
jpayne@68: 			return getSketchesMap(a, params);
jpayne@68: 		}else{
jpayne@68: 			return getSketchesList(a, params);
jpayne@68: 		}
jpayne@68: 	}
jpayne@68: 	
jpayne@68: 	/** Return true if added. */
jpayne@68: 	private boolean addToTaxSet(int sketchID, IntHashSetList taxSet, int taxLevelExtended){
jpayne@68: 		Sketch sk=refSketches.get(sketchID);
jpayne@68: 		int taxID=sk.taxID;
jpayne@68: 		if(taxID<0 || taxID>=minFakeID){return false;}
jpayne@68: 		taxID=taxtree.getIdAtLevelExtended(taxID, taxLevelExtended);
jpayne@68: 		return taxSet.add(taxID);
jpayne@68: 	}
jpayne@68: 	
jpayne@68: 	public SketchResults getSketchesList(final Sketch a, DisplayParams params){
jpayne@68: 		final int minHits=params.minHits, contamLevel=params.contamLevel();
jpayne@68: 		final boolean countContamHits=params.needContamCounts();//, metaFilter=params.hasMetaFilters(), taxFilter=params.hasTaxFilters();
jpayne@68: 		final Timer t=(printTime ? new Timer() : null);
jpayne@68: 		
jpayne@68: 		final int[] singleton=new int[1];
jpayne@68: 		final IntList idList=new IntList(Tools.min(targetSketchSize, indexLimit, 1000));
jpayne@68: 		AbstractBitSet abs=a.indexBitSet();
jpayne@68: 		assert((abs==null)!=countContamHits);
jpayne@68: 		
jpayne@68: 		final IntHashSetList taxSet;
jpayne@68: 		final int[][] taxHits;
jpayne@68: 		if(contamLevel>=0){
jpayne@68: 			taxSet=new IntHashSetList(31);
jpayne@68: 			taxHits=new int[a.length()][];
jpayne@68: 			assert(taxtree!=null) : "A TaxTree is required for this operation.";
jpayne@68: 		}else{
jpayne@68: 			taxSet=null;
jpayne@68: 			taxHits=null;
jpayne@68: 		}
jpayne@68: 		
jpayne@68: 		for(int i=0; i<a.keys.length; i++){
jpayne@68: 			long key=a.keys[i];
jpayne@68: 			AbstractKmerTable set=tableArray[(int)(key%WAYS)];
jpayne@68: //			System.err.println(set.getValue(key));
jpayne@68: 			final int[] ids=set.getValues(key, singleton);
jpayne@68: //			System.err.println(Arrays.toString(ids));
jpayne@68: 			if(ids!=null && ids[0]>=0){
jpayne@68: 				int incr=0;
jpayne@68: 				for(int id : ids){
jpayne@68: 					if(id>=0){
jpayne@68: 						final int trueID=id-1;//Minimum id is 1, indicating sketch 0.
jpayne@68: 						idList.add(trueID);//Minimum id is 1, indicating sketch 0.
jpayne@68: 						incr++;
jpayne@68: 						if(taxSet!=null){addToTaxSet(trueID, taxSet, contamLevel);}
jpayne@68: 					}
jpayne@68: 				}
jpayne@68: 				if(countContamHits && incr>0){abs.increment(i, incr);}
jpayne@68: 				if(taxSet!=null && taxSet.size()>0){
jpayne@68: 					taxHits[i]=taxSet.toArray();
jpayne@68: 					taxSet.clear();
jpayne@68: 				}
jpayne@68: 			}
jpayne@68: 			
jpayne@68: 		}
jpayne@68: 		
jpayne@68: //		assert(abs!=null);
jpayne@68: //		assert(false) : abs.cardinality();
jpayne@68: 		
jpayne@68: 		if(printTime){
jpayne@68: 			t.stop("\nTime for searching index: \t");
jpayne@68: 			t.start();
jpayne@68: 		}
jpayne@68: 		
jpayne@68: //		System.err.println("idList.size:"+idList.size);
jpayne@68: 		if(idList.size<minHits){return new SketchResults(a);}//null breaks some things
jpayne@68: 		idList.sort();
jpayne@68: 		
jpayne@68: 		if(printTime){
jpayne@68: 			t.stop("Time for sorting "+idList.size()+" hits:\t");
jpayne@68: 			t.start();
jpayne@68: 		}
jpayne@68: 		
jpayne@68: 		ArrayList<Sketch> list=new ArrayList<Sketch>(Tools.min(8, idList.size));
jpayne@68: 
jpayne@68: 		int last=-1;
jpayne@68: 		int hits=0;
jpayne@68: 		for(int i=0; i<idList.size; i++){
jpayne@68: 			int id=idList.get(i);
jpayne@68: 			if(id==last){
jpayne@68: //				System.err.println("A: "+last+", "+id+", "+count+", "+minHits);
jpayne@68: 				hits++;
jpayne@68: 			}else{
jpayne@68: //				System.err.println("B: "+last+", "+id+", "+count+", "+minHits);
jpayne@68: 				if(last>-1 && (hits>=minHits)){
jpayne@68: 					final Sketch ref=refSketches.get(last);
jpayne@68: 					list.add(ref);
jpayne@68: 				}
jpayne@68: 				last=id;
jpayne@68: 				hits=0;
jpayne@68: 			}
jpayne@68: 		}
jpayne@68: 		if(last>-1 && (hits>=minHits)){
jpayne@68: 			final Sketch ref=refSketches.get(last);
jpayne@68: //			if((!metaFilter || ref.passesMeta(params)) && (!taxFilter || params.passesFilter(ref))){list.add(ref);}
jpayne@68: 			list.add(ref);
jpayne@68: 		}
jpayne@68: 		if(printTime){
jpayne@68: 			t.stop("Time for fetching sketches: \t");
jpayne@68: 		}
jpayne@68: 		return list.isEmpty() ? new SketchResults(a) : new SketchResults(a, list, taxHits);
jpayne@68: 	}
jpayne@68: 	
jpayne@68: //	static ThreadLocal<IntHashMap> intMapHolder=new ThreadLocal<IntHashMap>();
jpayne@68: 	
jpayne@68: 	public final int[] getSketchIdsMap(long key, int[] singleton){
jpayne@68: 		AbstractKmerTable set=tableArray[(int)(key%WAYS)];
jpayne@68: 		final int[] ids=set.getValues(key, singleton);
jpayne@68: 		return ids;
jpayne@68: 	}
jpayne@68: 	
jpayne@68: 	public SketchResults getSketchesMap(final Sketch a, DisplayParams params){
jpayne@68: 		final int minHits=params.minHits, contamLevel=params.contamLevel();
jpayne@68: 		final boolean countContamHits=params.needContamCounts();//, metaFilter=params.hasMetaFilters(), taxFilter=params.hasTaxFilters();
jpayne@68: 		final Timer t=(printTime ? new Timer() : null);
jpayne@68: 		final int[] singleton=new int[1];
jpayne@68: 		
jpayne@68: 		final IntHashMap idMap=new IntHashMap(Tools.min(targetSketchSize, indexLimit, intMapSize), 0.7f);
jpayne@68: 		
jpayne@68: 		AbstractBitSet abs=a.indexBitSet();
jpayne@68: 		assert((abs==null)!=countContamHits);
jpayne@68: 		
jpayne@68: 		final IntHashSetList taxSet;
jpayne@68: 		final int[][] taxHits;
jpayne@68: 		if(contamLevel>=0){
jpayne@68: 			taxSet=new IntHashSetList(31);
jpayne@68: 			taxHits=new int[a.length()][];
jpayne@68: 			assert(taxtree!=null) : "A TaxTree is required for this operation.";
jpayne@68: 		}else{
jpayne@68: 			taxSet=null;
jpayne@68: 			taxHits=null;
jpayne@68: 		}
jpayne@68: //		assert(false) : (taxHits==null)+", "+contamLevel;
jpayne@68: 		
jpayne@68: 		if(printTime){
jpayne@68: 			t.stop("\nTime for allocation:      \t");
jpayne@68: 			t.start();
jpayne@68: 		}
jpayne@68: 		
jpayne@68: 		int[] refHitCounts;
jpayne@68: 		if(params.printRefHits){
jpayne@68: 			refHitCounts=new int[a.keys.length];
jpayne@68: 			a.setRefHitCounts(refHitCounts);
jpayne@68: 		}else{refHitCounts=null;}
jpayne@68: 		
jpayne@68: 		for(int i=0; i<a.keys.length; i++){
jpayne@68: 			long key=a.keys[i];
jpayne@68: 			final int[] ids=getSketchIdsMap(key, singleton);
jpayne@68: 			
jpayne@68: 			if(ids!=null && ids[0]>=0){
jpayne@68: 				int incr=0;
jpayne@68: 				for(int id : ids){
jpayne@68: 					if(id>=0){
jpayne@68: 						final int trueID=id-1;//Minimum id is 1, indicating sketch 0.
jpayne@68: 						idMap.increment(trueID);
jpayne@68: 						if(!allToAll || compareSelf){incr++;}
jpayne@68: 						if(taxSet!=null){addToTaxSet(trueID, taxSet, contamLevel);}
jpayne@68: 						if(refHitCounts!=null && trueID!=a.sketchID){refHitCounts[i]++;}
jpayne@68: 					}
jpayne@68: 				}
jpayne@68: 				if(countContamHits && incr>0){abs.increment(i, incr);}
jpayne@68: 				if(taxSet!=null && taxSet.size()>0){
jpayne@68: 					taxHits[i]=taxSet.toArray();
jpayne@68: 					taxSet.clear();
jpayne@68: 				}
jpayne@68: 			}
jpayne@68: 		}
jpayne@68: 		
jpayne@68: 		if(printTime){
jpayne@68: 			t.stop("Time for searching index: \t");
jpayne@68: 			System.err.println("Size:                   \t"+idMap.size());
jpayne@68: 			t.start();
jpayne@68: 		}
jpayne@68: 		
jpayne@68: //		System.err.println("idList.size:"+idList.size);
jpayne@68: 		final int size=idMap.size();
jpayne@68: 		if(size==0){return new SketchResults(a);}//null breaks some things
jpayne@68: 		
jpayne@68: 		ArrayList<Sketch> list=new ArrayList<Sketch>(Tools.min(8, size));
jpayne@68: 
jpayne@68: 		final int[] keys=idMap.keys();
jpayne@68: 		final int[] values=idMap.values();
jpayne@68: 		
jpayne@68: 		for(int i=0; i<keys.length; i++){
jpayne@68: 			int value=values[i];
jpayne@68: 			if(value>=minHits){
jpayne@68: 				int id=keys[i];
jpayne@68: 				final Sketch ref=refSketches.get(id);
jpayne@68: //				if((!metaFilter || ref.passesMeta(params)) && (!taxFilter || params.passesFilter(ref))){list.add(ref);}
jpayne@68: 				list.add(ref);
jpayne@68: 			}
jpayne@68: 		}
jpayne@68: 		if(printTime){
jpayne@68: 			t.stop("Time for fetching sketches: \t");
jpayne@68: 		}
jpayne@68: 		return list.isEmpty() ? new SketchResults(a) : new SketchResults(a, list, taxHits);
jpayne@68: 	}
jpayne@68: 	
jpayne@68: 	/*--------------------------------------------------------------*/
jpayne@68: 	
jpayne@68: 	public class IndexThread extends Thread {
jpayne@68: 		
jpayne@68: 		public IndexThread(AtomicInteger nextIndex_, AtomicLong keyCount_, AtomicLong uniqueKeyCount_){
jpayne@68: 			buffer=new HashBuffer(tableArray, 1000, 31, true, false);
jpayne@68: 			nextIndex=nextIndex_;
jpayne@68: 			keyCount=keyCount_;
jpayne@68: 			uniqueKeyCount=uniqueKeyCount_;
jpayne@68: 		}
jpayne@68: 		
jpayne@68: 		@Override
jpayne@68: 		public void run(){
jpayne@68: //			System.err.println("Thread running.");
jpayne@68: 			int id=nextIndex.getAndIncrement();
jpayne@68: 			final int numSketches=refSketches.size();
jpayne@68: 			final int limit0=Tools.min((AUTOSIZE || AUTOSIZE_LINEAR ? Integer.MAX_VALUE : targetSketchSize), indexLimit);
jpayne@68: //			System.err.println("numSketches="+numSketches);
jpayne@68: 			while(id<numSketches){
jpayne@68: 				final Sketch sk=refSketches.get(id);
jpayne@68: 				final long[] array=sk.keys;
jpayne@68: 				final int limit=Tools.min(array.length, limit0);
jpayne@68: //				System.err.println("limit="+limit);
jpayne@68: 				for(int i=0; i<limit; i++){
jpayne@68: 					long key=array[i];
jpayne@68: 					buffer.set(key, id+1);//Must be greater than zero
jpayne@68: 					codesProcessedT++;
jpayne@68: 				}
jpayne@68: 				id=nextIndex.getAndIncrement();
jpayne@68: 			}
jpayne@68: 			long temp=buffer.flush();
jpayne@68: 			
jpayne@68: 			synchronized(this){
jpayne@68: 				codesProcessedT+=0;
jpayne@68: 				success=true;
jpayne@68: 				keyCount.getAndAdd(codesProcessedT);
jpayne@68: 				uniqueKeyCount.getAndAdd(buffer.uniqueAdded);
jpayne@68: //				if(codesProcessedT>0){System.err.println(codesProcessedT);}
jpayne@68: 			}
jpayne@68: 		}
jpayne@68: 		
jpayne@68: 		AtomicInteger nextIndex;
jpayne@68: 		AtomicLong keyCount;
jpayne@68: 		AtomicLong uniqueKeyCount;
jpayne@68: 		long codesProcessedT=0;
jpayne@68: 		HashBuffer buffer;
jpayne@68: 		boolean success=false;
jpayne@68: 		
jpayne@68: 	}
jpayne@68: 	
jpayne@68: 	/*--------------------------------------------------------------*/
jpayne@68: 	
jpayne@68: 	public final KmerTableSet tables;
jpayne@68: 	public final AbstractKmerTable[] tableArray;
jpayne@68: 	public final ArrayList<Sketch> refSketches;
jpayne@68: 	
jpayne@68: 	public boolean errorState=false;
jpayne@68: 
jpayne@68: 	private static final boolean printTime=false;
jpayne@68: 	public static boolean useIntMap=true;
jpayne@68: //	public static boolean useIntMapBinary=false;
jpayne@68: 	public static int intMapSize=1000;
jpayne@68: 	public static int indexLimit=Integer.MAX_VALUE;
jpayne@68: 	public static final int WAYS=31;
jpayne@68: 	
jpayne@68: }