package edu.ohsu.sonmezsysbio.cloudbreak.command;

import com.beust.jcommander.Parameter;
import com.beust.jcommander.Parameters;
import com.google.common.base.Joiner;
import edu.ohsu.sonmezsysbio.cloudbreak.Cloudbreak;
import edu.ohsu.sonmezsysbio.cloudbreak.io.HDFSWriter;
import net.sf.samtools.SAMFileReader;
import net.sf.samtools.SAMRecord;
import net.sf.samtools.SAMRecordIterator;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FSDataOutputStream;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.io.SequenceFile;
import org.apache.hadoop.io.Text;
import org.apache.hadoop.io.compress.SnappyCodec;
import org.apache.log4j.Logger;

import java.io.BufferedWriter;
import java.io.File;
import java.io.IOException;
import java.io.OutputStreamWriter;
import java.util.ArrayList;
import java.util.List;
import java.util.zip.GZIPOutputStream;

/**
 * Created by IntelliJ IDEA.
 * User: cwhelan
 * Date: 10/15/12
 * Time: 10:50 AM
 */
@Parameters(separators = "=", commandDescription = "Load paired fastq files into HDFS")
public class CommandReadSAMFileIntoHDFS implements CloudbreakCommand {

    private static org.apache.log4j.Logger logger = Logger.getLogger(CommandReadSAMFileIntoHDFS.class);

    @Parameter(names = {"--HDFSDataDir"}, required = true)
    String hdfsDataDir;

    @Parameter(names = {"--samFile"}, required = true)
    String samFile;

    @Parameter(names = {"--outFileName"})
    String outFileName = "reads.txt";

    @Parameter(names = {"--compress"})
    String compress = "snappy";


    public void run(Configuration conf) throws Exception {
        Configuration config = new Configuration();

        FileSystem hdfs = FileSystem.get(config);
        Path p = new Path(hdfsDataDir + "/" + outFileName);

        HDFSWriter writer = new HDFSWriter();
        if ("snappy".equals(compress)) {
            writer.seqFileWriter = SequenceFile.createWriter(hdfs, config, p, Text.class, Text.class, SequenceFile.CompressionType.BLOCK, new SnappyCodec());
        } else {
            FSDataOutputStream outputStream = hdfs.create(p);
            BufferedWriter bufferedWriter = null;
            if ("gzip".equals(compress)) {
                bufferedWriter = new BufferedWriter(new OutputStreamWriter(new GZIPOutputStream(outputStream)));
            } else {
                bufferedWriter = new BufferedWriter(new OutputStreamWriter(outputStream));
            }
            writer.textFileWriter = bufferedWriter;
        }
        try {
            readFile(writer, samFile);
        } finally {
            writer.close();
        }

    }
    private void readFile(HDFSWriter writer, String samFile) throws IOException {
        SAMFileReader samFileReader = new SAMFileReader(new File(samFile));
        samFileReader.setValidationStringency(SAMFileReader.ValidationStringency.LENIENT);
        SAMRecordIterator it = samFileReader.iterator();
        String currentReadName = "";
        List<String> read1Records = new ArrayList<String>();
        List<String> read2Records = new ArrayList<String>();
        Joiner readAlignmentJoiner = Joiner.on(Cloudbreak.ALIGNMENT_SEPARATOR);
        long i = 0;
        while (it.hasNext()) {
            SAMRecord samRecord = it.next();
            // todo check sorting against previous record
            String readName = samRecord.getReadName();
            if (! readName.equals(currentReadName) && ! currentReadName.equals("")) {
                logger.debug("writing " + readName);
                writeRecords(currentReadName, writer, read1Records, read2Records, readAlignmentJoiner, i);
                currentReadName = readName;
                i++;
            }

            if (currentReadName.equals("")) {
                currentReadName = readName;
            }
            if (samRecord.getReadPairedFlag() && ! samRecord.getReadUnmappedFlag()) {
                if (samRecord.getFirstOfPairFlag()) {
                    read1Records.add(samRecord.getSAMString().trim());
                } else {
                    read2Records.add(samRecord.getSAMString().trim());
                }
            }
        }
        writeRecords(currentReadName, writer, read1Records, read2Records, readAlignmentJoiner, i);
    }

    private void writeRecords(String currentReadName, HDFSWriter writer, List<String> read1Records, List<String> read2Records, Joiner readAlignmentJoiner, long i) throws IOException {
        logger.debug("r1 records: " + read1Records.size());
        logger.debug("r2 records: " + read1Records.size());
        // todo: not importing OEA mappings yet
        if (read1Records.size() > 0 && read2Records.size() > 0) {
            writer.write(new Text(currentReadName), readAlignmentJoiner.join(read1Records) + Cloudbreak.READ_SEPARATOR + readAlignmentJoiner.join(read2Records) + "\n");
        }
        read1Records.clear();
        read2Records.clear();
    }
}

@@ -2944,24 +2944,25 @@
 iterator();%0A
+%0A
         Stri
@@ -2990,1965 +2990,234 @@
 %22%22;%0A
-        List%3CString%3E read1Records = new ArrayList%3CString%3E();%0A        List%3CString%3E read2Records = new ArrayList%3CString%3E();%0A        Joiner readAlignmentJoiner = Joiner.on(Cloudbreak.ALIGNMENT_SEPARATOR);%0A        long i = 0;%0A        while (it.hasNext()) %7B%0A            SAMRecord samRecord = it.next();%0A            // todo check sorting against previous record%0A            String readName = samRecord.getReadName();%0A            if (! readName.equals(currentReadName) && ! currentReadName.equals(%22%22)) %7B%0A                logger.debug(%22writing %22 + readName);%0A                writeRecords(currentReadName, writer, read1Records, read2Records, readAlignmentJoiner, i);%0A                currentReadName = readName;%0A                i++;%0A            %7D%0A%0A            if (currentReadName.equals(%22%22)) %7B%0A                currentReadName = readName;%0A            %7D%0A            if (samRecord.getReadPairedFlag() && ! samRecord.getReadUnmappedFlag()) %7B%0A                if (samRecord.getFirstOfPairFlag()) %7B%0A                    read1Records.add(samRecord.getSAMString().trim());%0A                %7D else %7B%0A                    read2Records.add(samRecord.getSAMString().trim());%0A                %7D%0A            %7D%0A        %7D%0A        writeRecords(currentReadName, writer, read1Records, read2Records, readAlignmentJoiner, i);%0A    %7D%0A%0A    private void writeRecords(String currentReadName, HDFSWriter writer, List%3CString%3E read1Records, List%3CString%3E read2Records, Joiner readAlignmentJoiner, long i) throws IOException %7B%0A        logger.debug(%22r1 records: %22 + read1Records.size());%0A        logger.debug(%22r2 records: %22 + read1Records.size());%0A        // todo: not importing OEA mappings yet%0A        if (read1Records.size() %3E 0 && read2Records.size() %3E 0) %7B%0A            writer.write(new Text(currentReadName), readAlignmentJoiner.join(read1Records) + Cloudbreak.READ_SEPARATOR + readAlignmentJoiner.join(read2Records) + %22%5Cn%22);%0A        %7D%0A        read1Records.clear();%0A        read2Records.clear();%0A    %7D
+%0A        while (it.hasNext()) %7B%0A            SAMRecord samRecord = it.next();%0A            String readName = samRecord.getReadName();%0A            writer.write(new Text(currentReadName), samRecord.getSAMString());%0A        %7D%0A    %7D%0A
 %0A%7D%0A

make reader just dump raw SAM records into HDFS to be sorted later