// ==========================================================================
//                 SeqAn - The Library for Sequence Analysis
// ==========================================================================
// Copyright (c) 2006-2013, Knut Reinert, FU Berlin
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:
//
//     * Redistributions of source code must retain the above copyright
//       notice, this list of conditions and the following disclaimer.
//     * Redistributions in binary form must reproduce the above copyright
//       notice, this list of conditions and the following disclaimer in the
//       documentation and/or other materials provided with the distribution.
//     * Neither the name of Knut Reinert or the FU Berlin nor the names of
//       its contributors may be used to endorse or promote products derived
//       from this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
// ARE DISCLAIMED. IN NO EVENT SHALL KNUT REINERT OR THE FU BERLIN BE LIABLE
// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
// LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
// OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH
// DAMAGE.
//
// ==========================================================================
// Author: Tobias Rausch <rausch@embl.de>
// Author: Manuel Holtgrewe <manuel.holtgrewe@fu-berlin.de>
// ==========================================================================
// Shows how to use the bam_io module to estimate the insert size from a SAM
// or BAM file.
// ==========================================================================

#include <iostream>
#include <vector>

#include <seqan/basic.h>
#include <seqan/sequence.h>
#include <seqan/stream.h>
#include <seqan/bam_io.h>

#if SEQAN_HAS_ZLIB

using namespace seqan;

// Stores information about a library.

struct LibraryInfo
{
    // This is the notation introduced by Illumina.
    //
    // F+: >R1 R2>
    // F-: >R2 R1>
    // R+: >R1 R2<
    // R-: <R2 R1>
    enum Orientation
    {
        F_PLUS  = 0,
        F_MINUS = 1,
        R_PLUS  = 2,
        R_MINUS = 3
    };

    unsigned median;
    double   stdDev;
    unsigned maxNormalISize;
    Orientation defaultOrient;

    LibraryInfo() : median(0), stdDev(0), maxNormalISize(0), defaultOrient(F_PLUS)
    {}
};

inline int
getStrandIndependentOrientation(BamAlignmentRecord const& rec)
{
    if (!hasFlagRC(rec)) {
        if (!(hasFlagNextRC(rec))) {
            return (rec.beginPos < rec.pNext) ? LibraryInfo::F_PLUS : LibraryInfo::F_MINUS;
        } else {
            return (rec.beginPos < rec.pNext) ? LibraryInfo::R_PLUS : LibraryInfo::R_MINUS;
        }
    } else {
        if (!hasFlagNextRC(rec)) {
            return (rec.beginPos > rec.pNext) ? LibraryInfo::R_PLUS : LibraryInfo::R_MINUS;
        } else {
            return (rec.beginPos > rec.pNext) ? LibraryInfo::F_PLUS : LibraryInfo::F_MINUS;
        }
    }
}

bool endsWith(CharString const & str, CharString const & x)
{
    typedef Size<CharString>::Type TSize;

    TSize len = std::min(length(str), length(x));
    TSize pos = length(str) - len;
    return suffix(str, pos) == x;
}

template <typename TStreamOrReader, typename TFormatTag>
bool performEstimation(LibraryInfo & libInfo, TStreamOrReader & streamOrReader, TFormatTag const & formatTag)
{
    typedef StringSet<String<char> >   TNameStore;
    typedef NameStoreCache<TNameStore> TNameStoreCache;
    
    // Vector of all insert sizes.
    typedef std::vector<unsigned int> TVecISize;
    TVecISize vecISize;


    // BAM Meta Info.
    TNameStore      nameStore;
    TNameStoreCache nameStoreCache(nameStore);
    BamIOContext<TNameStore> context(nameStore, nameStoreCache);

    // Read Header.
    BamHeader header;
    readRecord(header, context, streamOrReader, formatTag);

    // Orientations.
    unsigned orientationCounters[4] = {0, 0, 0, 0};

    // Running mean.
    __uint64 runningMean = 0;
    unsigned pairCount = 0;

    // Count.
    BamAlignmentRecord record;
    while (!atEnd(streamOrReader))
    {
        if (readRecord(record, context, streamOrReader, formatTag) != 0)
        {
            std::cerr << "Error reading SAM/BAM file.\n";
            return false;
        }

        // Skip all records except for first mate of properly mapped pairs.
        if (!hasFlagMultiple(record) || hasFlagUnmapped(record) || hasFlagNextUnmapped(record) ||
            hasFlagSecondary(record) || hasFlagFirst(record))
            continue;

        vecISize.push_back(abs(record.tLen));
        runningMean += abs(record.tLen);
        pairCount += 1;

        orientationCounters[getStrandIndependentOrientation(record)] += 1;
    }

    // Get default orientation.
    unsigned orientMax = 0;
    for (unsigned i = 0; i < 4; ++i)
        if (orientationCounters[i] > orientationCounters[orientMax])
            orientMax = i;
    LibraryInfo::Orientation orientations[4] = { LibraryInfo::F_PLUS, LibraryInfo::F_MINUS, LibraryInfo::R_PLUS, LibraryInfo::R_MINUS };
    libInfo.defaultOrient = orientations[orientMax];

    // Trim off the chimera peak in mate-pair libraries.
    if (libInfo.defaultOrient == LibraryInfo::R_MINUS && (runningMean / pairCount) >= 1000u)
    {
        typedef typename TVecISize::const_iterator TVecISizeIter;
        TVecISize vecISizeTmp;
        for(TVecISizeIter it = vecISize.begin(); it < vecISize.end(); ++it)
            if (*it > 1000)
                vecISizeTmp.push_back(*it);
        std::swap(vecISize, vecISizeTmp);
    }

    
    // Check that this is a proper paired-end library
    if (vecISize.empty())
        return true;

    // Get library stats.
    //
    // Start with median.
    typedef typename TVecISize::iterator TVecISizeIter;
    TVecISizeIter begin = vecISize.begin();
    TVecISizeIter end = vecISize.end();
    std::nth_element(begin, begin + (end - begin) / 2, end);
    begin = vecISize.begin();
    end = vecISize.end();
    libInfo.median = *(begin + (end - begin) / 2);
    // Standard deviation is next.
    //
    // SD calculation cutoffs are 7 SDs to the left and right assuming 10% library deviation.
    libInfo.stdDev = 0;
    double cutoffMax = libInfo.median + 7 * 0.1 * libInfo.median;
    double cutoffMin = libInfo.median - 7 * 0.1 * libInfo.median;
    if ((cutoffMin < 0) || (cutoffMax < cutoffMin))
        cutoffMin = 0; 
    unsigned int count = 0;
    for(;begin < end; ++begin)
    {
        if ((*begin >= cutoffMin) && (*begin <= cutoffMax))
        {
            libInfo.stdDev += (*begin - libInfo.median) * (*begin - libInfo.median);
            ++count;
        }
    }
    libInfo.stdDev = sqrt(libInfo.stdDev / count);
    libInfo.maxNormalISize = static_cast<unsigned>(libInfo.median + 3 * libInfo.stdDev);

    return true;
}

int main(int argc, char const ** argv)
{
    if (argc != 2)
    {
        std::cerr << "Invalid arguments!\n"
                  << "USAGE: bam_library_size {IN.sam,IN.bam}\n";
        return 1;
    }

    LibraryInfo libInfo;

    // Guess file type from extension.
    if (endsWith(argv[1], ".sam"))
    {
        std::ifstream inSamStream(toCString(argv[1]), std::ios::binary | std::ios::in);
        if (!inSamStream.good())
        {
            std::cerr << "Could not open input SAM file " << argv[1] << "\n";
            return 1;
        }

        RecordReader<std::ifstream, SinglePass<> > reader(inSamStream);

        if (!performEstimation(libInfo, reader, Sam()))
            return 1;
    }
    else
    {
        Stream<Bgzf> inBamStream;
        if (!open(inBamStream, toCString(argv[1]), "r"))
        {
            std::cerr << "Could not open input BAM file " << argv[1] << "\n";
            return 1;
        }

        if (!performEstimation(libInfo, inBamStream, Bam()))
            return 1;
    }

    // Print result.
    std::cout << "Library Information\n\n"
              << "path:                       " << argv[1] << "\n"
              << "median:                     " << libInfo.median << "\n"
              << "standard deviation:         " << libInfo.stdDev << "\n"
              << "maximum normal insert size: " << libInfo.maxNormalISize << "\n";

    std::cout << "orientation:                ";
    switch (libInfo.defaultOrient)
    {
        case LibraryInfo::F_PLUS:
            std::cout << "F+ R1 ---> ---> R2\n";
            break;
        case LibraryInfo::F_MINUS:
            std::cout << "F- R1 ---> ---> R2\n";
            break;
        case LibraryInfo::R_PLUS:
            std::cout << "R+ R1 ---> <--- R2\n";
            break;
        case LibraryInfo::R_MINUS:
            std::cout << "R- R1 <--- ---> R2\n";
            break;
    }

    return 0;
}

#else

int main(int, char const **)
{
    std::cerr << "bam_library_size can only be compiled correctly with zlib." << std::endl;
    return 0;
}

#endif  // #if SEQAN_HAS_ZLIB