package cz.vutbr.feec.utko.xonder02.lda;

import java.io.*;
import java.nio.file.Files;
import java.nio.file.StandardCopyOption;
import java.text.DecimalFormat;
import java.util.*;

public class Analyzer {
    private Dictionary dictionary = new Dictionary();
    private Loader loader = new Loader(dictionary);

    private TopicStats[] analyzeCorpus(Corpus corpus, Lda lda) throws Exception {
        int numberOfEmotions = loader.getEmotionDictionary().getNumberOfTokens();
        TopicStats[] topicStatses = TopicStats.generate(lda.getNumberOfTopics(), numberOfEmotions);

        for (Document document : corpus) {
            int realEmotionId = loader.getDocumentEmotionId(document);

            List<DocumentTopicDistance> topics = lda.findTopic(document, 0.01);
            topicStatses[topics.get(0).getTopic().getTopicId()].addCount(realEmotionId);

            for (DocumentTopicDistance documentTopicDistance : topics) {
                int topicId = documentTopicDistance.getTopic().getTopicId();
                topicStatses[topicId].addGamma(realEmotionId, documentTopicDistance.getDistance());
            }

            /*if (document.getId() % 100 == 0) {
                System.out.print("Document processed: ");
                System.out.println(document.getId());
            }*/
        }

        return topicStatses;
    }

    private void analyzeTestCorpus(Corpus corpus, Lda lda, TopicStats[] topicStatses) throws Exception {
        double completeAvg = 0;

        for (Document document : corpus) {
            int realEmotionId = loader.getDocumentEmotionId(document);

            List<DocumentTopicDistance> topics = lda.findTopic(document, 0.01);

            double avg = 0;
            int topicId = 0;
            for (DocumentTopicDistance documentTopicDistance : topics) {
                topicId = documentTopicDistance.getTopic().getTopicId();
                double[] stats = topicStatses[topicId].statsByBestWithDocumentCount(topicStatses);
                double percent = stats[realEmotionId];
                avg = percent;
                break;
                //avg += percent * documentTopicDistance.getDistance();
            }

            //avg /= topics.size();
            completeAvg += avg;

            /*System.out.print("Document processed: quality ");
            System.out.print(avg * 100);
            System.out.print(", emotion ");
            System.out.print(realEmotionId);
            System.out.print(", topic ID ");
            System.out.println(topicId);*/
        }

        completeAvg /= (double)corpus.documentCount();

        System.out.print("Overall quality: ");
        System.out.print(completeAvg * 100);
    }

    public void test() {
        Corpus trainingCorpus = new Corpus();
        String[] folders = new String[] { "afraid", "anger", "neutral", "sadness", "satisfaction", "surprise" };
        HashMap<String, String> folderEmotion = new HashMap<String, String>();
        folderEmotion.put("afraid", "afraid");
        folderEmotion.put("anger", "anger");
        folderEmotion.put("neutral", "neutral");
        folderEmotion.put("sadness", "sadness");
        folderEmotion.put("satisfaction", "satisfaction");
        folderEmotion.put("surprise", "surprise");

        try {
            for (String folder : folders) {
                loader.readDirectory(trainingCorpus, new File("data/train/" + folder + "/YES/"), folderEmotion.get(folder), false);
            }

            System.out.print("Loaded ");
            System.out.print(trainingCorpus.documentCount());
            System.out.println(" documents from training corpus");

            System.out.print("Number of unique tokens: ");
            System.out.println(dictionary.getNumberOfTokens());

            System.out.print("Number of average unique tokens per document: ");
            System.out.println(trainingCorpus.averageNumberOfTokensPerDocument());
            System.out.println();

            TfIdf tfIdf = new TfIdf(dictionary);
            tfIdf.computeForCorpus(trainingCorpus);

            int topicCount = 6;
            Lda lda = new Lda(trainingCorpus, topicCount, 5000, 0.00001, 15);
            //lda.getState().saveToFile(new File("state.ser"));

            //Lda lda = new Lda(State.loadFromFile(new File("state.ser")), 50, 0.01, 10);

            TopicStats[] trainingTopicStatus = analyzeCorpus(trainingCorpus, lda);

            for (int i = 0; i < topicCount; i++) {
                TopicStats topicStats = trainingTopicStatus[i];
                System.out.print("Topic ");
                System.out.print(i + 1);
                System.out.print(" with ");
                System.out.print(trainingTopicStatus[i].documentCount());
                System.out.print(" documents, ");
                System.out.println(TopicStats.statsToString(topicStats.statsByCount(), loader.getEmotionDictionary()));

                System.out.print("By count: ");
                System.out.println(TopicStats.statsToString(topicStats.statsByCountNormalize(), loader.getEmotionDictionary()));
                //System.out.println(TopicStats.statsToString(topicStats.statsByCountWithDocumentCountNormalize(trainingTopicStatus), loader.getEmotionDictionary()));

                System.out.print("By gamma: ");
                System.out.println(TopicStats.statsToString(topicStats.statsByGamma(), loader.getEmotionDictionary()));

                //System.out.print("By gamma with document count: ");
               // System.out.println(TopicStats.statsToString(topicStats.statsByGammaWithDocumentCount(trainingTopicStatus), loader.getEmotionDictionary()));
                //lda.getTopic(i).printPairs(dictionary, 10);
                System.out.println();
            }

            Corpus testCorpus = new Corpus();

            for (int i = 0; i < 300; i++) {
                int folderId = (int) (Math.random() * folders.length);
                String folder  = folders[folderId];

                File file = loader.readRandomFileFromDirectory(testCorpus, new File("data/test/" + folder + "/YES/"), folderEmotion.get(folder), true);

                /*File output = new File("data/test2/" + folder + "/" + i + '-' + file.getName());
                Files.copy(file.toPath(), output.toPath(), StandardCopyOption.REPLACE_EXISTING);*/
            }


            /*for (String folder : folders) {
                loader.readDirectory(testCorpus, new File("data/test2/" + folder + "/"), folderEmotion.get(folder), true);
            }*/

            System.out.print("Loaded ");
            System.out.print(testCorpus.documentCount());
            System.out.println(" from test corpus");

            System.out.print("Number of orphan tokens: ");
            System.out.println(dictionary.getNumberOfTokensOnlyInOneDocument());
            System.out.println();

            tfIdf.computeForCorpus(testCorpus);

            analyzeTestCorpus(testCorpus, lda, trainingTopicStatus);

        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    public void test2() {
        Corpus trainingCorpus = new Corpus();
        Corpus trainingTestCorpus = new Corpus();
        Corpus testCorpus = new Corpus();


        String[] folders = new String[] { "afraid", "anger", "neutral", "sadness", "satisfaction", "surprise" };
        HashMap<String, String> folderEmotion = new HashMap<String, String>();
        folderEmotion.put("afraid", "afraid");
        folderEmotion.put("anger", "anger");
        folderEmotion.put("neutral", "neutral");
        folderEmotion.put("sadness", "sadness");
        folderEmotion.put("satisfaction", "satisfaction");
        folderEmotion.put("surprise", "surprise");

        try {
            for (String folder : folders) {
                loader.readDirectory(trainingCorpus, new File("data/train/" + folder + "/YES/"), folderEmotion.get(folder), false);
            }

            for (String folder : folders) {
                loader.readDirectory(trainingTestCorpus, new File("data/test/" + folder + "/YES/"), folderEmotion.get(folder), false);
            }

            trainingTestCorpus.merge(trainingCorpus);

            System.out.print("Number of unique tokens: ");
            System.out.println(dictionary.getNumberOfTokens());

            System.out.print("Number of unique tokens per document: ");
            System.out.println(trainingCorpus.averageNumberOfTokensPerDocument());

            TfIdf tfIdf = new TfIdf(dictionary);
            tfIdf.computeForCorpus(trainingTestCorpus);

            int topicCount = 50;
            Lda lda = new Lda(trainingTestCorpus, topicCount, 5000, 0.00001, 15);
            //lda.getState().saveToFile(new File("state.ser"));

            //Lda lda = new Lda(State.loadFromFile(new File("state.ser")), 50, 0.01, 10);

            TopicStats[] trainingTopicStatus = analyzeCorpus(trainingCorpus, lda);

            for (int i = 0; i < topicCount; i++) {
                TopicStats topicStats = trainingTopicStatus[i];
                System.out.print("Topic ");
                System.out.print(i);
                System.out.print(" with ");
                System.out.print(trainingTopicStatus[i].documentCount());
                System.out.print(" documents, ");
                System.out.println(TopicStats.statsToString(topicStats.statsByCount(), loader.getEmotionDictionary()));

                System.out.print("By count: ");
                System.out.println(TopicStats.statsToString(topicStats.statsByCountNormalize(), loader.getEmotionDictionary()));
                //System.out.println(TopicStats.statsToString(topicStats.statsByCountWithDocumentCountNormalize(trainingTopicStatus), loader.getEmotionDictionary()));

                System.out.print("By gamma: ");
                System.out.println(TopicStats.statsToString(topicStats.statsByGamma(), loader.getEmotionDictionary()));

                //System.out.print("By gamma with document count: ");
                // System.out.println(TopicStats.statsToString(topicStats.statsByGammaWithDocumentCount(trainingTopicStatus), loader.getEmotionDictionary()));
                //lda.getTopic(i).printPairs(dictionary, 10);
                System.out.println();
            }



            /*for (int i = 0; i < 300; i++) {
                int folderId = (int) (Math.random() * folders.length);
                String folder  = folders[folderId];

                File file = loader.readRandomFileFromDirectory(testCorpus, new File("data/test/" + folder + "/YES/"), folderEmotion.get(folder), true);

                File output = new File("data/test2/" + folder + "/" + i + '-' + file.getName());
                Files.copy(file.toPath(), output.toPath(), StandardCopyOption.REPLACE_EXISTING);
            }*/


            for (String folder : folders) {
                loader.readDirectory(testCorpus, new File("data/test2/" + folder + "/"), folderEmotion.get(folder), true);
            }

            System.out.print("Number of orphan tokens: ");
            System.out.println(dictionary.getNumberOfTokensOnlyInOneDocument());

            tfIdf.computeForCorpus(testCorpus);

            analyzeTestCorpus(testCorpus, lda, trainingTopicStatus);

        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    public static void main(String[] args) {
        Analyzer analyzer = new Analyzer();
        analyzer.test();
    }

    static class TopicStats {
        public Stat[] stats;
        private int totalDocumentCount = 0;

        public static TopicStats[] generate(int topicCount, int emotionCount) {
            TopicStats[] topicStatses = new TopicStats[topicCount];
            for (int i = 0; i < topicCount; i++) {
                topicStatses[i] = new TopicStats(emotionCount);
            }
            return topicStatses;
        }

        public TopicStats(int emotionCount) {
            stats = new Stat[emotionCount];
            for (int i = 0; i < emotionCount; i++) {
                stats[i] = new Stat();
            }
        }

        public void addCount(int emotionId) {
            stats[emotionId].realCount++;
        }

        public void addGamma(int emotionId, double gamma) {
            Stat stat = stats[emotionId];
            stat.gamma += gamma;
            stat.count++;
            totalDocumentCount++;
        }

        public int documentCount() {
            int sum = 0;
            for (Stat stat : stats) {
                sum += stat.realCount;
            }

            return sum;
        }

        public double[] statsByBest() {
            int bestId = 0;
            double best = stats[bestId].realCount;

            for (int i = 1; i < stats.length; i++) {
                 if (best <  stats[i].realCount) {
                     best = stats[i].realCount;
                     bestId = i;
                 }
            }

            double[] output = new double[stats.length];
            output[bestId] = 1;
            return output;
        }

        public double[] statsByBestWithDocumentCount(TopicStats[] topicStats) {
            double[] count = statsByCountWithDocumentCount(topicStats);
            return computeBest(count);
        }

        public double[] statsByBestGamma() {
            double[] gamma = statsByGamma();
            return computeBest(gamma);
        }

        private double[] computeBest(double[] input) {
            int bestId = 0;
            double best = input[bestId];

            for (int i = 1; i < input.length; i++) {
                if (best < input[i]) {
                    best = input[i];
                    bestId = i;
                }
            }

            double[] output = new double[stats.length];
            output[bestId] = 1;
            return output;
        }

        public double[] statsByCount() {
            double[] output = new double[stats.length];
            for (int i = 0; i < stats.length; i++) {
                output[i] = stats[i].realCount;
            }

            return output;
        }

        public double[] statsByCountNormalize() {
            return normalize(statsByCount());
        }

        public double[] statsByCountWithDocumentCount(TopicStats[] topicStats) {
            double[] multiplier = perDocument(topicStats);

            double[] output = new double[stats.length];
            for (int i = 0; i < stats.length; i++) {
                output[i] = stats[i].realCount * multiplier[i];
            }

            return output;
        }

        public double[] statsByCountWithDocumentCountNormalize(TopicStats[] topicStats) {
            return normalize(statsByCountWithDocumentCount(topicStats));
        }

        public double[] statsByGamma() {
            double[] output = new double[stats.length];

            for (int i = 0; i < stats.length; i++) {
                if (stats[i].count == 0) {
                    output[i] = 0;
                } else {
                    output[i] = stats[i].gamma / (double) stats[i].count;
                }
            }

            return normalize(output);
        }

        public double[] statsByGammaWithDocumentCount(TopicStats[] topicStats) {
            double[] multiplier = perDocument(topicStats);
            double[] output = new double[stats.length];

            for (int i = 0; i < stats.length; i++) {
                if (stats[i].count == 0) {
                    output[i] = 0;
                } else {
                    output[i] = (stats[i].gamma / (double) stats[i].count) / multiplier[i];
                }
            }

            return normalize(output);
        }

        private double[] normalize(double[] input) {
            double sum = 0;
            for (double anInput : input) {
                sum += anInput;
            }

            if (sum == 0) {
                return input;
            }

            for (int i = 0; i < input.length; i++) {
                input[i] /= sum;
            }

            return input;
        }

        private double[] perDocument(TopicStats[] topicStats) {
            double[][] statsPerTopic = new double[topicStats.length][];
            int documentCount = 0;

            for (int i = 0; i < topicStats.length; i++) {
                statsPerTopic[i] = topicStats[i].statsByCount();
                documentCount += topicStats[i].documentCount();
            }

            double[] output = new double[stats.length];

            for (int i = 0; i < stats.length; i++) {
                int documentWithEmotion = 0;
                for (double[] aStatsPerTopic : statsPerTopic) {
                    documentWithEmotion += aStatsPerTopic[i];
                }

                output[i] = documentCount / (double) documentWithEmotion;
            }

            return output;
        }

        public static String statsToString(double[] stats, Dictionary emotionDictionary) {
            ArrayList<DoubleArrayPair> sorted = new ArrayList<DoubleArrayPair>(stats.length);

            for (int i = 0; i < stats.length; i++) {
                sorted.add(new DoubleArrayPair(i, stats[i]));
            }

            /*Collections.sort(sorted, new Comparator<DoubleArrayPair>() {
                @Override
                public int compare(DoubleArrayPair o1, DoubleArrayPair o2) {
                    return o2.value.compareTo(o1.value);
                }
            });*/

            List<String> stringList = new ArrayList<String>();

            DecimalFormat decimalFormat = new DecimalFormat("#0.00");
            for (DoubleArrayPair pair : sorted) {
                stringList.add(decimalFormat.format(pair.value) + "x " + emotionDictionary.getTokenById(pair.key));
            }

            if (stringList.size() > 0) {
                StringBuilder builder = new StringBuilder();
                builder.append(stringList.remove(0));

                for (String s : stringList) {
                    builder.append(", ");
                    builder.append(s);
                }

                return builder.toString();
            }

            return "";

        }

        static class DoubleArrayPair {
            public Integer key;
            public Double value;

            public DoubleArrayPair(Integer key, Double value) {
                this.key = key;
                this.value = value;
            }
        }

        class Stat {
            public Integer count = 0;
            public Integer realCount = 0;
            public Double gamma = .0;
        }
    }
}
