// ***************************************************************************
// BamRandomAccessController_p.cpp (c) 2011 Derek Barnett
// Marth Lab, Department of Biology, Boston College
// ---------------------------------------------------------------------------
// Last modified: 25 October 2011(DB)
// ---------------------------------------------------------------------------
// Manages random access operations in a BAM file
// **************************************************************************

#include "api/internal/bam/BamRandomAccessController_p.h"
#include "api/BamIndex.h"
#include "api/internal/bam/BamReader_p.h"
#include "api/internal/index/BamIndexFactory_p.h"
#include "api/internal/utils/BamException_p.h"
using namespace BamTools;
using namespace BamTools::Internal;

#include <cassert>
#include <sstream>

BamRandomAccessController::BamRandomAccessController()
    : m_index(0)
    , m_hasAlignmentsInRegion(true)
{}

BamRandomAccessController::~BamRandomAccessController()
{
    Close();
}

void BamRandomAccessController::AdjustRegion(const int& referenceCount)
{

    // skip if no index available
    if (m_index == 0) {
        return;
    }

    // see if any references in region have alignments
    m_hasAlignmentsInRegion = false;
    int currentId = m_region.LeftRefID;
    const int rightBoundRefId =
        (m_region.isRightBoundSpecified() ? m_region.RightRefID : referenceCount - 1);
    while (currentId <= rightBoundRefId) {
        m_hasAlignmentsInRegion = m_index->HasAlignments(currentId);
        if (m_hasAlignmentsInRegion) {
            break;
        }
        ++currentId;
    }

    // if no data found on any reference in region
    if (!m_hasAlignmentsInRegion) {
        return;
    }

    // if left bound of desired region had no data, use first reference that had data
    // otherwise, leave requested region as-is
    if (currentId != m_region.LeftRefID) {
        m_region.LeftRefID = currentId;
        m_region.LeftPosition = 0;
    }
}

// returns alignments' "RegionState": { Before|Overlaps|After } current region
BamRandomAccessController::RegionState BamRandomAccessController::AlignmentState(
    const BamAlignment& alignment) const
{

    // if region has no left bound at all
    if (!m_region.isLeftBoundSpecified()) {
        return OverlapsRegion;
    }

    // handle unmapped reads - return AFTER region to halt processing
    if (alignment.RefID == -1) {
        return AfterRegion;
    }

    // if alignment is on any reference before left bound reference
    if (alignment.RefID < m_region.LeftRefID) {
        return BeforeRegion;

        // if alignment is on left bound reference
    } else if (alignment.RefID == m_region.LeftRefID) {

        // if alignment starts at or after left bound position
        if (alignment.Position >= m_region.LeftPosition) {

            if (m_region.isRightBoundSpecified() &&  // right bound is specified AND
                m_region.LeftRefID ==
                    m_region.RightRefID &&  // left & right bounds on same reference AND
                alignment.Position >=
                    m_region.RightPosition) {  // alignment starts on or after right bound position
                return AfterRegion;

                // otherwise, alignment overlaps region
            } else {
                return OverlapsRegion;
            }
        }

        // alignment starts before left bound position
        else {

            // if alignment overlaps left bound position
            if (alignment.GetEndPosition() > m_region.LeftPosition) {
                return OverlapsRegion;
            } else {
                return BeforeRegion;
            }
        }
    }

    // otherwise alignment is on a reference after left bound reference
    else {

        // if region has a right bound
        if (m_region.isRightBoundSpecified()) {

            // alignment is on any reference between boundaries
            if (alignment.RefID < m_region.RightRefID) {
                return OverlapsRegion;

                // alignment is on any reference after right boundary
            } else if (alignment.RefID > m_region.RightRefID) {
                return AfterRegion;

                // alignment is on right bound reference
            } else {

                // if alignment starts before right bound position
                if (alignment.Position < m_region.RightPosition) {
                    return OverlapsRegion;
                } else {
                    return AfterRegion;
                }
            }
        }

        // otherwise, alignment starts after left bound and there is no right bound given
        else {
            return OverlapsRegion;
        }
    }
}

void BamRandomAccessController::Close()
{
    ClearIndex();
    ClearRegion();
}

void BamRandomAccessController::ClearIndex()
{
    if (m_index) {
        delete m_index;
        m_index = 0;
    }
}

void BamRandomAccessController::ClearRegion()
{
    m_region.clear();
    m_hasAlignmentsInRegion = true;
}

bool BamRandomAccessController::CreateIndex(BamReaderPrivate* reader,
                                            const BamIndex::IndexType& type)
{
    // skip if reader is invalid
    assert(reader);
    if (!reader->IsOpen()) {
        SetErrorString("BamRandomAccessController::CreateIndex",
                       "cannot create index for unopened reader");
        return false;
    }

    // create new index of requested type
    BamIndex* newIndex = BamIndexFactory::CreateIndexOfType(type, reader);
    if (newIndex == 0) {
        std::stringstream s;
        s << "could not create index of type: " << type;
        SetErrorString("BamRandomAccessController::CreateIndex", s.str());
        return false;
    }

    // attempt to build index from current BamReader file
    if (!newIndex->Create()) {
        const std::string indexError = newIndex->GetErrorString();
        const std::string message = "could not create index: \n\t" + indexError;
        SetErrorString("BamRandomAccessController::CreateIndex", message);
        return false;
    }

    // save new index & return success
    SetIndex(newIndex);
    return true;
}

std::string BamRandomAccessController::GetErrorString() const
{
    return m_errorString;
}

bool BamRandomAccessController::HasIndex() const
{
    return (m_index != 0);
}

bool BamRandomAccessController::HasRegion() const
{
    return (!m_region.isNull());
}

bool BamRandomAccessController::IndexHasAlignmentsForReference(const int& refId)
{
    return m_index->HasAlignments(refId);
}

bool BamRandomAccessController::LocateIndex(BamReaderPrivate* reader,
                                            const BamIndex::IndexType& preferredType)
{
    // look up index filename, deferring to preferredType if possible
    assert(reader);
    const std::string& indexFilename =
        BamIndexFactory::FindIndexFilename(reader->Filename(), preferredType);

    // if no index file found (of any type)
    if (indexFilename.empty()) {
        const std::string message =
            std::string("could not find index file for:") + reader->Filename();
        SetErrorString("BamRandomAccessController::LocateIndex", message);
        return false;
    }

    // otherwise open & use index file that was found
    return OpenIndex(indexFilename, reader);
}

bool BamRandomAccessController::OpenIndex(const std::string& indexFilename,
                                          BamReaderPrivate* reader)
{

    // attempt create new index of type based on filename
    BamIndex* index = BamIndexFactory::CreateIndexFromFilename(indexFilename, reader);
    if (index == 0) {
        const std::string message = std::string("could not open index file: ") + indexFilename;
        SetErrorString("BamRandomAccessController::OpenIndex", message);
        return false;
    }

    // attempt to load data from index file
    if (!index->Load(indexFilename)) {
        const std::string indexError = index->GetErrorString();
        const std::string message = std::string("could not load index data from file: ") +
                                    indexFilename + "\n\t" + indexError;
        SetErrorString("BamRandomAccessController::OpenIndex", message);
        return false;
    }

    // save new index & return success
    SetIndex(index);
    return true;
}

bool BamRandomAccessController::RegionHasAlignments() const
{
    return m_hasAlignmentsInRegion;
}

void BamRandomAccessController::SetErrorString(const std::string& where, const std::string& what)
{
    m_errorString = where + ": " + what;
}

void BamRandomAccessController::SetIndex(BamIndex* index)
{
    if (m_index) {
        ClearIndex();
    }
    m_index = index;
}

bool BamRandomAccessController::SetRegion(const BamRegion& region, const int& referenceCount)
{

    // store region
    m_region = region;

    // cannot jump when no index is available
    if (!HasIndex()) {
        SetErrorString("BamRandomAccessController", "cannot jump if no index data available");
        return false;
    }

    // adjust region as necessary to reflect where data actually begins
    AdjustRegion(referenceCount);

    // if no data present, return true
    //   * Not an error, but future attempts to access alignments in this region will not return data
    //     Returning true is useful in a BamMultiReader setting where some BAM files may
    //     lack alignments in regions where other files still have data available.
    if (!m_hasAlignmentsInRegion) {
        return true;
    }

    // return success/failure of jump to specified region,
    //
    //  * Index::Jump() is allowed to modify the m_hasAlignmentsInRegion flag
    //    This covers 'corner case' where a region is requested that lies beyond the last
    //    alignment on a reference. If this occurs, any subsequent calls to GetNextAlignment[Core]
    //    will not return data. BamMultiReader will still be able to successfully pull alignments
    //    from a region from other files even if this one has no data.
    if (!m_index->Jump(m_region, &m_hasAlignmentsInRegion)) {
        const std::string indexError = m_index->GetErrorString();
        const std::string message = std::string("could not set region\n\t") + indexError;
        SetErrorString("BamRandomAccessController::OpenIndex", message);
        return false;
    } else {
        return true;
    }
}
