#!/usr/bin/perl

use strict;

# Parse VCF files to extract haplotype dosage information into text files suitable for reading into R;
# Searches $globdir for files with extension vcf.gz.
# For each VCF it finds, it produces a corresponding haplotype dosage file with same basename.
# Thins the output by only printing every 10th line (or whatever $thin is set to)
# Outputfiles are written to $outdir
# Assumes vcf files have been generated by stitch and contain haplotype dosage information

my @id;
my $thin = @ARGV[0];
my $n = 0;
my $globdir = "/lustre/projects/MAGIC_WHEAT//5STITCH/STITCH_haplotypes/LD_pruned_vcfs/*.vcf.gz";
my @vcf = glob( $globdir ); #"/SAN/mottlab/chickpea_raw/Founders_Rice_MAGIC/STITCH/imputation-out_2/chr01_1_locus1_K8/*vcf.gz");
my $outdir = "/lustre/projects/MAGIC_WHEAT//5STITCH/STITCH_haplotypes/LD_pruned_haplotypes";

foreach my $vcf ( @vcf ) {
    my ($basename) = $vcf =~ /([^\/]+)\.vcf\.gz$/;
    my $outfile = "$outdir/$basename.hap.txt";
    my $cmd = "zcat < $vcf |";
    open(CMD, $cmd) || die "Could not $cmd\n";
    open(OUT, ">$outfile") || die "Could not write to $outfile\n";
    print "$vcf -> $outfile\n";

    while(<CMD>) {
	if (/\#CHROM\s+POS\s+ID\s+REF\s+ALT\s+QUAL\s+FILTER\s+INFO\s+FORMAT/ ){
	    chomp;
	    @id = split;
	    splice @id, 0, 9 ;
	    print OUT join( "\t", "chr","bp","K", @id). "\n";
	}
	elsif ( /PASS/) {
	    if ( ($n++)% $thin == 0 ) {
		chomp;
		my @a = split;
		my $chr = $a[0];
		my $bp = $a[1];
		my @h;
		my $F=0;
		for( my $k=9;$k<=$#a;$k++) {
		    my @d = split /[\,\:]/, $a[$k] ;
		    splice @d, 0, 5;
		    $h[$k-9] = \@d;
		    $F = $#d;
		}
		for(my $f=0;$f<=$F;$f++) {
		    print OUT join( "\t", $chr, $bp, $f+1) ;
		    foreach my $d (@h) {
			print OUT "\t$d->[$f]";
		    }
		    print OUT "\n";
		}
	    }
	}
    }
    close(OUT);
}
