package org.broadinstitute.sting.playground.gatk.walkers.fasta;

import org.broadinstitute.sting.gatk.contexts.AlignmentContext;
import org.broadinstitute.sting.gatk.contexts.ReferenceContext;
import org.broadinstitute.sting.gatk.refdata.*;
import org.broadinstitute.sting.gatk.walkers.*;
import org.broadinstitute.sting.utils.*;
import org.broadinstitute.sting.utils.cmdLine.Argument;

import java.util.Iterator;

// create a fasta sequence file from a reference, intervals, and rod(s) of variants
// if there are multiple variants at a site, we take the first one seen

@WalkerName("FastaAlternateReferenceMaker")
@Requires(value={DataSource.REFERENCE})
public class FastaAlternateReferenceWalker extends RefWalker<Pair<GenomeLoc, String>, Pair<GenomeLoc, String>> {

    @Argument(fullName="maskSNPs", shortName="mask", doc="print 'N' at SNP sites instead of the alternate allele", required=false)
    private Boolean MASK_SNPS = false;

	private StringBuffer sb = new StringBuffer();
    int deletionBasesRemaining = 0;

    public Pair<GenomeLoc, String> map(RefMetaDataTracker rodData, ReferenceContext ref, AlignmentContext context) {
        if ( deletionBasesRemaining > 0 ) {
            deletionBasesRemaining--;
            return new Pair<GenomeLoc, String>(context.getLocation(), "");
        }

        Iterator<ReferenceOrderedDatum> rods = rodData.getAllRods().iterator();
        while ( rods.hasNext() ) {
            ReferenceOrderedDatum rod = rods.next();
            if ( !(rod instanceof AllelicVariant) )
                continue;

            // if we have multiple variants at a locus, just take the first damn one we see for now
            AllelicVariant variant = (AllelicVariant)rod;
            if ( variant.isDeletion() ) {
                deletionBasesRemaining = variant.length();
                // delete the next n bases, not this one
                return new Pair<GenomeLoc, String>(context.getLocation(), String.valueOf(ref.getBase()));
            } else if ( variant.isInsertion() ) {
                return new Pair<GenomeLoc, String>(context.getLocation(), String.valueOf(ref.getBase()).concat(variant.getAltBasesFWD()));
            } else if ( variant.isSNP() ) {
                if ( MASK_SNPS )
                    return new Pair<GenomeLoc, String>(context.getLocation(), "N");
                else
                    return new Pair<GenomeLoc, String>(context.getLocation(), variant.getAltBasesFWD());
            }
        }

        // if we got here then we're just ref
        return new Pair<GenomeLoc, String>(context.getLocation(), String.valueOf(ref.getBase()));
	}

    public Pair<GenomeLoc, String> reduceInit() {
        return new Pair<GenomeLoc, String>(null, "");
    }

	public Pair<GenomeLoc, String> reduce(Pair<GenomeLoc, String> value, Pair<GenomeLoc, String> sum) {
        // if there is no interval to the left, then this is the first one
        if ( sum.first == null ) {
            sum.first = value.first;
            sum.second = value.second;
        }
        // if the intervals don't overlap, print out the leftmost one and start a new one
        // (end of contig or new interval)
        else if ( value.first.getStart() != sum.first.getStop() + 1 ) {
            printFasta(sum.first, sum.second);
            sum.first = value.first;
            sum.second = value.second;
        }
        // otherwise, merge them
        else {
            sum.first = GenomeLocParser.setStop(sum.first,value.first.getStop());
            sum.second = sum.second.concat(value.second);
        }
		return sum;
	}

    public void onTraversalDone(Pair<GenomeLoc, String> sum) {
        if (sum.second != null)
            printFasta(sum.first, sum.second);
    }

    private void printFasta(GenomeLoc loc, String s) {
        out.println(">" + loc);
        int lines = s.length() / 60;
        int currentStart = 0;
        for (int i=0; i < lines; i++) {
            out.println(s.substring(currentStart, currentStart+60));
            currentStart += 60;
        }
        out.println(s.substring(currentStart));
    }
}
@@ -932,16 +932,212 @@
 = false;
+%0A    @Argument(fullName=%22outputSequenomFormat%22, shortName=%22sequenom%22, doc=%22output results in sequenom format (overrides 'maskSNPs' argument)%22, required=false)%0A    private Boolean SEQUENOM = false;
 %0A%0A%09priva
@@ -1327,24 +1327,81 @@
  context) %7B%0A
+        String refBase = String.valueOf(ref.getBase());%0A%0A
         if (
@@ -1539,20 +1539,97 @@
 ation(),
+ (SEQUENOM ? (deletionBasesRemaining == 0 ? refBase.concat(%22%5D%22) : refBase) :
  %22%22)
+)
 ;%0A      
@@ -2274,36 +2274,51 @@
 (), 
-String.valueOf(ref.get
+(SEQUENOM ? refBase.concat(%22%5B-/%22) : ref
 Base
-()
 ));%0A
@@ -2445,37 +2445,80 @@
 (), 
-String.valueOf(ref
+(SEQUENOM ? refBase.concat(%22%5B+/%22+variant
 .get
+Alt
 Base
-())
+sFWD()+%22%5D%22) : refBase
 .con
@@ -2547,16 +2547,17 @@
 sFWD()))
+)
 ;%0A      
@@ -2598,45 +2598,8 @@
 ) %7B%0A
-                if ( MASK_SNPS )%0A    
     
@@ -2672,112 +2672,38 @@
 (), 
-%22N%22);%0A                else%0A                    return new Pair%3CGenomeLoc, String%3E(context.getLocation(),
+(SEQUENOM %7C%7C MASK_SNPS ? %22N%22 :
  var
@@ -2724,16 +2724,17 @@
 esFWD())
+)
 ;%0A      
@@ -2868,37 +2868,15 @@
 (), 
-String.valueOf(ref.get
+ref
 Base
-())
 );%0A%09

option to output in sequenom input format