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

import org.broad.tribble.util.variantcontext.Allele;
import org.broad.tribble.util.variantcontext.Genotype;
import org.broad.tribble.util.variantcontext.VariantContext;
import org.broad.tribble.vcf.VCFConstants;
import org.broadinstitute.sting.commandline.Argument;
import org.broadinstitute.sting.gatk.contexts.AlignmentContext;
import org.broadinstitute.sting.gatk.contexts.ReferenceContext;
import org.broadinstitute.sting.gatk.contexts.variantcontext.VariantContextUtils;
import org.broadinstitute.sting.gatk.refdata.RefMetaDataTracker;
import org.broadinstitute.sting.gatk.walkers.RodWalker;
import org.broadinstitute.sting.gatk.walkers.RMD;
import org.broadinstitute.sting.gatk.walkers.Requires;
import org.broadinstitute.sting.utils.GenomeLoc;
import org.broadinstitute.sting.utils.MathUtils;
import org.broadinstitute.sting.utils.SampleUtils;

import java.io.PrintStream;
import java.util.*;

/**
 * Produces an input file to Beagle imputation engine, listing genotype likelihoods for each sample in input variant file
 */
@Requires(value={},referenceMetaData=@RMD(name=ProduceBeagleInputWalker.ROD_NAME, type=VariantContext.class))
public class ProduceBeagleInputWalker extends RodWalker<Integer, Integer> {

    public static final String ROD_NAME = "variant";

    @Argument(fullName = "beagle_file", shortName = "beagle", doc = "File to print BEAGLE-specific data for use with imputation", required = true)
    public PrintStream  beagleWriter = null;
    @Argument(fullName = "genotypes_file", shortName = "genotypes", doc = "File to print reference genotypes for error analysis", required = false)
    public PrintStream beagleGenotypesWriter  = null;
    @Argument(fullName = "inserted_nocall_rate", shortName = "nc_rate", doc = "Rate (0-1) at which genotype no-calls will be randomly inserted, for testing", required = false)
    public double insertedNoCallRate  = 0;

    private Set<String> samples = null;
    private Random generator;

    public void initialize() {
        generator = new Random();

        samples = SampleUtils.getSampleListWithVCFHeader(getToolkit(), Arrays.asList(ROD_NAME));

        beagleWriter.print("marker alleleA alleleB");
        for ( String sample : samples )
            beagleWriter.print(String.format(" %s %s %s", sample, sample, sample));

        beagleWriter.println();
        if (beagleGenotypesWriter != null)
            beagleGenotypesWriter.println("dummy header");
    }

    public Integer map( RefMetaDataTracker tracker, ReferenceContext ref, AlignmentContext context ) {

         if( tracker != null ) {
            GenomeLoc loc = context.getLocation();
            VariantContext vc_eval;

            vc_eval = tracker.getVariantContext(ref, ROD_NAME, null, loc, false);
            if ( vc_eval == null || vc_eval.isFiltered() )
                return 0;

             if (!vc_eval.isSNP())
                     return 0;

             if (!vc_eval.isBiallelic())
                     return 0;
             
            // output marker ID to Beagle input file
            beagleWriter.print(String.format("%s ", VariantContextUtils.getLocation(vc_eval).toString()));

            if (beagleGenotypesWriter != null)
                beagleGenotypesWriter.print(String.format("%s ", VariantContextUtils.getLocation(vc_eval).toString()));

            for (Allele allele: vc_eval.getAlleles()) {
                // TODO -- check whether this is really needed by Beagle
                String bglPrintString;
                if (allele.isNoCall() || allele.isNull())
                    bglPrintString = "0";
                else
                    bglPrintString = allele.toString().substring(0,1);  // get rid of * in case of reference allele

                beagleWriter.print(String.format("%s ", bglPrintString));
            }

            if ( !vc_eval.hasGenotypes() )
                return 0;

            Map<String, Genotype> genotypes = vc_eval.getGenotypes();
            for ( String sample : samples ) {
                // use sample as key into genotypes structure
                Genotype genotype = genotypes.get(sample);
                if (genotype.isCalled() && genotype.hasAttribute(VCFConstants.GENOTYPE_LIKELIHOODS_KEY)) {
                    String[] glArray = genotype.getAttributeAsString(VCFConstants.GENOTYPE_LIKELIHOODS_KEY).split(",");

                    double[] likeArray = new double[glArray.length];

                    // convert to double array so we can normalize
                    int k=0;
                    for (String gl : glArray) {
                        likeArray[k++] = Double.valueOf(gl);
                    }

                    double[] normalizedLikelihoods = MathUtils.normalizeFromLog10(likeArray);
                    // see if we need to randomly mask out genotype in this position.
                    Double d = generator.nextDouble();
                    if (d > insertedNoCallRate ) {
//                    System.out.format("%5.4f ", d);
                        for (Double likeVal: normalizedLikelihoods)
                            beagleWriter.print(String.format("%5.4f ",likeVal));
                    }
                    else {
                        // we are masking out this genotype
                        beagleWriter.print("0.33 0.33 0.33 ");
                    }

                    if (beagleGenotypesWriter != null) {
                        char a = genotype.getAllele(0).toString().charAt(0);
                        char b = genotype.getAllele(0).toString().charAt(0);

                        beagleGenotypesWriter.format("%c %c ", a, b);
                    }
                }
                else if (genotype.isCalled() && !genotype.hasAttribute(VCFConstants.GENOTYPE_LIKELIHOODS_KEY)) {
                    // hack to deal with input VCFs with no genotype likelihoods.  Just assume the called genotype
                    // is confident.  This is useful for Hapmap and 1KG release VCFs.
                    double AA = 0.02;
                    double AB = 0.02;
                    double BB = 0.02;

                    if (genotype.isHomRef()) { AA = 0.96; }
                    else if (genotype.isHet()) { AB = 0.96; }
                    else if (genotype.isHomVar()) { BB = 0.96; }

                    beagleWriter.printf("%.2f %.2f %.2f ", AA, AB, BB);

                    if (beagleGenotypesWriter != null) {
                        char a = genotype.getAllele(0).toString().charAt(0);
                        char b = genotype.getAllele(0).toString().charAt(0);

                        beagleGenotypesWriter.format("%c %c ", a, b);
                    }
                }
                else  {
                    beagleWriter.print("0.33 0.33 0.33 "); // write 1/3 likelihoods for uncalled genotypes.
                    if (beagleGenotypesWriter != null)
                         beagleGenotypesWriter.print(". . ");
                }
            }

            beagleWriter.println();
            if (beagleGenotypesWriter != null)
                 beagleGenotypesWriter.println();
        }
        return 1;
        
    }

    public Integer reduceInit() {
        return 0; // Nothing to do here
    }

    public Integer reduce( Integer value, Integer sum ) {
        return 0; // Nothing to do here
    }

    public void onTraversalDone( Integer sum ) {

   }
    
}
