package com.tzld.piaoquan.recommend.similarity.word2vec;

import com.ansj.vec.Word2VEC;
import com.ansj.vec.domain.WordEntry;

import java.io.IOException;
import java.util.Collections;
import java.util.Iterator;
import java.util.List;
import java.util.Set;
import java.util.TreeSet;
import java.util.Map.Entry;

/**
 * @author dyp
 */


public class Word2Vec {
    private Word2VEC vec = new Word2VEC();
    private boolean loadModel = false;

    public Word2Vec() {
    }

    public void loadGoogleModelFromOss(String endpoint, String bucketName, String path, String accessKeyId, String accessKetSecret) throws IOException {
        this.vec.loadGoogleModelFromOss(endpoint, bucketName, path, accessKeyId, accessKetSecret);
        this.loadModel = true;
    }

    public void loadGoogleModel(String path) throws IOException {
        this.vec.loadGoogleModel(path);
        this.loadModel = true;
    }

    public float[] getWordVector(String word) {
        return !this.loadModel ? null : this.vec.getWordVector(word);
    }

    private float calDist(float[] vec1, float[] vec2) {
        float dist = 0.0F;

        for (int i = 0; i < vec1.length; ++i) {
            dist += vec1[i] * vec2[i];
        }

        return dist;
    }

    private void calSum(float[] sum, float[] vec) {
        for (int i = 0; i < sum.length; ++i) {
            sum[i] += vec[i];
        }

    }

    public float wordSimilarity(String word1, String word2) {
        if (!this.loadModel) {
            return 0.0F;
        } else {
            float[] word1Vec = this.getWordVector(word1);
            float[] word2Vec = this.getWordVector(word2);
            return word1Vec != null && word2Vec != null ? this.calDist(word1Vec, word2Vec) : 0.0F;
        }
    }

    public Set<WordEntry> getSimilarWords(String word, int maxReturnNum) {
        if (!this.loadModel) {
            return null;
        } else {
            float[] center = this.getWordVector(word);
            if (center == null) {
                return Collections.emptySet();
            } else {
                int resultSize = this.vec.getWords() < maxReturnNum ? this.vec.getWords() : maxReturnNum;
                TreeSet<WordEntry> result = new TreeSet();
                double min = 4.9E-324D;
                Iterator var9 = this.vec.getWordMap().entrySet().iterator();

                while (var9.hasNext()) {
                    Entry<String, float[]> entry = (Entry) var9.next();
                    float[] vector = (float[]) entry.getValue();
                    float dist = this.calDist(center, vector);
                    if (result.size() <= resultSize) {
                        result.add(new WordEntry((String) entry.getKey(), dist));
                        min = (double) ((WordEntry) result.last()).score;
                    } else if ((double) dist > min) {
                        result.add(new WordEntry((String) entry.getKey(), dist));
                        result.pollLast();
                        min = (double) ((WordEntry) result.last()).score;
                    }
                }

                result.pollFirst();
                return result;
            }
        }
    }

    private float calMaxSimilarity(String centerWord, List<String> wordList) {
        float max = -1.0F;
        if (wordList.contains(centerWord)) {
            return 1.0F;
        } else {
            Iterator var5 = wordList.iterator();

            while (var5.hasNext()) {
                String word = (String) var5.next();
                float temp = this.wordSimilarity(centerWord, word);
                if (temp != 0.0F && temp > max) {
                    max = temp;
                }
            }

            if (max == -1.0F) {
                return 0.0F;
            } else {
                return max;
            }
        }
    }

    public float fastSentenceSimilarity(List<String> sentence1Words, List<String> sentence2Words) {
        if (!this.loadModel) {
            return 0.0F;
        } else if (!sentence1Words.isEmpty() && !sentence2Words.isEmpty()) {
            float[] sen1vector = new float[this.vec.getSize()];
            float[] sen2vector = new float[this.vec.getSize()];
            double len1 = 0.0D;
            double len2 = 0.0D;

            int i;
            float[] tmp;
            for (i = 0; i < sentence1Words.size(); ++i) {
                tmp = this.getWordVector(sentence1Words.get(i));
                if (tmp != null) {
                    this.calSum(sen1vector, tmp);
                }
            }

            for (i = 0; i < sentence2Words.size(); ++i) {
                tmp = this.getWordVector(sentence2Words.get(i));
                if (tmp != null) {
                    this.calSum(sen2vector, tmp);
                }
            }

            for (i = 0; i < this.vec.getSize(); ++i) {
                len1 += (double) (sen1vector[i] * sen1vector[i]);
                len2 += (double) (sen2vector[i] * sen2vector[i]);
            }

            if (len1 * len2 == 0) {
                return 0.0F;
            }
            return (float) ((double) this.calDist(sen1vector, sen2vector) / Math.sqrt(len1 * len2));
        } else {
            return 0.0F;
        }
    }

    public float sentenceSimilarity(List<String> sentence1Words, List<String> sentence2Words) {
        if (!this.loadModel) {
            return 0.0F;
        } else if (!sentence1Words.isEmpty() && !sentence2Words.isEmpty()) {
            float sum1 = 0.0F;
            float sum2 = 0.0F;
            int count1 = 0;
            int count2 = 0;

            int i;
            for (i = 0; i < sentence1Words.size(); ++i) {
                if (this.getWordVector(sentence1Words.get(i)) != null) {
                    ++count1;
                    sum1 += this.calMaxSimilarity(sentence1Words.get(i), sentence2Words);
                }
            }

            for (i = 0; i < sentence2Words.size(); ++i) {
                if (this.getWordVector(sentence2Words.get(i)) != null) {
                    ++count2;
                    sum2 += this.calMaxSimilarity(sentence2Words.get(i), sentence1Words);
                }
            }

            return (sum1 + sum2) / (float) (count1 + count2);
        } else {
            return 0.0F;
        }
    }

    public float sentenceSimilarity(List<String> sentence1Words, List<String> sentence2Words, float[] weightVector1, float[] weightVector2) throws Exception {
        if (!this.loadModel) {
            return 0.0F;
        } else if (!sentence1Words.isEmpty() && !sentence2Words.isEmpty()) {
            if (sentence1Words.size() == weightVector1.length && sentence2Words.size() == weightVector2.length) {
                float sum1 = 0.0F;
                float sum2 = 0.0F;
                float divide1 = 0.0F;
                float divide2 = 0.0F;

                int i;
                float wordMaxSimi;
                for (i = 0; i < sentence1Words.size(); ++i) {
                    if (this.getWordVector(sentence1Words.get(i)) != null) {
                        wordMaxSimi = this.calMaxSimilarity(sentence1Words.get(i), sentence2Words);
                        sum1 += wordMaxSimi * weightVector1[i];
                        divide1 += weightVector1[i];
                    }
                }

                for (i = 0; i < sentence2Words.size(); ++i) {
                    if (this.getWordVector(sentence2Words.get(i)) != null) {
                        wordMaxSimi = this.calMaxSimilarity(sentence2Words.get(i), sentence1Words);
                        sum2 += wordMaxSimi * weightVector2[i];
                        divide2 += weightVector2[i];
                    }
                }

                return (sum1 + sum2) / (divide1 + divide2);
            } else {
                throw new Exception("length of word list and weight vector is different");
            }
        } else {
            return 0.0F;
        }
    }
}
