package com.uas.search.sort;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

import org.apache.lucene.index.FieldInvertState;
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.index.NumericDocValues;
import org.apache.lucene.search.CollectionStatistics;
import org.apache.lucene.search.Explanation;
import org.apache.lucene.search.TermStatistics;
import org.apache.lucene.search.similarities.Similarity;
import org.apache.lucene.util.BytesRef;
import org.apache.lucene.util.SmallFloat;

/**
 * @author sunyj
 * @since 2017年7月3日 下午5:08:11
 */
public class CustomSimilarity extends Similarity {

	/** Cache of decoded bytes. */
	private static final float[] NORM_TABLE = new float[256];

	static {
		for (int i = 0; i < 256; i++) {
			NORM_TABLE[i] = SmallFloat.byte315ToFloat((byte) i);
		}
	}

	@Override
	public long computeNorm(FieldInvertState state) {
		float normValue = lengthNorm(state);
		return encodeNormValue(normValue);
	}

	@Override
	public SimWeight computeWeight(float queryBoost, CollectionStatistics collectionStats,
			TermStatistics... termStats) {
		final Explanation idf = termStats.length == 1 ? idfExplain(collectionStats, termStats[0])
				: idfExplain(collectionStats, termStats);
		return new IDFStats(collectionStats.field(), idf, queryBoost);
	}

	@Override
	public SimScorer simScorer(SimWeight weight, LeafReaderContext context) throws IOException {
		IDFStats idfstats = (IDFStats) weight;
		return new TFIDFSimScorer(idfstats, context.reader().getNormValues(idfstats.field));
	}

	/**
	 * True if overlap tokens (tokens with a position of increment of zero) are
	 * discounted from the document's length.
	 */
	protected boolean discountOverlaps = true;

	public float lengthNorm(FieldInvertState state) {
		System.out.println("lengthNorm... ");
		final int numTerms;
		if (discountOverlaps)
			numTerms = state.getLength() - state.getNumOverlap();
		else
			numTerms = state.getLength();
		return state.getBoost() * ((float) (1.0 / Math.sqrt(numTerms)));
	}

	public final long encodeNormValue(float f) {
		return SmallFloat.floatToByte315(f);
	}

	/**
	 * Computes a score factor for a simple term and returns an explanation for
	 * that score factor.
	 * 
	 * <p>
	 * The default implementation uses:
	 * 
	 * <pre class="prettyprint">
	 * idf(docFreq, searcher.maxDoc());
	 * </pre>
	 * 
	 * Note that {@link CollectionStatistics#maxDoc()} is used instead of
	 * {@link org.apache.lucene.index.IndexReader#numDocs()
	 * IndexReader#numDocs()} because also {@link TermStatistics#docFreq()} is
	 * used, and when the latter is inaccurate, so is
	 * {@link CollectionStatistics#maxDoc()}, and in the same direction. In
	 * addition, {@link CollectionStatistics#maxDoc()} is more efficient to
	 * compute
	 * 
	 * @param collectionStats
	 *            collection-level statistics
	 * @param termStats
	 *            term-level statistics for the term
	 * @return an Explain object that includes both an idf score factor and an
	 *         explanation for the term.
	 */
	public Explanation idfExplain(CollectionStatistics collectionStats, TermStatistics termStats) {
		final long df = termStats.docFreq();
		final long max = collectionStats.maxDoc();
		final float idf = idf(df, max);
		return Explanation.match(idf, "idf(docFreq=" + df + ", maxDocs=" + max + ")");
	}

	/**
	 * Computes a score factor for a phrase.
	 * 
	 * <p>
	 * The default implementation sums the idf factor for each term in the
	 * phrase.
	 * 
	 * @param collectionStats
	 *            collection-level statistics
	 * @param termStats
	 *            term-level statistics for the terms in the phrase
	 * @return an Explain object that includes both an idf score factor for the
	 *         phrase and an explanation for each term.
	 */
	public Explanation idfExplain(CollectionStatistics collectionStats, TermStatistics termStats[]) {
		final long max = collectionStats.maxDoc();
		float idf = 0.0f;
		List<Explanation> subs = new ArrayList<>();
		for (final TermStatistics stat : termStats) {
			final long df = stat.docFreq();
			final float termIdf = idf(df, max);
			subs.add(Explanation.match(termIdf, "idf(docFreq=" + df + ", maxDocs=" + max + ")"));
			idf += termIdf;
		}
		return Explanation.match(idf, "idf(), sum of:", subs);
	}

	/** Implemented as <code>log(numDocs/(docFreq+1)) + 1</code>. */
	public float idf(long docFreq, long numDocs) {
		System.out.println("idf... ");
		return (float) (Math.log(numDocs / (double) (docFreq + 1)) + 1.0);
	}

	/** Implemented as <code>sqrt(freq)</code>. */
	public float tf(float freq) {
		System.out.println("tf... ");
		return (float) Math.sqrt(freq);
	}

	/**
	 * Decodes the norm value, assuming it is a single byte.
	 * 
	 * @see #encodeNormValue(float)
	 */
	public final float decodeNormValue(long norm) {
		return NORM_TABLE[(int) (norm & 0xFF)]; // & 0xFF maps negative bytes to
												// positive above 127
	}

	/** Implemented as <code>1 / (distance + 1)</code>. */
	public float sloppyFreq(int distance) {
		return 1.0f / (distance + 1);
	}

	/** The default implementation returns <code>1</code> */
	public float scorePayload(int doc, int start, int end, BytesRef payload) {
		return 1;
	}

	private Explanation explainQuery(IDFStats stats) {
		List<Explanation> subs = new ArrayList<>();

		Explanation boostExpl = Explanation.match(stats.queryBoost, "boost");
		if (stats.queryBoost != 1.0f)
			subs.add(boostExpl);
		subs.add(stats.idf);

		Explanation queryNormExpl = Explanation.match(stats.queryNorm, "queryNorm");
		subs.add(queryNormExpl);

		return Explanation.match(boostExpl.getValue() * stats.idf.getValue() * queryNormExpl.getValue(),
				"queryWeight, product of:", subs);
	}

	private Explanation explainField(int doc, Explanation freq, IDFStats stats, NumericDocValues norms) {
		Explanation tfExplanation = Explanation.match(tf(freq.getValue()),
				"tf(freq=" + freq.getValue() + "), with freq of:", freq);
		Explanation fieldNormExpl = Explanation.match(norms != null ? decodeNormValue(norms.get(doc)) : 1.0f,
				"fieldNorm(doc=" + doc + ")");

		return Explanation.match(tfExplanation.getValue() * stats.idf.getValue() * fieldNormExpl.getValue(),
				"fieldWeight in " + doc + ", product of:", tfExplanation, stats.idf, fieldNormExpl);
	}

	private Explanation explainScore(int doc, Explanation freq, IDFStats stats, NumericDocValues norms) {
		Explanation queryExpl = explainQuery(stats);
		Explanation fieldExpl = explainField(doc, freq, stats, norms);
		if (queryExpl.getValue() == 1f) {
			return fieldExpl;
		}
		return Explanation.match(queryExpl.getValue() * fieldExpl.getValue(),
				"score(doc=" + doc + ",freq=" + freq.getValue() + "), product of:", queryExpl, fieldExpl);
	}

	/**
	 * Collection statistics for the TF-IDF model. The only statistic of
	 * interest to this model is idf.
	 */
	private static class IDFStats extends SimWeight {
		private final String field;
		/** The idf and its explanation */
		private final Explanation idf;
		private float queryNorm;
		private float queryWeight;
		private final float queryBoost;
		private float value;

		public IDFStats(String field, Explanation idf, float queryBoost) {
			// TODO: Validate?
			this.field = field;
			this.idf = idf;
			this.queryBoost = queryBoost;
			this.queryWeight = idf.getValue() * queryBoost; // compute query
															// weight
		}

		@Override
		public float getValueForNormalization() {
			// TODO: (sorta LUCENE-1907) make non-static class and expose this
			// squaring via a nice method to subclasses?
			return queryWeight * queryWeight; // sum of squared weights
		}

		@Override
		public void normalize(float queryNorm, float topLevelBoost) {
			this.queryNorm = queryNorm * topLevelBoost;
			queryWeight *= this.queryNorm; // normalize query weight
			value = queryWeight * idf.getValue(); // idf for document
		}
	}

	private final class TFIDFSimScorer extends SimScorer {
		private final IDFStats stats;
		private final float weightValue;
		private final NumericDocValues norms;

		TFIDFSimScorer(IDFStats stats, NumericDocValues norms) throws IOException {
			this.stats = stats;
			this.weightValue = stats.value;
			this.norms = norms;
		}

		@Override
		public float score(int doc, float freq) {
			final float raw = tf(freq) * weightValue; // compute tf(f)*weight

			return norms == null ? raw : raw * decodeNormValue(norms.get(doc)); // normalize
																				// for
																				// field
		}

		@Override
		public float computeSlopFactor(int distance) {
			return sloppyFreq(distance);
		}

		@Override
		public float computePayloadFactor(int doc, int start, int end, BytesRef payload) {
			return scorePayload(doc, start, end, payload);
		}

		@Override
		public Explanation explain(int doc, Explanation freq) {
			return explainScore(doc, freq, stats, norms);
		}
	}
}
