/*
 * 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.playground.gatk.walkers.vcftools;

import org.broad.tribble.vcf.*;
import org.broadinstitute.sting.gatk.contexts.AlignmentContext;
import org.broadinstitute.sting.gatk.contexts.ReferenceContext;
import org.broadinstitute.sting.gatk.refdata.RefMetaDataTracker;
import org.broadinstitute.sting.gatk.refdata.utils.GATKFeature;
import org.broadinstitute.sting.gatk.walkers.RodWalker;
import org.broadinstitute.sting.commandline.Argument;
import org.broadinstitute.sting.utils.genotype.vcf.*;
import org.broadinstitute.sting.utils.genotype.vcf.VCFWriter;

import java.io.File;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.Set;

/**
 * Extracts subsets of a VCF file like one or more samples, all or only variant loci, all or filtered loci.
 */
public class VCFSubsetWalker extends RodWalker<ArrayList<VCFRecord>, VCFWriter> {
    @Argument(fullName="sample", shortName="SN", doc="Sample to include (or nothing to specify all samples)", required=false)
    private HashSet<String> SAMPLES;

    @Argument(fullName="vcfsubset", shortName="O", doc="File to write VCF subset to", required=false)
    private File VPATH = null;

    @Argument(fullName="includeNonVariants", shortName="INV", doc="Include non-variant loci", required=false)
    private boolean INCLUDE_NON_VARIANTS = false;

    @Argument(fullName="includeFiltered", shortName="IF", doc="Include filtered loci", required=false)
    private boolean INCLUDE_FILTERED = false;

    private VCFHeader vheader = null;
    private VCFWriter vwriter = null;

    public void initializeWriter() {
        Set<VCFHeaderLine> metaData = new HashSet<VCFHeaderLine>();
        metaData.add(new VCFHeaderLine("source", "VariantsToVCF"));
        metaData.add(new VCFHeaderLine("reference", this.getToolkit().getArguments().referenceFile.getAbsolutePath()));

        Set<String> additionalColumns = new HashSet<String>();
        additionalColumns.add("FORMAT");
        additionalColumns.addAll(SAMPLES);

        vheader = new VCFHeader(metaData, additionalColumns);
        if (VPATH != null) {
            vwriter = new VCFWriter(VPATH);
            vwriter.writeHeader(vheader);
        }
    }

    public ArrayList<VCFRecord> map(RefMetaDataTracker tracker, ReferenceContext ref, AlignmentContext context) {
        ArrayList<VCFRecord> records = new ArrayList<VCFRecord>();

        if (tracker != null) {
            for (GATKFeature feature : tracker.getAllRods()) {
                Object rod = feature.getUnderlyingObject();
                if (rod instanceof VCFRecord) {
                    VCFRecord vcfrod = (VCFRecord) rod;

                    if (SAMPLES == null) {
                        SAMPLES = new HashSet<String>();
                        SAMPLES.addAll(vcfrod.getHeader().getGenotypeSamples());
                    }

                    if (VPATH != null && vwriter == null) {
                        initializeWriter();
                    }

                    //out.println(record.toStringEncoding(vcfrod.getHeader()));

                    records.add(vcfrod);
                }
            }
        }
        
        return records;
    }

    public VCFWriter reduceInit() {
        return vwriter;
    }

    private VCFRecord subsetRecord(VCFRecord record) {
        ArrayList<VCFGenotypeRecord> genotypeRecords = new ArrayList<VCFGenotypeRecord>();
        HashSet<VCFGenotypeEncoding> genotypeEncodingSet = new HashSet<VCFGenotypeEncoding>();

        for ( VCFGenotypeRecord gr : record.getVCFGenotypeRecords() ) {
            if (SAMPLES.contains(gr.getSampleName())) {
                genotypeRecords.add(gr);

                for (VCFGenotypeEncoding allele : gr.getAlleles()) {
                    if (!allele.getBases().equalsIgnoreCase(record.getReference())) {
                        genotypeEncodingSet.add(allele);
                    }
                }
            }
        }

        ArrayList<VCFGenotypeEncoding> genotypeEncodings = new ArrayList<VCFGenotypeEncoding>();
        for (VCFGenotypeEncoding allele : genotypeEncodingSet) {
            genotypeEncodings.add(allele);
        }

        VCFRecord subset = new VCFRecord(record.getReference(),
                                         record.getChr(),
                                         record.getStart(),
                                         record.getID(),
                                         genotypeEncodings,
                                         record.getQual(),
                                         record.getFilterString(),
                                         record.getInfoValues(),
                                         record.getGenotypeFormatString(),
                                         genotypeRecords);

        return subset;
    }

    public VCFWriter reduce(ArrayList<VCFRecord> records, VCFWriter writer) {
        for (VCFRecord record : records) {
            VCFRecord subset = subsetRecord(record);

            boolean isVariant = false;

            if (subset.getVCFGenotypeRecords().size() > 0) {
                for ( VCFGenotypeEncoding ge : subset.getVCFGenotypeRecords().get(0).getAlleles() ) {
                    if (!record.getReference().equals(ge.getBases())) {
                        isVariant = true;
                    }
                }
            }

            if ((isVariant || INCLUDE_NON_VARIANTS) && (!subset.isFiltered() || INCLUDE_FILTERED)) {
                if (vwriter != null) {
                    vwriter.addRecord(subset);
                } else {
                    out.println(subset.toStringEncoding(vheader));
                }
            }
        }

        return writer;
    }

    public void onTraversalDone(VCFWriter writer) {
        if (vwriter != null) {
            vwriter.close();
        }
    }
}

@@ -5219,24 +5219,88 @@
 odingSet) %7B%0A
+            if (!allele.getBases().equalsIgnoreCase(%22.%22)) %7B%0A    
             
@@ -5322,32 +5322,46 @@
 gs.add(allele);%0A
+            %7D%0A
         %7D%0A%0A     

Don't list '.' as an alternate allele, dummy!