Added GP field

public/biopet-tools/src/main/scala/nl/lumc/sasc/biopet/tools/ImputeToVcf.scala → public/biopet-tools/src/main/scala/nl/lumc/sasc/biopet/tools/GensToVcf.scala

@@ -3,23 +3,26 @@ package nl.lumc.sasc.biopet.tools
 import java.io.File
 import java.util
 
-import htsjdk.tribble.index.IndexCreator
+import htsjdk.variant.variantcontext.{GenotypeBuilder, Allele, VariantContextBuilder}
 import htsjdk.variant.variantcontext.writer.{Options, VariantContextWriterBuilder, AsyncVariantContextWriter}
 import htsjdk.variant.vcf._
 import nl.lumc.sasc.biopet.utils.ToolCommand
 
+import scala.collection.JavaConversions._
+
 import scala.io.Source
 
 /**
   * Created by pjvanthof on 15/03/16.
   */
-object ImputeToVcf extends ToolCommand {
+object GensToVcf extends ToolCommand {
 
   case class Args(inputGenotypes: File = null,
                   inputInfo: Option[File] = None,
                   outputVcf: File = null,
                   sampleFile: File = null,
-                  referenceFasta: Option[File] = None) extends AbstractArgs
+                  referenceFasta: Option[File] = None,
+                  contig: String = null) extends AbstractArgs
 
   class OptParser extends AbstractOptParser {
     opt[File]('g', "inputGenotypes") required () maxOccurs 1 valueName "<file>" action { (x, c) =>
@@ -37,6 +40,9 @@ object ImputeToVcf extends ToolCommand {
     opt[File]('R', "referenceFasta") maxOccurs 1 valueName "<file>" action { (x, c) =>
       c.copy(referenceFasta = Some(x))
     } text "reference fasta file"
+    opt[String]('c', "contig") required () maxOccurs 1 valueName "<file>" action { (x, c) =>
+      c.copy(contig = x)
+    } text "contig of impute file"
   }
 
   def main(args: Array[String]): Unit = {
@@ -47,7 +53,7 @@ object ImputeToVcf extends ToolCommand {
     Source.fromFile(cmdArgs.sampleFile).getLines().toArray.drop(2).map(_.split("\t").head).foreach(samples.add(_))
 
     val metaLines = new util.HashSet[VCFHeaderLine]()
-    metaLines.add(new VCFFormatHeaderLine("PL", VCFHeaderLineCount.UNBOUNDED, VCFHeaderLineType.Float, ""))
+    metaLines.add(new VCFFormatHeaderLine("GP", VCFHeaderLineCount.UNBOUNDED, VCFHeaderLineType.String, ""))
 
     //TODO: Add reference dict
     val header = new VCFHeader(metaLines, samples)
@@ -58,7 +64,30 @@ object ImputeToVcf extends ToolCommand {
       .build)
     writer.writeHeader(header)
 
+    val genotypeIt = Source.fromFile(cmdArgs.inputGenotypes).getLines()
+    val infoIt = cmdArgs.inputInfo.map(Source.fromFile(_).getLines())
+
+    for (genotypeLine <- genotypeIt) {
+      val genotypeValues = genotypeLine.split(" ")
+      val start = genotypeValues(2).toInt
+      val end = genotypeValues(3).length - 1 + start
+      val genotypes = samples.toList.zipWithIndex.map { case (sampleName, index) =>
+        new GenotypeBuilder()
+          .name(sampleName)
+          .attribute("GP", Array(genotypeValues(5 + (index * 3)), genotypeValues(5 + (index * 3) + 1), genotypeValues(5 + (index * 3) + 2)).map(_.toDouble))
+          .make()
+      }
 
+      val builder = (new VariantContextBuilder)
+        .chr(cmdArgs.contig)
+        .alleles(genotypeValues(3), genotypeValues(4))
+        .start(start)
+        .stop(end)
+        .genotypes(genotypes)
+      val id = genotypeValues(1)
+      if (id.startsWith(cmdArgs.contig + ":")) writer.add(builder.make())
+      else writer.add(builder.id(id).make())
+    }
 
     writer.close()
   }