<commit_before>﻿using Microsoft.Extensions.Logging;
using Sir.VectorSpace;
using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.IO;
using System.Threading.Tasks;

namespace Sir.Search
{
    /// <summary>
    /// Read session targeting multiple collections.
    /// </summary>
    public class SearchSession : DocumentStreamSession, IDisposable, ISearchSession
    {
        private readonly Database _sessionFactory;
        private readonly IModel _model;
        private readonly ILogger _logger;

        public SearchSession(
            string directory, 
            Database sessionFactory,
            IModel model,
            ILogger logger = null) : base(directory, sessionFactory)
        {
            _sessionFactory = sessionFactory;
            _model = model;
            _logger = logger;
        }

        public SearchResult Search(Query query, int skip, int take)
        {
            var result = Execute(query, skip, take);

            if (result != null)
            {
                var numOfTerms = query.TotalNumberOfTerms();
                var scoreMultiplier = (double)1 / numOfTerms;
                var docs = ReadDocs(result.SortedDocuments, query.Select, scoreMultiplier);

                return new SearchResult(query, result.Total, docs.Count, docs);
            }

            return new SearchResult(query, 0, 0, new Document[0]);
        }

        public Document SearchScalar(Query query)
        {
            var result = Execute(query, 0, 1);

            if (result != null)
            {
                var numOfTerms = query.TotalNumberOfTerms();
                var scoreMultiplier = (double)1 / numOfTerms;
                var docs = ReadDocs(result.SortedDocuments, query.Select, scoreMultiplier);

                return docs.Count > 0 ? docs[0] : null;
            }

            return null;
        }

        private ScoredResult Execute(Query query, int skip, int take)
        {
            var timer = Stopwatch.StartNew();

            // Scan
            Scan(query);
            LogDebug($"scanning took {timer.Elapsed}");
            timer.Restart();

            // Map each query term to a list of document IDs
            PostingsMapper.Map(query, _sessionFactory);
            LogDebug($"mapping took {timer.Elapsed}");
            timer.Restart();

            // Reduce
            IDictionary<(ulong, long), double> scoredResult = new Dictionary<(ulong, long), double>();
            Reducer.Reduce(query, ref scoredResult);
            LogDebug($"reducing took {timer.Elapsed}");
            timer.Restart();

            // Sort
            var sorted = Sort(scoredResult, skip, take);
            LogDebug($"sorting took {timer.Elapsed}");

            return sorted;
        }

        /// <summary>
        /// Score each term and find their posting list locations.
        /// </summary>
        private void Scan(Query query)
        {
            if (query == null)
                return;

            Parallel.ForEach(query.AllTerms(), term =>
            //foreach (var term in query.AllTerms())
            {
                var columnReader = GetColumnReader(term.Directory, term.CollectionId, term.KeyId);

                if (columnReader != null)
                {
                    using (columnReader)
                    {
                        var hit = columnReader.ClosestMatch(term.Vector, _model);

                        if (hit != null)
                        {
                            term.Score = hit.Score;
                            term.PostingsOffsets = hit.Node.PostingsOffsets ?? new List<long> { hit.Node.PostingsOffset };
                        }
                    }
                }
            });
        }

        private static ScoredResult Sort(IDictionary<(ulong, long), double> documents, int skip, int take)
        {
            var sortedByScore = new List<KeyValuePair<(ulong, long), double>>(documents);

            sortedByScore.Sort(
                delegate (KeyValuePair<(ulong, long), double> pair1,
                KeyValuePair<(ulong, long), double> pair2)
                {
                    return pair2.Value.CompareTo(pair1.Value);
                }
            );

            var index = skip > 0 ? skip : 0;
            int count;

            if (take == 0)
                count = sortedByScore.Count - (index);
            else
                count = Math.Min(sortedByScore.Count - (index), take);

            return new ScoredResult 
            { 
                SortedDocuments = sortedByScore.GetRange(index, count), 
                Total = sortedByScore.Count 
            };
        }

        private IList<Document> ReadDocs(
            IEnumerable<KeyValuePair<(ulong collectionId, long docId), double>> docIds, <commit_msg>search single threaded<commit_after>﻿using Microsoft.Extensions.Logging;
using Sir.VectorSpace;
using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.IO;
using System.Threading.Tasks;

namespace Sir.Search
{
    /// <summary>
    /// Read session targeting multiple collections.
    /// </summary>
    public class SearchSession : DocumentStreamSession, IDisposable, ISearchSession
    {
        private readonly Database _sessionFactory;
        private readonly IModel _model;
        private readonly ILogger _logger;

        public SearchSession(
            string directory, 
            Database sessionFactory,
            IModel model,
            ILogger logger = null) : base(directory, sessionFactory)
        {
            _sessionFactory = sessionFactory;
            _model = model;
            _logger = logger;
        }

        public SearchResult Search(Query query, int skip, int take)
        {
            var result = Execute(query, skip, take);

            if (result != null)
            {
                var numOfTerms = query.TotalNumberOfTerms();
                var scoreMultiplier = (double)1 / numOfTerms;
                var docs = ReadDocs(result.SortedDocuments, query.Select, scoreMultiplier);

                return new SearchResult(query, result.Total, docs.Count, docs);
            }

            return new SearchResult(query, 0, 0, new Document[0]);
        }

        public Document SearchScalar(Query query)
        {
            var result = Execute(query, 0, 1);

            if (result != null)
            {
                var numOfTerms = query.TotalNumberOfTerms();
                var scoreMultiplier = (double)1 / numOfTerms;
                var docs = ReadDocs(result.SortedDocuments, query.Select, scoreMultiplier);

                return docs.Count > 0 ? docs[0] : null;
            }

            return null;
        }

        private ScoredResult Execute(Query query, int skip, int take)
        {
            var timer = Stopwatch.StartNew();

            // Scan
            Scan(query);
            LogDebug($"scanning took {timer.Elapsed}");
            timer.Restart();

            // Map each query term to a list of document IDs
            PostingsMapper.Map(query, _sessionFactory);
            LogDebug($"mapping took {timer.Elapsed}");
            timer.Restart();

            // Reduce
            IDictionary<(ulong, long), double> scoredResult = new Dictionary<(ulong, long), double>();
            Reducer.Reduce(query, ref scoredResult);
            LogDebug($"reducing took {timer.Elapsed}");
            timer.Restart();

            // Sort
            var sorted = Sort(scoredResult, skip, take);
            LogDebug($"sorting took {timer.Elapsed}");

            return sorted;
        }

        /// <summary>
        /// Find the angle in vector space between the term and the closest matching node and record its posting list address.
        /// </summary>
        private void Scan(Query query)
        {
            if (query == null)
                return;

            var readers = new Dictionary<(string, ulong, long), IColumnReader>();

            try
            {
                foreach (var term in query.AllTerms())
                {
                    IColumnReader reader;
                    var key = (term.Directory, term.CollectionId, term.KeyId);

                    if (!readers.TryGetValue(key, out reader))
                    {
                        reader = GetColumnReader(term.Directory, term.CollectionId, term.KeyId);

                        if (reader != null)
                        {
                            readers.Add(key, reader);
                        }
                    }

                    if (reader != null)
                    {
                        var hit = reader.ClosestMatch(term.Vector, _model);

                        if (hit != null)
                        {
                            term.Score = hit.Score;
                            term.PostingsOffsets = hit.Node.PostingsOffsets ?? new List<long> { hit.Node.PostingsOffset };
                        }
                    }
                }
            }
            finally
            {
                foreach (var reader in readers.Values)
                {
                    reader.Dispose();
                }
            }
        }

        private static ScoredResult Sort(IDictionary<(ulong, long), double> documents, int skip, int take)
        {
            var sortedByScore = new List<KeyValuePair<(ulong, long), double>>(documents);

            sortedByScore.Sort(
                delegate (KeyValuePair<(ulong, long), double> pair1,
                KeyValuePair<(ulong, long), double> pair2)
                {
                    return pair2.Value.CompareTo(pair1.Value);
                }
            );

            var index = skip > 0 ? skip : 0;
            int count;

            if (take == 0)
                count = sortedByScore.Count - (index);
            else
                count = Math.Min(sortedByScore.Count - (index), take);

            return new ScoredResult 
            { 
                SortedDocuments = sortedByScore.GetRange(index, count), 
                Total = sortedByScore.Count 
            };
        }

        private IList<Document> ReadDocs(
            IEnumerable<KeyValuePair<(ulong collectionId, long docId), double>> docIds, 
{'fraction_non_alphanumeric': 0.09141301176699407, 'fraction_numerical': 0.0031119323154721384, 'mean_word_length': 1.4497379704621247, 'pattern_counts': {'":': 0, '<': 34, '<?xml version=': 0, '>': 39, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3616542', 'n_tokens_mistral': 2649, 'n_tokens_neox': 2485, 'n_words': 653}