/*
 * BWTBinarySearch.h
 *
 *  Created on: 17. 02. 2021
 *      Author: Martin Prokopic
 */

#pragma once

#include <sstream>
#include <alib/set>
#include <string/LinearString.h>
#include <string/LinearStringTerminatingSymbolBWT.h>
#include <stringology/indexing/BWTAuxiliaryDataConstruct.h>
#include <indexes/stringology/BWTAuxiliaryData.h>
#include <alib/measure>

namespace stringology {

namespace exact {

class BWTBinarySearch {
public:
    /**
     * Use binary search to find occurrences of pattern in subject (a BWT encoded string).
     *
     * @tparam SymbolType
     * @param subject the input BWT-encoded string
     * @param pattern the pattern to search for
     * @return indices of occurrences of pattern in the original (decoded) text, starting at 0
     */
    template < class SymbolType >
    static ext::set < unsigned > match ( const string::LinearStringTerminatingSymbolBWT < SymbolType > & subject,
                                         const string::LinearString < SymbolType > & pattern ) {
        return match ( subject, pattern,
                       stringology::indexing::BWTAuxiliaryDataConstruct::construct ( subject,
                                                                                     indexes::stringology::BWTAuxiliaryDataType::M_OCC_PREFIX_SUM |
                                                                                     indexes::stringology::BWTAuxiliaryDataType::W_FORWARD_DECODE |
                                                                                     indexes::stringology::BWTAuxiliaryDataType::I_FIRST_TO_TEXT ));
    }

    /**
     * Use binary search to find occurrences of pattern in subject (a BWT encoded string), reusing precomputed auxiliary data.
     * Required data types are M_OCC_PREFIX_SUM, W_FORWARD_DECODE and I_FIRST_TO_TEXT.
     *
     * @tparam SymbolType
     * @param subject the input BWT-encoded string
     * @param pattern the pattern to search for
     * @param data precomputed auxiliary data
     * @return indices of occurrences of pattern in the original (decoded) text, starting at 0
     */
    template < class SymbolType >
    static ext::set < unsigned > match ( const string::LinearStringTerminatingSymbolBWT < SymbolType > & subject,
                                         const string::LinearString < SymbolType > & pattern,
                                         const indexes::stringology::BWTAuxiliaryData < SymbolType > & data );
};

template < class SymbolType >
ext::set < unsigned > BWTBinarySearch::match ( const string::LinearStringTerminatingSymbolBWT < SymbolType > & subject,
                                               const string::LinearString < SymbolType > & pattern,
                                               const indexes::stringology::BWTAuxiliaryData < SymbolType > & data ) {
    if ( ! data.isDataTypePresent ( indexes::stringology::BWTAuxiliaryDataType::M_OCC_PREFIX_SUM ) ||
         ! data.isDataTypePresent ( indexes::stringology::BWTAuxiliaryDataType::W_FORWARD_DECODE ) ||
         ! data.isDataTypePresent ( indexes::stringology::BWTAuxiliaryDataType::I_FIRST_TO_TEXT ))
        throw exception::CommonException ( "BWT Binary Search requires M, W and I arrays to be present." );

    if ( pattern.getAlphabet ( ).contains ( subject.getTerminatingSymbol ( ))) {
        std::ostringstream msg;
        msg << "Pattern cannot contain subject's terminating symbol (" << subject.getTerminatingSymbol ( ) << ").";
        throw exception::CommonException ( msg.str ( ));
    }
    if ( pattern.empty ( ))
        throw exception::CommonException ( "Pattern cannot be empty." );

    alphabet::SymbolCompare < SymbolType > symbolCompare;
    auto binarySearchCompare = [ & ] ( size_t patternIndex, size_t subjectIndex ) {
        subjectIndex = data.getWForwardDecode ( )[ subjectIndex ]; // wrap around from last character to first character of shift
        while ( patternIndex < pattern.size ( ) &&
                subject.getContent ( )[ subjectIndex ] == pattern.getContent ( )[ patternIndex ] ) {
            subjectIndex = data.getWForwardDecode ( )[ subjectIndex ];
            patternIndex ++;
        }
        if ( patternIndex == pattern.size ( ))
            return 0;
        if ( symbolCompare ( pattern.getContent ( )[ patternIndex ], subject.getContent ( )[ subjectIndex ] ))
            return - 1;
        return 1;
    };

    measurements::start ( "BWTBinarySearch", measurements::Type::ALGORITHM );

    size_t low, high, lowerBound = 0, upperBound = subject.size ( );

    {
        // optimization: we can limit initial binary search range using M and K arrays
        const auto & firstSymbol = pattern.getContent ( )[ 0 ];
        auto it = data.getMOccPrefixSum ( ).find ( firstSymbol );
        if ( it == data.getMOccPrefixSum ( ).end ( ))
            lowerBound = subject.size ( ); // first character of pattern is not present in subject's alphabet
        else {
            lowerBound = it->second;
            upperBound = lowerBound + data.getKOcc ( ).at ( firstSymbol );
        }
    }

    // lower bound - first index greater or equal to pattern
    low = lowerBound, high = upperBound;
    while ( low < high ) {
        size_t mid = ( low + high ) / 2;
        if ( binarySearchCompare ( 0, mid ) > 0 ) // pattern > current
            low = mid + 1;
        else
            high = mid;
    }
    lowerBound = low;

    if ( lowerBound == upperBound || binarySearchCompare ( 0, lowerBound ) != 0 ) {
        measurements::end ( );

        return ext::set < unsigned int > ( ); // no matches
    }

    // upper bound - first index greater than pattern
    low = lowerBound, high = upperBound;
    while ( low < high ) {
        size_t mid = ( low + high ) / 2;
        if ( binarySearchCompare ( 0, mid ) >= 0 ) // pattern >= current
            low = mid + 1;
        else
            high = mid;
    }
    upperBound = low;

    ext::set < unsigned int > matches;
    for ( size_t i = lowerBound; i < upperBound; i ++ )
        matches.insert ( data.getIFirstToText ( )[ i ] );

    measurements::end ( );

    return matches;
}

} /* namespace exact */

} /* namespace stringology */
