package net.ziech.testing;

import java.io.IOException;
import java.io.Reader;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.nio.FloatBuffer;
import java.nio.IntBuffer;
import java.util.Map;
import java.util.Random;

import org.apache.lucene.analysis.Analyzer;
import org.apache.lucene.analysis.Tokenizer;
import org.apache.lucene.analysis.icu.segmentation.ICUTokenizer;
import org.apache.lucene.document.Document;
import org.apache.lucene.document.DocumentStoredFieldVisitor;
import org.apache.lucene.document.FloatDocValuesField;
import org.apache.lucene.document.NumericDocValuesField;
import org.apache.lucene.index.AtomicReaderContext;
import org.apache.lucene.index.DirectoryReader;
import org.apache.lucene.index.FieldInfo;
import org.apache.lucene.index.IndexWriter;
import org.apache.lucene.index.IndexWriterConfig;
import org.apache.lucene.index.NumericDocValues;
import org.apache.lucene.queries.function.FunctionQuery;
import org.apache.lucene.queries.function.FunctionValues;
import org.apache.lucene.queries.function.ValueSource;
import org.apache.lucene.queries.function.valuesource.FloatFieldSource;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.TopDocs;
import org.apache.lucene.store.RAMDirectory;
import org.apache.lucene.util.Version;

public class LuceneFloatSourceTest {
	
	private static class LuceneNumericFieldBasedGeometricDistance extends ValueSource {
		private final boolean withoutSquareRoot;
		private final String xField;
		private final String yField;

		public LuceneNumericFieldBasedGeometricDistance(final String xField, final String yField, boolean withoutSquareRoot) {
			this.xField = xField;
			this.yField = yField;
			this.withoutSquareRoot = withoutSquareRoot;
		}

		@Override
		public FunctionValues getValues(Map context,
				AtomicReaderContext readerContext) throws IOException {
			final NumericDocValues xValues = readerContext.reader().getNumericDocValues(this.xField);
			final NumericDocValues yValues = readerContext.reader().getNumericDocValues(this.yField);
			if (withoutSquareRoot) {
				return new FunctionValues() {
					
					@Override
					public float floatVal(int doc) {
						final long y = yValues.get(doc);
						final long x = xValues.get(doc);
						return (float) (y*y + x*x);
					}
	
					@Override
					public String toString(int doc) {
						return "geometric length function value source";
					}
				};
			}
			return new FunctionValues() {
				
				@Override
				public float floatVal(int doc) {
					final long y = yValues.get(doc);
					final long x = xValues.get(doc);
					return (float) Math.sqrt(y*y + x*x);
				}

				@Override
				public String toString(int doc) {
					return "geometric length function value source";
				}
			};
		}

		@Override
		public boolean equals(Object o) {
			return false;
		}

		@Override
		public int hashCode() {
			return 0;
		}

		@Override
		public String description() {
			return "";
		}
	}
	
	private static class LuceneValueSourceBasedGeometricDistance extends ValueSource {
		private FloatFieldSource xSource;
		private FloatFieldSource ySource;
		private final boolean withoutSquareRoot;

		public LuceneValueSourceBasedGeometricDistance(final FloatFieldSource xSource, final FloatFieldSource ySource, boolean withoutSquareRoot) {
			this.xSource = xSource;
			this.ySource = ySource;
			this.withoutSquareRoot = withoutSquareRoot;
		}

		@Override
		public FunctionValues getValues(Map context,
				AtomicReaderContext readerContext) throws IOException {
			final FunctionValues yValues = ySource.getValues(context, readerContext);
			final FunctionValues xValues = xSource.getValues(context, readerContext);
			if (withoutSquareRoot) {
				return new FunctionValues() {
					
					@Override
					public float floatVal(int doc) {
						final float y = yValues.floatVal(doc);
						final float x = xValues.floatVal(doc);
						return (float) (y*y + x*x);
					}
	
					@Override
					public String toString(int doc) {
						return "geometric length function value source";
					}
				};
			}
			return new FunctionValues() {
				
				@Override
				public float floatVal(int doc) {
					final float y = yValues.floatVal(doc);
					final float x = xValues.floatVal(doc);
					return (float) Math.sqrt(y*y + x*x);
				}

				@Override
				public String toString(int doc) {
					return "geometric length function value source";
				}
			};
		}

		@Override
		public boolean equals(Object o) {
			return false;
		}

		@Override
		public int hashCode() {
			return 0;
		}

		@Override
		public String description() {
			return "";
		}
	}
	
	private static class DirectIntBufferBasedGeometricDistance extends ValueSource {
		
		private final IntBuffer xbuf;
		private final IntBuffer ybuf;

		public DirectIntBufferBasedGeometricDistance(final IntBuffer x, final IntBuffer y) {
			this.xbuf = x;
			this.ybuf = y;
		}

		@Override
		public FunctionValues getValues(Map context,
				AtomicReaderContext readerContext) throws IOException {
			return new FunctionValues() {
				
				@Override
				public float floatVal(int doc) {
					final long y = ybuf.get(doc);
					final long x = xbuf.get(doc);
					
					return (float) (y*y + x*x);
				}

				@Override
				public String toString(int doc) {
					return "geometric length function value source";
				}
			};
		}

		@Override
		public boolean equals(Object o) {
			return false;
		}

		@Override
		public int hashCode() {
			return 0;
		}

		@Override
		public String description() {
			return "";
		}
	}
	/**
	 * In theory one would need one direct buffer pair per segment
	 * @author Christian_2
	 *
	 */
	private static class DirectBufferBasedGeometricDistance extends ValueSource {
		
		private final FloatBuffer xbuf;
		private final FloatBuffer ybuf;
		private final boolean withoutSquareRoot;

		public DirectBufferBasedGeometricDistance(final FloatBuffer x, final FloatBuffer y, boolean withoutSquareRoot) {
			this.xbuf = x;
			this.ybuf = y;
			this.withoutSquareRoot = withoutSquareRoot;
		}

		@Override
		public FunctionValues getValues(Map context,
				AtomicReaderContext readerContext) throws IOException {
			if (withoutSquareRoot) {
				return new FunctionValues() {
					
					@Override
					public float floatVal(int doc) {
						final float y = ybuf.get(doc);
						final float x = xbuf.get(doc);
						
						return (float) (y*y + x*x);
					}

					@Override
					public String toString(int doc) {
						return "geometric length function value source";
					}
				};
			}
			return new FunctionValues() {
				
				@Override
				public float floatVal(int doc) {
					final float y = ybuf.get(doc);
					final float x = xbuf.get(doc);
					
					return (float) Math.sqrt(y*y + x*x);
				}

				@Override
				public String toString(int doc) {
					return "geometric length function value source";
				}
			};
		}

		@Override
		public boolean equals(Object o) {
			return false;
		}

		@Override
		public int hashCode() {
			return 0;
		}

		@Override
		public String description() {
			return "";
		}
	}
	
	public static class Counter extends DocumentStoredFieldVisitor {
		
		public float res = 0.0f;

		@Override
		public void floatField(FieldInfo fieldInfo, float value) {
			res += value;
		}
	};

	public static void main(String[] args) throws IOException {
		ByteBuffer dbb = ByteBuffer.allocateDirect(100000000);
		ByteBuffer dbb2 = ByteBuffer.allocateDirect(1000000000);
		ByteBuffer dbbInt = ByteBuffer.allocateDirect(100000000);
		ByteBuffer dbbInt2 = ByteBuffer.allocateDirect(1000000000);
		dbb.order(ByteOrder.nativeOrder());
		FloatBuffer dfb = dbb.asFloatBuffer();
		FloatBuffer dfb2 = dbb2.asFloatBuffer();
		IntBuffer dib = dbbInt.asIntBuffer();
		IntBuffer dib2 = dbbInt2.asIntBuffer();
		Random rnd = new Random();
		final int size = dfb.limit();
		for (int i=0; i<size; i++) {
			float f = rnd.nextFloat();
			dfb.put(f);
			dib.put((int) (f*100000000.0f));
			f = rnd.nextFloat();
			dfb2.put(f);
			dib2.put((int) (f*100000000.0f));
		}
		
		IndexWriterConfig conf = new IndexWriterConfig(Version.LUCENE_44, new Analyzer(){

			@Override
			protected TokenStreamComponents createComponents(String fieldName,
					Reader reader) {
				Tokenizer tokenizer = new ICUTokenizer(reader);
				return new TokenStreamComponents(tokenizer);
			}
			
		});
		RAMDirectory dir = new RAMDirectory();
		/*File f = new File("temp");
		if (f.exists()) {
			for (File child : f.listFiles()) {
				child.delete();
			}
		} else {
			f.mkdir();
		}
		FSDirectory dir = FSDirectory.open(f);*/
		IndexWriter iw = new IndexWriter(dir, conf);
		for (int i=0; i<size; i++) {
			Document doc = new Document();
			doc.add(new FloatDocValuesField("x", dfb.get(i)));
			doc.add(new FloatDocValuesField("y", dfb2.get(i)));
			doc.add(new NumericDocValuesField("xInt", dib.get(i)));
			doc.add(new NumericDocValuesField("yInt", dib2.get(i)));
			iw.addDocument(doc);
			if (i % 1000000 == 0) {
				System.out.println("Indexed "+i);
			}
		}
		iw.commit();
		iw.forceMerge(1);
		iw.close();
		
		System.out.println("now reading numbers");
		DirectoryReader dr = DirectoryReader.open(dir);
		
		final FloatFieldSource ySource = new FloatFieldSource("y");
		final FloatFieldSource xSource = new FloatFieldSource("x");
		IndexSearcher searcher = new IndexSearcher(dr);

		for (int i=0; i<3; i++) {
			long start = System.currentTimeMillis();
			TopDocs topDocs = searcher.search(new FunctionQuery(new DirectBufferBasedGeometricDistance(dfb, dfb2, true)), 1);
			long end = System.currentTimeMillis();
			System.out.println("Scoring "+topDocs.totalHits+" documents with direct float buffers (without square root) took "+(end-start));
		}
		
		for (int i=0; i<3; i++) {
			long start = System.currentTimeMillis();
			TopDocs topDocs = searcher.search(new FunctionQuery(new DirectBufferBasedGeometricDistance(dfb, dfb2, false)), 1);
			long end = System.currentTimeMillis();
			System.out.println("Scoring "+topDocs.totalHits+" documents with direct float buffers (and a square root) took "+(end-start));
		}
		
		for (int i=0; i<3; i++) {
			long start = System.currentTimeMillis();
			TopDocs topDocs = searcher.search(new FunctionQuery(new LuceneValueSourceBasedGeometricDistance(xSource, ySource, true)), 1);
			long end = System.currentTimeMillis();
			System.out.println("Scoring "+topDocs.totalHits+" documents with a lucene float value source (without square root) took "+(end-start));
		}
		
		for (int i=0; i<3; i++) {
			long start = System.currentTimeMillis();
			TopDocs topDocs = searcher.search(new FunctionQuery(new LuceneValueSourceBasedGeometricDistance(xSource, ySource, false)), 1);
			long end = System.currentTimeMillis();
			System.out.println("Scoring "+topDocs.totalHits+" documents with a lucene float value source (and a square root) took "+(end-start));
		}
		
		for (int i=0; i<3; i++) {
			long start = System.currentTimeMillis();
			TopDocs topDocs = searcher.search(new FunctionQuery(new DirectIntBufferBasedGeometricDistance(dib, dib2)), 1);
			long end = System.currentTimeMillis();
			System.out.println("Scoring "+topDocs.totalHits+" documents with direct int buffers (without square root) took "+(end-start));
		}
		
		for (int i=0; i<3; i++) {
			long start = System.currentTimeMillis();
			TopDocs topDocs = searcher.search(new FunctionQuery(new LuceneNumericFieldBasedGeometricDistance("xInt", "yInt", true)), 1);
			long end = System.currentTimeMillis();
			System.out.println("Scoring "+topDocs.totalHits+" documents with a lucene numeric values (without square root) source took "+(end-start));
		}
		
		for (int i=0; i<3; i++) {
			long start = System.currentTimeMillis();
			TopDocs topDocs = searcher.search(new FunctionQuery(new LuceneNumericFieldBasedGeometricDistance("xInt", "yInt", false)), 1);
			long end = System.currentTimeMillis();
			System.out.println("Scoring "+topDocs.totalHits+" documents with a lucene numeric values (and a square root) source took "+(end-start));
		}
	}
	
}