package uk.ac.ebi.quickgo.annotation.service.statistics;

import uk.ac.ebi.quickgo.annotation.model.*;
import uk.ac.ebi.quickgo.rest.search.AggregateSearchQueryTemplate;
import uk.ac.ebi.quickgo.rest.search.BasicSearchQueryTemplate;
import uk.ac.ebi.quickgo.rest.search.RetrievalException;
import uk.ac.ebi.quickgo.rest.search.SearchService;
import uk.ac.ebi.quickgo.rest.search.query.QueryRequest;
import uk.ac.ebi.quickgo.rest.search.query.QuickGOQuery;
import uk.ac.ebi.quickgo.rest.search.request.converter.ConvertedFilter;
import uk.ac.ebi.quickgo.rest.search.request.converter.FilterConverterFactory;
import uk.ac.ebi.quickgo.rest.search.results.AggregateResponse;
import uk.ac.ebi.quickgo.rest.search.results.AggregationBucket;
import uk.ac.ebi.quickgo.rest.search.results.AggregationResult;
import uk.ac.ebi.quickgo.rest.search.results.QueryResult;

import java.util.*;
import java.util.stream.Collectors;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import static uk.ac.ebi.quickgo.rest.search.AggregateFunction.UNIQUE;

/**
 * Service that collects distribution statistics of annotations and gene products throughout a given set of annotation
 * fields.
 *
 * @author Ricardo Antunes
 */
@Service
public class AnnotationStatisticsService implements StatisticsService {
    private static final long NO_COUNT_FOR_GROUP_ERROR = -1L;

    private static final int FIRST_PAGE = 1;
    private static final int RESULTS_PER_PAGE = 0;

    private final Logger logger = LoggerFactory.getLogger(this.getClass());

    private final FilterConverterFactory converterFactory;
    private final SearchService<Annotation> searchService;
    private final StatsRequestConverter converter;

    @Autowired
    public AnnotationStatisticsService(FilterConverterFactory converterFactory,
            SearchService<Annotation> searchService,
            StatsRequestConverter converter) {
        this.converterFactory = converterFactory;
        this.searchService = searchService;
        this.converter = converter;
    }

    @Override public QueryResult<StatisticsGroup> calculate(AnnotationRequest request) {
        List<AnnotationRequest.StatsRequest> statsRequest = request.createStatsRequests();

        QueryRequest queryRequest = buildQueryRequest(request);

        QueryResult<Annotation> annotationQueryResult = searchService.findByQuery(queryRequest);

        AggregateResponse globalAggregation = annotationQueryResult.getAggregation();

        List<StatisticsGroup> statsGroups = statsRequest.stream()
                .map(req -> convertResponse(globalAggregation, req))
                .collect(Collectors.toList());

        return new QueryResult.Builder<>(statsGroups.size(), statsGroups).build();
    }

    private QueryRequest buildQueryRequest(AnnotationRequest request) {
        BasicSearchQueryTemplate basicTemplate = new BasicSearchQueryTemplate(Collections.emptyList());
        AggregateSearchQueryTemplate aggregateTemplate = new AggregateSearchQueryTemplate();

        BasicSearchQueryTemplate.Builder basicBuilder = basicTemplate.newBuilder()
                .setQuery(QuickGOQuery.createAllQuery())
                .setFilters(request.createFilterRequests().stream()
                        .map(converterFactory::convert)
                        .map(ConvertedFilter::getConvertedValue)
                        .collect(Collectors.toSet()))
                .setPage(FIRST_PAGE)
                .setPageSize(RESULTS_PER_PAGE);

        return aggregateTemplate.newBuilder(basicBuilder)
                .setAggregate(converter.convert(request.createStatsRequests()))
                .build();
    }

    private StatisticsGroup convertResponse(AggregateResponse globalAggregation, AnnotationRequest.StatsRequest statsRequest) {
        StatisticsConverter converter =
                new StatisticsConverter(statsRequest.getGroupName(), statsRequest.getGroupField());

        long totalHits = extractCount(globalAggregation, statsRequest.getGroupField());

        if (totalHits == NO_COUNT_FOR_GROUP_ERROR) {
            throw new RetrievalException("Unable to calculate statistics for group: " + statsRequest.getGroupName());
        }

        return converter.convert(globalAggregation.getNestedAggregations(), totalHits);
    }

    /**
     * Extracts the counts made on the whole data set for a given group.
     *
     * @see AnnotationRequest.StatsRequest#getGroupName()
     *
     * @param globalAggregation the aggregation object containing the group count values
     * @param groupField the name of the groupField the count was made upon
     * @return an object containing the global counts of things that are of interest
     */
    private long extractCount(AggregateResponse globalAggregation, String groupField) {
        return globalAggregation.getAggregationResult(UNIQUE, groupField)
                .map(agg -> (long) agg.getResult()).orElse(NO_COUNT_FOR_GROUP_ERROR);
    }

    /**
     * Converts a collection of {@link AggregateResponse} data retrieved from a {@link QueryResult}, into a
     * {@link StatisticsGroup}, that can be presented to the client.
     * </p>
     * This class is capable of creating a single {@link StatisticsGroup} per call to the
     * {@link StatisticsConverter#convert(Collection, long)} method.
     */
    private class StatisticsConverter {
        private final String groupField;
        private final String groupName;

        StatisticsConverter(String groupName, String groupField) {
            this.groupName = groupName;
            this.groupField = groupField;
        }

        StatisticsGroup convert(Collection<AggregateResponse> aggregations, long totalHits) {
            StatisticsGroup statsGroup = new StatisticsGroup(groupName, totalHits);

            aggregations.stream()
                    .map(agg -> createStatsType(agg, totalHits))
                    .forEach(statsGroup::addStatsType);

            return statsGroup;
        }

        private StatisticsByType createStatsType(AggregateResponse aggregation, long totalHits) {
            StatisticsByType type = new StatisticsByType(aggregation.getName());

            Set<AggregationBucket> buckets = aggregation.getBuckets();

            buckets.stream()
                    .map(bucket -> createStatsValue(bucket, totalHits))
                    .forEach(value -> {
                        if (value.isPresent()) {
                            type.addValue(value.get());
                        } else {
                            logger.warn("No stats for groupField {}, for type {}", groupField, type.getType());
                        }
                    });

            return type;
        }

        private Optional<StatisticsValue> createStatsValue(AggregationBucket bucket, long totalHits) {
            Optional<AggregationResult> resultOpt = bucket.getAggregationResult(UNIQUE, groupField);

            return resultOpt.
                    map(aggResult -> new StatisticsValue(bucket.getValue(), (long) aggResult.getResult(), totalHits));
        }
    }
}
@@ -137,80 +137,15 @@
 rch.
-AggregateSearchQueryTemplate;%0Aimport uk.ac.ebi.quickgo.rest.search.Basic
+Default
 Sear
@@ -270,16 +270,70 @@
 ervice;%0A
+import uk.ac.ebi.quickgo.rest.search.SearchableField;%0A
 import u
@@ -857,17 +857,98 @@
 va.util.
-*
+Collection;%0Aimport java.util.List;%0Aimport java.util.Optional;%0Aimport java.util.Set
 ;%0Aimport
@@ -1872,16 +1872,77 @@
 erter;%0A%0A
+    private final DefaultSearchQueryTemplate queryTemplate;%0A%0A
     @Aut
@@ -2124,16 +2124,62 @@
 onverter
+,%0A            SearchableField searchableFields
 ) %7B%0A    
@@ -2303,16 +2303,91 @@
 nverter;
+%0A%0A        queryTemplate = new DefaultSearchQueryTemplate(searchableFields);
 %0A    %7D%0A%0A
@@ -3174,259 +3174,20 @@
     
-BasicSearchQueryTemplate basicTemplate = new BasicSearchQueryTemplate(Collections.emptyList());%0A        AggregateSearchQueryTemplate aggregateTemplate = new AggregateSearchQueryTemplate();%0A%0A        BasicSearchQueryTemplate.Builder basicBuilder = basic
+return query
 Temp
@@ -3278,19 +3278,19 @@
        .
-set
+add
 Filters(
@@ -3591,68 +3591,8 @@
 AGE)
-;%0A%0A        return aggregateTemplate.newBuilder(basicBuilder)
 %0A   
@@ -3781,16 +3781,28 @@
 egation,
+%0A           
  Annotat

Modify the stats service to use the DefaultSearchQueryTemplate instead of the combination of BasicSearchQueryTemplate and AggregationQueryTemplate.