/*
 * Copyright (c) 2010 The Broad Institute
 *
 * Permission is hereby granted, free of charge, to any person
 * obtaining a copy of this software and associated documentation
 * files (the "Software"), to deal in the Software without
 * restriction, including without limitation the rights to use,
 * copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the
 * Software is furnished to do so, subject to the following
 * conditions:
 *
 * The above copyright notice and this permission notice shall be
 * included in all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
 * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
 * OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
 * NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
 * HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
 * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
 * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR
 * THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */

package org.broadinstitute.sting.oneoffprojects.walkers;

import org.broadinstitute.sting.gatk.refdata.ReadMetaDataTracker;
import org.broadinstitute.sting.gatk.walkers.ReadWalker;
import org.broadinstitute.sting.gatk.contexts.ReferenceContext;
import org.broadinstitute.sting.commandline.Argument;
import org.broadinstitute.sting.commandline.Output;
import org.broadinstitute.sting.utils.GenomeLoc;
import org.broadinstitute.sting.utils.GenomeLocParser;
import net.sf.samtools.SAMRecord;
import net.sf.samtools.SAMFileWriter;
import org.broadinstitute.sting.utils.exceptions.UserException;

import java.io.*;

/**
 * Created by IntelliJ IDEA.
 * User: rpoplin
 * Date: Oct 23, 2009
 *
 * This walker is designed to work as the second pass in a two-pass processing step.
 * It does a by-read traversal calculating a read quality score based on a number of factors:
 *  1.) Average neighborhood quality score for the length of the read (data generated by NeighborhoodQualityWalker)
 *  2.) Is this read's mate mapped to a different chromosome?
 *  3.) The mapping quality for this read.
 *  4.) Number of reference mismatches in this read.
 * This walker creates a new bam file in which each read is annotated by this read quality score
 *  in addition if the read quality score is below the given threshold, the read is flagged.
 *
 * This walker requires as input the file of (GenomeLoc QualityScore)'s generated by NeighborhoodQualityWalker.
 * This walker accepts as input a threshold in order to flag reads which are of unacceptable read quality.
 *
 * This walker is designed to be used in conjunction with NeighborhoodQualityWalker.
 */

public class ReadQualityScoreWalker extends ReadWalker<SAMRecord, SAMFileWriter> {
    @Output
    protected PrintStream out;
    @Argument(fullName = "inputQualityFile", shortName = "if", doc = "Input quality score file generated by NeighborhoodQualityWalker", required = true)
    protected String inputQualityFile = null;
    @Argument(fullName = "outputBamFile", shortName = "of", doc = "Write output to this BAM filename instead of STDOUT", required = false)
    protected SAMFileWriter outputBamFile = null;
    @Argument(fullName = "threshold", shortName = "th", doc="Flag reads whose read quality score is below this threshold", required = false)
    protected int qualityThreshold = 13;

    private BufferedReader inputReader = null;
    private static String line = null;

    public SAMRecord map( ReferenceContext ref, SAMRecord read, ReadMetaDataTracker metaDataTracker ) {
        return read; // all the work is done in the reduce step for this walker
    }

    public SAMFileWriter reduceInit() {
        try {
            inputReader = new BufferedReader( new FileReader ( inputQualityFile ) );
        } catch ( FileNotFoundException e) {
            throw new UserException.CouldNotReadInputFile(new File(inputQualityFile), e);
		} catch (IOException e) {
            throw new UserException.CouldNotReadInputFile(new File(inputQualityFile), e);
		}
        return outputBamFile;
    }

    public SAMFileWriter reduce( SAMRecord read, SAMFileWriter output ) {

		int readQualityScore = 0;
		float meanNeighborhoodQuality = 0.0f;

        // The large block of code below is parsing through the input file and calculating the meanNeighborhoodQuality over the length of the read
        //  It does this by first skipping in the file to where the current read starts and marking that location
        //  Next it continues reading lines for the length of the read generating a sum of neighborhood quality
        //  When it reaches the end of the read it jumps back to the marker so that it can be used by the next read
        // BUGBUG: This assumes reads will be sorted by start location
        float sumNeighborhoodQuality = 0.0f;
        int numLines = 0;
        GenomeLoc readLoc = GenomeLocParser.createGenomeLoc( read );
        if( readLoc.size() > 0 ) { // only calculate mean NQS if the read has a well formed GenomeLoc, if not NQS will be zero
            try {
                if( line == null ) {
                    line = inputReader.readLine();
                    if( line == null ) { throw new UserException.MalformedFile(new File(inputQualityFile), "Input file is empty" ); }
                }
                String[] halves = line.split( " ", 2 );
                GenomeLoc curLoc = GenomeLocParser.parseGenomeLoc( halves[0] );
                while( curLoc.isBefore( readLoc ) ) { // Loop until the beginning of the read
                    line = inputReader.readLine();
                    if( line == null ) { throw new UserException.MalformedFile(new File(inputQualityFile), "Input file doesn't encompass all reads. Can't find beginning of read: " + readLoc ); }
                    halves = line.split( " ", 2 );
                    curLoc = GenomeLocParser.parseGenomeLoc( halves[0] );
                }
                // now we have skipped ahead in the input file to where this read starts
                logger.debug( "Starting: " + curLoc + ", read: " + readLoc + "\t size: " + readLoc.size() );
                inputReader.mark( 30 * ( (int)readLoc.size() + 3 ) ); // BUGBUG: Is this a sufficient buffer size?
                String savedLine = line;

                while( !curLoc.isPast( readLoc ) ) {  // Loop until just past the end of the read
                    sumNeighborhoodQuality += Float.parseFloat( halves[1] );
                    numLines++;
                    line = inputReader.readLine();
                    if( line == null ) { throw new UserException.MalformedFile(new File(inputQualityFile), "Input file doesn't encompass all reads. Can't find end of read: " + readLoc ); }
                    halves = line.split( " ", 2 );
                    curLoc = GenomeLocParser.parseGenomeLoc( halves[0] );
                }
                // now we have parsed the input file up to where the read ends
                // reset back to the mark in order to parse the next read in the next call to the reduce function
                inputReader.reset();
                line = savedLine;

            } catch ( FileNotFoundException e ) {
                throw new UserException.CouldNotReadInputFile(new File(inputQualityFile), e);
            } catch (IOException e ) {
                throw new UserException.CouldNotReadInputFile(new File(inputQualityFile), e);
            }

            meanNeighborhoodQuality = sumNeighborhoodQuality / ((float) numLines);
        }

		
        // Find out if this read's mate mapped to a different chromosome
        //boolean isGoodPair = ( read.getReadPairedFlag() ? read.getProperPairFlag() : true );
        boolean isGoodPair = ( !read.getReadPairedFlag() || read.getProperPairFlag() ); // optimized version of above line

        // Get the mapping quality for this read
        int mappingQuality = read.getMappingQuality();

        // Get the number of reference mismatches in this read
        assert read.getReadLength() > 0 : "Read length must be greater than zero.";
        float mismatchRate = 1.0f;
        if( read.getAttribute("NM") != null ) {
            mismatchRate = ((float) Integer.parseInt(read.getAttribute("NM").toString())) / ((float) read.getReadLength());
        }

		
		// Calculate the three additional metrics that go into a read quality score
		// BUGBUG: some analysis is needed to determine reasonable quality values and rates for the exponentials
		float scoreMate = ( isGoodPair ? 40.0f : 2.0f );
		float scoreMapping = 40.0f * (float) Math.exp( -0.02f * Math.max( 99.0f - mappingQuality, 0.0f ) );
									// exp decay with rate 0.02, scaled to Q=40 when mapping quality is 99
        float scoreMismatch = 40.0f * (float) Math.exp( -27.0f * mismatchRate );
									// exp decay with rate 27.0, scaled to Q=40 when the mismatch rate is 0% for this read
		
		// BUGBUG: some analysis is needed to determine reasonable weights for each metric
		readQualityScore = Math.round( 0.6f * meanNeighborhoodQuality + 0.1f * scoreMate + 0.05f * scoreMapping + 0.25f * scoreMismatch );
        if( readQualityScore == 0 ) { readQualityScore = 1; }
        assert readQualityScore > 0 : "Read quality score must be positive and nonzero.";

		// Add the read quality score to the read in the new bam file and flag it if quality is below the given threshold
        // BUGBUG: which attributes should be set here?
        read.setAttribute( "XR", readQualityScore );
        if( readQualityScore < qualityThreshold ) { 
            read.setAttribute( "ZR", 1 );
        }

        // verbose debug printing lines
        logger.debug( read.getReadName() + " " + readQualityScore );
        logger.debug( "neighborhood quality =\t" + meanNeighborhoodQuality );
        logger.debug( "mate mismatch? =\t" + isGoodPair + " --> " + scoreMate );
        logger.debug( "mapping quality =\t" + mappingQuality + " --> " + scoreMapping );
        logger.debug( "ref mismatch rate =\t" + mismatchRate + " --> " + scoreMismatch );

        // This printout useful for making histograms of scores in Matlab
        //out.println( readQualityScore + " " + meanNeighborhoodQuality + " " + scoreMate + " " + scoreMapping + " " + scoreMismatch );

        // Add the read to the output bam file or output to STDOUT
        if ( output != null ) {
            output.addAlignment( read );
        } else {
            out.println( read.format() );
        }

        return output;
    }

    public void onTraversalDone( SAMFileWriter reduceResult ) {
    }
    
}
