package htsjdk.samtools;

import htsjdk.samtools.seekablestream.SeekablePathStream;
import htsjdk.samtools.seekablestream.SeekableStream;
import htsjdk.samtools.util.RuntimeIOException;

import java.io.File;
import java.io.IOException;
import java.nio.file.Path;
import java.util.*;

/**
 * Implementation of the CSI index for BAM files.
 * The CSI index extends the BAI index by allowing a more flexible
 * binning scheme, with variable depth (number of levels) and
 * bin sizes, thus allowing for genomic regions longer than 2^29-1.
 */
public class CSIIndex extends AbstractBAMFileIndex implements BrowseableBAMIndex {

    /**
     * Don't add default values to these variables, as they will override the
     * value assign to them by the {@link #initParameters()} method, called from
     * the superclass constructor.
     */
    private int binDepth;
    private int minShift;
    private int maxBins;
    private int maxSpan;
    private byte[] auxData;
    private int nReferences;
    private long metaDataPos;

    /**
     * Constructors
     */

    public CSIIndex(final SeekableStream stream, final SAMSequenceDictionary dictionary) {
        this(IndexFileBufferFactory.getBuffer(stream), stream.getSource(), dictionary);
    }

    public CSIIndex(final Path path, final SAMSequenceDictionary dictionary) throws IOException {
        this(new SeekablePathStream(path), dictionary);
    }

    public CSIIndex(final File file, boolean enableMemoryMapping, final SAMSequenceDictionary dictionary) {
        this(IndexFileBufferFactory.getBuffer(file, enableMemoryMapping), file.getName(), dictionary);
    }

    private CSIIndex(final IndexFileBuffer indexFileBuffer, final String source, final SAMSequenceDictionary dictionary) {
        super(indexFileBuffer, source, dictionary);
    }

    /**
     * Getters and setters
     */

    /**
     * Bin depth is the number of levels of the index. By default,
     * BAI has 6 levels. CSI makes this variable.
     */
    public int getBinDepth() {
        return binDepth;
    }

    private void setBinDepth(int binDepth) { this.binDepth = binDepth; }

    /**
     * 2^(min shift) is the smallest width of a bin
     */
    public int getMinShift() {
        return minShift;
    }

    private void setMinShift(int minShift) {
        this.minShift = minShift;
    }

    public int getMaxBins() {
        return maxBins;
    }

    private void setMaxBins(int binDepth) { this.maxBins = ((1<<3*binDepth) - 1)/7; }

    public int getMaxSpan() {
        return maxSpan;
    }

    private void setMaxSpan(int binDepth, int minShift) {
        this.maxSpan = 1<<(minShift + 3*(binDepth - 1));
    }

    public byte[] getAuxData() { return auxData; }

    private void setAuxData(byte[] auxData) { this.auxData = auxData; }

    @Override
    public int getNumberOfReferences() { return nReferences; }

    private void setNumberOfReferences(int nReferences) { this.nReferences = nReferences; }

    /**
     * Computes the number of bins on the given level.
     * @param levelNumber Level for which to compute the size.
     * @return
     */
    @Override
    public int getLevelSize(final int levelNumber) {
        if (levelNumber >= getBinDepth()) {
            throw new SAMException("Level number (" + levelNumber + ") is greater than or equal to maximum (" + getBinDepth() + ").");
        }
        return 1<<3*(levelNumber);
    }

    /**
     * Extends the functionality of {@link AbstractBAMFileIndex#getFirstBinInLevel(int)} ,
     * which cannot be overridden due to its static nature.
     */
    public int getFirstBinInLevelForCSI(final int levelNumber) {
        if (levelNumber >= getBinDepth()) {
            throw new SAMException("Level number (" + levelNumber + ") is greater than or equal to maximum (" + getBinDepth() + ").");
        }
        return ((1<<3*levelNumber) - 1)/7;
    }

    @Override
    public int getLevelForBin(Bin bin) {
        if(bin == null || bin.getBinNumber() > getMaxBins()) {
            throw new SAMException("Tried to get level for invalid bin: " +  bin);
        }
        for (int i = getBinDepth()-1; i > -1 ; i--) {
              if (bin.getBinNumber() >= getFirstBinInLevelForCSI(i)) {
                  return i;
              }
        }
        throw new SAMException("Unable to find correct level for bin: " + bin);
    }

    @Override
    public int getFirstLocusInBin(Bin bin) {
        if(bin == null || bin.getBinNumber() > getMaxBins()) {
            throw new SAMException("Tried to get first locus for invalid bin: " + bin);
        }
        int level = getLevelForBin(bin);
        int firstBinOnLevel = getFirstBinInLevelForCSI(level);
        int levelSize = getLevelSize(level);

        return (bin.getBinNumber() - firstBinOnLevel)*(getMaxSpan()/levelSize) + 1;
    }

    @Override
    public int getLastLocusInBin(Bin bin) {
        if(bin == null || bin.getBinNumber() > getMaxBins()) {
            throw new SAMException("Tried to get last locus for invalid bin: " + bin);
        }
        int level = getLevelForBin(bin);
        int firstBinOnLevel = getFirstBinInLevelForCSI(level);
        int levelSize = getLevelSize(level);

        return (bin.getBinNumber() - firstBinOnLevel + 1)*(getMaxSpan()/levelSize);
    }

    @Override
    public BinList getBinsOverlapping(int referenceIndex, int startPos, int endPos) {
        final BitSet regionBins = GenomicIndexUtil.regionToBins(startPos, endPos, getMinShift(), getBinDepth());
        if (regionBins == null) {
            return null;
        }
        return new BinList(referenceIndex,regionBins);
    }

    @Override
    public BAMFileSpan getSpanOverlapping(int referenceIndex, int startPos, int endPos) {
        final BAMIndexContent queryResults = query(referenceIndex, startPos, endPos);
        int initialBinNumber = getFirstBinInLevelForCSI(getBinDepth() - 1) + (startPos - 1 >> getMinShift());
        long minimumOffset = 0L;
        Bin targetBin;

        if(queryResults == null) {
            return null;
        }

        /** Compute 'minimumOffset' by searching the lowest level bin containing 'startPos'.
            If the computed bin is not in the index, try the next bin to the left, belonging
            to the same parent. If it is the first sibling bin, try the parent bin.
         */

        do {
            int firstBinNumber;
            targetBin = queryResults.getBins().getBin(initialBinNumber);
            if (targetBin != null) {
                break;
            }
            firstBinNumber = (getParentBinNumber(initialBinNumber)<<3) + 1;
            if (initialBinNumber > firstBinNumber) {
                initialBinNumber--;
            } else {
                initialBinNumber = getParentBinNumber(initialBinNumber);
            }
        } while(initialBinNumber != 0);

        if (initialBinNumber == 0) {
            targetBin = queryResults.getBins().getBin(initialBinNumber);
        }

        if (targetBin != null && targetBin instanceof BinWithOffset) {
            minimumOffset = ((BinWithOffset) targetBin).getlOffset();
        }

        List<Chunk> chunkList = new ArrayList<Chunk>();
        for(final Chunk chunk: queryResults.getAllChunks()) {
            chunkList.add(chunk.clone());
        }

        chunkList = Chunk.optimizeChunkList(chunkList, minimumOffset);
        return new BAMFileSpan(chunkList);
    }

    @Override
    public BAMFileSpan getSpanOverlapping(final Bin bin) {
        if(bin == null) {
            return null;
        }

        final int referenceSequence = bin.getReferenceSequence();
        final BAMIndexContent queryResults = getQueryResults(referenceSequence);

        if(queryResults == null) {
            return null;
        }

        final int binLevel = getLevelForBin(bin);
        final int firstLocusInBin = getFirstLocusInBin(bin);
        long minimumOffset = bin instanceof BinWithOffset ? ((BinWithOffset)bin).getlOffset() : 0L;

        // Add the specified bin to the tree if it exists.
        final List<Bin> binTree = new ArrayList<Bin>();
        if(queryResults.containsBin(bin)) {
            binTree.add(queryResults.getBins().getBin(bin.getBinNumber()));
        }

        int currentBinLevel = binLevel;
        while(--currentBinLevel >= 0) {
            final int binStart = getFirstBinInLevelForCSI(currentBinLevel);
            final int binWidth = getMaxSpan()/getLevelSize(currentBinLevel);
            final int parentBinNumber = firstLocusInBin/binWidth + binStart;
            final Bin parentBin = queryResults.getBins().getBin(parentBinNumber);
            if(parentBin != null && queryResults.containsBin(parentBin)) {
                binTree.add(parentBin);
            }
        }

        List<Chunk> chunkList = new ArrayList<Chunk>();
        for(final Bin coveringBin: binTree) {
            for(final Chunk chunk: coveringBin.getChunkList())
                chunkList.add(chunk.clone());
        }

        chunkList = Chunk.optimizeChunkList(chunkList, minimumOffset);
        return new BAMFileSpan(chunkList);
    }

    @Override
    public long getStartOfLastLinearBin() {
        if (metaDataPos > 0 && position() != metaDataPos) {
            seek(metaDataPos);
        }

        final int sequenceIndex = getNumberOfReferences();
        long loffset = -1L;
        for (int i = 0; i < sequenceIndex; i++) {

            final int nBins = readInteger(); // n_bin

            for (int j = 0; j < nBins; j++) {
                readInteger(); // bin
                loffset = readLong(); // loffset
                final int nChunks = readInteger(); // n_chunk
                skipBytes(BAMFileConstants.CSI_CHUNK_SIZE * nChunks);
            }
        }

        return loffset;
    }

    @Override
    protected void verifyIndexMagicNumber(final String sourceName) {
        // Verify the magic number.
        if (BAMFileConstants.CSI_MAGIC_OFFSET != position()) {
            seek(BAMFileConstants.CSI_MAGIC_OFFSET);
        }
        final byte[] buffer = new byte[BAMFileConstants.CSI_MINSHIFT_OFFSET];
        readBytes(buffer); // magic
        if (!Arrays.equals(buffer, BAMFileConstants.CSI_INDEX_MAGIC)) {
            throw new RuntimeIOException("Invalid file header in BAM CSI index " + sourceName +
                    ": " + new String(buffer));
        }
    }

    private void readMinShiftAndBinDepth() {
        if (BAMFileConstants.CSI_MINSHIFT_OFFSET != position()) {
            seek(BAMFileConstants.CSI_MINSHIFT_OFFSET);
        }
        setMinShift(readInteger()); // min_shift
        setBinDepth(readInteger() + 1); // depth - HTSlib doesn't count the first level (bin 0)
        setMaxBins(binDepth);
        setMaxSpan(binDepth, minShift);
    }

    private void readAuxDataAndNRef() {
        if (BAMFileConstants.CSI_AUXDATA_OFFSET != position()) {
            seek(BAMFileConstants.CSI_AUXDATA_OFFSET);
        }
        //set the aux data length first
        byte[] auxData = new byte[readInteger()]; // l_aux
        readBytes(auxData); // aux
        setAuxData(auxData);
        setNumberOfReferences(readInteger()); // n_ref
        metaDataPos = position(); // save the metadata position for delayed reading
    }

    @Override
    protected final void initParameters() {
        readMinShiftAndBinDepth();
        readAuxDataAndNRef();
        setSequenceIndexes(getNumberOfReferences());
    }

    public int getParentBinNumber(int binNumber) {
        if (binNumber >= getMaxBins()) {
            throw new SAMException("Tried to get parent bin for invalid bin (" + binNumber + ").");
        }
        if (binNumber == 0) {
            return 0;
        }
        return (binNumber - 1) >> 3;
    }

    public int getParentBinNumber(Bin bin) {
        if (bin == null) {
            throw new SAMException("Tried to get parent bin for null bin.");
        }
        return getParentBinNumber(bin.getBinNumber());
    }

    /**
     * The maximum possible bin number for this reference sequence.
     * This is based on the maximum coordinate position of the reference
     * which is based on the size of the reference
     */
    private int getMaxBinNumberForReference(final int reference) {
        try {
            final int sequenceLength = getBamDictionary().getSequence(reference).getSequenceLength();
            return getFirstBinInLevelForCSI(getBinDepth() - 1) + (sequenceLength >> getMinShift());
        } catch (final Exception e) {
            return getMaxBins();
        }
    }

    @Override
    protected BAMIndexContent query(final int referenceSequence, final int startPos, final int endPos) {
        if (metaDataPos > 0 && position() != metaDataPos) {
            seek(metaDataPos);
        }

        final List<Chunk> metaDataChunks = new ArrayList<Chunk>();

        final int sequenceCount = getNumberOfReferences();

        if (referenceSequence >= sequenceCount) {
            return null;
        }

        final BitSet regionBins = GenomicIndexUtil.regionToBins(startPos, endPos, getMinShift(), getBinDepth());
        if (regionBins == null) {
            return null;
        }

        skipToSequence(referenceSequence);

        final int binCount = readInteger(); // n_bin
        boolean metaDataSeen = false;
        final Bin[] bins = new BinWithOffset[getMaxBinNumberForReference(referenceSequence) +1];
        for (int binNumber = 0; binNumber < binCount; binNumber++) {
            final int indexBin = readInteger(); // bin
            final long lOffset = readLong(); // l_offset
            final int nChunks = readInteger();  // n_chunk
            List<Chunk> chunks;

            Chunk lastChunk = null;
            if (regionBins.get(indexBin)) {
                chunks = new ArrayList<Chunk>(nChunks);
                readChunks(nChunks, chunks);
            } else if (indexBin == getMaxBins() + 1) {
                // meta data - build the bin so that the count of bins is correct;
                // but don't attach meta chunks to the bin, or normal queries will be off
                readChunks(nChunks, metaDataChunks);
                metaDataSeen = true;
                continue; // don't create a Bin
            } else {
                skipBytes(BAMFileConstants.CSI_CHUNK_SIZE * nChunks);
                chunks = Collections.emptyList();
            }
            final BinWithOffset bin = new BinWithOffset(referenceSequence, indexBin, lOffset);
            bin.setChunkList(chunks);
            bin.setLastChunk(lastChunk);
            bins[indexBin] = bin;
        }

        return new BAMIndexContent(referenceSequence, bins, binCount - (metaDataSeen? 1 : 0), new BAMIndexMetaData(metaDataChunks), null);
    }

    /**
     * Return meta data for the given reference including information about number of aligned, unaligned, and noCoordinate records
     *
     * @param reference the reference of interest
     * @return meta data for the reference
     */
    @Override
    public BAMIndexMetaData getMetaData(final int reference) {
        if (metaDataPos > 0 && position() != metaDataPos) {
            seek(metaDataPos);
        }

        final List<Chunk> metaDataChunks = new ArrayList<Chunk>();

        final int sequenceCount = getNumberOfReferences();

        if (reference >= sequenceCount) {
            return null;
        }

        skipToSequence(reference);

        final int binCount = readInteger(); // n_bin
        for (int binNumber = 0; binNumber < binCount; binNumber++) {
            final int indexBin = readInteger(); // bin
            final long lOffset = readLong(); // loffset
            final int nChunks = readInteger(); // n_chunk
            if (indexBin == getMaxBins() + 1) {
                readChunks(nChunks, metaDataChunks);
            } else {
                skipBytes(BAMFileConstants.CSI_CHUNK_SIZE * nChunks);
            }
        }
        return new BAMIndexMetaData(metaDataChunks);
    }

    /**
     * Returns count of records unassociated with any reference. Call before the index file is closed
     *
     * @return meta data at the end of the bam index that indicates count of records holding no coordinates
     * or null if no meta data (old index format)
     */
    @Override
    public Long getNoCoordinateCount() {
        if (metaDataPos > 0 && position() != metaDataPos) {
            seek(metaDataPos);
        }

        skipToSequence(getNumberOfReferences());
        try { // in case of old index file without meta data
            return readLong();
        } catch (final Exception e) {
            return null;
        }
    }

    @Override
    public BAMIndexContent getQueryResults(final int referenceSequence) {
        return query(referenceSequence, 1, -1);
    }

    @Override
    protected void skipToSequence(final int sequenceIndex) {
        if(sequenceIndex > getNumberOfReferences()) {
            throw new SAMException("Sequence index (" + sequenceIndex + ") is greater than maximum (" + getNumberOfReferences() + ").");
        }

        //Use sequence position cache if available
        if(sequenceIndexes[sequenceIndex] != -1){
            seek(sequenceIndexes[sequenceIndex]);
            return;
        }

        if (metaDataPos > 0 && position() != metaDataPos) {
            seek(metaDataPos);
        }
        for (int i = 0; i < sequenceIndex; i++) {

            final int nBins = readInteger(); // n_bin

            for (int j = 0; j < nBins; j++) {
                readInteger(); // bin
                readLong(); // loffset
                final int nChunks = readInteger(); // n_chunk
                skipBytes(BAMFileConstants.CSI_CHUNK_SIZE * nChunks);
            }
        }

        //Update sequence position cache
        sequenceIndexes[sequenceIndex] = position();
    }
}
