#!/usr/bin/perl

use strict;
use warnings;

#
# usage
#   hapl_generator.pl WORDLEN MIN_MAF VCF_FILE FASTA_FILE
#

my $wlen = $ARGV[0];
my $minmaf = $ARGV[1];
my $vcf = $ARGV[2];
my $fasta = $ARGV[3];

# IUPAC codes
my %iupac = (
  "AC" => "M",
  "AG" => "R",
  "AT" => "W",
  "CG" => "S",
  "CT" => "Y",
  "GT" => "K",
  "ACG" => "V",
  "ACT" => "H",
  "AGT" => "D",
  "CGT" => "B",
  "ACGT" => "N"
);

# Read chromosome
my @seqs;
open (my $ifs, $fasta) or die $!;
while (my $line = <$ifs>) {
  next if (substr ($line, 0, 1) eq ">");
  chomp ($line);
  push (@seqs, $line);
}
close ($ifs);
my $seq = join ("", @seqs);
                                                                                                                        
my %htypes;

my $lastpos = 0;
my $prev = -1000;
my @codes = ();
open ($ifs, $vcf) or die $!;
while (my $line = <$ifs>) {
  next if (substr ($line, 0, 1) eq "#");
  chomp ($line);
  my @t = split ("\t", $line);
  my $chr = $t[0];
  my $pos = $t[1];
  next if ($pos <= $lastpos);
  my $id = $t[2];
  my $ref = $t[3];
  my $alt = $t[4];
  my $dsc = $t[7];
  my $maf = 0;
  @t = split (";", $dsc);
  foreach (@t) {
    my @tt = split ("=", $_);
    if ($tt[0] eq "CAF") {
      @tt = split (",", $tt[1]);
      $maf = 1 - $tt[0];
      last;
    }
  }
  #if ($maf >= $minmaf) {
  #  printf (STDOUT "%s\n", $line);
  #}
  my $first = $pos - ($wlen - 2);
  if (($first > $prev) || (scalar (@codes) >= 10)) {
    process (@codes);
    @codes = ();
  }
  my $last = $first + ($wlen - 2) + length ($ref) + ($wlen - 1) - 1;
  $lastpos = $pos;
  $prev = $last;
  if (length ($ref) > 1) {
    # Deletion
    my $left = substr ($seq, $first - 1, $wlen - 2);
    my $mid = $ref;
    my $right = substr ($seq, $pos - 1 + length ($ref), $wlen - 1);
    next if ((index ($left, "N") >= 0) || (index ($right, "N") >= 0));
    my $code = "$chr $first $last $id $mid";
    my $ht = "$left$mid$right";
    #push (@htypes, $ht);
    #$htypes{$code} = $ht;
    #printf (STDOUT "%s\n", $line);
    #printf (STDOUT "%s %s\n", $ht, $code);
    my @alt = split (",", $alt);
    foreach (@alt) {
      $mid = $_;
      #my $code = "$chr $first $last $id $mid";
      $ht = "$ht,$left$mid$right";
    }
    push (@codes, $code);
    $htypes{$code} = $ht;
    #printf (STDOUT "%s %s\n", $ht, $code);
  } else {
    my $left = substr ($seq, $first - 1, $wlen - 1);
    my $mid = $ref;
    my $right = substr ($seq, $pos - 1 + length ($ref), $wlen - 1);
    next if ((index ($left, "N") >= 0) || (index ($right, "N") >= 0));
    if (length ($alt) == 1) {
      #SNP
      my @l = ($ref, $alt);
      @l = sort (@l);
      my $c = join ("", @l);
      my $mid = $iupac{$c};
      my $code = "$chr $first $last $id $mid";
      my $ht = "$left$mid$right";
      push (@codes, $code);
      $htypes{$code} = $ht;
      #printf (STDOUT "%s\n", $line);
      #printf (STDOUT "%s %s\n", $ht, $code);
    } else {
      # Insertion and/or multi
    }
  }
}
close ($ifs);
process (@codes);

sub process {
  if (!@_) {
    return;
  }
  #printf (STDOUT "Processing\n");
  #foreach (@_) {
  #  printf (STDOUT "  %s - %s\n", $_, $htypes{$_});
  #}
  my $code = $_[0];
  my @t = split (" ", $code);
  my $first = $t[1];
  my $seq = $htypes{$code};
  @t = split (",", $seq);
  foreach (@t) {
    rec ($code, $first, $_, \@_, 1);
  }
}

sub rec {
  my $code = $_[0];
  my $first = $_[1];
  my $seq = $_[2];
  my $l = $_[3];
  my @l = @$l;
  my $pos = $_[4];
  my @t = split (" ", $code);
  my $last = $t[2];
  #printf (STDERR "%s %s %s %s %s\n", $code, $first, $seq, scalar (@l), $pos);
  if ($pos >= @l) {
    printf (STDOUT ">%s\n%s\n", $code, $seq);
  } else {
    my $scode = $l[$pos];
    my @st = split (" ", $scode);
    my $sfirst = $st[1];
    my $sseq = $htypes{$scode};
    @st = split (",", $sseq);
    foreach (@st) {
      my $stail = substr ($_, $wlen - 1);
      my $rel = $last - $sfirst + 1 - ($wlen - 1);
      my $left = substr ($seq, 0, length ($seq) - $rel);
      my $str = "$left$stail";
      rec ($code, $first, $str, $l, $pos + 1);
    }
  }
}

