package com.github.sparkcaller;

import htsjdk.samtools.*;
import picard.sam.MergeSamFiles;
import scala.Tuple2;

import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;

public class SAMFileUtils {

    public static File mergeBAMFiles(List<File> samFiles, String outputFileName) {
        MergeSamFiles samMerger = new MergeSamFiles();
        ArrayList<String> args = new ArrayList<>();

        for (File samFile : samFiles) {
            args.add("I=");
            args.add(samFile.getPath());
        }

        args.add("O=");
        args.add(outputFileName);

        samMerger.instanceMain(args.toArray(new String[0]));

        return new File(outputFileName);
    }

    public static int getSAMNumberOfLines(SamReader samReader) {
        SAMRecordIterator samIterator = samReader.iterator();

        int samRecordSize = 0;
        // Get the size of the SAM record.
        while(samIterator.hasNext()) {
            SAMRecord record = samIterator.next();
            samRecordSize++;
        }
        samIterator.close();

        return samRecordSize;
    }

    public static List<File> splitBAMFile(File bamFile, int numSplits) throws IOException {
        List<File> outputFiles = new ArrayList<>();

        SamReader samFile = SamReaderFactory.makeDefault().open(bamFile);

        int SAMNumRecords = getSAMNumberOfLines(samFile);

        SAMFileWriterFactory samWriterFactory =  new SAMFileWriterFactory();
        SAMRecordIterator samIterator = samFile.iterator();

        // We rather want to have fewer records per file and more files, as this makes it easier to utilize all nodes.
        int recordsPerFile = SAMNumRecords / numSplits;

        int currNumRecords = 0;
        int currFileNum = 0;

        SAMFileHeader samHeader = samFile.getFileHeader();

        String splitBaseFilename = Utils.removeExtenstion(bamFile.getPath(), "bam");
        Tuple2<SAMFileWriter, File> currSplitTuple = getSplitSAMWriter(splitBaseFilename, currFileNum, true,
                                                                       samHeader, samWriterFactory);
        SAMFileWriter currSplitSAMWriter = currSplitTuple._1;
        outputFiles.add(currSplitTuple._2);

        while(samIterator.hasNext()) {
            if (currNumRecords == recordsPerFile) {
                currSplitSAMWriter.close();

                currNumRecords = 0;
                currFileNum++;
                currSplitTuple = getSplitSAMWriter(splitBaseFilename, currFileNum, true, samHeader, samWriterFactory);

                currSplitSAMWriter = currSplitTuple._1;
                outputFiles.add(currSplitTuple._2);
            }

            currSplitSAMWriter.addAlignment(samIterator.next());
            currNumRecords++;
        }

        currSplitSAMWriter.close();
        samIterator.close();
        samFile.close();

        return outputFiles;
    }

    public static List<Tuple2<String, File>> splitBAMByChromosome(File bamFile) throws IOException {
        SamReader samFile = SamReaderFactory.makeDefault().open(bamFile);
        SAMFileHeader samFileHeader = samFile.getFileHeader();
        HashMap<String, SAMFileWriter> contigMapper = new HashMap<>();
        SAMFileWriterFactory samWriterFactory =  new SAMFileWriterFactory();

        SAMRecordIterator samIterator = samFile.iterator();
        String prevContig = "NONE";
        SAMFileWriter currWriter = null;

        List<Tuple2<String, File>> outputFiles = new ArrayList<>();
        while(samIterator.hasNext()) {
            SAMRecord record = samIterator.next();
            String currContig = record.getContig();

            if (currContig == null) {
                currContig = "unmapped";
            }

            // When the input is sorted, do not perform a loop up for every record.
            if (!prevContig.equals(currContig)) {
                currWriter = contigMapper.get(currContig);
            }

            // Create a new writer if a new contig was found.
            if (currWriter == null) {
                String newSAMFilename = Utils.removeExtenstion(bamFile.getPath(), "bam") + "-" + currContig + ".bam";
                File newSAMFile = new File(newSAMFilename);
                SAMFileWriter newWriter = samWriterFactory.makeSAMOrBAMWriter(samFileHeader, true, newSAMFile);
                newWriter.addAlignment(record);
                contigMapper.put(currContig, newWriter);

                currWriter = newWriter;
                outputFiles.add(new Tuple2<>(currContig, newSAMFile));
            } else {
                currWriter.addAlignment(record);
            }

            prevContig = currContig;
        }
        contigMapper.values().forEach(SAMFileWriter::close);

        return outputFiles;
    }

    public static Tuple2<SAMFileWriter, File> getSplitSAMWriter(String baseFilename, int id, boolean isSorted,
                                                                SAMFileHeader SAMheader,
                                                                SAMFileWriterFactory samFileWriterFactory) {

        File newSplitFile = new File(baseFilename + "-split-" + Integer.toString(id) + ".bam");
        SAMFileWriter currSAMWriter = samFileWriterFactory.makeSAMOrBAMWriter(SAMheader, isSorted, newSplitFile);
        return new Tuple2<>(currSAMWriter, newSplitFile);

    }
}

@@ -747,2201 +747,8 @@
  %7D%0A%0A
-    public static int getSAMNumberOfLines(SamReader samReader) %7B%0A        SAMRecordIterator samIterator = samReader.iterator();%0A%0A        int samRecordSize = 0;%0A        // Get the size of the SAM record.%0A        while(samIterator.hasNext()) %7B%0A            SAMRecord record = samIterator.next();%0A            samRecordSize++;%0A        %7D%0A        samIterator.close();%0A%0A        return samRecordSize;%0A    %7D%0A%0A    public static List%3CFile%3E splitBAMFile(File bamFile, int numSplits) throws IOException %7B%0A        List%3CFile%3E outputFiles = new ArrayList%3C%3E();%0A%0A        SamReader samFile = SamReaderFactory.makeDefault().open(bamFile);%0A%0A        int SAMNumRecords = getSAMNumberOfLines(samFile);%0A%0A        SAMFileWriterFactory samWriterFactory =  new SAMFileWriterFactory();%0A        SAMRecordIterator samIterator = samFile.iterator();%0A%0A        // We rather want to have fewer records per file and more files, as this makes it easier to utilize all nodes.%0A        int recordsPerFile = SAMNumRecords / numSplits;%0A%0A        int currNumRecords = 0;%0A        int currFileNum = 0;%0A%0A        SAMFileHeader samHeader = samFile.getFileHeader();%0A%0A        String splitBaseFilename = Utils.removeExtenstion(bamFile.getPath(), %22bam%22);%0A        Tuple2%3CSAMFileWriter, File%3E currSplitTuple = getSplitSAMWriter(splitBaseFilename, currFileNum, true,%0A                                                                       samHeader, samWriterFactory);%0A        SAMFileWriter currSplitSAMWriter = currSplitTuple._1;%0A        outputFiles.add(currSplitTuple._2);%0A%0A        while(samIterator.hasNext()) %7B%0A            if (currNumRecords == recordsPerFile) %7B%0A                currSplitSAMWriter.close();%0A%0A                currNumRecords = 0;%0A                currFileNum++;%0A                currSplitTuple = getSplitSAMWriter(splitBaseFilename, currFileNum, true, samHeader, samWriterFactory);%0A%0A                currSplitSAMWriter = currSplitTuple._1;%0A                outputFiles.add(currSplitTuple._2);%0A            %7D%0A%0A            currSplitSAMWriter.addAlignment(samIterator.next());%0A            currNumRecords++;%0A        %7D%0A%0A        currSplitSAMWriter.close();%0A        samIterator.close();%0A        samFile.close();%0A%0A        return outputFiles;%0A    %7D%0A%0A
     
@@ -2621,587 +2621,8 @@
  %7D%0A%0A
-    public static Tuple2%3CSAMFileWriter, File%3E getSplitSAMWriter(String baseFilename, int id, boolean isSorted,%0A                                                                SAMFileHeader SAMheader,%0A                                                                SAMFileWriterFactory samFileWriterFactory) %7B%0A%0A        File newSplitFile = new File(baseFilename + %22-split-%22 + Integer.toString(id) + %22.bam%22);%0A        SAMFileWriter currSAMWriter = samFileWriterFactory.makeSAMOrBAMWriter(SAMheader, isSorted, newSplitFile);%0A        return new Tuple2%3C%3E(currSAMWriter, newSplitFile);%0A%0A
     

Remove several functions which are no longer needed.