aswi2020anonymous-gitlab/be/fulltextsearch/src/main/java/cz/zcu/kiv/aswi/fulltextsearch/SolrService.java @ 6d9b0887

package cz.zcu.kiv.aswi.fulltextsearch;

import cz.zcu.kiv.aswi.fulltextsearch.document.*;
import cz.zcu.kiv.aswi.fulltextsearch.model.SearchResponse;
import javaxt.io.Image;
import org.apache.solr.client.solrj.SolrQuery;
import org.apache.solr.client.solrj.SolrServerException;
import org.apache.solr.client.solrj.impl.HttpSolrClient;
import org.apache.solr.client.solrj.request.schema.SchemaRequest;
import org.apache.solr.client.solrj.response.QueryResponse;
import org.apache.solr.client.solrj.response.schema.SchemaResponse;
import org.apache.solr.common.SolrDocument;
import org.apache.solr.common.SolrDocumentList;
import org.apache.solr.common.SolrInputDocument;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.http.MediaType;
import org.springframework.web.multipart.MultipartFile;

import javax.imageio.ImageIO;
import javax.xml.bind.JAXBException;
import java.awt.image.BufferedImage;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.util.*;
import java.util.stream.Collectors;

public class SolrService {

    // for docker
    private static final String URL = "http://fulltextsearch-solr";

    // for localhost
    //private static final String URL = "http://localhost";

    private static final String CORE_NAME_DOCUMENTS = "documents";
    private static final String CORE_NAME_FILES = "files";

    private static final String PREFIX_TEXT_LINE = "text_line_";
    private static final String PREFIX_TEXT_WORD = "text_word_";
    private static final String SUFFIX_COORDS = "_coords";

    private static final String FIELD_DOC_FILENAME = "doc_filename";
    private static final String FIELD_IMG_FILENAME = "img_filename";
    private static final String FIELD_TEXT_REGION = "text_region";
    private static final String FIELD_TEXT_REGION_COORDS = "text_region" + SUFFIX_COORDS;
    private static final String FIELD_DOC_CONTENT = "doc_content";
    private static final String FIELD_IMG_CONTENT = "img_content";
    private static final float IMG_OUTPUT_QUALITY = 0.5f;
    private static final Logger LOG = LoggerFactory.getLogger(IndexController.class);

    private static final int SEARCH_QUERY_ROWS = 100;

    private HttpSolrClient solrClientDocuments;
    private HttpSolrClient solrClientFiles;
    private List<String> fieldNamesDocumentCore = new ArrayList<>();

    public SolrService() {
        connect();
    }

    private void connect() {
        LOG.info("Connecting to Solr");
        String urlString = URL + ":8983/solr/" + CORE_NAME_DOCUMENTS;
        solrClientDocuments = new HttpSolrClient.Builder(urlString).build();

        urlString = URL + ":8983/solr/" + CORE_NAME_FILES;
        solrClientFiles = new HttpSolrClient.Builder(urlString).build();
    }

    public int uploadFiles(List<MultipartFile> files) throws IOException, SolrServerException, JAXBException {
        LOG.trace("SolrService method uploadFiles called");
        Map<String, List<MultipartFile>> pairs = new HashMap<>();

        LOG.debug("Preprocessing files for upload to Solr");
        for (MultipartFile file: files) {
            String filename = file.getOriginalFilename();
            if (filename == null) {
                LOG.error("Skipping file with empty filename");
                continue;
            }

            String type = file.getContentType();
            if (type == null) {
                LOG.error("Skipping file " + filename + " (unknown type)");
                continue;
            }

            int index = filename.lastIndexOf('.');
            String filenameWithoutExtension = filename;
            if (index >= 0) {
                filenameWithoutExtension = filename.substring(0, index);
            }

            if (type.equals(MediaType.IMAGE_JPEG_VALUE) || type.equals(MediaType.IMAGE_PNG_VALUE)
                    || type.equals(MediaType.APPLICATION_XML_VALUE) || type.equals(MediaType.TEXT_XML_VALUE)) {

                List<MultipartFile> list = pairs.get(filenameWithoutExtension);
                if (list == null) {
                    list = new LinkedList<>();
                    list.add(file);
                    pairs.put(filenameWithoutExtension, list);
                } else {
                    list.add(file);
                }
                LOG.debug("File " + filename + " added");
            } else {
                LOG.error("Not supported type of file! File " + filename + " is type " + type);
            }
        }
        return addFiles(pairs);
    }

    private int addFiles(Map<String, List<MultipartFile>> pairs) throws IOException, SolrServerException, JAXBException {
        LOG.trace("SolrService method addFiles called");
        LOG.debug("Uploading files to Solr");
        getFieldNamesFromDocumentCore();
        addBasicFieldsToDocumentsCore();

        addAllFieldsToFilesCore();

        int count = 0;
        Set<String> keys = pairs.keySet();
        for (String key: keys) {
            List<MultipartFile> pair = pairs.get(key);
            if (pair.size() != 2) continue;

            MultipartFile first = pair.get(0);
            if (first.getContentType().equals(MediaType.APPLICATION_XML_VALUE) || first.getContentType().equals(MediaType.TEXT_XML_VALUE)) {
                addToFilesCore(first, pair.get(1));
            } else {
                addToFilesCore(pair.get(1), first);
            }
            count += 2;
        }

        LOG.debug("Commiting Solr Documents and Files core changes");
        solrClientFiles.commit();
        solrClientDocuments.commit();

        LOG.info("Files uploaded");
        return count;
    }

    private void addToFilesCore(MultipartFile xmlFile, MultipartFile imageFile) throws IOException, SolrServerException, JAXBException {
        LOG.trace("SolrService method addToFilesCore called");
        SolrInputDocument solrInputDocument = new SolrInputDocument();

        LOG.debug("Adding a file pair to File core");
        solrInputDocument.addField(FIELD_DOC_FILENAME, xmlFile.getOriginalFilename());
        solrInputDocument.addField(FIELD_IMG_FILENAME, imageFile.getOriginalFilename());
        solrInputDocument.addField(FIELD_DOC_CONTENT, xmlFile.getBytes());
        solrInputDocument.addField(FIELD_IMG_CONTENT, Base64.getEncoder().encode(imageFile.getBytes()));

        solrClientFiles.add(solrInputDocument);

        LOG.debug("Loading XML file");
        XMLLoader xmlLoader = new XMLLoader();
        PcGts doc = xmlLoader.loadFile(xmlFile.getInputStream(), xmlFile.getOriginalFilename());
        addToDocumentsCore(doc, imageFile.getOriginalFilename());
    }

    private void addAllFieldsToFilesCore() throws IOException, SolrServerException {
        LOG.trace("SolrService method addAllFieldsToFilesCore called");
        List<String> fieldNamesFilesCore = new ArrayList<>();

        SchemaRequest.Fields fieldRequest = new SchemaRequest.Fields();
        SchemaResponse.FieldsResponse fieldsResponse = fieldRequest.process(solrClientFiles);
        List<Map<String, Object>> fields = fieldsResponse.getFields();

        for (Map<String, Object> field: fields) {
            fieldNamesFilesCore.add(field.get("name").toString());
        }

        String[] filenames = new String[]{ FIELD_DOC_FILENAME, FIELD_IMG_FILENAME };
        String[] contents = new String[]{ FIELD_DOC_CONTENT, FIELD_IMG_CONTENT };

        LOG.debug("Adding all fields to File core");
        for (String name: filenames) {
            if (fieldNamesFilesCore.contains(name)) continue;

            try {
                SchemaRequest.Update request = getAddTextFieldRequest(name);
                request.process(solrClientFiles);
            } catch (Exception e) {
                LOG.error("Couldn't add field = " + name);
                LOG.error(e.toString());
            }
        }

        for (String name: contents) {
            if (fieldNamesFilesCore.contains(name)) continue;

            try {
                Map<String, Object> fieldAttributes = new LinkedHashMap<>();
                fieldAttributes.put("name", name);
                fieldAttributes.put("type", "binary");
                fieldAttributes.put("indexed", false);
                fieldAttributes.put("stored", true);
                fieldAttributes.put("multiValued", false);
                SchemaRequest.Update request = new SchemaRequest.AddField(fieldAttributes);
                request.process(solrClientFiles);
            } catch (Exception e) {
                LOG.error("Couldn't add field = " + name);
                LOG.error(e.toString());
            }
        }
    }

    private void addBasicFieldsToDocumentsCore() {
        LOG.trace("SolrService method addBasicFieldsToDocumentsCore called");
        LOG.debug("Adding basic fields to Documents core");
        addFieldToDocumentCore(FIELD_DOC_FILENAME);
        addFieldToDocumentCore(FIELD_IMG_FILENAME);
        addFieldToDocumentCore(FIELD_TEXT_REGION);
        addFieldToDocumentCore(FIELD_TEXT_REGION_COORDS);
    }

    private void getFieldNamesFromDocumentCore() throws IOException, SolrServerException {
        LOG.trace("SolrService method getFieldNamesFromDocumentCore called");
        LOG.debug("Getting field names from Documents core");
        SchemaRequest.Fields fieldRequest = new SchemaRequest.Fields();
        SchemaResponse.FieldsResponse fieldsResponse = fieldRequest.process(solrClientDocuments);
        List<Map<String, Object>> fields = fieldsResponse.getFields();

        for (Map<String, Object> field: fields) {
            fieldNamesDocumentCore.add(field.get("name").toString());
        }
    }


    private void addFieldToDocumentCore(String name) {
        LOG.trace("SolrService method addFieldToDocumentCore called");

        if (fieldNamesDocumentCore.contains(name)) return;
        LOG.debug("Adding field to Documents core");
        try {
            SchemaRequest.Update request = getAddTextFieldRequest(name);
            request.process(solrClientDocuments);
            fieldNamesDocumentCore.add(name);
        } catch (Exception e) {
            LOG.error("Couldn't add field = " + name);
            LOG.error(e.toString());
        }
    }

    private SchemaRequest.Update getAddTextFieldRequest(String name) {
        LOG.trace("SolrService method getAddTextFieldRequest called");
        Map<String, Object> fieldAttributes = new LinkedHashMap<>();
        fieldAttributes.put("name", name);
        fieldAttributes.put("type", "text_general");
        fieldAttributes.put("stored", true);
        fieldAttributes.put("multiValued", false);
        return new SchemaRequest.AddField(fieldAttributes);
    }

    public void deleteAll() throws IOException, SolrServerException {
        LOG.trace("SolrService method getAddTextFieldRequest called");
        LOG.info("Deleting ALL from Documents core");
        solrClientDocuments.deleteByQuery("*:*");
        solrClientDocuments.commit();

        LOG.info("Deleting ALL from Files core");
        solrClientFiles.deleteByQuery("*:*");
        solrClientFiles.commit();
    }

    public byte[] getImageFromFilesCore(String filename) throws IOException, SolrServerException {
        LOG.trace("SolrService method getImageFromFilesCore called");
        LOG.debug("Getting an image from Files core");
        SolrQuery solrQuery = new SolrQuery();
        solrQuery.set("q", FIELD_IMG_FILENAME + ":" + "\"" + filename + "\"");
        SolrDocumentList results = solrClientFiles.query(solrQuery).getResults();
        if (results.getNumFound() < 1) {
            LOG.error("Image " + filename + " not found!");
        } else if (results.getNumFound() > 1) {
            LOG.error("Image " + filename + " has multiple files. Using first one...");
        }

        SolrDocument result = results.get(0);
        return (byte[]) result.getFieldValue(FIELD_IMG_CONTENT);
    }

    public SearchResponse query(String query) throws IOException, SolrServerException  {
        LOG.trace("SolrService method query called");
        LOG.debug("Processing query - creating Solr query");
        SolrQuery solrQuery = new SolrQuery();
        solrQuery.set("q", FIELD_TEXT_REGION + ":" + query);
        solrQuery.setHighlight(true);
        solrQuery.addHighlightField(PREFIX_TEXT_LINE +  "*");
        solrQuery.addHighlightField(PREFIX_TEXT_WORD +  "*");
        solrQuery.setRows(SEARCH_QUERY_ROWS);
        LOG.debug("Processing query - sending Solr query");
        QueryResponse response = solrClientDocuments.query(solrQuery);

        SolrDocumentList docList = response.getResults();
        Map<String, Map<String, List<String>>> highlight = response.getHighlighting();

        SearchResponse searchResponse = new SearchResponse();
        searchResponse.setExpression(query);

        LOG.debug("Processing query - creating Search response");
        List<TextRegion> textRegions = new LinkedList<>();
        for (SolrDocument solrDocument: docList) {
            TextRegion textRegion = new TextRegion();

            // text, coords
            textRegion.setRegionText((String) solrDocument.getFieldValue(FIELD_TEXT_REGION));
            String pointsString = (String) solrDocument.getFieldValue(FIELD_TEXT_REGION_COORDS);
            Point[] points = Coords.parsePointString(pointsString);
            textRegion.setRegionCoords(new Coords(points));

            // filename
            String documentName = ((String) solrDocument.getFieldValue(FIELD_DOC_FILENAME)).replaceAll(".xml", "");
            textRegion.setDocumentName(documentName);

            // text lines and text words
            Map<String, List<String>> map = highlight.get(solrDocument.getFieldValue("id").toString());
            Set<String> linesAndWordsNames = map.keySet();

            Set<String> linesNames = linesAndWordsNames.stream().filter(name -> name.startsWith(PREFIX_TEXT_LINE)).collect(Collectors.toSet());
            Set<String> wordNames = linesAndWordsNames.stream().filter(name -> name.startsWith(PREFIX_TEXT_WORD)).collect(Collectors.toSet());

            Map<String, TextLine> textLines = new HashMap<>();
            for (String lineName: linesNames) {
                TextLine textLine = new TextLine();

                textLine.setLineText((String) solrDocument.getFieldValue(lineName));
                String coordsValue = (String) solrDocument.getFieldValue(lineName + SUFFIX_COORDS);
                Point[] linePoints = Coords.parsePointString(coordsValue);

                //-----
                for (Point p: linePoints){
                    p.setX(p.getX() - textRegion.getRegionCoords().getPoints()[0].getX());
                    p.setY(p.getY() - textRegion.getRegionCoords().getPoints()[0].getY());
                }
                //-----

                textLine.setLineCoords(new Coords(linePoints));

                textLines.put(lineName, textLine);
            }

            Map<String, List<TextWord>> textWords = new HashMap<>();
            for (String wordName: wordNames) {
                TextWord textWord = new TextWord();

                textWord.setWordText((String) solrDocument.getFieldValue(wordName));
                String coordsValue = (String) solrDocument.getFieldValue(wordName + SUFFIX_COORDS);
                Point[] wordPoints = Coords.parsePointString(coordsValue);

                //-----
                for (Point p: wordPoints){
                    p.setX(p.getX() - textRegion.getRegionCoords().getPoints()[0].getX());
                    p.setY(p.getY() - textRegion.getRegionCoords().getPoints()[0].getY());
                }
                //-----

                textWord.setWordCoords(new Coords(wordPoints));

                String lineName = wordName.substring(wordName.indexOf(PREFIX_TEXT_LINE));

                List<TextWord> lineWords = textWords.get(lineName);
                if (lineWords == null) {
                    lineWords = new LinkedList<>();
                }

                lineWords.add(textWord);
                textWords.put(lineName, lineWords);
            }

            for (String lineName: textWords.keySet()) {
                TextLine line = textLines.get(lineName);
                line.setTextWords(textWords.get(lineName).toArray(new TextWord[0]));
            }

            textRegion.setTextLines(textLines.values().toArray(new TextLine[0]));

            // cropping image
            String imgFilename = (String) solrDocument.getFieldValue(FIELD_IMG_FILENAME);
            byte[] imageBytes = getImageFromFilesCore(imgFilename);

            BufferedImage img = ImageIO.read(new ByteArrayInputStream(Base64.getDecoder().decode(imageBytes)));
            Image cropImg = new Image(img);
            cropImg.crop(textRegion.getRegionCoords().getPoints()[0].getX(),textRegion.getRegionCoords().getPoints()[0].getY(),
                    textRegion.getRegionCoords().getPoints()[1].getX()-textRegion.getRegionCoords().getPoints()[0].getX(),
                    textRegion.getRegionCoords().getPoints()[2].getY()-textRegion.getRegionCoords().getPoints()[0].getY());

            cropImg.setOutputQuality(IMG_OUTPUT_QUALITY);

            imageBytes = cropImg.getByteArray();
            imageBytes = Base64.getEncoder().encode(imageBytes);

            // recalculating Coords - during init for words and lines
            for (int k = 1; k<textRegion.getRegionCoords().getPoints().length; k++){
                textRegion.getRegionCoords().getPoints()[k].setX(textRegion.getRegionCoords().getPoints()[k].getX() - textRegion.getRegionCoords().getPoints()[0].getX());
                textRegion.getRegionCoords().getPoints()[k].setY(textRegion.getRegionCoords().getPoints()[k].getY() - textRegion.getRegionCoords().getPoints()[0].getY());
            }
            textRegion.getRegionCoords().getPoints()[0].setX(0);
            textRegion.getRegionCoords().getPoints()[0].setY(0);

            textRegion.setImageCut(new String(imageBytes));

            // image size
            BufferedImage image = ImageIO.read(new ByteArrayInputStream(Base64.getDecoder().decode(imageBytes)));
            DocumentSize documentSize = new DocumentSize(image.getHeight(), image.getWidth());
            textRegion.setDocumentSize(documentSize);

            // generate random string id for response of given length
            textRegion.setRandomId(Utils.generateRandomStringId(10));

            textRegions.add(textRegion);
        }

        searchResponse.setTextRegions(textRegions.toArray(new TextRegion[0]));
        LOG.debug("Processing query - sending Search response");
        LOG.info("Query received and processed");
        return searchResponse;
    }

    private void addToDocumentsCore(PcGts document, String imageFilename) throws IOException, SolrServerException {
        LOG.trace("SolrService method addToDocumentsCore called");
        LOG.debug("Adding to Documents core");
        List<TextRegion> regions = document.getPage().getTextRegions();
        if (regions == null || regions.size() == 0) return;

        for (int i = 0; i < regions.size(); i++) {
            TextRegion textRegion = regions.get(i);
            SolrInputDocument solrInputDocument = new SolrInputDocument();

            solrInputDocument.addField(FIELD_DOC_FILENAME, document.getFilename());
            solrInputDocument.addField(FIELD_IMG_FILENAME, imageFilename);
            solrInputDocument.addField(FIELD_TEXT_REGION, textRegion.getTextEquiv().getUnicode());
            solrInputDocument.addField(FIELD_TEXT_REGION_COORDS, textRegion.getRegionCoords().getPointsString());

            TextLine[] textLines = textRegion.getTextLines();
            if (textLines == null || textLines.length == 0) continue;
            for (int j = 0; j < textLines.length; j++) {
                TextLine textLine = textLines[j];
                String prefix = PREFIX_TEXT_LINE + j;

                addFieldToDocumentCore(prefix);
                solrInputDocument.addField(prefix, textLine.getTextEquiv().getUnicode());

                addFieldToDocumentCore(prefix + SUFFIX_COORDS);
                solrInputDocument.addField(prefix + SUFFIX_COORDS, textLine.getLineCoords().getPointsString());

                TextWord[] textWords = textLine.getTextWords();
                if (textWords == null || textWords.length == 0) continue;
                for (int k = 0; k < textWords.length; k++) {
                    TextWord textWord = textWords[k];
                    String wordPrefix = PREFIX_TEXT_WORD + k + "_" + prefix;

                    addFieldToDocumentCore(wordPrefix);
                    solrInputDocument.addField(wordPrefix, textWord.getTextEquiv().getUnicode());

                    addFieldToDocumentCore(wordPrefix + SUFFIX_COORDS);
                    solrInputDocument.addField(wordPrefix + SUFFIX_COORDS, textWord.getWordCoords().getPointsString());
                }
            }
            solrClientDocuments.add(solrInputDocument);
        }
    }

    public String info() throws IOException, SolrServerException {
        LOG.trace("SolrService method info called");
        LOG.debug("Displaying information about files in Solr");
        SolrQuery q = new SolrQuery("*:*");
        q.setRows(0);  // don't actually request any data
        return "Number of documents in " + CORE_NAME_DOCUMENTS + " core: " + solrClientDocuments.query(q).getResults().getNumFound() +
                "\nNumber of documents in " + CORE_NAME_FILES + " core: " + solrClientFiles.query(q).getResults().getNumFound();
    }

    public List<String> listAllFiles() throws IOException, SolrServerException {
        LOG.trace("SolrService method listAllFiles called");
        LOG.debug("Displaying ALL files in Solr");
        SolrQuery solrQuery = new SolrQuery();
        solrQuery.set("q", FIELD_DOC_FILENAME + ":*");
        solrQuery.setRows(500); // todo not a magic number!
        SolrDocumentList docList = solrClientFiles.query(solrQuery).getResults();
        return docList.stream().map(doc -> doc.getFieldValue(FIELD_DOC_FILENAME).toString().replaceAll(".xml", "")).collect(Collectors.toList());
    }

    public String listSingleFile(String documentName) throws IOException, SolrServerException {
        LOG.trace("SolrService method listSingleFile called");
        LOG.debug("Displaying single file in Solr");
        SolrQuery solrQuery = new SolrQuery();
        solrQuery.set("q", FIELD_DOC_FILENAME + ":*");
        //solrQuery.set("q", FIELD_DOC_FILENAME + ":" + documentName + ".xml");
        solrQuery.setRows(500); // todo not a magic number!
        SolrDocumentList docList = solrClientFiles.query(solrQuery).getResults();
        String ret = "";

        for (SolrDocument solrDocument : docList){
            if (solrDocument.getFieldValue(FIELD_DOC_FILENAME).toString().equals(documentName + ".xml")){
                String imgFilename = (String) solrDocument.getFieldValue(FIELD_IMG_FILENAME);
                byte[] imageBytes = getImageFromFilesCore(imgFilename);

                BufferedImage img = ImageIO.read(new ByteArrayInputStream(Base64.getDecoder().decode(imageBytes)));
                Image cropImg = new Image(img);
                cropImg.setOutputQuality(IMG_OUTPUT_QUALITY);

                imageBytes = cropImg.getByteArray();
                imageBytes = Base64.getEncoder().encode(imageBytes);
                ret = new String(imageBytes);
                break;
            }
        }
        if (ret.equals("")){
            LOG.error("File was NOT found in Solr");
        } else{
            LOG.debug("File was found in Solr");
        }
        return ret;
    }

}