/**
 * Copyright (c) 2011, The University of Southampton and the individual contributors.
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without modification,
 * are permitted provided that the following conditions are met:
 *
 *   *  Redistributions of source code must retain the above copyright notice,
 *      this list of conditions and the following disclaimer.
 *
 *   *  Redistributions in binary form must reproduce the above copyright notice,
 *      this list of conditions and the following disclaimer in the documentation
 *      and/or other materials provided with the distribution.
 *
 *   *  Neither the name of the University of Southampton nor the names of its
 *      contributors may be used to endorse or promote products derived from this
 *      software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
 * ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
 * ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */
package org.openimaj.text.nlp.namedentity;

import info.bliki.wiki.filter.PlainTextConverter;
import info.bliki.wiki.model.WikiModel;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.File;
import java.io.FileNotFoundException;
import java.io.FileReader;
import java.io.FileWriter;
import java.io.IOException;
import java.util.HashMap;

import javax.xml.parsers.DocumentBuilder;
import javax.xml.parsers.DocumentBuilderFactory;
import javax.xml.parsers.ParserConfigurationException;

import org.apache.commons.lang.StringEscapeUtils;
import org.apache.lucene.document.FieldType;
import org.apache.lucene.store.SimpleFSDirectory;
import org.openimaj.text.nlp.namedentity.YagoEntityCandidateFinderFactory.YagoEntityCandidateFinder;
import org.openimaj.text.nlp.namedentity.YagoEntityContextScorerFactory.YagoEntityContextScorer;
import org.openimaj.text.nlp.namedentity.YagoEntityExactMatcherFactory.YagoEntityExactMatcher;
import org.w3c.dom.Document;
import org.w3c.dom.NodeList;
import org.xml.sax.SAXException;

/**
 * This class has various methods that can be used to build the resources
 * required by {@link YagoEntityCandidateFinder},
 * {@link YagoEntityContextScorer} and {@link YagoEntityExactMatcher}. These
 * resources are a text File of entity aliases, and a lucene index of contextual
 * data.
 *
 * The directory of the stripped down Yago tsv files is required. This directory
 * can be built with {@link SeedBuilder}.
 *
 * @author Laurence Willmore (lgw1e10@ecs.soton.ac.uk)
 * @author Sina Samangooei (ss@ecs.soton.ac.uk)
 *
 */
public class EntityExtractionResourceBuilder {

        /**
         * Default file name for the alias text file.
         */
        public static String DEFAULT_ALIAS_NAME = "AliasMapFile.txt";
        /**
         * Default directory name for the lucene index.
         */
        public static String DEFAULT_CONTEXT_NAME = "YagoLucene";
        private static String DEFAULT_ROOT_NAME = ".YagoEntityExtraction";
        private static String wikiApiPrefix = "http://en.wikipedia.org/w/api.php?format=xml&action=query&titles=";
        private static String wikiApiSuffix = "&prop=revisions&rvprop=content";
        private boolean verbose = true;
        // This will build for location entities. There are too many for memory.
        // Leave false.
        private boolean locations = false;
        private static BufferedWriter logOut;

        /**
         * Builds the alias text file in the default location.
         *
         * @param seedDirectoryPath
         *            = path location of the stripped down Yago .tsv files.
         */
        public void buildCandidateAliasFile(String seedDirectoryPath) {
                buildCandidateAliasFile(seedDirectoryPath, getDefaultRootPath()
                                + File.separator + DEFAULT_ALIAS_NAME);
        }

        /**
         * Builds the alias text file in the specified location.
         *
         * @param seedDirectoryPath
         *            = path location of the stripped down Yago .tsv files.
         * @param destinationPath
         *            = path to build the alias text file.
         */
        public void buildCandidateAliasFile(String seedDirectoryPath,
                        String destinationPath)
        {
                writeAliasFile(getEntities(seedDirectoryPath), destinationPath,
                                seedDirectoryPath);
        }

        /**
         * Builds the lucene index in the default path.
         *
         * @param seedDirectoryPath
         *            = path location of the stripped down Yago .tsv files.
         */
        public void buildContextLuceneIndex(String seedDirectoryPath) {
                buildContextLuceneIndex(seedDirectoryPath, getDefaultRootPath()
                                + File.separator + DEFAULT_CONTEXT_NAME);
        }

        /**
         * Builds the lucene index at the specified path.
         *
         * @param seedDirectoryPath
         * @param destinationPath
         */
        public void buildContextLuceneIndex(String seedDirectoryPath,
                        String destinationPath)
        {
                try {
                        buildIndex(getEntities(seedDirectoryPath), destinationPath,
                                        seedDirectoryPath);
                } catch (final IOException e) {
                        e.printStackTrace();
                }
        }

        /**
         * Builds the alias text file and the lucene index in the default root
         * directory.
         *
         * @param seedDirectoryPath
         */
        public void buildAll(String seedDirectoryPath) {
                validateFileStructure();
                createLogging(getDefaultRootPath() + File.separator + "log.txt");
                buildAll(seedDirectoryPath, getDefaultRootPath());
                try {
                        logOut.flush();
                        logOut.close();
                } catch (final IOException e) {
                        e.printStackTrace();
                }
        }

        /**
         * Builds the alias text file and the lucene index in the specified root
         * directory.
         *
         * @param seedDirectoryPath
         * @param destinationPath
         */
        public void buildAll(String seedDirectoryPath, String destinationPath) {
                // Get the entities as people and organisations
                print("Building All...");
                final HashMap<String, YagoNamedEntity> entities = getEntities(seedDirectoryPath);
                writeAliasFile(entities, destinationPath + File.separator
                                + DEFAULT_ALIAS_NAME, seedDirectoryPath);
                try {
                        buildIndex(entities, destinationPath + File.separator
                                        + DEFAULT_CONTEXT_NAME, seedDirectoryPath);
                } catch (final IOException e) {
                        e.printStackTrace();
                }
                print("Done");
        }

        /**
         * @return default root directory path for all YagoEntity resources.
         */
        public static String getDefaultRootPath() {
                return System.getProperty("user.home") + File.separator
                                + DEFAULT_ROOT_NAME;
        }

        /**
         * @return default alias text file path.
         */
        public static String getDefaultAliasFilePath() {
                return getDefaultRootPath() + File.separator + DEFAULT_ALIAS_NAME;
        }

        /**
         * @return defualt lucene directory path.
         */
        public static String getDefaultIndexDirectoryPath() {
                return getDefaultRootPath() + File.separator + DEFAULT_CONTEXT_NAME;
        }

        public static String getAliasFrom(String rootName) {
                String result;
                String noGeo = null;
                if (rootName.startsWith("geoent_")) {
                        noGeo = rootName.substring(rootName.indexOf('_') + 1,
                                        rootName.lastIndexOf('_'));
                } else
                        noGeo = rootName;
                final String spaces = noGeo.replaceAll("_", " ");
                String noParen;
                if (spaces.contains("("))
                        noParen = spaces.substring(0, spaces.indexOf("("));
                else
                        noParen = spaces;
                String dropComma;
                if (noParen.contains(","))
                        dropComma = noParen.substring(0, spaces.indexOf(","));
                else
                        dropComma = noParen;
                result = dropComma;
                return result;
        }

        private void validateFileStructure() {
                final File rootDir = new File(getDefaultRootPath());
                if (!rootDir.isDirectory()) {
                        rootDir.mkdir();
                }
                final File indexDir = new File(getDefaultRootPath() + File.separator
                                + DEFAULT_CONTEXT_NAME);
                if (!indexDir.isDirectory()) {
                        indexDir.mkdir();
                } else {
                        for (final File f : indexDir.listFiles())
                                f.delete();
                }
        }

        private static void createLogging(String logFilePath) {
                final File f = new File(logFilePath);
                if (!f.isFile()) {
                        try {
                                f.createNewFile();
                        } catch (final IOException e) {
                                e.printStackTrace();
                        }
                } else {
                }
                FileWriter fstream = null;
                try {
                        fstream = new FileWriter(logFilePath);
                        logOut = new BufferedWriter(fstream);
                        logOut.write("");
                } catch (final IOException e) {
                        // TODO Auto-generated catch block
                        e.printStackTrace();
                }

        }

        private void buildIndex(HashMap<String, YagoNamedEntity> entities,
                        String destinationPath, String seedDirectoryPath)
                                        throws IOException
        {
                print("Building Index...");
                setEntityContextValues(entities, seedDirectoryPath);
                print("Initializing Lucene objects...");

                // initialize lucene objects
                final String[] names = { "uri", "context", "type" };
                FieldType[] types;
                final FieldType ti = new FieldType();
                ti.setIndexed(true);
                ti.setTokenized(true);
                ti.setStored(true);
                final FieldType n = new FieldType();
                n.setStored(true);
                n.setIndexed(true);
                types = new FieldType[3];
                types[0] = n;
                types[1] = ti;
                types[2] = n;
                final File f = new File(destinationPath);
                final QuickIndexer qi = new QuickIndexer(new SimpleFSDirectory(f));

                // Initialize wiki objects
                final DocumentBuilderFactory docBuilderFactory = DocumentBuilderFactory
                                .newInstance();
                DocumentBuilder docBuilder = null;
                Document doc;
                try {
                        docBuilder = docBuilderFactory.newDocumentBuilder();
                } catch (final ParserConfigurationException e) {
                        e.printStackTrace();
                }
                doc = null;
                final WikiModel wikiModel = new WikiModel(
                                "http://www.mywiki.com/wiki/${image}",
                                "http://www.mywiki.com/wiki/${title}");
                int count = 0;
                print("Building Lucene Index...");
                for (final YagoNamedEntity entity : entities.values()) {
                        count++;
                        if (count % 5000 == 0)
                                print("Processed " + count);
                        // if wikiURL, add wiki to context
                        if (entity.wikiURL != null) {
                                final String title = entity.wikiURL.substring(entity.wikiURL
                                                .lastIndexOf("/") + 1);
                                try {
                                        doc = docBuilder.parse(wikiApiPrefix + title
                                                        + wikiApiSuffix);
                                } catch (final SAXException e) {
                                        e.printStackTrace();
                                } catch (final IOException e) {
                                        e.printStackTrace();
                                }
                                doc.getDocumentElement().normalize();
                                final NodeList revisions = doc.getElementsByTagName("rev");
                                if (revisions.getLength() > 0) {
                                        final String markup = revisions.item(0).getTextContent();

                                        // convert markup dump to plaintext.
                                        final String plainStr = wikiModel.render(
                                                        new PlainTextConverter(), markup);
                                        // add it to the context.
                                        entity.addContext(plainStr);
                                }
                        }
                        final String[] values = { entity.rootName, entity.getContext(),
                                        entity.type.toString() };
                        qi.addDocumentFromFields(names, values, types);
                }
                qi.finalise();
        }

        private void setEntityContextValues(
                        final HashMap<String, YagoNamedEntity> entities,
                        String seedDirectoryPath)
        {
                print("Setting Context Values...");
                BufferedReader in = null;
                // Created
                try {
                        in = openFileAsReadStream(seedDirectoryPath + File.separator
                                        + "created_stripped.tsv");
                } catch (final FileNotFoundException e) {
                        e.printStackTrace();
                }
                StreamLooper sl = new StreamLooper(in) {
                        @Override
                        protected void doWork(String s) {
                                final String[] values = s.split("\\s+");
                                final String rootName = values[1];
                                final String context = convertResource(values[2]);
                                if (entities.keySet().contains(rootName)) {
                                        entities.get(rootName).addContext(context);
                                }
                        }
                };
                sl.loop();

                // wikiAnchorText
                try {
                        in = openFileAsReadStream(seedDirectoryPath + File.separator
                                        + "hasWikipediaAnchorText_stripped.tsv");
                } catch (final FileNotFoundException e) {
                        e.printStackTrace();
                }
                sl = new StreamLooper(in) {
                        @Override
                        protected void doWork(String s) {
                                final String[] values = s.split("\\s+");
                                final String rootName = values[1];
                                final String context = convertLiteral(values[2]);
                                if (entities.keySet().contains(rootName)) {
                                        entities.get(rootName).addContext(context);
                                }
                        }
                };
                sl.loop();

                // wikiUrl

                try {
                        in = openFileAsReadStream(seedDirectoryPath + File.separator
                                        + "hasWikipediaUrl_stripped.tsv");
                } catch (final FileNotFoundException e) {
                        e.printStackTrace();
                }
                sl = new StreamLooper(in) {
                        @Override
                        protected void doWork(String s) {
                                final String[] values = s.split("\\s+");
                                final String rootName = values[1];
                                if (entities.keySet().contains(rootName)) {
                                        entities.get(rootName).wikiURL = values[2].replaceAll("\"",
                                                        "");
                                }
                        }
                };
                sl.loop();
                // validate
                print("Validating Context...");
                int noContext = 0;
                for (final YagoNamedEntity ne : entities.values()) {
                        for (final String alias : ne.aliasList) {
                                ne.addContext(alias);
                        }
                        if ((ne.getContext() == null || ne.getContext().equals(""))
                                        && ne.wikiURL == null)
                        {
                                noContext++;
                        }
                }
                print("No Context: " + noContext);
        }

        private void setEntityAliasValues(
                        final HashMap<String, YagoNamedEntity> entities,
                        String seedDirectoryPath)
        {
                print("Setting Alias Values...");
                // Populate 'isCalled'
                BufferedReader in = null;
                try {
                        in = openFileAsReadStream(seedDirectoryPath + File.separator
                                        + "isCalled_stripped.tsv");
                } catch (final FileNotFoundException e) {
                        e.printStackTrace();
                }
                StreamLooper sl = new StreamLooper(in) {
                        @Override
                        protected void doWork(String s) {
                                final String[] values = s.split("\\s+");
                                final String rootName = values[1];
                                final String alias = convertLiteral(values[2]);
                                if (entities.keySet().contains(rootName)) {
                                        entities.get(rootName).addAlias(alias);
                                }
                        }
                };
                sl.loop();

                // populate 'means'

                try {
                        in = openFileAsReadStream(seedDirectoryPath + File.separator
                                        + "means_stripped.tsv");
                } catch (final FileNotFoundException e) {
                        e.printStackTrace();
                }
                sl = new StreamLooper(in) {
                        @Override
                        protected void doWork(String s) {
                                final String[] values = s.split("\\s+");
                                final String rootName = values[2];
                                final String alias = convertLiteral(values[1]);
                                // System.out.println(alias);
                                if (entities.keySet().contains(rootName)) {
                                        entities.get(rootName).addAlias(alias);
                                }
                        }
                };
                sl.loop();
                print("Validating Aliases...");
                for (final YagoNamedEntity ne : entities.values()) {
                        final String alias = getAliasFrom(ne.rootName);
                        ne.addAlias(alias);
                }
        }

        private void writeAliasFile(HashMap<String, YagoNamedEntity> entities,
                        String destinationPath, String seedDirectoryPath)
        {
                setEntityAliasValues(entities, seedDirectoryPath);

                BufferedWriter w;
                try {
                        w = openFileAsWriteStream(destinationPath);
                        w.write("");
                        for (final YagoNamedEntity ne : entities.values()) {
                                if (ne.aliasList.size() > 0) {
                                        w.append("+" + ne.rootName + "\n");
                                        for (final String alias : ne.aliasList) {
                                                w.append("." + alias + "\n");
                                        }
                                }
                        }
                } catch (final IOException e) {
                        e.printStackTrace();
                }
        }

        private HashMap<String, YagoNamedEntity> getEntities(
                        String seedDirectoryPath)
                        {
                print("Getting Entities...");
                final HashMap<String, YagoNamedEntity> result = new HashMap<String, YagoNamedEntity>();
                BufferedReader in = null;
                try {
                        in = openFileAsReadStream(seedDirectoryPath + File.separator
                                        + "wordnet_person_100007846.txt");
                } catch (final FileNotFoundException e2) {
                        e2.printStackTrace();
                }
                // get People
                StreamLooper sl = new StreamLooper(in) {
                        @Override
                        protected void doWork(String s) {
                                final String[] values = s.split("\\s+");
                                final String rootName = convertLiteral(values[1]);
                                if (!rootName.startsWith("Category:")) {
                                        final YagoNamedEntity ne = new YagoNamedEntity(rootName,
                                                        NamedEntity.Type.Person);
                                        result.put(rootName, ne);
                                }
                        }
                };
                sl.loop();

                // get Organisations
                try {
                        in = openFileAsReadStream(seedDirectoryPath + File.separator
                                        + "wordnet_organization_108008335.txt");
                } catch (final FileNotFoundException e1) {
                        e1.printStackTrace();
                }
                sl = new StreamLooper(in) {
                        @Override
                        protected void doWork(String s) {
                                final String[] values = s.split("\\s+");
                                final String rootName = convertLiteral(values[1]);
                                if (!(rootName.startsWith("Category:") || rootName
                                                .startsWith("geoent_")))
                                {
                                        final YagoNamedEntity ne = new YagoNamedEntity(rootName,
                                                        NamedEntity.Type.Organisation);
                                        result.put(rootName, ne);
                                }
                        }
                };
                sl.loop();

                if (locations) {
                        // get Locations
                        try {
                                in = openFileAsReadStream(seedDirectoryPath + File.separator
                                                + "wordnet_location_100027167.txt");
                        } catch (final FileNotFoundException e1) {
                                e1.printStackTrace();
                        }
                        sl = new StreamLooper(in) {
                                @Override
                                protected void doWork(String s) {
                                        final String[] values = s.split("\\s+");
                                        final String rootName = convertLiteral(values[1]);
                                        if (!rootName.startsWith("Category:")) {
                                                final YagoNamedEntity ne = new YagoNamedEntity(rootName,
                                                                NamedEntity.Type.Location);
                                                result.put(rootName, ne);
                                        }
                                }
                        };
                        sl.loop();
                }
                print("Total Entities: " + result.size());
                return result;
                        }

        public static BufferedReader openFileAsReadStream(String path)
                        throws FileNotFoundException
        {
                FileReader fr = null;
                fr = new FileReader(path);
                final BufferedReader br = new BufferedReader(fr);
                return br;
        }

        public static BufferedWriter openFileAsWriteStream(String path)
                        throws IOException
        {
                FileWriter fw = null;
                fw = new FileWriter(path);
                final BufferedWriter bw = new BufferedWriter(fw);
                return bw;
        }

        private static String convertLiteral(String literal) {
                final String escaped = StringEscapeUtils.unescapeJava(literal);
                String first = null;
                if (escaped.startsWith("\""))
                        first = escaped.substring(1);
                else
                        first = escaped;
                if (first.endsWith("\""))
                        return first.substring(0, first.length() - 1);
                else
                        return first;
        }

        private static String convertResource(String literal) {
                final String escaped = StringEscapeUtils.unescapeJava(literal);
                return escaped.replaceAll("_", " ");
        }

        private void print(String message) {
                if (verbose)
                        System.out.println(message);
                if (logOut != null) {
                        log(message);
                }
        }

        private void log(String message) {
                try {
                        logOut.append(message + "\n");
                } catch (final IOException e) {
                        e.printStackTrace();
                }
        }

        /**
         * Defualt main.
         *
         * @param args
         *            = path to the seed directory.
         */
        public static void main(String[] args) {
                new EntityExtractionResourceBuilder().buildCandidateAliasFile(args[0]);
        }

        /**
         * Helper class to iterate through the lines of a Reader to do a bit of work
         * on each.
         *
         * @author Laurence Willmore (lgw1e10@ecs.soton.ac.uk)
         *
         */
        public static abstract class StreamLooper {
                BufferedReader reader;

                public StreamLooper(BufferedReader reader) {
                        this.reader = reader;
                }

                /**
                 * Iterates through each line to do the work.
                 */
                public void loop() {
                        String s = null;
                        try {
                                while ((s = reader.readLine()) != null) {
                                        doWork(s);
                                }
                                reader.close();
                        } catch (final IOException e) {
                                e.printStackTrace();
                        }
                }

                /**
                 * Do what you want to each line here.
                 *
                 * @param s
                 */
                protected abstract void doWork(String s);
        }

}